Skip to content

Commit bb0fe53

Browse files
authored
fix(deps): align pyproject DeepEP pin with the container build (NVIDIA-NeMo#3608)
DeepEP is pinned in two places that nothing keeps in sync: docker/Dockerfile ARG DEEPEP_COMMIT = 42144303 (DeepEP NVIDIA-NeMo#638, 2026-05-20) pyproject.toml deep_ep.rev = 7febc6e2 (DeepEP NVIDIA-NeMo#580, 2026-02-09) The Dockerfile pin was bumped to pick up DeepEP NVIDIA-NeMo#638 (pads max_num_of_tokens_per_rank to the combine-kernel chunk size, fixing the HybridEP JIT static_assert), but the uv source was left behind. Since the `moe` extra is not part of `all`, the container installs DeepEP only from the Dockerfile source build -- so CI never exercised the stale pin, while anyone running `uv sync --extra moe` or `pip install nemo_automodel[moe]` outside the container got a DeepEP three months older than the one every recipe was validated against. Bump the uv source to 42144303 and update the declared `[[tool.uv.dependency-metadata]]` version to `v1.2.1+4214430` to match (deep_ep builds as `1.2.1+$(git rev-parse --short HEAD)`; uv trusts this string instead of building during resolution). Both lock files are refreshed accordingly. Add tests/unit_tests/test_deepep_pin_consistency.py so the next bump cannot land half-applied: it fails when the two pins disagree, and when the declared version stops naming the pinned commit. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
1 parent 1e40293 commit bb0fe53

4 files changed

Lines changed: 126 additions & 10 deletions

File tree

‎docker/common/uv-pytorch.lock‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ overrides = [
116116

117117
[[manifest.dependency-metadata]]
118118
name = "deep-ep"
119-
version = "1.2.1+7febc6e"
119+
version = "1.2.1+4214430"
120120
requires-dist = ["torch", "packaging", "ninja"]
121121

122122
[[manifest.dependency-metadata]]
@@ -1661,8 +1661,8 @@ wheels = [
16611661

16621662
[[package]]
16631663
name = "deep-ep"
1664-
version = "1.2.1+7febc6e"
1665-
source = { git = "https://github.com/deepseek-ai/DeepEP.git?rev=7febc6e25660af0f54d95dd781ecdcd62265ecca#7febc6e25660af0f54d95dd781ecdcd62265ecca" }
1664+
version = "1.2.1+4214430"
1665+
source = { git = "https://github.com/deepseek-ai/DeepEP.git?rev=42144303752422ade37f24bca9e2dde12df70e09#42144303752422ade37f24bca9e2dde12df70e09" }
16661666
dependencies = [
16671667
{ name = "ninja" },
16681668
{ name = "packaging" },
@@ -4265,7 +4265,7 @@ requires-dist = [
42654265
{ name = "databricks-sql-connector", marker = "extra == 'delta-databricks'", specifier = ">=4.4.0" },
42664266
{ name = "datasets", specifier = ">=4.0.0" },
42674267
{ name = "decord", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and extra == 'vlm-media'" },
4268-
{ name = "deep-ep", marker = "extra == 'moe'", git = "https://github.com/deepseek-ai/DeepEP.git?rev=7febc6e25660af0f54d95dd781ecdcd62265ecca" },
4268+
{ name = "deep-ep", marker = "extra == 'moe'", git = "https://github.com/deepseek-ai/DeepEP.git?rev=42144303752422ade37f24bca9e2dde12df70e09" },
42694269
{ name = "deltalake", marker = "extra == 'delta-databricks'", specifier = ">=1.0.0" },
42704270
{ name = "diffusers", marker = "extra == 'diffusion'", specifier = ">=0.39.0" },
42714271
{ name = "ffpa-attn", marker = "extra == 'ffpa'", specifier = ">=0.2.2" },

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ prerelease = "allow"
301301

302302
[tool.uv.sources]
303303
megatron-fsdp = { git = "https://github.com/yuhezhang-ai/Megatron-LM.git", rev = "455389c480af6b3acdca74c7830c68b3274eb083", subdirectory = "megatron/core/distributed/fsdp/src" }
304-
deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "7febc6e25660af0f54d95dd781ecdcd62265ecca" }
304+
deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", rev = "42144303752422ade37f24bca9e2dde12df70e09" }
305305
cut-cross-entropy = { git = "https://github.com/apple/ml-cross-entropy.git", rev = "87a86aba72cfd2f0d8abecaf81c13c4528ea07d8" }
306306
magi_attention = { git = "https://github.com/SandAI-org/MagiAttention.git", rev = "d7ea8afd44c790b65fab68a04a6a0fdd5adbf182" }
307307
liger-kernel = { git = "https://github.com/linkedin/Liger-Kernel.git", rev = "1f1a4b8d6a6c3c1ddb3573a46480c41cef25bde6" }
@@ -360,7 +360,7 @@ nv_grouped_gemm = ["setuptools"]
360360
[[tool.uv.dependency-metadata]]
361361
name = "deep_ep"
362362
# This version has to match the version in the commit/rev/tag used
363-
version = "v1.2.1+7febc6e"
363+
version = "v1.2.1+4214430"
364364
requires-dist = ["torch", "packaging", "ninja"]
365365

366366
[[tool.uv.dependency-metadata]]
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""DeepEP is pinned twice: once for the container build and once for uv/pip installs.
16+
17+
`docker/Dockerfile` builds DeepEP from source at `ARG DEEPEP_COMMIT`, which is what every
18+
CI job and every published image actually runs. `pyproject.toml` pins the same dependency
19+
under `[tool.uv.sources]` for anyone installing `nemo_automodel[moe]` outside the container.
20+
Nothing links the two, so bumping one and forgetting the other silently ships a different
21+
DeepEP to users than the one the tests validated. These tests fail on that drift.
22+
"""
23+
24+
import re
25+
from pathlib import Path
26+
27+
try:
28+
import tomllib # ty: ignore[unresolved-import]
29+
except ModuleNotFoundError:
30+
import tomli as tomllib # ty: ignore[unresolved-import]
31+
32+
REPO_ROOT = Path(__file__).resolve().parents[2]
33+
PYPROJECT_PATH = REPO_ROOT / "pyproject.toml"
34+
DOCKERFILE_PATH = REPO_ROOT / "docker" / "Dockerfile"
35+
36+
_FULL_SHA = r"[0-9a-f]{40}"
37+
_DOCKERFILE_ARG = re.compile(rf"^ARG\s+DEEPEP_COMMIT=({_FULL_SHA})\s*$", re.MULTILINE)
38+
# `git rev-parse --short` never abbreviates below 7 characters.
39+
_MIN_ABBREV_LEN = 7
40+
41+
42+
def _load_pyproject() -> dict:
43+
with PYPROJECT_PATH.open("rb") as handle:
44+
return tomllib.load(handle)
45+
46+
47+
def _pyproject_deepep_rev() -> str:
48+
"""The commit `uv`/`pip` resolve `deep_ep` from, per `[tool.uv.sources]`."""
49+
sources = _load_pyproject()["tool"]["uv"]["sources"]
50+
assert "deep_ep" in sources, f"[tool.uv.sources] in {PYPROJECT_PATH.name} no longer declares `deep_ep`"
51+
52+
rev = sources["deep_ep"]["rev"]
53+
assert re.fullmatch(_FULL_SHA, rev), (
54+
f"[tool.uv.sources] deep_ep.rev must be a full 40-character commit SHA so it can be compared "
55+
f"against DEEPEP_COMMIT in docker/Dockerfile, got {rev!r}"
56+
)
57+
return rev
58+
59+
60+
def _pyproject_deepep_declared_version() -> str:
61+
"""The hand-written `[[tool.uv.dependency-metadata]]` version used to skip a build during resolution."""
62+
entries = _load_pyproject()["tool"]["uv"]["dependency-metadata"]
63+
matches = [entry for entry in entries if entry["name"] == "deep_ep"]
64+
assert len(matches) == 1, (
65+
f"expected exactly one [[tool.uv.dependency-metadata]] entry named `deep_ep` in "
66+
f"{PYPROJECT_PATH.name}, found {len(matches)}"
67+
)
68+
return matches[0]["version"]
69+
70+
71+
def _dockerfile_deepep_commit() -> str:
72+
"""The commit the container builds DeepEP from, per `ARG DEEPEP_COMMIT`."""
73+
matches = _DOCKERFILE_ARG.findall(DOCKERFILE_PATH.read_text())
74+
assert len(matches) == 1, (
75+
f"expected exactly one `ARG DEEPEP_COMMIT=<40-char sha>` line in docker/Dockerfile, "
76+
f"found {len(matches)}; this test cannot tell which pin is authoritative otherwise"
77+
)
78+
return matches[0]
79+
80+
81+
def test_pyproject_deepep_rev_matches_dockerfile_commit():
82+
"""The uv/pip pin and the container build must resolve to the same DeepEP commit."""
83+
pyproject_rev = _pyproject_deepep_rev()
84+
dockerfile_commit = _dockerfile_deepep_commit()
85+
86+
assert pyproject_rev == dockerfile_commit, (
87+
"DeepEP pin drift: the container and the uv/pip install path would ship different DeepEP builds.\n"
88+
f" pyproject.toml [tool.uv.sources] deep_ep.rev = {pyproject_rev}\n"
89+
f" docker/Dockerfile ARG DEEPEP_COMMIT = {dockerfile_commit}\n"
90+
"Bump whichever is stale, then refresh both lock files (see CONTRIBUTING.md) so "
91+
"`uv.lock` and `docker/common/uv-pytorch.lock` agree."
92+
)
93+
94+
95+
def test_deepep_dependency_metadata_version_matches_pinned_rev():
96+
"""The declared `1.2.1+<short-sha>` local version must name the commit that is actually pinned.
97+
98+
`uv` trusts this string instead of building DeepEP during resolution; if it names a different
99+
commit than `deep_ep.rev`, the lock file records a version that the built wheel never produces.
100+
"""
101+
rev = _pyproject_deepep_rev()
102+
declared = _pyproject_deepep_declared_version()
103+
104+
assert "+" in declared, (
105+
f"[[tool.uv.dependency-metadata]] deep_ep.version must carry a `+<short-sha>` local segment "
106+
f"identifying the pinned commit, got {declared!r}"
107+
)
108+
short_sha = declared.rsplit("+", 1)[1]
109+
110+
assert len(short_sha) >= _MIN_ABBREV_LEN and rev.startswith(short_sha), (
111+
"DeepEP declared version does not match the pinned commit.\n"
112+
f" [tool.uv.sources] deep_ep.rev = {rev}\n"
113+
f" [[tool.uv.dependency-metadata]] version = {declared}\n"
114+
f"`deep_ep` builds as `1.2.1+$(git rev-parse --short HEAD)`, so the local segment must be "
115+
f"an abbreviation of the pinned rev (expected `+{rev[:_MIN_ABBREV_LEN]}`, got `+{short_sha}`)."
116+
)

‎uv.lock‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)