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

[rugpanov]

Why

The per-environment constraint artifacts mirror the cluster image faithfully — right for the image, but some ML pins can't install on a developer machine.

What

envgen.py, applied to both pyproject.toml and constraints.txt:

• Strips +local version segments and pins the base release (torch 2.9.0+cu129 → torch~=2.9.0, flask 1.1.2+db1 → flask~=1.1.2). ~= is invalid with a local segment and the build exists only off-index, while its base is an ordinary PyPI version.
• Drops GPU-only distributions — nvidia-* / cuda-* CUDA wheels & tooling, plus triton, flash-attn, deepspeed, horovod, pynvml: all need an NVIDIA GPU (and a CUDA/MPI toolchain) a dev machine lacks.
• Drops Ubuntu system packages — python-apt, distro-info (base version not on PyPI).
• Adds .github/workflows/test.yml running the unit tests on push / PR touching .github/scripts/**, plus a new test_envgen.py and a README "What is intentionally not included" section.

Note on horovod: it was previously kept (a plain pin) in 13.3/14.3 ML, so dropping it is an independent policy change, not a consequence of the strip fix — it's the same source-only, GPU/MPI-only class as deepspeed/flash-attn, which can't pip install cleanly on a dev machine.

torch/torchvision are now pinned consistently across CPU and GPU ML envs. All affected artifacts are regenerated as a mechanical transform of the release-notes pins (no unrelated version drift).

Verification

• python -m unittest test_envgen — 11 passing; CI test-envgen green.
• No +local / GPU / system pins remain under python/; all 33 pyproject.toml re-parse as valid TOML.

The pandas-on-Python-3.12 fix for DBR 16.4 / serverless-v3 is split into #22.

This pull request and its description were written by Isaac.

[anton-107]

Verified against a fresh clone of the PR branch.

What it does

envgen.py gains one rule — drop any pin whose version contains + — plus GPU-only names (triton, flash-attn, deepspeed) in DROP and an nvidia- entry in DROP_PREFIX. Because _filtered feeds both builders, this cleans pyproject.toml and constraints.txt together. 38 artifact files are regenerated as a pure deletion, a new test_envgen.py covers the rule, and the README gains a "what is intentionally not included" section.

The central premise is correct and worth affirming: constraint-dependencies / -c bind a version only if the package enters resolution — they never force an install. So removing these pins cannot break a resolution that previously worked. The blast radius is genuinely low.

The delta-strip execution is also clean. I confirmed:

• 0 additions / 376 deletions under python/ — no unrelated version drift
• all 118 unique removed lines match the new rule, nothing unexplained
• all 33 pyproject.toml still parse as valid TOML
• no residual +local or GPU-only pins anywhere under python/

Issues

1. The blanket + drop removes packages whose base version is on PyPI

The rationale is "resolves nowhere off the cluster." That holds for torch==2.9.0+cu129. It does not hold for flask==1.1.2+db1 or horovod==0.28.1+db1 — +db1 means Databricks rebuilt it, but flask 1.1.2 and horovod 0.28.1 are ordinary PyPI releases. The result ends up decided by incidental docs formatting:

	docs list	after this PR
13.3 ML	horovod 0.28.1	horovod~=0.28.1 ✅
15.4 ML	horovod 0.28.1+db1	dropped
13.3 ML	flask 1.1.2+db1	dropped
15.4 ML	flask 2.2.5	flask~=2.2.5 ✅

Same package, same version in the horovod case, opposite treatment. For a repo whose stated purpose is "reproduce the runtime's Python environment locally," 13.3 ML now carries no flask constraint while 15.4 does.

Stripping the segment (flask~=1.1.2) preserves fidelity and is valid PEP 440.

One caveat: a blanket strip is also wrong — it would resurrect python-apt~=2.7.7 and distro-info~=1.7, Ubuntu system packages that are currently dropped only as a side effect of the + rule. Those belong in DROP by name; then +local can be stripped rather than dropped. Worth adding them to DROP regardless of what you decide here, so the intent is explicit and survives a future change to this rule.

2. torch is now unpinned in 17 of 19 ML environments, and the two survivors are the CPU/GPU inverse of what the description implies

The body says a plain torch 2.7.0 pin is kept because it resolves to a macOS/CPU wheel. True, but it's kept only where the docs happened to list it plain — which is the GPU pages:

17.3.x-cpu-ml-scala2.13   torch UNPINNED   (main: torch==2.7.0+cpu)
17.3.x-gpu-ml-scala2.13   torch~=2.7.0     <- survives
19.x-cpu-ml-scala2.13     torch UNPINNED   (main: torch==2.12.0+cpu)
19.x-gpu-ml-scala2.13     torch~=2.12.0    <- survives

Every other ML env, CPU and GPU alike, loses torch entirely. So the GPU artifact pins torch and its CPU sibling doesn't — backwards from the reasoning given. Fixing item 1 resolves this along with it.

3. Test bug: substring assertions will misfire on real inputs

assertNotIn("torch", out) passes only because the fixture contains no other torch* package. Real ML envs ship torchmetrics, torcheval, torch-optimizer, torchvision. I confirmed the failure — adding torchmetrics to the fixture makes the assertion fail even though the torch pin was correctly dropped:

numpy~=2.1.3
torcheval~=0.0.7
torchmetrics~=1.6.0
--- assertNotIn("torch", out) -> 'torch' in out == True  => test FAILS

Assert on the rendered line ("torch~=", or the exact list entry) instead of the bare name. Same pattern applies to triton — no collision today, but the same trap.

4. The tests don't run anywhere

.github/workflows/ contains only sync.yml, which runs sync.py and never the tests. Nine tests that execute only when someone types the command will rot. A ~10-line job on push/PR would close this, and it's the natural place to also assert the invariant in item 5.

5. req()'s new docstring states an invariant the code no longer enforces

The docstring says "Every version passed in is therefore a plain release" and the defensive + branch is gone. That's true for the two in-tree call sites, but req() is module-level and importable, and if a local version does reach it the output is silently invalid — confirmed with packaging:

torch~=2.9.0+cu129   INVALID -> Local version label can only be used with `==` or `!=`

That yields a corrupt-but-plausible artifact instead of an error. assert "+" not in version keeps it cheap and loud.

Related: dbconnect_pin is the one function reading raw pkgs rather than _filtered — good catch adding a test for it, and .split(".")[0] does make it safe. Worth a comment saying that's why.

6. The nvidia-ml-py rationale doesn't fit

The prefix comment calls it "GPU tooling, no local use" and the README says these "have no macOS wheel at all." nvidia-ml-py is a pure-Python NVML binding — it installs anywhere; it just can't do anything without a driver. Dropping it is defensible, but not for the stated reason, and it's grouped with genuinely platform-locked wheels like nvidia-cublas-cu12. I couldn't reach PyPI from my network to confirm the wheel tag, so please verify before adjusting the wording.

7. Nit

Body says "7 passing"; it's 9 (FilterNonLocal 3, BuildArtifacts 2, Req 2, DbconnectPin 2).

Verdict

The GPU-only name drops (nvidia-*, triton, flash-attn, deepspeed) are clearly right and I'd merge those as-is.

The blanket + rule is the part I'd push back on: it conflates "unpublishable build" with "Databricks-patched build of a real PyPI release," and the visible cost is torch disappearing from 17 ML environments plus the flask/horovod split. Fixing items 1 and 3 is a small change to the same lines this PR already touches. Item 4 is worth doing while the file is fresh.

_{Reviewed by Isaac on behalf of Anton.}

[rugpanov]

Thanks @anton-107 — all addressed (commits 04534b5 + e095775):

• 1 & 2 (blanket + drop too broad; torch unpinned / CPU-GPU inverted): switched from dropping +local pins to stripping the segment and pinning the base release (torch 2.9.0+cu129 → torch~=2.9.0, flask 1.1.2+db1 → flask~=1.1.2). torch/torchvision are now pinned consistently across CPU and GPU ML envs, and the flask/horovod split is gone. The two genuine Ubuntu system packages (python-apt, distro-info) are added to DROP by name so stripping doesn't resurrect them.
• 3 (fragile substring assertion): tests now assert on the rendered pin (torch~=2.9.0), with torchmetrics in the fixture as the neighbor guard you described — no bare-name checks.
• 5 (req() invariant): req() now strips the segment itself, so it can't emit an invalid ~=…+local; docstring updated accordingly.
• 6 (nvidia-ml-py rationale): README/comment corrected — it's pure-Python and installs anywhere, just useless without a driver; dropped by the same nvidia- prefix.
• 4 (tests didn't run): added .github/workflows/test.yml running the unit tests on push / PR touching .github/scripts/**.
• 7: count corrected (11).

Also checked the open issues (#18 pandas cp312, #19 grpcio-status conflict, #20 grpcio macOS wheel): this PR doesn't touch those packages, so they're unaffected — same spirit, separate follow-ups.

[anton-107]

Approving. Both substantive problems from the first pass are properly fixed, and the fix is better than what I suggested — putting the strip in req() makes the invalid output unreachable rather than merely asserted against.

Verified against a fresh clone at 73ec1a0 rather than taking the summary at face value:

• Tests pass exactly as CI invokes them. python -m unittest discover -s .github/scripts -p "test_*.py" → Ran 11 tests … OK. Count matches the corrected 11.
• Items 1 and 2 (the CPU/GPU inversion) — resolved. Siblings now agree: 17.3.x both torch~=2.7.0 torchvision~=0.22.0; 19.x both torch~=2.12.0 torchvision~=0.27.0. flask is pinned in both 13.3 and 15.4 (1.1.2 / 2.2.5) — no split.
• Item 3 (substring trap) — fixed properly. _check asserts on rendered pins (torch~=2.9.0, nvidia-cublas-cu12~=, triton~=, python-apt~=) with torchmetrics in the fixture as a live neighbor guard, and both build paths share it so pyproject and constraints can't drift apart.
• Item 5 — invariant now enforced by construction. I validated all 9,441 constraint lines across all 33 environments as PEP 508 with packaging: 0 invalid, no residual +, no single-component ~= pins.
• Artifacts are reproducible, not hand-edited. I re-derived all 33 constraints.txt from main's data using only the new rules: 0 mismatches, byte-for-byte. This is what makes me comfortable approving 38 regenerated files.
• Item 4 — CI job is correct. linux-ubuntu-latest with the runner-group rationale carried from sync.yml, SHA-pinned actions, permissions: contents: read, and the follow-up commit rightly adds test.yml to its own paths so the workflow gates changes to itself.

Blast radius of the strip is narrower than it sounds: only torch, torchvision, flask are actually affected. The other former +local packages (horovod, python-apt, distro-info) are all name-dropped now.

Two things to settle (non-blocking)

A. horovod lost its pin in 4 environments, and the PR body doesn't mention it. The comment justifies it as source-only, needing MPI/NCCL + a compiler — reasonable on the merits. But main carried horovod~=0.28.1 in 13.3.x and 14.3.x (cpu and gpu), where it was a plain pin the old + rule never touched. So this isn't a consequence of the strip fix; it's an independent policy change riding along. Worth a line in the description, since it's the one behavior change here that isn't traceable to a review comment.

B. pynvml and cuda-* escape the GPU-only rule, so the README now asserts a cleaner taxonomy than the code implements. Still pinned after this PR:

pynvml~=11.5.0          (14.3.x cpu + gpu ML)
cuda-bindings~=13.2.0   (19.x gpu ML)
cuda-pathfinder~=1.5.4
cuda-toolkit~=13.0.2

pynvml is the same NVML binding as nvidia-ml-py — same purpose, one dropped by prefix and one kept, purely because of how it's named. cuda-toolkit is squarely in the class the README describes. All four are pre-existing on main, so this is not a regression from this PR and I'm not blocking on it — but either extend DROP or soften the README wording so the two agree.

Follow-ups, not for this PR

• norm, parse_requirements, and _label are still untested. parse_requirements is the entry point for every artifact and carries real logic (the * footnote-marker strip, and a regex that silently skips non-matching lines — a malformed docs table would drop packages quietly). Obvious next test now that the harness exists.
• Nothing guards artifacts against script drift. sync.py --check exists for exactly this and isn't in CI; only the weekly sync.yml regenerates, so a future envgen.py change can land with stale python/ artifacts and CI stays green. Wiring --check into the test job would need the paths filter widened, so it's a separate change.

_{Reviewed by Isaac on behalf of Anton.}