From 57aca9799da64775cd899cc67e544424ec339432 Mon Sep 17 00:00:00 2001
From: Scott Todd <scott.todd0@gmail.com>
Date: Mon, 9 Dec 2024 09:05:15 -0800
Subject: [PATCH] Switch workflows to use new nightly iree-turbine packages.
 (#657)

This simplification will help with
https://github.com/nod-ai/shark-ai/issues/584.

Nightly releases of iree-turbine are now being built thanks to
https://github.com/iree-org/iree-turbine/pull/314 and published at the
same index as the other IREE packages thanks to
https://github.com/iree-org/iree/pull/19391.
---
 .github/workflows/ci-llama-large-tests.yaml   | 12 +++--
 .github/workflows/ci-llama-quick-tests.yaml   | 12 +++--
 .github/workflows/ci-sglang-benchmark.yml     |  9 ++--
 .../workflows/ci-sglang-integration-tests.yml |  5 +--
 .github/workflows/ci-shark-ai.yml             | 12 ++---
 .github/workflows/ci-sharktank.yml            | 45 +++++++++----------
 .github/workflows/ci_eval.yaml                | 19 ++++----
 .github/workflows/ci_eval_short.yaml          | 12 ++---
 docs/nightly_releases.md                      | 19 +++-----
 9 files changed, 59 insertions(+), 86 deletions(-)

diff --git a/.github/workflows/ci-llama-large-tests.yaml b/.github/workflows/ci-llama-large-tests.yaml
index 376f93938..379a5bac3 100644
--- a/.github/workflows/ci-llama-large-tests.yaml
+++ b/.github/workflows/ci-llama-large-tests.yaml
@@ -59,14 +59,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-turbine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Test with nightly releases, not what iree-turbine uses.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          # Install nightly IREE packages.
+          # We could also pin to a known working or stable version.
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
           pip freeze
 
diff --git a/.github/workflows/ci-llama-quick-tests.yaml b/.github/workflows/ci-llama-quick-tests.yaml
index ddbcc204a..418dbea2e 100644
--- a/.github/workflows/ci-llama-quick-tests.yaml
+++ b/.github/workflows/ci-llama-quick-tests.yaml
@@ -59,14 +59,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-turbine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Test with nightly releases, not what iree-turbine uses.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          # Install nightly IREE packages.
+          # We could also pin to a known working or stable version.
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
           pip freeze
 
diff --git a/.github/workflows/ci-sglang-benchmark.yml b/.github/workflows/ci-sglang-benchmark.yml
index d53189483..c549fc25a 100644
--- a/.github/workflows/ci-sglang-benchmark.yml
+++ b/.github/workflows/ci-sglang-benchmark.yml
@@ -65,14 +65,11 @@ jobs:
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --no-compile -r pytorch-cpu-requirements.txt
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
+          pip install -f https://iree.dev/pip-release-links.html --pre iree-turbine
           pip install --no-compile -r requirements.txt -e sharktank/ shortfin/
 
-          # Try with the latest nightly releases, not what iree-turbine pins.
-          # We could also pin to a known working or stable version.
-          # This should eventually stabilize. Do the best we can for now.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade \
+          # Pin to known-working versions.
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler==3.1.0rc20241204 \
             iree-base-runtime==3.1.0rc20241204 \
             "numpy<2.0"
diff --git a/.github/workflows/ci-sglang-integration-tests.yml b/.github/workflows/ci-sglang-integration-tests.yml
index 36a59779a..c3627226e 100644
--- a/.github/workflows/ci-sglang-integration-tests.yml
+++ b/.github/workflows/ci-sglang-integration-tests.yml
@@ -54,15 +54,14 @@ jobs:
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --no-compile -r pytorch-cpu-requirements.txt
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
           pip install --no-compile -r requirements.txt -e sharktank/ shortfin/
 
           # Use newest possible releases to be able to track commits that may
           # cause errors.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
             iree-base-runtime \
+            iree-turbine \
             "numpy<2.0"
 
           # Install SGLang and sentence_transformers
diff --git a/.github/workflows/ci-shark-ai.yml b/.github/workflows/ci-shark-ai.yml
index e662d125b..0e8ae0c87 100644
--- a/.github/workflows/ci-shark-ai.yml
+++ b/.github/workflows/ci-shark-ai.yml
@@ -55,16 +55,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -e sharktank/ shortfin/
 
-          # Install latest iree-tubrine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Try with the latest IREE nightly releases, not what iree-turbine pins.
+          # Install nightly IREE packages.
           # We could also pin to a known working or stable version.
-          # This should eventually stabilize. Do the best we can for now.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
           pip freeze
 
diff --git a/.github/workflows/ci-sharktank.yml b/.github/workflows/ci-sharktank.yml
index fc3d50bb6..4fdc5775a 100644
--- a/.github/workflows/ci-sharktank.yml
+++ b/.github/workflows/ci-sharktank.yml
@@ -60,10 +60,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Update to the latest iree packages.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
-            iree-base-compiler iree-base-runtime --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
+          # Install nightly IREE packages.
+          # We could also pin to a known working or stable version.
+          pip install -f https://iree.dev/pip-release-links.html --pre \
+            iree-base-compiler \
+            iree-base-runtime \
+            iree-turbine
 
       - name: Run sharktank tests
         if: ${{ !cancelled() }}
@@ -83,7 +85,7 @@ jobs:
       run:
         shell: bash
     env:
-      PIP_CACHE_DIR: "${{ github.workspace }}/.pip-cache"
+      VENV_DIR: ${{ github.workspace }}/.venv
       HF_HOME: "/data/huggingface"
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -94,15 +96,12 @@ jobs:
         with:
           python-version: ${{matrix.version}}
 
-      - name: Cache Pip Packages
-        uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2
-        id: cache-pip
-        with:
-          path: ${{ env.PIP_CACHE_DIR }}
-          key: pip-${{ steps.setup_python.outputs.python-version }}-${{ hashFiles('*requirements*.txt','sharktank/requirements*.txt') }}
+      - name: Create Python venv
+        run: python -m venv ${VENV_DIR}
 
       - name: Install sharktank deps
         run: |
+          source ${VENV_DIR}/bin/activate
           python -m pip install --no-compile --upgrade pip
           # Note: We install in three steps in order to satisfy requirements
           # from non default locations first. Installing the PyTorch CPU
@@ -110,21 +109,18 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-tubrine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Try with the latest IREE nightly releases, not what iree-turbine pins.
+          # Install nightly IREE packages.
           # We could also pin to a known working or stable version.
-          # This should eventually stabilize. Do the best we can for now.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
       - name: Run tests
       # TODO: unify with-t5-data and with-clip-data flags into a single flag
       # and make it possible to run only tests that require data.
         run: |
+          source ${VENV_DIR}/bin/activate
           pytest \
           --with-clip-data \
             --with-t5-data \
@@ -162,10 +158,13 @@ jobs:
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
-          # Update to the latest iree packages.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
-            iree-base-compiler iree-base-runtime --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
+
+          # Install nightly IREE packages.
+          # We could also pin to a known working or stable version.
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
+            iree-base-compiler \
+            iree-base-runtime \
+            iree-turbine
       - name: Run punet tests
         run: |
           pytest -v sharktank/ -m punet_quick \
diff --git a/.github/workflows/ci_eval.yaml b/.github/workflows/ci_eval.yaml
index 565bf3352..c41709f38 100644
--- a/.github/workflows/ci_eval.yaml
+++ b/.github/workflows/ci_eval.yaml
@@ -57,16 +57,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-tubrine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Try with the latest IREE nightly releases, not what iree-turbine pins.
+          # Install nightly IREE packages.
           # We could also pin to a known working or stable version.
-          # This should eventually stabilize. Do the best we can for now.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          pip install -f https://iree.dev/pip-release-links.html --pre \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
           pip freeze
 
@@ -119,9 +115,10 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-tubrine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
+          # Install nightly iree-turbine.
+          # We could also pin to a known working or stable version.
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
+            iree-turbine
 
       - name: Run perplexity test with Torch
         run: |
diff --git a/.github/workflows/ci_eval_short.yaml b/.github/workflows/ci_eval_short.yaml
index 6331e0709..fd3b5e5d7 100644
--- a/.github/workflows/ci_eval_short.yaml
+++ b/.github/workflows/ci_eval_short.yaml
@@ -56,16 +56,12 @@ jobs:
           pip install --no-compile -r pytorch-cpu-requirements.txt
           pip install --no-compile -r requirements.txt -r sharktank/requirements-tests.txt -e sharktank/
 
-          # Install latest iree-tubrine.
-          pip install --no-compile -f https://iree.dev/pip-release-links.html --src deps \
-            -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-
-          # Try with the latest IREE nightly releases, not what iree-turbine pins.
+          # Install nightly IREE packages.
           # We could also pin to a known working or stable version.
-          # This should eventually stabilize. Do the best we can for now.
-          pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
+          pip install -f https://iree.dev/pip-release-links.html --pre --upgrade \
             iree-base-compiler \
-            iree-base-runtime
+            iree-base-runtime \
+            iree-turbine
 
           pip freeze
 
diff --git a/docs/nightly_releases.md b/docs/nightly_releases.md
index f41374445..a7c182cd4 100644
--- a/docs/nightly_releases.md
+++ b/docs/nightly_releases.md
@@ -86,26 +86,19 @@ deactivate
 
 ## Installing newer versions of dependencies
 
-To install the `iree-turbine` package from the latest source:
-
-```bash
-pip install --src deps \
-  -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
-```
-
-To install the `iree-base-compiler` and `iree-base-runtime` packages from
-nightly releases:
+To install all IREE packages from nightly releases:
 
 ```bash
 pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
-  iree-base-compiler iree-base-runtime
+  iree-base-compiler \
+  iree-base-runtime \
+  iree-turbine
 ```
 
-To install all three packages together:
+To install an editable `iree-turbine` package from the latest source:
 
 ```bash
-pip install -f https://iree.dev/pip-release-links.html --upgrade --pre \
-  iree-base-compiler iree-base-runtime --src deps \
+pip install --src deps \
   -e "git+https://github.com/iree-org/iree-turbine.git#egg=iree-turbine"
 ```