From 0b54a370393a1e6bb691ad21c5608c74eb9035a1 Mon Sep 17 00:00:00 2001
From: Long Wu <wulong@zetyun.com>
Date: Tue, 3 Nov 2026 10:03:26 +0800
Subject: [PATCH] fix: KVBM wheel source

---
 container/templates/wheel_builder.Dockerfile  | 11 +++++++++-
 .../test_aisimulate_consistency.py            | 20 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/container/templates/wheel_builder.Dockerfile b/container/templates/wheel_builder.Dockerfile
index f04c168ddb5f..fe2ca221caa9 100644
--- a/container/templates/wheel_builder.Dockerfile
+++ b/container/templates/wheel_builder.Dockerfile
@@ -930,8 +930,17 @@ COPY pyproject.toml README.md LICENSE Cargo.toml Cargo.lock rust-toolchain.toml
 COPY lib/ /opt/dynamo/lib/
 COPY dingo/ /opt/dynamo/dingo/
 
-# Build kvbm wheel (with nixl linkage via auditwheel repair)
+# This stage does not inherit runtime_wheel_builder's Cargo source patch.
+# Stage the same pinned AISimulate source before KVBM runs Cargo metadata.
 ARG ENABLE_KVBM
+COPY container/deps/stage_aisimulate.py container/deps/aisimulate-sources.json /opt/dynamo/container/deps/
+RUN if [ "$ENABLE_KVBM" = "true" ]; then \
+        source ${VIRTUAL_ENV}/bin/activate && \
+        python /opt/dynamo/container/deps/stage_aisimulate.py \
+            --rust-source-only --output /opt/build-deps/aisimulate; \
+    fi
+
+# Build kvbm wheel (with nixl linkage via auditwheel repair)
 RUN --mount=type=secret,id=aws-web-identity-token,target=/run/secrets/aws-token \
     --mount=type=secret,id=aws-role-arn,env=AWS_ROLE_ARN \
     --mount=type=secret,id=actions-results-url,env=ACTIONS_RESULTS_URL \
diff --git a/tests/dependencies/test_aisimulate_consistency.py b/tests/dependencies/test_aisimulate_consistency.py
index d95648745293..037fc19d01ca 100644
--- a/tests/dependencies/test_aisimulate_consistency.py
+++ b/tests/dependencies/test_aisimulate_consistency.py
@@ -150,6 +150,26 @@ def _staging_module():
     return module
 
 
+@pytest.mark.parametrize(
+    "stage",
+    ["runtime_wheel_builder", "wheel_builder"],
+)
+def test_each_rust_builder_stages_aisimulate_before_maturin(stage):
+    template = (ROOT / "container/templates/wheel_builder.Dockerfile").read_text()
+    stages = re.split(r"^FROM .* AS ", template, flags=re.MULTILINE)
+    builder = next(part for part in stages if part.startswith(stage + "\n"))
+    copy = "COPY container/deps/stage_aisimulate.py container/deps/aisimulate-sources.json"
+    staging = "--rust-source-only --output /opt/build-deps/aisimulate"
+    assert builder.index("COPY lib/") < builder.index(copy)
+    assert builder.index(copy) < builder.index(staging) < builder.index("maturin build")
+    if stage == "wheel_builder":
+        assert builder.index("ARG ENABLE_KVBM") < builder.index(staging)
+        instructions = builder.replace("\\\n", "").splitlines()
+        run = next(line for line in instructions if staging in line)
+        assert run.startswith('RUN if [ "$ENABLE_KVBM" = "true" ]; then ')
+        assert run.endswith("fi")
+
+
 def test_rust_source_staging_preserves_transitive_lock(tmp_path, monkeypatch):
     staging = _staging_module()
     root = tmp_path / "build"