Drop non-local (GPU/CUDA and +local) pins from generated artifacts

.github/scripts/envgen.py

@@ -20,6 +20,15 @@
     packages dropped     be pip-installed locally or that ship vendored inside
                          setuptools (see DROP / DROP_PREFIX). py4j is kept; pyspark
                          is dropped so DB Connect supplies its own bundled build.
+  * Local segments     - a PEP 440 local version segment (``+cu129`` / ``+cpu`` /
+    stripped             ``+db1``) is stripped and the base release pinned
+                         (``torch 2.9.0+cu129`` -> ``torch~=2.9.0``); ``~=`` is
+                         invalid with a local segment and the build exists only
+                         off-index, while its base is an ordinary PyPI release. See ``req``.
+  * GPU-only dropped   - ``nvidia-*`` / ``cuda-*`` CUDA wheels and tooling, plus
+                         triton, flash-attn, deepspeed, horovod and pynvml — all need
+                         a GPU (and CUDA/MPI toolchain) a dev machine lacks (see
+                         DROP / DROP_PREFIX).
   * requires-python    - taken from the runtime's Python version (major.minor).
 
 This module is imported by ``sync.py`` (the weekly discovery + reconciliation Action).
@@ -30,11 +39,24 @@
 # libs, the spark client, pip itself, and deps vendored inside setuptools.
 DROP = {
     "pyspark", "dbus-python", "pygobject", "pip", "unattended-upgrades",
+    # Ubuntu system packages carried in the image but not pip-installable. Named
+    # explicitly so they stay dropped even though req() now strips local segments
+    # (their +ubuntu / +build local builds would otherwise resurrect as a base pin).
+    "python-apt", "distro-info",
     # setuptools-vendored
     "autocommand", "inflect", "typeguard", "backports-tarfile",
     "importlib-resources", "more-itertools",
+    # GPU-only: need an NVIDIA GPU + CUDA toolchain a dev machine does not have.
+    # (nvidia-* / cuda-* are dropped by prefix below.) horovod is the same class — a
+    # source-only distribution needing MPI/NCCL + a compiler to build. pynvml is the
+    # NVML binding (same as nvidia-ml-py); useless without a driver.
+    "triton", "flash-attn", "deepspeed", "horovod", "pynvml",
 }
-DROP_PREFIX = ("jaraco-",)        # jaraco.collections / jaraco.context / ...
+DROP_PREFIX = (
+    "jaraco-",        # jaraco.collections / jaraco.context / ... (setuptools-vendored)
+    "nvidia-",        # nvidia-* CUDA wheels (and nvidia-ml-py): GPU tooling, no local use
+    "cuda-",          # cuda-toolkit / cuda-bindings / cuda-pathfinder: CUDA tooling, GPU-only
+)
 
 
 def norm(name):
@@ -44,10 +66,15 @@ def norm(name):
 
 
 def req(name, version):
-    """Render one requirement. Compatible-release ``~=`` allows patch bumps, but it
-    is invalid with a local version segment (PEP 440), and a local build like
-    ``+cpu`` / ``+cu118`` / ``+db1`` is exactly what distinguishes CPU vs GPU ML
-    images and Databricks-patched packages — so those are pinned exactly with ``==``.
+    """Render one requirement. Compatible-release ``~=`` allows patch bumps.
+
+    A PEP 440 local version segment (``+cpu`` / ``+cu118`` / ``+db1``) is stripped:
+    ``~=`` is invalid with a local segment, and the segment names a build published
+    only off-index (``download.pytorch.org``) or rebuilt inside the image, while its
+    base release is an ordinary PyPI version. So the base is pinned instead
+    (``torch 2.9.0+cu129`` -> ``torch~=2.9.0``, ``flask 1.1.2+db1`` -> ``flask~=1.1.2``).
+    Ubuntu system builds like ``python-apt 2.7.7+ubuntu5.2`` would strip to a bogus
+    PyPI pin, so those are dropped by name in ``DROP`` before reaching here.
 
     ``databricks-sdk`` is a special case. It moves in lockstep with
     ``databricks-connect``, which is installed from PyPI in the dev group and declares
