aws-neuron · mehvaibh · May 22, 2026 · May 22, 2026
@@ -0,0 +1,249 @@
+ARG BUILD_STAGE=prod
+
+FROM public.ecr.aws/docker/library/ubuntu:24.04 AS base
+
+LABEL dlc_major_version="1"
+LABEL maintainer="Amazon AI"
+
+# This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 24
+ARG DEBIAN_FRONTEND=noninteractive
+ARG PYTHON=python3.13
+ARG PYTHON_VERSION=3.13.13
+ARG PIP=pip3
+ARG OMPI_VERSION=4.1.5
+ARG PYPI_SIMPLE_URL="https://pypi.org/simple/"
+
+# Python won’t try to write .pyc or .pyo files on the import of source modules
+# Force stdin, stdout and stderr to be totally unbuffered. Good for logging
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+ENV PYTHONIOENCODING=UTF-8
+ENV LANG=C.UTF-8
+ENV LC_ALL=C.UTF-8
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/aws/neuron/lib"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib64"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/openmpi/lib64"
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib"
+ENV PATH="/opt/aws/neuron/bin:${PATH}"
+
+RUN apt-get update \
+ && apt-get upgrade -y \
+ && apt-get install -y --no-install-recommends \
+    build-essential \
+    ca-certificates \
+    cmake \
+    curl \
+    emacs \
+    git \
+    gnupg2 \
+    gpg-agent \
+    jq \
+    libopencv-dev \
+    libglib2.0-0 \
+    libgl1-mesa-dri \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libssl-dev \
+    libsqlite3-dev \
+    libgdbm-dev \
+    libc6-dev \
+    libbz2-dev \
+    libncurses-dev \
+    libffi-dev \
+    libcap-dev \
+    libhwloc-dev \
+    openjdk-8-jdk-headless \
+    openjdk-8-jdk \
+    openjdk-8-jre \
+    openjdk-11-jdk \
+    openssl \
+    software-properties-common \
+    tk-dev \
+    unzip \
+    wget \
+    vim \
+    zlib1g-dev \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install Open MPI and configure SSH for MPI operator in k8s
+RUN mkdir -p /tmp/openmpi \
+ && cd /tmp/openmpi \
+ && wget --quiet https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-${OMPI_VERSION}.tar.gz \
+ && tar zxf openmpi-${OMPI_VERSION}.tar.gz \
+ && cd openmpi-${OMPI_VERSION} \
+ && ./configure --enable-orterun-prefix-by-default \
+ && make -j $(nproc) all \
+ && make install \
+ && ldconfig \
+ && rm -rf /tmp/openmpi
+
+# Install packages and configure SSH for MPI operator in k8s
+RUN apt-get update \
+ && apt-get install -y openmpi-bin openssh-server \
+ && mkdir -p /var/run/sshd \
+ && echo "    UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config \
+ && echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config \
+ && sed -i 's/#\(StrictModes \).*/\1no/g' /etc/ssh/sshd_config \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install Python
+RUN wget -q https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \
+ && tar -xzf Python-$PYTHON_VERSION.tgz \
+ && cd Python-$PYTHON_VERSION \
+ && ./configure --enable-shared --prefix=/usr/local \
+ && make -j $(nproc) && make install \
+ && cd .. && rm -rf ../Python-$PYTHON_VERSION* \
+ && ln -s /usr/local/bin/pip3 /usr/bin/pip \
+ && ln -s /usr/local/bin/$PYTHON /usr/local/bin/python \
+ && ${PIP} --no-cache-dir install --upgrade \
+    pip \
+    requests \
+    setuptools \
+    uv \
+ && rm -rf ~/.cache/pip/*
+
+# Install AWS CLI
+RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \
+ && unzip awscliv2.zip \
+ && ./aws/install \
+ && rm -rf awscliv2.zip aws
+
+# U24 will not allow installation of pip packages outside of venv without this flag
+# This is because U24 ships with Python 3.12 by default and installation into the Python
+# interpreter’s directory are disabled outside of a virtual environment. 
+# https://peps.python.org/pep-0668/
+RUN ${PIP} config set global.break-system-packages true
+
+# Install EFA
+RUN apt-get update \
+ && cd $HOME \
+ && curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz \
+ && wget https://efa-installer.amazonaws.com/aws-efa-installer.key && gpg --import aws-efa-installer.key \
+ && cat aws-efa-installer.key | gpg --fingerprint \
+ && wget https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz.sig && gpg --verify ./aws-efa-installer-latest.tar.gz.sig \
+ && tar -xf aws-efa-installer-latest.tar.gz \
+ && cd aws-efa-installer \
+ && ./efa_installer.sh -y -g --skip-kmod --skip-limit-conf --no-verify \
+ && cd $HOME \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+WORKDIR /
+
+# The ENV variables declared below are changed in the previous section
+# Grouping these ENV variables in the first section causes
+# ompi_info to fail. This is only observed in CPU containers
+ENV PATH="$PATH:/home/.openmpi/bin"
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/.openmpi/lib/"
+RUN ompi_info --parsable --all | grep mpi_built_with_cuda_support:value
+
+RUN mkdir -p /etc/pki/tls/certs && cp /etc/ssl/certs/ca-certificates.crt /etc/pki/tls/certs/ca-bundle.crt
+
+# Copy workaround script for incorrect hostname
+COPY changehostname.c /
+COPY --chmod=755 start_with_right_hostname.sh deep_learning_container.py /usr/local/bin/
+
+RUN HOME_DIR=/root \
+ && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
+ && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
+ && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
+ && chmod +x /usr/local/bin/testOSSCompliance \
+ && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
+ && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
+ && rm -rf ${HOME_DIR}/oss_compliance* \
+ && rm -rf /tmp/tmp*
+
+# Setting up APT and PIP repo for neuron artifacts
+ARG NEURON_APT_REPO=apt.repos.neuron.amazonaws.com
+ARG NEURON_APT_REPO_KEY
+ARG NEURON_PIP_REPO=pip.repos.neuron.amazonaws.com
+ARG NEURON_PIP_REPO_KEY
+RUN mkdir -p /etc/apt/keyrings \
+ && APT_REPO_PREFIX=$([ -n "${NEURON_APT_REPO_KEY}" ] && echo "${NEURON_APT_REPO_KEY}@" || echo "") \
+ && echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://${APT_REPO_PREFIX}${NEURON_APT_REPO} jammy main" > /etc/apt/sources.list.d/neuron.list \
+ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg
+
+# Neuron SDK components version numbers
+ARG NEURONX_RUNTIME_LIB_VERSION=2.32.31.0-0234f5ed2
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.32.28.0-452cba8de
+ARG NEURONX_TOOLS_VERSION=2.30.10.0-39feaa591
+
+ARG NEURONX_CC_VERSION=2.25.3371.0+f524f7f8
+ARG NEURONX_JAX_TRAINING_VERSION=0.10.0.1.0.9913+41e8ced2
+ARG NKI_VERSION=0.4.0+25940409122.gd30719f9
+ARG NEURON_AGENTIC_DEVELOPMENT_VERSION=1.1
+
+FROM base AS repo
+
+# Install Neuron components from the apt and pip repos (latest versions)
+RUN apt-get update \
+ && apt-get install -y \
+    aws-neuronx-tools \
+    aws-neuronx-collectives \
+    aws-neuronx-runtime-lib \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \
+ && ${PIP} install --no-cache-dir --force-reinstall \
+   --index-url ${PIP_REPO_URL} \
+   --extra-index-url ${PYPI_SIMPLE_URL} \
+   --trusted-host ${NEURON_PIP_REPO} \
+   neuronx-cc \
+   jax-neuronx \
+   nki \
+   neuron_agentic_development \
+ && rm -rf ~/.cache/pip/*
+
+# Override jax version brought in by jax-neuronx
+RUN ${PIP} install --no-cache-dir jax==0.9.0
+
+FROM base AS prod
+
+# Install Neuron components
+# Install Neuron Driver, Runtime and Tools
+RUN apt-get update \
+ && apt-get install -y \
+    aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
+    aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
+    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# Install JAX & Neuron CC
+RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \
+ && ${PIP} install --no-cache-dir --force-reinstall \
+   --index-url ${PIP_REPO_URL} \
+   --trusted-host ${NEURON_PIP_REPO} \
+   --extra-index-url ${PYPI_SIMPLE_URL} \
+   neuronx-cc==$NEURONX_CC_VERSION \
+   jax-neuronx==$NEURONX_JAX_TRAINING_VERSION \
+   nki==$NKI_VERSION \
+   neuron_agentic_development==$NEURON_AGENTIC_DEVELOPMENT_VERSION \
+ && rm -rf ~/.cache/pip/*
+
+# Override jax version brought in by jax-neuronx
+RUN ${PIP} install --no-cache-dir jax==0.9.0
+
+FROM ${BUILD_STAGE} AS final
+
+# Upgrade OS packages to latest versions
+RUN apt-get update \
+ && apt-get upgrade -y \
+ && rm -rf /var/lib/apt/lists/* \
+ && apt-get clean
+
+# Starts framework
+ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
+CMD ["/bin/bash"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
@@ -16,7 +16,7 @@

 # See http://bugs.python.org/issue19846
 ENV LANG=C.UTF-8
 ENV LD_LIBRARY_PATH=/opt/aws/neuron/lib:/lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
 ENV PATH=/opt/conda/bin:/opt/aws/neuron/bin:$PATH

 RUN apt-get update \
@@ -71,15 +71,13 @@
  && chmod +x ~/miniforge.sh \
  && ~/miniforge.sh -b -p /opt/conda \
  && rm ~/miniforge.sh \
- && /opt/conda/bin/conda update -y conda \
  && /opt/conda/bin/mamba install -c conda-forge -y \
     python=$PYTHON_VERSION \
     pyopenssl \
     cython \
     mkl-include \
     mkl \
     parso \
-    typing \
     # Below 2 are included in miniconda base, but not mamba so need to install
     conda-content-trust \
     charset-normalizer \
@@ -172,24 +170,25 @@

 # Setting up APT and PIP repo for neuron artifacts
 ARG NEURON_APT_REPO=apt.repos.neuron.amazonaws.com
 ARG NEURON_APT_REPO_KEY
 ARG NEURON_PIP_REPO=pip.repos.neuron.amazonaws.com
 ARG NEURON_PIP_REPO_KEY
 RUN mkdir -p /etc/apt/keyrings \
 && APT_REPO_PREFIX=$([ -n "${NEURON_APT_REPO_KEY}" ] && echo "${NEURON_APT_REPO_KEY}@" || echo "") \
 && echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://${APT_REPO_PREFIX}${NEURON_APT_REPO} jammy main" > /etc/apt/sources.list.d/neuron.list \
  && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg
 
 # Neuron SDK components version numbers
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.31.24.0-1a31ba186
-ARG NEURONX_RUNTIME_LIB_VERSION=2.31.24.0-0b044f4ce
-ARG NEURONX_TOOLS_VERSION=2.29.22.0-b486b0ade
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.32.28.0-452cba8de
+ARG NEURONX_RUNTIME_LIB_VERSION=2.32.31.0-0234f5ed2
+ARG NEURONX_TOOLS_VERSION=2.30.10.0-39feaa591
 
-ARG NEURONX_CC_VERSION=2.24.8799.0+6f62ff7c
-ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.13.26312+8e870898
-ARG NEURONX_DISTRIBUTED_VERSION=0.18.27753+1cafd54f
-ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.9.17334+ced6ae4e
-ARG NKI_VERSION=0.3.0+23928721754.g18aa1271
+ARG NEURONX_CC_VERSION=2.25.3371.0+f524f7f8
+ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.14.27725+e2ff0410
+ARG NEURONX_DISTRIBUTED_VERSION=0.19.28093+fc70b593
+ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.10.17970+8548ba25
+ARG NKI_VERSION=0.4.0+25940409122.gd30719f9
+ARG NEURON_AGENTIC_DEVELOPMENT_VERSION=1.1
 
 FROM base AS repo
 
@@ -213,6 +212,7 @@
    neuronx_distributed \
    neuronx_distributed_inference \
    nki \
+   neuron_agentic_development \
  && rm -rf ~/.cache/pip/*
 
 FROM base AS prod
@@ -237,6 +237,7 @@
    neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
    neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION \
    nki==$NKI_VERSION \
+   neuron_agentic_development==$NEURON_AGENTIC_DEVELOPMENT_VERSION \
  && rm -rf ~/.cache/pip/*
 
 FROM ${BUILD_STAGE} AS final

@@ -21,7 +21,7 @@
 ENV PYTHONIOENCODING=UTF-8
 ENV LANG=C.UTF-8
 ENV LC_ALL=C.UTF-8
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/aws/neuron/lib"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/efa/lib64"
 ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/opt/amazon/openmpi/lib64"
@@ -198,24 +198,25 @@

 # Setting up APT and PIP repo for neuron artifacts
 ARG NEURON_APT_REPO=apt.repos.neuron.amazonaws.com
 ARG NEURON_APT_REPO_KEY
 ARG NEURON_PIP_REPO=pip.repos.neuron.amazonaws.com
 ARG NEURON_PIP_REPO_KEY
 RUN mkdir -p /etc/apt/keyrings \
 && APT_REPO_PREFIX=$([ -n "${NEURON_APT_REPO_KEY}" ] && echo "${NEURON_APT_REPO_KEY}@" || echo "") \
 && echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://${APT_REPO_PREFIX}${NEURON_APT_REPO} jammy main" > /etc/apt/sources.list.d/neuron.list \
  && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") --retry 3 --retry-delay 1 --retry-all-errors -fSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg
 
 # Neuron SDK components
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.31.24.0-1a31ba186
-ARG NEURONX_RUNTIME_LIB_VERSION=2.31.24.0-0b044f4ce
-ARG NEURONX_TOOLS_VERSION=2.29.22.0-b486b0ade
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.32.28.0-452cba8de
+ARG NEURONX_RUNTIME_LIB_VERSION=2.32.31.0-0234f5ed2
+ARG NEURONX_TOOLS_VERSION=2.30.10.0-39feaa591
 
-ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.13.26312+8e870898
-ARG NEURONX_CC_VERSION=2.24.8799.0+6f62ff7c
-ARG NEURONX_DISTRIBUTED_VERSION=0.18.27753+1cafd54f
+ARG NEURONX_FRAMEWORK_VERSION=2.9.0.2.14.27725+e2ff0410
+ARG NEURONX_CC_VERSION=2.25.3371.0+f524f7f8
+ARG NEURONX_DISTRIBUTED_VERSION=0.19.28093+fc70b593
 ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.7.0
-ARG NKI_VERSION=0.3.0+23928721754.g18aa1271
+ARG NKI_VERSION=0.4.0+25940409122.gd30719f9
+ARG NEURON_AGENTIC_DEVELOPMENT_VERSION=1.1
 
 FROM base AS repo
 
@@ -238,6 +239,7 @@
    neuronx-cc \
    neuronx_distributed \
    nki \
+   neuron_agentic_development \
  && rm -rf ~/.cache/pip/*
 
 # Always install neuronx_distributed_training from the prod repo
@@ -269,6 +271,7 @@
    neuronx-cc==$NEURONX_CC_VERSION \
    neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
    nki==$NKI_VERSION \
+   neuron_agentic_development==$NEURON_AGENTIC_DEVELOPMENT_VERSION \
  && rm -rf ~/.cache/pip/*
 
 # Always install neuronx_distributed_training from the prod repo