aws-neuron · mounchin · Sep 19, 2025 · Sep 19, 2025
@@ -1,9 +1,8 @@
-hydra-core>=1.2.0,<1.3
+hydra-core>=1.3.0
 omegaconf>=2.2,<2.3
 pyyaml==6.0.1
 torchmetrics>=0.4.1rc0,<=0.10.3
-trl==0.10.1
-transformers==4.48.0
+transformers==4.52.4
 wandb
 webdataset>=0.1.48,<=0.1.62
 pandas
@@ -31,11 +30,11 @@ pytorch-lightning==2.5.0
 ipadic
 mecab-python3
 protobuf==3.20.3
-datasets==2.19.1
+datasets==3.0.0
 dill==0.3.8
 nemo_toolkit==2.1.0
 regex
-requests<2.32.0
+requests==2.32.4
 python-daemon
 huggingface_hub>=0.27.1
 multiprocess==0.70.16
@@ -44,4 +43,8 @@ numpy>=1.24.3,<=1.25.2
 rouge_score
 setuptools>=70.0
 lightning==2.5.0
-ml-dtypes==0.2.0
+ml-dtypes==0.5.0
+boto3==1.35.93
+botocore==1.35.93
+s3transfer==0.10.4
+s3fs
@@ -7,8 +7,8 @@
 
 # This arg required to stop docker build waiting for region configuration while installing tz data from ubuntu 22
 ARG DEBIAN_FRONTEND=noninteractive
-ARG PYTHON=python3.10
-ARG PYTHON_VERSION=3.10.12
+ARG PYTHON=python3.11
+ARG PYTHON_VERSION=3.11.13
 ARG PIP=pip3
 ARG OMPI_VERSION=4.1.5
 
@@ -19,7 +19,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"
@@ -149,9 +149,9 @@

 # 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} focal main" > /etc/apt/sources.list.d/neuron.list \
@@ -160,16 +160,11 @@
  && ${PIP} config set global.extra-index-url "${PIP_REPO_URL}"
 
 # Neuron SDK components version numbers
-ARG NEURON_ARTIFACT_PATH=/root/neuron_artifacts
-ARG IGNORE_MISSING_NEURON_COMPONENTS=false
-RUN IGNORE_MISSING_NEURON_COMPONENTS=$(echo ${IGNORE_MISSING_NEURON_COMPONENTS} | tr '[:upper:]' '[:lower:]')
-
-ARG NEURONX_RUNTIME_LIB_VERSION=2.27.23.0-8deec4dbf
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.27.34.0-ec8cd5e8b
-ARG NEURONX_TOOLS_VERSION=2.25.145.0
-
-ARG NEURONX_CC_VERSION=2.20.9961.0
-ARG NEURONX_JAX_TRAINING_VERSION=0.6.1.1.0.3499+2edccbed
+ARG NEURONX_RUNTIME_LIB_VERSION=2.28.23.0-dd5879008
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.28.27.0-bc30ece58
+ARG NEURONX_TOOLS_VERSION=2.26.14.0
+ARG NEURONX_CC_VERSION=2.21.18209.0+043b1bf7
+ARG NEURONX_JAX_TRAINING_VERSION=0.6.2.1.0.6446+d8c0de77
 
 FROM base AS repo
 
@@ -213,4 +208,4 @@
 ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
 CMD ["/bin/bash"]
 
-HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1
@@ -0,0 +1,204 @@
+ARG BUILD_STAGE=prod
+
+FROM public.ecr.aws/docker/library/ubuntu:22.04 AS base
+
+LABEL dlc_major_version="1"
+LABEL maintainer="Amazon AI"
+LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG PIP=pip3
+ARG PYTHON=python3.11
+ARG PYTHON_VERSION=3.11.13
+ARG TORCHSERVE_VERSION=0.11.0
+ARG SM_TOOLKIT_VERSION=2.0.25
+ARG MINIFORGE_VERSION=25.3.1-0
+
+# 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 \
+ && apt-get upgrade -y \
+ && apt-get install -y --no-install-recommends \
+    apt-transport-https \
+    build-essential \
+    ca-certificates \
+    cmake \
+    curl \
+    emacs \
+    git \
+    gnupg2 \
+    gpg-agent \
+    jq \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libcap-dev \
+    libhwloc-dev \
+    openjdk-11-jdk \
+    unzip \
+    vim \
+    wget \
+    zlib1g-dev \
+ && rm -rf /var/lib/apt/lists/* \
+ && rm -rf /tmp/tmp* \
+ && apt-get clean
+
+# https://github.com/docker-library/openjdk/issues/261 https://github.com/docker-library/openjdk/pull/263/files
+RUN keytool -importkeystore -srckeystore /etc/ssl/certs/java/cacerts -destkeystore /etc/ssl/certs/java/cacerts.jks -deststoretype JKS -srcstorepass changeit -deststorepass changeit -noprompt; \
+    mv /etc/ssl/certs/java/cacerts.jks /etc/ssl/certs/java/cacerts; \
+    /var/lib/dpkg/info/ca-certificates-java.postinst configure;
+
+RUN curl -L -o ~/miniforge.sh https://github.com/conda-forge/miniforge/releases/download/${MINIFORGE_VERSION}/Miniforge3-${MINIFORGE_VERSION}-Linux-x86_64.sh \
+ && 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 \
+ && /opt/conda/bin/conda clean -ya
+
+RUN /opt/conda/bin/mamba install -c conda-forge \
+    python=$PYTHON_VERSION \
+    scikit-learn \
+    h5py \
+    requests \
+ && conda clean -ya \
+ && pip install --upgrade pip \
+    --trusted-host pypi.org --trusted-host files.pythonhosted.org \
+ && ln -s /opt/conda/bin/pip /usr/local/bin/pip3 \
+ && pip install \
+    packaging \
+    enum-compat \
+    ipython \
+ && rm -rf ~/.cache/pip/*
+
+RUN ${PIP} install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org \
+ && ${PIP} install --no-cache-dir -U \
+    opencv-python>=4.8.1.78 \
+    "numpy<1.24,>1.21" \
+    "scipy>=1.8.0" \
+    six \
+    "awscli<2" \
+    pandas==1.* \
+    boto3 \
+    cryptography \
+    "protobuf>=3.18.3,<4" \
+    torchserve==${TORCHSERVE_VERSION} \
+    torch-model-archiver==${TORCHSERVE_VERSION} \
+ && rm -rf ~/.cache/pip/*
+
+ENV SAGEMAKER_SERVING_MODULE=sagemaker_pytorch_serving_container.serving:main
+ENV TEMP=/home/model-server/tmp
+
+RUN useradd -m model-server \
+ && mkdir -p /home/model-server/tmp /opt/ml/model \
+ && chown -R model-server /home/model-server /opt/ml/model
+
+COPY --chmod=755 neuron-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
+COPY --chmod=755 neuron-monitor.sh deep_learning_container.py /usr/local/bin/
+COPY --chmod=755 torchserve-neuron.sh /usr/local/bin/entrypoint.sh
+COPY config.properties /home/model-server
+
+RUN ${PIP} install --no-cache-dir "sagemaker-pytorch-inference==${SM_TOOLKIT_VERSION}" \
+ # patch default_pytorch_inference_handler.py to import torch_neuronx
+ && DEST_DIR=$(python -c "import os.path, sagemaker_pytorch_serving_container; print(os.path.dirname(sagemaker_pytorch_serving_container.__file__))") \
+ && DEST_FILE=${DEST_DIR}/default_pytorch_inference_handler.py \
+ && sed -i "s/import torch/import torch, torch_neuronx/" ${DEST_FILE} \
+ && rm -rf ~/.cache/pip/*
+
+# Compliance 
+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* \
+ # conda leaves an empty /root/.cache/conda/notices.cache file which is not removed by conda clean -ya
+ && rm -rf ${HOME_DIR}/.cache/conda
+
+RUN curl -o /license.txt  https://aws-dlc-licenses.s3.amazonaws.com/pytorch-2.8/license.txt
+
+# 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} focal main" > /etc/apt/sources.list.d/neuron.list \
+ && curl $([ -n "${NEURON_APT_REPO_KEY}" ] && echo "-u ${NEURON_APT_REPO_KEY}") -sSL "https://${NEURON_APT_REPO}/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB" | gpg --dearmor > /etc/apt/keyrings/neuron.gpg \
+ && PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \
+ && ${PIP} config set global.extra-index-url "${PIP_REPO_URL}"
+
+# Neuron SDK components version numbers
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.28.27.0-bc30ece58
+ARG NEURONX_RUNTIME_LIB_VERSION=2.28.23.0-dd5879008
+ARG NEURONX_TOOLS_VERSION=2.26.14.0
+
+ARG NEURONX_CC_VERSION=2.21.18209.0+043b1bf7
+ARG NEURONX_FRAMEWORK_VERSION=2.8.0.2.10.13553+1e4dd6ca
+ARG NEURONX_DISTRIBUTED_VERSION=0.15.22404+1f27bddf
+ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.6.10598+a59fdc00
+
+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} install --no-cache-dir \
+   neuronx-cc \
+   torch-neuronx \
+   neuronx_distributed \
+   neuronx_distributed_inference \
+ && rm -rf ~/.cache/pip/*
+
+FROM base AS prod
+
+# Install Neuron components with specific versions
+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
+
+RUN ${PIP} install --no-cache-dir \
+   neuronx-cc==$NEURONX_CC_VERSION \
+   torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
+   neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
+   neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION \
+ && rm -rf ~/.cache/pip/*
+
+FROM ${BUILD_STAGE} AS final
+
+EXPOSE 8080 8081
+
+ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]
+CMD ["/usr/local/bin/entrypoint.sh"]
+
+HEALTHCHECK CMD curl --fail http://localhost:8080/ping || exit 1