ci: run smoke tests across a python x torch version matrix (apple#32)

guru-desh · web-flow · commit 651345549045 · 2026-07-20T15:56:07.000-07:00
* ci: run smoke tests across a python x torch version matrix Smoke tests previously ran on one implicit torch version (whatever pyproject.toml's unpinned bound resolved to) across all supported python versions. Add torch_2_8/2_9/2_10/2_11 dependency groups (build-matched torchao/torchvision per pytorch/ao issue 2919) and a matrix.torch_group axis on the test-smoke CI job so each of the 4 torch minor versions is smoke tested independently, still covering all 3 python versions per job via the existing nox session. lowest_tested_torch/highest_tested_torch are renamed to torch_2_8/ torch_2_11 for a consistent naming scheme; the linux-tests/macos-tests full-suite jobs and their Makefile targets are unchanged. * fix: allow digits in AVAILABLE_GROUPS regex for torch_2_* group names setup_env.sh's group discovery used grep -E '^[a-z_-]+ = \[', which excludes digits and made torch_2_8/2_9/2_10/2_11 invisible to validate_groups — breaking test-highest-pytorch/test-lowest-pytorch with 'Unknown dependency group', caught by the fork CI run. * feat: make TORCH_GROUP the single source of truth for torch pinning setup_env.sh now reads TORCH_GROUP directly and folds it into every sync it runs, so make env/test/test-fast/docs/etc. all pin torch instead of floating on whatever pyproject.toml's plain bound resolves to. test-highest-pytorch/test-lowest-pytorch/env-highest-torch lock their group via an inline `export TORCH_GROUP=... &&` at the top of the recipe (rather than a hardcoded --with-torch_2_X flag), so their name stays a promise regardless of a command-line override — verified that GNU Make's own `override` directive achieves the same lock but trips up the mbake Makefile formatter's duplicate-target detector, which doesn't recognize that syntax. --with-torch_2_X stays as a guarded alternate path for direct, non-Make invocations of setup_env.sh; a guard now errors clearly if it disagrees with an explicitly-set TORCH_GROUP instead of surfacing a raw uv conflicting-groups resolver error. * refactor: read TORCH_GROUP directly in smoke tests, drop SMOKE_TEST_TORCH_GROUP Now that the Makefile exports TORCH_GROUP unconditionally, the smoke-test-specific SMOKE_TEST_TORCH_GROUP env var was redundant indirection. ci/nox/noxfile.py reads TORCH_GROUP directly, and `make test-smoke` no longer needs to translate the Makefile variable into a differently-named env var for the nox subprocess. * refactor(setup_env): drop bespoke TORCH_GROUP conflict guard The custom bash check that detected TORCH_GROUP disagreeing with an explicit --with-torch_2_X flag duplicated logic uv already provides: `uv sync --group torch_2_11 --group torch_2_8` fails on its own with a clear "Groups ... are incompatible with the conflicts: ..." error and non-zero exit. Simplifies the --all-groups exclusion loop and the plain --group append accordingly, removing the GROUP_REQUESTED/ TORCH_GROUP_ALREADY_SYNCED bookkeeping that only existed to feed it. * refactor(nox): drop unused default for TORCH_GROUP The Makefile always exports TORCH_GROUP before invoking nox, so the Python-side default was dead code for every make test-smoke run; kept only for a bare, non-Make nox invocation which isn't a supported entrypoint. * docs(Makefile): trim TORCH_GROUP comment The expanded rationale duplicated what's now self-evident from reading the three lock-target recipes directly. * fix(Makefile): narrow TORCH_GROUP override to the setup_env.sh call only `export TORCH_GROUP=torch_2_X && ...` set the variable for the whole recipe's shell, which happened to be harmless since nothing else in these three recipes reads TORCH_GROUP, but the intent was always to scope the override to just the one setup_env.sh invocation that needs it. Using `TORCH_GROUP=torch_2_X <command>` as an inline prefix does that precisely: it applies only to the first command in the `&&` chain, reverting to the Makefile's own default/command-line value for everything after. * refactor(Makefile): alias highest/lowest torch groups for maintainability Add HIGHEST_TORCH_GROUP/LOWEST_TORCH_GROUP variables and route the 4 places that hardcoded torch_2_11/torch_2_8 as a stand-in for "the extreme" (TORCH_GROUP's default, env-highest-torch, test-highest- pytorch, test-lowest-pytorch) through them instead. Addresses PR review feedback: bumping the project's torch version bounds now means changing 2 lines instead of manually finding every hardcoded reference. apple#32 (comment)
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -38,11 +38,23 @@ jobs:
         run: make check
 
   test-smoke:
-    name: Linux / make test-smoke
+    name: Linux / make test-smoke / torch=${{ matrix.torch_version }}
     runs-on: ubuntu-latest
     timeout-minutes: 60
     env:
       INSTALL_PRECOMMIT: "false"
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - torch_group: torch_2_8
+            torch_version: "2.8"
+          - torch_group: torch_2_9
+            torch_version: "2.9"
+          - torch_group: torch_2_10
+            torch_version: "2.10"
+          - torch_group: torch_2_11
+            torch_version: "2.11"
     steps:
       - name: Check out repository
         uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
@@ -51,7 +63,7 @@ jobs:
         with:
           enable-cache: false
       - name: Run `make test-smoke`
-        run: make test-smoke
+        run: make test-smoke TORCH_GROUP=${{ matrix.torch_group }}
 
   test-tutorials:
     name: Linux / make test-tutorials
diff --git a/Makefile b/Makefile
@@ -53,6 +53,17 @@ VENV_HIGHEST_TORCH ?= .venv-highest-torch
 VENV_LOWEST_TORCH ?= .venv-lowest-torch
 VENV_TUTORIAL ?= .venv-tutorial
 
+# The torch_2_* groups (pyproject.toml [dependency-groups]) currently at
+# each end of the supported range. Bump these two lines — nothing else —
+# when the project's torch version bounds change.
+HIGHEST_TORCH_GROUP := torch_2_11
+LOWEST_TORCH_GROUP := torch_2_8
+
+# Torch dependency group (pyproject.toml [dependency-groups]) that every
+# environment-building target (env, test, test-smoke, docs, ...) pins to.
+TORCH_GROUP ?= $(HIGHEST_TORCH_GROUP)
+export TORCH_GROUP
+
 # Documentation directory. Defaults to $(MAKEFILE_DIR)docs so the same recipe
 # works in both contexts:
 #
@@ -176,7 +187,7 @@ env: _maybe_patch_pyproject
 
 # Set up development environment with latest supported PyTorch version
 env-highest-torch: _maybe_patch_pyproject
-	@$(SETUP_ENV) --venv $(VENV_HIGHEST_TORCH) --python-version $(PYTHON_VERSION) --with-highest_tested_torch
+	@TORCH_GROUP=$(HIGHEST_TORCH_GROUP) $(SETUP_ENV) --venv $(VENV_HIGHEST_TORCH) --python-version $(PYTHON_VERSION)
 	@$(call write_active_venv,$(VENV_HIGHEST_TORCH))
 
 # Set up environment for running tutorials (quantization notebook)
@@ -235,6 +246,7 @@ test-slow:
 	@$(MAKE) test PYTEST_ARGS="--marker slow"
 
 # Run smoke tests only (pass PYTEST_ARGS for custom flags, e.g., make test-smoke PYTEST_ARGS="--junitxml=results.xml").
+# Pass TORCH_GROUP to smoke test against a specific torch version (default: HIGHEST_TORCH_GROUP).
 test-smoke:
 	@$(call use_env,VENV) && \
 	echo "Running smoke tests..." && \
@@ -244,20 +256,20 @@ test-smoke:
 # Run tests on lowest supported PyTorch version (pass PYTEST_ARGS for custom flags)
 test-lowest-pytorch:
 	@echo "Running tests on lowest PyTorch version supported..."
-	@$(call use_env,VENV_LOWEST_TORCH,--with-lowest_tested_torch) && \
-		echo "Testing with lowest supported PyTorch versions" && \
-		uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
-		$(RUN_TESTS) $(PYTEST_ARGS) && \
-		echo "All tests passed!"
+	@TORCH_GROUP=$(LOWEST_TORCH_GROUP) $(call use_env,VENV_LOWEST_TORCH) && \
+	echo "Testing with lowest supported PyTorch versions" && \
+	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
+	$(RUN_TESTS) $(PYTEST_ARGS) && \
+	echo "All tests passed!"
 
 # Run tests on highest supported PyTorch version (pass PYTEST_ARGS for custom flags)
 test-highest-pytorch:
 	@echo "Running tests on highest PyTorch version supported..."
-	@$(call use_env,VENV_HIGHEST_TORCH,--with-highest_tested_torch) && \
-		echo "Testing with latest supported PyTorch versions" && \
-		uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
-		$(RUN_TESTS) $(PYTEST_ARGS) && \
-		echo "All tests passed!"
+	@TORCH_GROUP=$(HIGHEST_TORCH_GROUP) $(call use_env,VENV_HIGHEST_TORCH) && \
+	echo "Testing with latest supported PyTorch versions" && \
+	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
+	$(RUN_TESTS) $(PYTEST_ARGS) && \
+	echo "All tests passed!"
 
 # Run tutorial notebook tests
 test-tutorials:
diff --git a/ci/nox/noxfile.py b/ci/nox/noxfile.py
@@ -34,8 +34,14 @@
 options.default_venv_backend = "uv"
 options.error_on_missing_interpreters = True
 
+TORCH_GROUP = os.environ.get("TORCH_GROUP")
 
-@session(python=get_supported_python_versions(), uv_extras=["coreai"], uv_groups=["test"])
+
+@session(
+    python=get_supported_python_versions(),
+    uv_extras=["coreai"],
+    uv_groups=["test", TORCH_GROUP],
+)
 def smoke_tests(session: Session) -> None:
     """Smoke test the package build and coreai_opt imports and basic functionality.
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -46,7 +46,7 @@ dependencies = [
     # based on other package versions. We can either 1) add a stricter check to require the newer torchao
     # version for everyone, or 2) add a runtime check in src/coreai_opt/init.py to ensure torchao
     # version >= 0.15.0 for torch version >= 2.9.0. Opting option 1.
-    # These torch versions must be in bounds of lowest_tested_torch and highest_tested_torch
+    # These torch versions must be in bounds of torch_2_8, torch_2_9, torch_2_10, and torch_2_11
     "torch>=2.8.0,<=2.11.0",
     "torchao>=0.15.0,<=0.17.0",
     "tqdm>=4.65",
@@ -107,20 +107,6 @@ docs = [
 # This applies to the `coreml` group as well.
 coreai = [ "coreai-opt[coreai]" ]
 coreml = [ "coreai-opt[coreml]" ]
-# Used in CI to force latest mimimum supported torch version
-# These torch versions must be in bounds of torch versions listed in project dependencies
-highest_tested_torch = [
-    "torch==2.11.0",
-    "torchao==0.17.0",
-    "torchvision==0.26.0",
-]
-# Used in CI to force lowest mimimum supported torch version
-# These torch versions must be in bounds of torch versions listed in project dependencies
-lowest_tested_torch = [
-    "torch==2.8.0",
-    "torchao==0.15.0",
-    "torchvision==0.23.0",
-]
 pre-commit = [
     "bashate>=2.1.1",
     "darker>=3.0.0",
@@ -144,6 +130,27 @@ pre-commit = [
     "tomli-w>=1.0.0",
 ]
 tamm-export = []
+# These torch versions must be in bounds of torch versions listed in project dependencies
+torch_2_10 = [
+    "torch==2.10.0",
+    "torchao==0.16.0",
+    "torchvision==0.25.0",
+]
+torch_2_11 = [
+    "torch==2.11.0",
+    "torchao==0.17.0",
+    "torchvision==0.26.0",
+]
+torch_2_8 = [
+    "torch==2.8.0",
+    "torchao==0.15.0",
+    "torchvision==0.23.0",
+]
+torch_2_9 = [
+    "torch==2.9.1",
+    "torchao==0.15.0",
+    "torchvision==0.24.1",
+]
 # Since torch and torchao are project dependencies, we need to include torchvision in dev
 # This allows standard `make test` to find torchvision
 torchvision = [
@@ -190,8 +197,10 @@ exclude-newer-package = { coreai-core = false, coreai-torch = false }
 # Declare conflicting groups so uv does not error
 conflicts = [
     [
-        { group = "highest_tested_torch" },
-        { group = "lowest_tested_torch" },
+        { group = "torch_2_8" },
+        { group = "torch_2_9" },
+        { group = "torch_2_10" },
+        { group = "torch_2_11" },
     ],
 ]
 [tool.uv.sources]
diff --git a/scripts/make/setup_env.sh b/scripts/make/setup_env.sh
@@ -59,8 +59,8 @@ fi
 #    - in_section && /^\[/: If we're in the section AND hit another section header (line starting with [)
 #      * exit: Stop processing (we've left the dependency-groups section)
 #    - in_section {print}: If we're in the section, print the line
-# 2. grep -E '^[a-z_-]+ = \[': Filter to lines that define groups
-#    - ^[a-z_-]+: Group name at start of line (lowercase letters, hyphens, underscores)
+# 2. grep -E '^[a-z0-9_-]+ = \[': Filter to lines that define groups
+#    - ^[a-z0-9_-]+: Group name at start of line (lowercase letters, digits, hyphens, underscores)
 #    - = \[: Followed by space, equals sign, space, opening bracket
 # 3. cut -d' ' -f1: Extract just the group name
 #    - -d' ': Use space as delimiter
@@ -75,7 +75,7 @@ AVAILABLE_GROUPS=$(
         in_section && /^\[/ { exit }
         in_section { print }
     ' "$PYPROJECT_TOML" |
-        grep -E '^[a-z_-]+ = \[' |
+        grep -E '^[a-z0-9_-]+ = \[' |
         cut -d' ' -f1 |
         tr '\n' ' '
 )
@@ -109,7 +109,7 @@ ENSURE_MODE=false
 
 # Groups excluded from --all-groups due to mutual conflicts in pyproject.toml.
 # tamm-export is omitted because it's opt-in only (never in default-groups or --all-groups).
-CONFLICTING_GROUPS=("highest_tested_torch" "lowest_tested_torch")
+CONFLICTING_GROUPS=("torch_2_8" "torch_2_9" "torch_2_10" "torch_2_11")
 
 show_help() {
     echo "Usage: $0 [OPTIONS]"
@@ -127,6 +127,7 @@ show_help() {
     echo ""
     echo "Environment variables:"
     echo "  VENV                Virtual environment name, overrides --venv (default: .venv)"
+    echo "  TORCH_GROUP         Torch dependency group to pin; source of truth for torch pinning"
     echo ""
     echo "Available dependency groups: $AVAILABLE_GROUPS"
     echo "Conflicting groups (excluded from --all-groups): ${CONFLICTING_GROUPS[*]}"
@@ -135,8 +136,8 @@ show_help() {
     echo "  $0 --python-version 3.11                               # Setup with dev group only"
     echo "  $0 --python-version 3.11 --with-docs                   # Setup with dev and docs groups"
     echo "  $0 --python-version 3.11 --all-groups                                    # Setup with all non-conflicting groups"
-    echo "  $0 --python-version 3.11 --all-groups --with-highest_tested_torch        # Setup with all groups and highest torch"
-    echo "  $0 --python-version 3.11 --all-groups --with-lowest_tested_torch         # Setup with all groups and lowest torch"
+    echo "  TORCH_GROUP=torch_2_11 $0 --python-version 3.11 --all-groups             # Setup with all groups and torch 2.11"
+    echo "  TORCH_GROUP=torch_2_8 $0 --python-version 3.11 --all-groups              # Setup with all groups and torch 2.8"
     echo "  $0 --python-version 3.11 --venv .venv-exp              # Setup with custom venv name"
     echo "  $0 --python-version 3.11 --with-docs --venv .venv-exp  # Setup with docs group and custom venv name"
     echo "  $0 --python-version 3.12                               # Setup with Python 3.12"
@@ -222,6 +223,7 @@ validate_groups() {
 # Validate dependency group names early
 [[ ${#EXTRA_GROUPS[@]} -gt 0 ]] && validate_groups "${EXTRA_GROUPS[@]}"
 [[ ${#EXCLUDE_GROUPS[@]} -gt 0 ]] && validate_groups "${EXCLUDE_GROUPS[@]}"
+[[ -n "${TORCH_GROUP:-}" ]] && validate_groups "$TORCH_GROUP"
 
 # Check for conflicts between --with-<group> and --without-<group>
 if [[ ${#EXTRA_GROUPS[@]} -gt 0 && ${#EXCLUDE_GROUPS[@]} -gt 0 ]]; then
@@ -267,24 +269,35 @@ if [[ "$ENSURE_MODE" == "true" ]] && [ -f "$VENV/bin/python" ]; then
         for GROUP in "${EXTRA_GROUPS[@]}"; do
             case "$GROUP" in
             docs) IMPORT_STMTS+="; import sphinx" ;;
-            highest_tested_torch | lowest_tested_torch)
-                IMPORT_STMTS+="; import torchao"
-                EXPECTED_TORCH="$(group_torch_pin "$GROUP")"
-                # These groups always pin torch, so an empty result means the
-                # pyproject parse regressed — fail loudly instead of silently
-                # skipping the version check (which would reintroduce the bug).
-                if [[ -z "$EXPECTED_TORCH" ]]; then
-                    echo "Error: could not parse a torch pin for group '$GROUP' in $PYPROJECT_TOML" >&2
-                    exit 1
-                fi
-                IMPORT_STMTS+="; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH'"
-                ;;
             rio) IMPORT_STMTS+="; import turi_lightning" ;;
             tamm-export) IMPORT_STMTS+="; import tamm_export" ;;
             esac
         done
     fi
 
+    # TORCH_GROUP is the single source of truth for torch pinning; fall back
+    # to a --with-torch_2_* flag for direct, non-Make invocations. Checked
+    # unconditionally rather than only when present in EXTRA_GROUPS, since
+    # TORCH_GROUP drives the sync regardless of what's in EXTRA_GROUPS.
+    EXPECTED_GROUP="${TORCH_GROUP:-}"
+    if [[ -z "$EXPECTED_GROUP" ]]; then
+        for GROUP in "${EXTRA_GROUPS[@]:-}"; do
+            [[ " ${CONFLICTING_GROUPS[*]} " == *" ${GROUP} "* ]] && EXPECTED_GROUP="$GROUP"
+        done
+    fi
+    if [[ -n "$EXPECTED_GROUP" ]]; then
+        IMPORT_STMTS+="; import torchao"
+        EXPECTED_TORCH="$(group_torch_pin "$EXPECTED_GROUP")"
+        # These groups always pin torch, so an empty result means the
+        # pyproject parse regressed — fail loudly instead of silently
+        # skipping the version check (which would reintroduce the bug).
+        if [[ -z "$EXPECTED_TORCH" ]]; then
+            echo "Error: could not parse a torch pin for '$EXPECTED_GROUP' in $PYPROJECT_TOML" >&2
+            exit 1
+        fi
+        IMPORT_STMTS+="; import torch; assert torch.__version__.split('+')[0] == '$EXPECTED_TORCH'"
+    fi
+
     if "$VENV/bin/python" -c "$IMPORT_STMTS" 2>/dev/null; then
         exit 0
     fi
@@ -324,9 +337,13 @@ echo "[2/3] Installing dependencies..."
 SYNC_CMD=(uv sync --active)
 if [[ "$ALL_GROUPS" == "true" ]]; then
     SYNC_CMD+=(--all-groups)
-    # Exclude conflicting groups unless explicitly requested via --with-*
+    # Exclude conflicting torch groups other than the one TORCH_GROUP or
+    # --with-<group> selects, so --all-groups doesn't try to sync every
+    # torch_2_* group at once. If both TORCH_GROUP and --with-<other-group>
+    # name different groups, neither gets excluded here and `uv sync` itself
+    # rejects the combination via its own conflicting-groups resolution.
     for GROUP in "${CONFLICTING_GROUPS[@]}"; do
-        if [[ ! " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]]; then
+        if [[ "$GROUP" != "${TORCH_GROUP:-}" ]] && [[ ! " ${EXTRA_GROUPS[*]:-} " == *" ${GROUP} "* ]]; then
             SYNC_CMD+=(--no-group "$GROUP")
         fi
     done
@@ -335,6 +352,14 @@ elif [[ ${#EXTRA_GROUPS[@]} -gt 0 ]]; then
         SYNC_CMD+=(--group "$GROUP")
     done
 fi
+# TORCH_GROUP is the single source of truth for torch pinning: append it
+# unconditionally (already covered under --all-groups by not being
+# excluded above). A disagreeing --with-<other-torch-group> ends up as a
+# second, different --group flag here, which `uv sync` itself rejects via
+# its conflicting-groups resolution rather than a bespoke check.
+if [[ -n "${TORCH_GROUP:-}" && "$ALL_GROUPS" != "true" ]]; then
+    SYNC_CMD+=(--group "$TORCH_GROUP")
+fi
 # Apply explicit group exclusions (e.g., --without-coreai)
 if [[ ${#EXCLUDE_GROUPS[@]} -gt 0 ]]; then
     for GROUP in "${EXCLUDE_GROUPS[@]}"; do
