docker: fix the ARM64 image build - #129
Conversation
1616dc5 to
250527c
Compare
amtellezfernandez
changed the title
250527c to
34922af
Compare
34922af to
d7a573d
Compare
amtellezfernandez
changed the title
|
Some context on why this is worth taking: the Dockerfile already selects an ARM64 base ( One of the four changes is a bug independent of ARM support: Rebased onto current main; single commit, Dockerfile only. Verification on GB10 hardware is in the description. |
d7a573d to
6c3879f
Compare
The Dockerfile selects an ARM64 base already, but the build fails past that point on
aarch64. Four changes, all keyed off TARGETARCH so the x86_64 path is untouched:
1. Re-declare `ARG TARGETARCH` inside the final stage. BuildKit's builtin
platform args are always available to FROM lines, so `FROM base-${TARGETARCH}`
resolves without any ARG at all -- but an ARG line after a FROM belongs to that
stage, and is visible to RUN only there and in stages derived from it. The
existing declaration sits inside the dcgm-exporter stage, which the final stage
does not derive from, so any arch conditional there compares against the empty
string on every architecture -- presenting as a step that completes suspiciously
fast rather than as an error. Verified with minimal-repro Dockerfiles covering
global, ancestor-stage, non-ancestor-stage and no-ARG placements.
2. Configure NVIDIA's CUDA apt repo on arm64: datacenter-gpu-manager-4-cuda12 comes from
it and the ARM64 base image does not have it. The repo path is derived from the base
image's own /etc/os-release rather than hardcoded, so it follows the base image.
3. Install an explicit package subset on arm64 instead of `--extra all --extra recipes`
(see the docker_local proposal for why each is excluded).
4. Skip the PyG compiled extensions on arm64: that step's wheel index is pinned to an x86
CUDA build (2.8.0+cu128) and publishes nothing for aarch64. Only alpasim-trafficsim
needs them, and it is excluded above, so they are unneeded rather than unavailable.
Verified by a real `docker build` on NVIDIA GB10 (Grace Blackwell) against main @ 1e801ca:
33.1GB image, zero errors, all seven installed packages import on aarch64, torch
2.13.0+cu130, and torch.cuda.is_available() true with --gpus all.
6c3879f to
c665775
Compare
1 participant
The Dockerfile already picks an ARM64 base (
FROM nvcr.io/nvidia/pytorch:25.08-py3 AS base-arm64), but the build fails further down, so that path is currently unreachable. This finishes it.The target is DGX Spark (GB10). With these changes the runtime, controller, physics and trafficsim services build and run natively there, and we've run complete closed-loop rollouts on that hardware using the video-model renderer (
runtime.renderer.kind=video_model, OmniDreams via FlashDreams), which suits GB10's unified memory. The one component still x86-only is the prebuilt NuRec/sensorsim image, which that renderer path doesn't use.What breaks on aarch64
Four separate failures past base-image selection. All four fixes key off
TARGETARCH, so the x86_64 path is untouched.1.
ARG TARGETARCHisn't in scope in the final stage. BuildKit's builtin platform args are always available toFROMlines, soFROM base-${TARGETARCH}resolves with noARGat all. But anARGafter aFROMbelongs to that stage, and only stages derived from it inherit the value. The existing declaration sits in the dcgm-exporter stage, which the final stage doesn't derive from, so any arch conditional there compares against the empty string on every architecture. It shows up as a step finishing suspiciously fast, not as an error, which is why static reading misses it. I checked this with minimal repro Dockerfiles covering global, ancestor-stage, non-ancestor-stage and no-ARG placements.2. No CUDA apt repo on the ARM64 base.
datacenter-gpu-manager-4-cuda12comes from it andnvcr.io/nvidia/pytorchships only plain Ubuntu ports repos, soapt-get installfails before anything else runs. The repo path is derived from the base image's own/etc/os-release, so it follows the base image across Ubuntu releases instead of pinning one.3.
--extra all --extra recipescan't resolve.alpasim-toolspulls PyQt5 (no aarch64 wheel),alpasim_driver/alpasim_pluginspull x86-only packages liketensordictthrough the learned-driver models,alpasim_wizardruns on the host, andalpasim-trafficsimpullstorch_geometric. The arm64 branch installs an explicit subset.Why an explicit list and not a declared extra consumed via
uv sync --extra <subset>: measured on GB10, going through the project resolution installs a+cputorch on aarch64 (2.8.0+cpu,torch.version.cuda == None), while the direct install resolves2.13.0+cu130with CUDA available. Until torch sourcing handles aarch64, the sync path quietly produces a CUDA-less runtime image.4. PyG extensions are x86-pinned. That step's wheel index points at
2.8.0+cu128and publishes nothing for aarch64. Onlyalpasim-trafficsimneeds those extensions and it's excluded above, so on arm64 they aren't needed at all.Verification
Full
docker buildon GB10 hardware against currentmain:if [ "arm64" = "arm64" ]andif [ "arm64" != "arm64" ], i.e.TARGETARCHresolving instead of empty (the failure mode from item 1).alpasim_controller,alpasim_runtime,alpasim_physics,alpasim_grpc,alpasim_utils,alpasim_plugins, andalpasim_evalvia its top-level moduleeval),torch 2.13.0+cu130,torch.cuda.is_available()true with--gpus all.import PyQt5andimport tensordictboth raiseModuleNotFoundError.Item 1 is arguably a bug independent of ARM support, since it makes any arch conditional in the final stage compare against an empty string. Happy to split it out if the platform scope here needs more discussion.
Rebased onto current
mainafter the August sync rewrote the Dockerfile, which is what the earlier conflict was. Single commit, Dockerfile only.