diff --git a/docker/art-gpu.Dockerfile b/docker/art-gpu.Dockerfile index 576ff902b..d2370b3e2 100644 --- a/docker/art-gpu.Dockerfile +++ b/docker/art-gpu.Dockerfile @@ -140,6 +140,7 @@ RUN if ! getent group messagebus >/dev/null; then groupadd -r messagebus; fi \ COPY --from=builder --chown=sky:sky /opt/uv-cache /opt/uv-cache COPY --from=builder --chown=sky:sky /opt/uv-python /opt/uv-python +COPY --from=builder --chown=sky:sky /opt/src/art/megatron_runtime/pyproject.toml /opt/src/art/megatron_runtime/uv.lock /opt/src/art/megatron_runtime/ USER sky WORKDIR /home/sky diff --git a/tests/unit/test_gpu_image_build_context.py b/tests/unit/test_gpu_image_build_context.py index a3279ada6..becea27b8 100644 --- a/tests/unit/test_gpu_image_build_context.py +++ b/tests/unit/test_gpu_image_build_context.py @@ -13,3 +13,23 @@ def test_gpu_image_build_context_copies_every_local_docker_source() -> None: continue for source in shlex.split(line)[1:-1]: assert f"${{repo_root}}/{source}" in build_script + + +def test_gpu_image_retains_only_managed_megatron_project_metadata() -> None: + dockerfile = (ROOT / "docker/art-gpu.Dockerfile").read_text() + final_stage = dockerfile.rsplit("\nFROM ", 1)[1] + copies = [ + shlex.split(line) + for line in final_stage.splitlines() + if line.startswith("COPY ") and "/opt/src/art" in line + ] + assert copies == [ + [ + "COPY", + "--from=builder", + "--chown=sky:sky", + "/opt/src/art/megatron_runtime/pyproject.toml", + "/opt/src/art/megatron_runtime/uv.lock", + "/opt/src/art/megatron_runtime/", + ] + ]