Build nightly shark-ai -package (#535)

Refactors the scripts so that a version suffix can be passed as an
argument. This is used to generate the suffix only once in the workflow
instead of computing it in multiple scripts and partly scattered across
different steps (`write_requirements.py` runs in a different step).
Further adds dependencies to IREE packages also for nightly shark-ai
builds.

Requires #519.

.github/workflows/build_packages.yml

@@ -40,15 +40,19 @@ jobs:
       - name: Generate release candidate versions
         id: version_rc
         run: |
-          sharktank_package_version=$(python3 build_tools/python_deploy/compute_local_version.py sharktank)
-          shortfin_package_version=$(python3 build_tools/python_deploy/compute_local_version.py shortfin)
+          version_suffix="$(printf 'rc%(%Y%m%d)T')"
+          echo "version_suffix=${version_suffix}" >> $GITHUB_ENV
+          sharktank_package_version=$(python3 build_tools/python_deploy/compute_local_version.py --version-suffix=${version_suffix} sharktank)
+          shortfin_package_version=$(python3 build_tools/python_deploy/compute_local_version.py --version-suffix=${version_suffix} shortfin)
+          sharkai_package_version=$(python3 build_tools/python_deploy/compute_common_version.py -rc --version-suffix=${version_suffix} --write-json)
       - name: Upload version_local.json
         uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
         with:
           name: version_local
           path: |
             sharktank/version_local.json
             shortfin/version_local.json
+            shark-ai/version_local.json
 
   build_packages:
     name: "${{ matrix.package }} :: ${{ matrix.platform }} :: ${{ matrix.python-version }}"
@@ -61,6 +65,10 @@ jobs:
       matrix:
         include:
           # Ubuntu packages.
+          - runs-on: ubuntu-24.04
+            platform: linux-x86_64
+            package: shark-ai
+            python-version: cp311-cp311  # Ignored (generic wheel), set for workflow naming
           - runs-on: ubuntu-24.04
             platform: linux-x86_64
             package: sharktank
@@ -103,6 +111,15 @@ jobs:
         path: ./c/
         merge-multiple: true
 
+    - name: Build shark-ai (Linux x86_64)
+      if: "matrix.package == 'shark-ai' && matrix.platform == 'linux-x86_64'"
+      env:
+        OUTPUT_DIR: "${{ github.workspace }}/bindist"
+      run: |
+        [ -e ./bindist/* ] && rm ./bindist/*
+        ./c/build_tools/python_deploy/write_requirements.py --version-suffix=${version_suffix}
+        ./c/shark-ai/build_tools/build_linux_package.sh
+
     - name: Build sharktank (Linux x86_64)
       if: "matrix.package == 'sharktank' && matrix.platform == 'linux-x86_64'"
       env:

build_tools/python_deploy/compute_common_version.py

@@ -24,6 +24,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--write-json", action="store_true")
+parser.add_argument("--version-suffix", action="store", type=str)
 
 release_type = parser.add_mutually_exclusive_group()
 release_type.add_argument("-stable", "--stable-release", action="store_true")  # default
@@ -37,6 +38,10 @@
     sys.stderr.write("error: A release type is required\n")
     sys.exit(1)
 
+if args.stable_release and args.version_suffix:
+    sys.stderr.write("error: A version suffix is only supported for stable releases\n")
+    sys.exit(1)
+
 THIS_DIR = Path(__file__).parent.resolve()
 REPO_ROOT = THIS_DIR.parent.parent
 
@@ -70,7 +75,12 @@ def write_version_info():
     COMMON_VERSION = SHORTFIN_BASE_VERSION
 
 if args.nightly_release:
-    COMMON_VERSION += "rc" + datetime.today().strftime("%Y%m%d")
+    if args.version_suffix:
+        VERSION_SUFFIX = args.version_suffix
+    else:
+        VERSION_SUFFIX = "rc" + datetime.today().strftime("%Y%m%d")
+
+    COMMON_VERSION += VERSION_SUFFIX
 
 if args.write_json:
     version_local = {"package-version": COMMON_VERSION}

build_tools/python_deploy/compute_local_version.py

@@ -19,6 +19,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("path", type=Path)
+parser.add_argument("--version-suffix", action="store", type=str)
 args = parser.parse_args()
 
 VERSION_FILE = args.path / "version.json"
@@ -38,11 +39,14 @@ def write_version_info():
 
 version_info = load_version_info()
 
+if args.version_suffix:
+    VERSION_SUFFIX = args.version_suffix
+else:
+    VERSION_SUFFIX = "rc" + datetime.today().strftime("%Y%m%d")
+
 PACKAGE_VERSION = version_info.get("package-version")
 PACKAGE_BASE_VERSION = Version(PACKAGE_VERSION).base_version
-PACKAGE_LOCAL_VERSION = (
-    PACKAGE_BASE_VERSION + "rc" + datetime.today().strftime("%Y%m%d")
-)
+PACKAGE_LOCAL_VERSION = PACKAGE_BASE_VERSION + VERSION_SUFFIX
 
 version_local = {"package-version": PACKAGE_LOCAL_VERSION}
 

build_tools/python_deploy/write_requirements.py

@@ -6,10 +6,11 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # This script writes the `packaging/shark-ai/requirements.txt` file and pins
-# the versions of the dependencies accordingly. For nighly releases,
+# the versions of the dependencies accordingly. For nightly releases,
 #  * sharktank
 #  * shortfin
-# get pinned to the corresponding nighly version. For stable releases,
+# get pinned to the corresponding nightly version. The IREE packages are
+# unpinned. For stable releases,
 # * iree-base-compiler
 # * iree-base-runtime
 # * iree-turbine
@@ -52,27 +53,28 @@ def write_requirements(requirements):
 metapackage_version = load_version_info(VERSION_FILE_LOCAL)
 PACKAGE_VERSION = metapackage_version.get("package-version")
 
-sharktank_version = load_version_info(VERSION_FILE_SHARKTANK)
-SHARKTANK_PACKAGE_VERSION = sharktank_version.get("package-version")
+# sharktank_version = load_version_info(VERSION_FILE_SHARKTANK)
+# SHARKTANK_PACKAGE_VERSION = sharktank_version.get("package-version")
 
 shortfin_version = load_version_info(VERSION_FILE_SHORTFIN)
 SHORTFIN_PACKAGE_VERSION = shortfin_version.get("package-version")
 
 stable_packages_list = ["iree-base-compiler", "iree-base-runtime", "iree-turbine"]
 
 if Version(PACKAGE_VERSION).is_prerelease:
+    requirements = ""
+    for package in stable_packages_list:
+        requirements += package + "\n"
     # TODO: Include sharktank as a dependencies of future releases
     # requirements = (
     #     "sharktank=="
     #     + Version(SHARKTANK_PACKAGE_VERSION).base_version
-    #     + "rc"
     #     + args.version_suffix
     #     + "\n"
     # )
     requirements += (
         "shortfin=="
         + Version(SHORTFIN_PACKAGE_VERSION).base_version
-        + "rc"
         + args.version_suffix
     )