build: OPS-1140: Refactor sglang dockerfile to support wideep (#3792)

Signed-off-by: Tushar Sharma <tusharma@nvidia.com>
Signed-off-by: Dan Aloni <dan.aloni@vastdata.com>
Signed-off-by: tzulingk@nvidia.com <tzulingk@nvidia.com>
Signed-off-by: Tzu-Ling Kan <tzulingk@nvidia.com>
Signed-off-by: Graham King <grahamk@nvidia.com>
Signed-off-by: Ziqi Fan <ziqif@nvidia.com>
Signed-off-by: Dan Gil <dagil@nvidia.com>
Signed-off-by: Dillon Cullinan <dcullinan@nvidia.com>
Signed-off-by: Harrison Saturley-Hall <hsaturleyhal@nvidia.com>
Signed-off-by: Anant Sharma <anants@nvidia.com>
Co-authored-by: Dan Aloni <dan.aloni@vastdata.com>
Co-authored-by: Ziqi Fan <ziqif@nvidia.com>
Co-authored-by: Elyas Mehtabuddin <emehtabuddin@nvidia.com>
Co-authored-by: Tzu-Ling Kan <tzulingk@nvidia.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Co-authored-by: Graham King <grahamk@nvidia.com>
Co-authored-by: dagil-nvidia <dagil@nvidia.com>
Co-authored-by: Rohan Varma <rohanv@nvidia.com>
Co-authored-by: Dillon Cullinan <dcullinan92@gmail.com>
Co-authored-by: Harrison Saturley-Hall <hsaturleyhal@nvidia.com>
Co-authored-by: Anant Sharma <anants@nvidia.com>

container/Dockerfile

@@ -240,10 +240,10 @@ RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
 RUN --mount=type=secret,id=aws-key-id,env=AWS_ACCESS_KEY_ID \
     --mount=type=secret,id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY \
     if [ "$ARCH" = "arm64" ]; then \
-        cd ${NIXL_SRC_DIR} && uv build . --out-dir /opt/dynamo/wheelhouse/nixl \
+        cd ${NIXL_SRC_DIR} && uv build . --out-dir /opt/dynamo/wheelhouse/nixl --python $PYTHON_VERSION \
         --config-settings=setup-args="-Ddisable_gds_backend=true"; \
     else \
-        cd ${NIXL_SRC_DIR} && uv build . --out-dir /opt/dynamo/wheelhouse/nixl; \
+        cd ${NIXL_SRC_DIR} && uv build . --out-dir /opt/dynamo/wheelhouse/nixl --python $PYTHON_VERSION; \
     fi
 
 ##################################
@@ -296,9 +296,8 @@ ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16} \
     PATH=/usr/local/cargo/bin:/opt/dynamo/venv/bin:$PATH
 
 # Install system dependencies
-ARG PYTHON_VERSION
 RUN dnf update -y \
-    && dnf install -y llvm-toolset protobuf-compiler python${PYTHON_VERSION}-devel wget unzip \
+    && dnf install -y llvm-toolset protobuf-compiler wget unzip \
     && dnf clean all \
     && rm -rf /var/cache/dnf
 
@@ -324,7 +323,14 @@ ENV PROTOC=/usr/local/bin/protoc
 COPY --from=base $RUSTUP_HOME $RUSTUP_HOME
 COPY --from=base $CARGO_HOME $CARGO_HOME
 COPY --from=base $NIXL_PREFIX $NIXL_PREFIX
-COPY --from=base $VIRTUAL_ENV $VIRTUAL_ENV
+
+ARG PYTHON_VERSION
+RUN mkdir -p /opt/dynamo/venv && \
+    uv venv /opt/dynamo/venv --python $PYTHON_VERSION
+
+ENV VIRTUAL_ENV=/opt/dynamo/venv \
+    PATH="/opt/dynamo/venv/bin:${PATH}"
+
 
 # Install SCCACHE if requested
 COPY container/use-sccache.sh /tmp/use-sccache.sh

container/build.sh

@@ -329,7 +329,6 @@ get_options() {
                 missing_requirement "$1"
             fi
             ;;
-
         --vllm-max-jobs)
             # Set MAX_JOBS for vLLM compilation (only used by Dockerfile.vllm)
             if [ "$2" ]; then
@@ -716,7 +715,10 @@ fi
 if [ -n "${MAX_JOBS}" ]; then
     BUILD_ARGS+=" --build-arg MAX_JOBS=${MAX_JOBS} "
 fi
-
+if [[ $FRAMEWORK == "SGLANG" ]]; then
+    echo "Forcing Python version to 3.10 for sglang image build"
+    BUILD_ARGS+=" --build-arg PYTHON_VERSION=3.10"
+fi
 # Add sccache build arguments
 if [ "$USE_SCCACHE" = true ]; then
     BUILD_ARGS+=" --build-arg USE_SCCACHE=true"
@@ -725,7 +727,10 @@ if [ "$USE_SCCACHE" = true ]; then
     BUILD_ARGS+=" --secret id=aws-key-id,env=AWS_ACCESS_KEY_ID"
     BUILD_ARGS+=" --secret id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY"
 fi
-
+if [[ "$PLATFORM" == *"linux/arm64"* && "${FRAMEWORK}" == "SGLANG" ]]; then
+    # Add arguments required for sglang blackwell build
+    BUILD_ARGS+=" --build-arg GRACE_BLACKWELL=true --build-arg BUILD_TYPE=blackwell_aarch64"
+fi
 LATEST_TAG="--tag dynamo:latest-${FRAMEWORK,,}"
 if [ -n "${TARGET}" ] && [ "${TARGET}" != "local-dev" ]; then
     LATEST_TAG="${LATEST_TAG}-${TARGET}"