set PYTHON_VERSION for python installations

Author: nv-tusharma

container/Dockerfile

@@ -238,8 +238,9 @@ RUN if [ "$ARCH" = "arm64" ]; then \
 ##################################
 
 # Create and activate virtual environment
+ARG PYTHON_VERSION
 RUN mkdir -p /opt/dynamo/venv && \
-    uv venv /opt/dynamo/venv --python 3.12
+    uv venv /opt/dynamo/venv --python $PYTHON_VERSION
 
 ENV VIRTUAL_ENV=/opt/dynamo/venv \
     PATH="/opt/dynamo/venv/bin:${PATH}"
@@ -280,8 +281,9 @@ 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 python3.12-devel wget \
+    && dnf install -y llvm-toolset protobuf-compiler python${PYTHON_VERSION}-devel wget \
     && dnf clean all \
     && rm -rf /var/cache/dnf
 

container/Dockerfile.vllm

@@ -42,6 +42,8 @@ ARG FLASHINF_REF="v0.2.11"
 # without adding if statements everywhere, so just define both as ARGs for now.
 ARG ARCH=amd64
 ARG ARCH_ALT=x86_64
+# Python configuration
+ARG PYTHON_VERSION=3.12
 
 ARG DYNAMO_BASE_IMAGE="dynamo:latest-none"
 FROM ${DYNAMO_BASE_IMAGE} AS dynamo_base
@@ -75,18 +77,19 @@ RUN apt-get update -y \
         ibverbs-utils \
         libibumad-dev \
         libibverbs-dev \
-        librdmacm-dev \
         libnuma-dev \
+        librdmacm-dev \
         rdma-core \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
 ### VIRTUAL ENVIRONMENT SETUP ###
 
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+ARG PYTHON_VERSION
 # Create virtual environment
 RUN mkdir -p /opt/dynamo/venv && \
-    uv venv /opt/dynamo/venv --python 3.12
+    uv venv /opt/dynamo/venv --python $PYTHON_VERSION
 
 # Activate virtual environment
 ENV VIRTUAL_ENV=/opt/dynamo/venv \
@@ -240,12 +243,13 @@ $LD_LIBRARY_PATH
 ### VIRTUAL ENVIRONMENT SETUP ###
 
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
-RUN uv venv ${VIRTUAL_ENV} --python 3.12
+ARG PYTHON_VERSION
+RUN uv venv ${VIRTUAL_ENV} --python $PYTHON_VERSION
 
 # Copy virtual environment from framework image to avoid re-installing framework + vllm dependencies
 COPY --from=framework \
-    /opt/dynamo/venv/lib/python3.12/site-packages \
-    /opt/dynamo/venv/lib/python3.12/site-packages
+    /opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages \
+    /opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages
 
 # Install dynamo, NIXL, and dynamo-specific dependencies
 COPY benchmarks/ /opt/dynamo/benchmarks/