@@ -60,8 +87,7 @@ def req(name, version):
     floor while letting databricks-connect's own metadata govern the exact version
     within that window. See issue #16.
     """
-    if "+" in version:
-        return f"{name}=={version}"
+    version = version.split("+", 1)[0]
     if name == "databricks-sdk":
         version = ".".join(version.split(".")[:2])
     return f"{name}~={version}"
@@ -81,6 +107,9 @@ def parse_requirements(text):
 
 
 def _filtered(pkgs):
+    # Inclusion only: drop the name-based DROP set and DROP_PREFIX. A PEP 440 local
+    # version segment is NOT a reason to drop — its base release is on PyPI — so
+    # those pins are kept here and req() strips the segment when rendering.
     return {n: v for n, v in pkgs.items()
             if n not in DROP and not n.startswith(DROP_PREFIX)}
 

.github/scripts/test_envgen.py

@@ -0,0 +1,107 @@
+"""Unit tests for envgen.py (run: python -m unittest test_envgen)."""
+import unittest
+
+from envgen import _filtered, build_constraints, build_pyproject, dbconnect_pin, req
+
+
+class FilterTest(unittest.TestCase):
+    def test_keeps_local_version_packages(self):
+        # A PEP 440 local version segment (+cu129 / +cpu / +db1) is NOT a reason to
+        # drop: the base release is on PyPI, so the pin is kept and the segment is
+        # stripped later by req(). Only GPU-only and system packages are dropped.
+        pkgs = {"torch": "2.9.0+cu129", "flask": "1.1.2+db1", "numpy": "2.1.3"}
+        self.assertEqual(_filtered(pkgs), pkgs)
+
+    def test_drops_gpu_only_by_name(self):
+        # Every nvidia-* distribution is a CUDA runtime component; triton /
+        # flash-attn / deepspeed are GPU-only as well. _filtered receives keys already
+        # PEP 503-normalized by parse_requirements/norm, so it compares lowercased.
+        pkgs = {
+            "nvidia-cublas-cu12": "12.6.4.1",
+            "nvidia-cudnn-cu12": "9.5.1.17",
+            "triton": "3.3.0",
+            "flash-attn": "2.7.4.post1",
+            "deepspeed": "0.16.5",
+            "horovod": "0.28.1",
+            "pynvml": "11.5.0",
+            "cuda-toolkit": "13.0.2",
+            "cuda-bindings": "13.2.0",
+            "numpy": "2.1.3",
+        }
+        self.assertEqual(_filtered(pkgs), {"numpy": "2.1.3"})
+
+    def test_drops_system_packages(self):
+        # Ubuntu system packages carried in the image but not pip-installable. They
+        # must be dropped by name so stripping local segments does not resurrect them.
+        pkgs = {"python-apt": "2.7.7+ubuntu5.2", "distro-info": "1.7+build1", "numpy": "2.1.3"}
+        self.assertEqual(_filtered(pkgs), {"numpy": "2.1.3"})
+
+    def test_keeps_installable_pins(self):
+        pkgs = {"ray": "2.37.0", "databricks-sdk": "0.67.0", "numpy": "2.1.3", "pyarrow": "21.0.0"}
+        self.assertEqual(_filtered(pkgs), pkgs)
+
+
+class ReqTest(unittest.TestCase):
+    def test_strips_local_version_segment(self):
+        # ~= is invalid with a local segment (PEP 440), and the segment names a build
+        # that only exists off-index; strip it so the base release is pinned instead.
+        self.assertEqual(req("torch", "2.9.0+cu129"), "torch~=2.9.0")
+        self.assertEqual(req("torch", "2.7.0+cpu"), "torch~=2.7.0")
+        self.assertEqual(req("flask", "1.1.2+db1"), "flask~=1.1.2")
+
+    def test_compatible_release_default(self):
+        self.assertEqual(req("numpy", "2.1.3"), "numpy~=2.1.3")
+
+    def test_databricks_sdk_widened_to_major_minor(self):
+        self.assertEqual(req("databricks-sdk", "0.67.0"), "databricks-sdk~=0.67")
+
+
+class BuildArtifactsTest(unittest.TestCase):
+    # torchmetrics is a real ML package that must survive; it also guards against a
+    # bare-substring assertion mistaking "torch~=..." for "torchmetrics".
+    pkgs = {
+        "numpy": "2.1.3",
+        "torch": "2.9.0+cu129",
+        "torchmetrics": "1.6.0",
+        "nvidia-cublas-cu12": "12.6.4.1",
+        "triton": "3.3.0",
+        "python-apt": "2.7.7+ubuntu5.2",
+        "pyarrow": "21.0.0",
+    }
+
+    def _check(self, out):
+        # +local stripped and pinned; ordinary pins kept.
+        self.assertIn("torch~=2.9.0", out)
+        self.assertIn("torchmetrics~=1.6.0", out)
+        self.assertIn("numpy~=2.1.3", out)
+        self.assertIn("pyarrow~=21.0.0", out)
+        self.assertNotIn("+cu129", out)
+        # Dropped by name — assert on the rendered pin, not a bare substring.
+        self.assertNotIn("nvidia-cublas-cu12~=", out)
+        self.assertNotIn("triton~=", out)
+        self.assertNotIn("python-apt~=", out)
+
+    def test_pyproject(self):
+        self._check(build_pyproject(self.pkgs, "serverless-v4", "3.12.3"))
+
+    def test_constraints(self):
+        self._check(build_constraints(self.pkgs, "serverless-v4"))
+
+
+class DbconnectPinTest(unittest.TestCase):
+    def test_strips_local_version_segment(self):
+        # databricks-connect is installed from the dev group as a plain PyPI release;
+        # the pin is normalized to ~=MAJOR.0, so a local segment in the release-notes
+        # version is discarded and never lands in an artifact. dbconnect_pin reads raw
+        # pkgs (not _filtered), so this guards that the normalization does the stripping.
+        self.assertEqual(
+            dbconnect_pin({"databricks-connect": "17.3.1+db1"}),
+            "databricks-connect~=17.0",
+        )
+
+    def test_none_when_absent(self):
+        self.assertIsNone(dbconnect_pin({"numpy": "2.1.3"}))
+
+
+if __name__ == "__main__":
+    unittest.main()

.github/workflows/test.yml

@@ -0,0 +1,34 @@
+name: test-envgen
+
+# Run the envgen unit tests so the artifact-transformation rules (drops, +local
+# stripping, databricks-sdk / databricks-connect pinning) can't silently regress.
+# envgen.py has no third-party dependencies, so this is a plain stdlib unittest run.
+
+on:
+  push:
+    paths:
+      - ".github/scripts/**"
+      - ".github/workflows/test.yml"
+  pull_request:
+    paths:
+      - ".github/scripts/**"
+      - ".github/workflows/test.yml"
+
+permissions:
+  contents: read
+
+jobs:
+  unittest:
+    # See sync.yml: GitHub-hosted ubuntu-latest is outside the databricks org IP
+    # allow list; linux-ubuntu-latest is the databricks-protected-runner-group.
+    runs-on: linux-ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@93cb6efe18208431cddfb8368fd83d5badbf9bfd # v5.0.1
+
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
+        with:
+          python-version: "3.12"
+
+      - name: Run envgen unit tests
+        run: python -m unittest discover -s .github/scripts -p "test_*.py" -v

README.md

@@ -51,7 +51,36 @@ serverless), so resolving a target to its artifact is a deterministic lookup.
   so the pip path is constraints-only unless DB Connect is installed explicitly.
 
 Both are a mechanical transform of the official package list published in the
-Databricks release notes — see `.github/scripts/envgen.py` for the rules.
+Databricks release notes — see [what is intentionally not included](#what-is-intentionally-not-included) and
+`.github/scripts/envgen.py` for the rules.
+
+## What is intentionally not included
+
+Not every package in the release-notes list is emitted verbatim. `envgen.py` drops
+the ones that can't — or shouldn't — install on a developer machine, and strips
+version markers that would make a pin unresolvable, so `uv sync` / `pip install -c`
+stay resolvable. Applied to **both** `pyproject.toml` and `constraints.txt`:
+
+- **System / OS packages** (dropped) — `pip`, `pyspark` (DB Connect supplies its own
+  bundled build; `py4j` is kept), `dbus-python`, `pygobject`, `unattended-upgrades`,
+  `python-apt`, `distro-info`.
+- **setuptools-vendored** (dropped) — `more-itertools`, `jaraco-*`, `inflect`,
+  `typeguard`, … — shipped inside setuptools, not installed standalone.
+- **GPU-only distributions** (dropped) — the `nvidia-*` and `cuda-*` CUDA runtime
+  components and tooling, plus `triton`, `flash-attn`, `deepspeed`, `horovod`,
+  `pynvml`: they need an NVIDIA GPU (and a CUDA toolchain / MPI to build) a dev machine
+  lacks. (`nvidia-ml-py` and `pynvml` are pure-Python NVML bindings — installable
+  anywhere but useless without a driver, so dropped by prefix/name.)
+- **Local version segments** (stripped, not dropped) — a `+cpu` / `+cuXXX` / `+db1`
+  segment names a build published only off-index (`download.pytorch.org`) or rebuilt
+  inside the image, and `~=` is invalid with a local segment. The segment is stripped
+  and the base release pinned (`torch 2.9.0+cu129` → `torch~=2.9.0`, `flask 1.1.2+db1`
+  → `flask~=1.1.2`), so `uv` resolves a platform-appropriate wheel. (Ubuntu system
+  builds like `python-apt 2.7.7+ubuntu5.2` are dropped by name above instead, since
+  their base version is not on PyPI.)
+
+The exact lists live in `DROP` / `DROP_PREFIX` and `_filtered()` / `req()` in
+`.github/scripts/envgen.py`.
 
 ## How it stays in sync
 
@@ -77,9 +106,10 @@ is best-effort. Nobody hand-edits the `python/` artifacts.
 - **DBR ML (CPU + GPU)** — for each `*-ml` runtime, a separate environment is produced
   per cluster type: `<ver>.x-cpu-ml-…` and `<ver>.x-gpu-ml-…`. Newer ML pages link
   downloadable `requirements-{cpu,gpu}-*.txt`; older ones render inline tables under
-  `python-libraries-on-{cpu,gpu}-clusters`. The GPU set carries the CUDA builds
-  (e.g. `torch==…+cu118`); the CPU set carries `…+cpu`. Local builds are pinned with
-  `==` (compatible-release `~=` is invalid with a `+local` segment).
+  `python-libraries-on-{cpu,gpu}-clusters`. The GPU set lists the CUDA builds
+  (e.g. `torch …+cu118`) and the CPU set lists `…+cpu`; the generated artifacts strip
+  the `+local` segment and pin the base release (see
+  [what is intentionally not included](#what-is-intentionally-not-included)).
 
 The Action runs it; you only need to run it locally to debug:
 
@@ -104,5 +134,7 @@ doc for the full rationale.
 - [x] Serverless (v1–vN) — auto-discovered + synced; ML base environment (`-ml`) when published (v5+)
 - [x] DBR standard runtimes — auto-discovered from the index + HTML-table parsing
 - [x] DBR ML runtimes (CPU + GPU) — downloadable requirements or inline tables
-- [ ] PyTorch index config in ML `pyproject.toml` (so `uv` fetches the matching
-      `+cpu` / `+cuXXX` torch build, not just pins it)
+- [ ] PyTorch index config in ML `pyproject.toml`. Today the `+cpu` / `+cuXXX`
+      torch/torchvision builds are stripped to a base pin (see [what is intentionally not included](#what-is-intentionally-not-included));
+      adding PyTorch's index would let `uv` fetch the exact `+cpu` / `+cuXXX` build the
+      runtime ships, rather than a base-version wheel.

python/dbr/13.3.x-cpu-ml-scala2.12/constraints.txt

@@ -68,7 +68,7 @@ fastapi~=0.98.0
 fastjsonschema~=2.18.0
 fasttext~=0.9.2
 filelock~=3.6.0
-flask==1.1.2+db1
+flask~=1.1.2
 flatbuffers~=23.5.26
 fonttools~=4.25.0
 frozenlist~=1.4.0
@@ -94,7 +94,6 @@ gviz-api~=1.10.0
 h11~=0.14.0
 h5py~=3.7.0
 holidays~=0.27.1
-horovod~=0.28.1
 htmlmin~=0.1.12
 httplib2~=0.20.2
 httptools~=0.6.0
@@ -276,8 +275,8 @@ tiktoken~=0.4.0
 tokenize-rt~=4.2.1
 tokenizers~=0.13.3
 tomli~=2.0.1
-torch==1.13.1+cpu
-torchvision==0.14.1+cpu
+torch~=1.13.1
+torchvision~=0.14.1
 tornado~=6.1
 tqdm~=4.64.1
 traitlets~=5.1.1

python/dbr/13.3.x-cpu-ml-scala2.12/pyproject.toml

@@ -80,7 +80,7 @@ constraint-dependencies = [
     "fastjsonschema~=2.18.0",
     "fasttext~=0.9.2",
     "filelock~=3.6.0",
-    "flask==1.1.2+db1",
+    "flask~=1.1.2",
     "flatbuffers~=23.5.26",
     "fonttools~=4.25.0",
     "frozenlist~=1.4.0",
@@ -106,7 +106,6 @@ constraint-dependencies = [
     "h11~=0.14.0",
     "h5py~=3.7.0",
     "holidays~=0.27.1",
-    "horovod~=0.28.1",
     "htmlmin~=0.1.12",
     "httplib2~=0.20.2",
     "httptools~=0.6.0",
@@ -288,8 +287,8 @@ constraint-dependencies = [
     "tokenize-rt~=4.2.1",
     "tokenizers~=0.13.3",
     "tomli~=2.0.1",
-    "torch==1.13.1+cpu",
-    "torchvision==0.14.1+cpu",
+    "torch~=1.13.1",
+    "torchvision~=0.14.1",
     "tornado~=6.1",
     "tqdm~=4.64.1",
     "traitlets~=5.1.1",

python/dbr/13.3.x-gpu-ml-scala2.12/constraints.txt

@@ -69,8 +69,7 @@ fastapi~=0.98.0
 fastjsonschema~=2.18.0
 fasttext~=0.9.2
 filelock~=3.6.0
-flash-attn~=1.0.7
-flask==1.1.2+db1
+flask~=1.1.2
 flatbuffers~=23.5.26
 fonttools~=4.25.0
 frozenlist~=1.4.0
@@ -96,7 +95,6 @@ gviz-api~=1.10.0
 h11~=0.14.0
 h5py~=3.7.0
 holidays~=0.27.1
-horovod~=0.28.1
 htmlmin~=0.1.12
 httplib2~=0.20.2
 httptools~=0.6.0
@@ -277,8 +275,8 @@ tiktoken~=0.4.0
 tokenize-rt~=4.2.1
 tokenizers~=0.13.3
 tomli~=2.0.1
-torch==1.13.1+cu117
-torchvision==0.14.1+cu117
+torch~=1.13.1
+torchvision~=0.14.1
 tornado~=6.1
 tqdm~=4.64.1
 traitlets~=5.1.1