Harden distributed optimizer integration and release gates

.github/workflows/ci.yml

@@ -31,11 +31,11 @@ jobs:
     name: Lint (ruff)
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
         with:
           persist-credentials: false
 
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: "3.12"
           cache: pip
@@ -52,11 +52,11 @@ jobs:
     name: Build sdist + wheel
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
         with:
           persist-credentials: false
 
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: "3.12"
           cache: pip
@@ -72,27 +72,43 @@ jobs:
         run: python -m twine check dist/*
 
   cpu:
-    name: CPU tests (py${{ matrix.python-version }})
+    name: CPU tests (py${{ matrix.python-version }}, ${{ matrix.torch-label }})
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        include:
+          - python-version: "3.10"
+            torch-spec: "torch==2.5.0"
+            torch-label: "torch floor 2.5.0"
+          - python-version: "3.11"
+            torch-spec: "torch"
+            torch-label: "latest torch"
+          - python-version: "3.12"
+            torch-spec: "torch"
+            torch-label: "latest torch"
+          - python-version: "3.13"
+            torch-spec: "torch"
+            torch-label: "latest torch"
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
 
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
         with:
           python-version: ${{ matrix.python-version }}
           cache: pip
 
       - name: Install CPU PyTorch + package
+        env:
+          TORCH_SPEC: ${{ matrix.torch-spec }}
         run: |
           python -m pip install --upgrade pip
           # CPU-only torch wheel: the CUDA wheel is ~2GB and pointless here.
-          pip install torch --index-url https://download.pytorch.org/whl/cpu
+          pip install "$TORCH_SPEC" --index-url https://download.pytorch.org/whl/cpu
           # torch is already satisfied, so this only pulls numba/numpy/ninja + gefen.
-          pip install .[test]
+          pip install '.[test]'
 
       - name: Byte-compile all sources
         run: python -m compileall -q src/gefen
@@ -124,12 +140,14 @@ jobs:
     if: github.event_name == 'workflow_dispatch' && inputs.run_gpu_tests
     runs-on: [self-hosted, gpu]
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
 
       - name: Install package (CUDA torch expected on the runner)
         run: |
           python -m pip install --upgrade pip
-          pip install .[test]
+          pip install '.[test]'
 
       - name: Run full test suite (CUDA)
         # Multi-GPU / FSDP2 cases skip themselves when the runner has one GPU.

.github/workflows/release.yml

@@ -5,8 +5,16 @@ name: Release
 #   git tag v0.2.1.dev1 && git push origin v0.2.1.dev1   -> build + TestPyPI only
 #   git tag v0.2.1       && git push origin v0.2.1        -> build + TestPyPI + PyPI
 #
-# The SAME artifacts built once in `build` are promoted through both indexes, so
-# what lands on PyPI is byte-identical to what you smoke-tested on TestPyPI.
+# `ci.yml` intentionally runs on branch pushes and pull requests, not tag pushes.
+# This workflow therefore carries installed-wheel CPU and Transformers Trainer
+# gates. Every gate tests the artifact built in `build`, and every publish job
+# downloads that same artifact.
+#
+# The two-GPU CUDA/JIT/distributed gate runs locally, not on a hosted runner:
+# `scripts/release_gpu_gate.sh <tag>` downloads this run's `dist` artifact and
+# runs the mandatory GPU test list against the installed wheel with zero skips
+# allowed. Approving the `testpypi` environment below is the release manager's
+# attestation that the local GPU gate passed for this exact artifact.
 #
 # Auth is OIDC trusted publishing (no API tokens stored). The manual approval
 # gates are GitHub Environment "required reviewers", configured in
@@ -33,6 +41,7 @@ jobs:
   build:
     name: Build & verify artifacts
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     outputs:
       version: ${{ steps.ver.outputs.version }}
       prerelease: ${{ steps.ver.outputs.prerelease }}
@@ -49,11 +58,36 @@ jobs:
           # could taint the wheel uploaded to PyPI. Build tooling installs fast.
 
       - name: Install build tooling
-        run: python -m pip install --upgrade pip build twine
+        run: |
+          python -m pip install 'pip==26.1.2'
+          python -m pip install 'build==1.5.1' 'twine==6.2.0'
+
+      - name: Pin artifact timestamps to the tagged commit
+        run: echo "SOURCE_DATE_EPOCH=$(git log -1 --format=%ct)" >> "$GITHUB_ENV"
 
       - name: Build sdist + wheel
         # Pure-Python build (kernels JIT at runtime), so no CUDA toolchain needed.
-        run: python -m build
+        run: |
+          normalize_sdist() {
+            local SDIST SCRATCH TOPDIR
+            SDIST=$(find dist -maxdepth 1 -type f -name '*.tar.gz' -print -quit)
+            SCRATCH=$(mktemp -d)
+            tar -xzf "$SDIST" -C "$SCRATCH"
+            TOPDIR=$(find "$SCRATCH" -mindepth 1 -maxdepth 1 -type d -printf '%f\n')
+            test -n "$TOPDIR"
+            tar --sort=name --mtime="@${SOURCE_DATE_EPOCH}" --owner=0 --group=0 --numeric-owner --format=posix --pax-option=delete=atime,delete=ctime -C "$SCRATCH" -cf - "$TOPDIR" | gzip -n > "${SDIST}.normalized"
+            mv "${SDIST}.normalized" "$SDIST"
+            rm -rf "$SCRATCH"
+          }
+          build_once() {
+            rm -rf dist
+            python -m build
+            normalize_sdist
+          }
+          build_once
+          sha256sum dist/* | sort -k2 > /tmp/gefen-dist.sha256
+          build_once
+          sha256sum -c /tmp/gefen-dist.sha256
 
       - name: twine check
         run: python -m twine check dist/*
@@ -72,24 +106,178 @@ jobs:
             exit 1
           fi
           echo "version=$PKG_VER" >> "$GITHUB_OUTPUT"
-          # PEP 440 prerelease markers (.devN / aN / bN / rcN) => TestPyPI only.
-          if echo "$PKG_VER" | grep -Eq '(\.dev|a|b|rc)[0-9]+$'; then
-            echo "prerelease=true" >> "$GITHUB_OUTPUT"
-          else
-            echo "prerelease=false" >> "$GITHUB_OUTPUT"
-          fi
+          # Any PEP 440 prerelease/dev release stops after TestPyPI.
+          PRERELEASE=$(python -c 'from packaging.version import Version; import sys; print(str(Version(sys.argv[1]).is_prerelease).lower())' "$PKG_VER")
+          echo "prerelease=$PRERELEASE" >> "$GITHUB_OUTPUT"
 
       - name: Upload artifacts
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
         with:
           name: dist
           path: dist/
 
+  cpu_wheel_tests:
+    name: Installed-wheel CPU tests (${{ matrix.torch-label }})
+    needs: build
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - python-version: "3.10"
+            torch-spec: "torch==2.5.0"
+            torch-label: "torch floor 2.5.0"
+          - python-version: "3.12"
+            torch-spec: "torch"
+            torch-label: "latest torch"
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
+        with:
+          name: dist
+          path: dist/
+
+      - name: Install the built wheel
+        env:
+          TORCH_SPEC: ${{ matrix.torch-spec }}
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install "$TORCH_SPEC" --index-url https://download.pytorch.org/whl/cpu
+          WHEEL=$(ls dist/*.whl)
+          python -m pip install "$WHEEL" pytest
+
+      - name: Verify installed metadata and packaged JIT resources
+        env:
+          EXPECTED_VERSION: ${{ needs.build.outputs.version }}
+        run: |
+          python - <<'PY'
+          import importlib.metadata as metadata
+          import importlib.resources as resources
+          import os
+          from pathlib import Path
+
+          import gefen
+
+          version = metadata.version("gefen-x")
+          assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"])
+          module_path = Path(gefen.__file__).resolve()
+          checkout_src = (Path.cwd() / "src").resolve()
+          assert checkout_src not in module_path.parents, module_path
+          for symbol in ("Gefen", "GefenMuon", "GefenMuonHybrid"):
+              assert hasattr(gefen, symbol), symbol
+
+          kernel_root = resources.files("gefen.kernels")
+          expected_sources = (
+              "automatic_gefen_fused_binding.cpp",
+              "automatic_gefen_fused_kernel.cu",
+              "automatic_vmean_binding.cpp",
+              "automatic_vmean_kernel.cu",
+              "exact_histogram_fused_binding.cpp",
+              "exact_histogram_fused_kernel.cu",
+              "period_variance_binding.cpp",
+              "period_variance_kernel.cu",
+          )
+          missing = [name for name in expected_sources if not kernel_root.joinpath(name).is_file()]
+          assert not missing, missing
+          assert resources.files("gefen").joinpath("py.typed").is_file()
+          requirements = metadata.requires("gefen-x") or []
+          normalized_requirements = [requirement.partition(";")[0].strip().lower() for requirement in requirements]
+          assert any(requirement == "ninja" for requirement in normalized_requirements)
+          assert any(requirement.startswith("setuptools>=") for requirement in normalized_requirements)
+          print("wheel OK:", module_path, version, len(expected_sources), "JIT sources")
+          PY
+          ninja --version
+
+      - name: Run CPU test suite against the wheel
+        env:
+          CUDA_VISIBLE_DEVICES: ""
+        run: python -m pytest tests -q -ra
+  framework_wheel_tests:
+    name: Installed-wheel Transformers Trainer resume
+    needs: build
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+      - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
+        with:
+          persist-credentials: false
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
+        with:
+          name: dist
+          path: dist/
+
+      - name: Install wheel and pinned framework surface
+        env:
+          EXPECTED_VERSION: ${{ needs.build.outputs.version }}
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          WHEEL=$(ls dist/*.whl)
+          python -m pip install "$WHEEL" pytest 'transformers==5.5.0' 'accelerate==1.14.0'
+          python - <<'PY'
+          import importlib.metadata as metadata
+          import os
+          from pathlib import Path
+
+          import gefen
+
+          version = metadata.version("gefen-x")
+          assert version == os.environ["EXPECTED_VERSION"], (version, os.environ["EXPECTED_VERSION"])
+          module_path = Path(gefen.__file__).resolve()
+          assert (Path.cwd() / "src").resolve() not in module_path.parents, module_path
+          print("installed wheel:", module_path, version)
+          PY
+
+      - name: Run Trainer checkpoint-continuation matrix
+        env:
+          CUDA_VISIBLE_DEVICES: ""
+        run: python -m pytest -q -ra --junitxml=release-framework.xml tests/test_transformers_trainer_resume.py
+
+      - name: Reject skipped framework-gate tests
+        run: |
+          python - <<'PY'
+          import xml.etree.ElementTree as ET
+
+          root = ET.parse("release-framework.xml").getroot()
+          cases = root.findall(".//testcase")
+          assert cases, "framework gate collected no tests"
+          skipped = [
+              "{}::{}".format(case.attrib.get("classname", ""), case.attrib.get("name", ""))
+              for case in cases
+              if case.find("skipped") is not None
+          ]
+          assert not skipped, "mandatory framework tests skipped:\n{}".format("\n".join(skipped))
+          print("mandatory framework tests:", len(cases), "passed with zero skips")
+          PY
+
+      - name: Upload framework release-gate report
+        if: always()
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+        with:
+          name: release-framework-junit-${{ github.run_attempt }}
+          path: release-framework.xml
+          if-no-files-found: ignore
+
   testpypi:
     name: Publish to TestPyPI (gate 1)
-    needs: build
+    needs: [build, cpu_wheel_tests, framework_wheel_tests]
     runs-on: ubuntu-latest
     # Approval gate #1: the `testpypi` environment's required reviewers.
+    # Approving attests that `scripts/release_gpu_gate.sh <tag>` passed locally
+    # against this run's built wheel (see the header comment).
     environment:
       name: testpypi
       url: https://test.pypi.org/project/gefen-x/${{ needs.build.outputs.version }}/
@@ -104,13 +292,10 @@ jobs:
         uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0
         with:
           repository-url: https://test.pypi.org/legacy/
-          # A real release also passes through here; if that version was already
-          # tested on TestPyPI, don't hard-fail on the duplicate.
-          skip-existing: true
 
   pypi:
     name: Publish to PyPI (gate 2)
-    needs: [build, testpypi]
+    needs: [build, cpu_wheel_tests, framework_wheel_tests, testpypi]
     # Prerelease tags stop at TestPyPI; only clean vX.Y.Z tags reach PyPI.
     if: needs.build.outputs.prerelease == 'false'
     runs-on: ubuntu-latest

.gitignore

@@ -18,6 +18,7 @@ benchmarks/data/
 # Generated optimizer-matrix results, logs, and checkpoints
 benchmarks/training_matrix/out/
 benchmarks/training_matrix/RESULTS_*.md
+benchmarks/trainer_resume/out/
 
 # Toy fine-tuning example outputs (generated data + saved checkpoints)
 examples/toy-finetune/data/