Skip to content

Updated Dockerfiles #213

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aws-ruhip merged 2 commits into from
Oct 20, 2026
Merged

Updated Dockerfiles #213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
04fed91
Updated Dockerfiles based on build artifacts uploaded to the Neuron D…
aws-ruhip Oct 20, 2026
4bb66ee
Update structure
aws-ruhip Oct 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 56 additions & 30 deletions jax/training/0.10.0/Dockerfile.neuronx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
ARG BUILD_STAGE=prod

FROM public.ecr.aws/docker/library/ubuntu:24.04 AS base
FROM public.ecr.aws/docker/library/ubuntu:26.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
# Keep apt non-interactive so tzdata's timezone/region prompt can't stall the build
ARG DEBIAN_FRONTEND=noninteractive
ARG PYTHON=python3.13
ARG PYTHON_VERSION=3.13.13
ARG PYTHON=python3.14
ARG PYTHON_VERSION=3.14.7
ARG PIP=pip3
ARG OMPI_VERSION=4.1.5
ARG PYPI_SIMPLE_URL="https://pypi.org/simple/"
Expand All @@ -20,7 +20,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"

Check warning on line 23 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / build (0.10.0, jax/training/0.10.0)

Variables should be defined before their use 

UndefinedVar: Usage of undefined variable '$LD_LIBRARY_PATH'
More info: https://docs.docker.com/go/dockerfile/rule/undefined-var/
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"
Expand All @@ -28,7 +28,7 @@
ENV PATH="/opt/aws/neuron/bin:${PATH}"
ENV PATH="${PATH}:/opt/amazon/efa/bin"

RUN apt-get update \

Check failure on line 31 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3008 warning: Pin versions in apt get install. Instead of `apt-get install <package>` use `apt-get install <package>=<version>`
&& apt-get upgrade -y \
&& apt-get install -y --no-install-recommends \
build-essential \
Expand Down Expand Up @@ -68,7 +68,7 @@
&& apt-get clean

# Install Open MPI and configure SSH for MPI operator in k8s
RUN mkdir -p /tmp/openmpi \

Check failure on line 71 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

SC2046 warning: Quote this to prevent word splitting.

Check failure on line 71 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3003 warning: Use WORKDIR to switch to a directory
&& 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 \
Expand All @@ -80,7 +80,7 @@
&& rm -rf /tmp/openmpi

# Install packages and configure SSH for MPI operator in k8s
RUN apt-get update \

Check failure on line 83 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3015 info: Avoid additional packages by specifying `--no-install-recommends`

Check failure on line 83 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3008 warning: Pin versions in apt get install. Instead of `apt-get install <package>` use `apt-get install <package>=<version>`
&& apt-get install -y openmpi-bin openssh-server \
&& mkdir -p /var/run/sshd \
&& echo " UserKnownHostsFile /dev/null" >> /etc/ssh/ssh_config \
Expand All @@ -91,7 +91,7 @@
&& apt-get clean

# Install Python
RUN wget -q https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \

Check failure on line 94 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

SC2046 warning: Quote this to prevent word splitting.

Check failure on line 94 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3003 warning: Use WORKDIR to switch to a directory
&& tar -xzf Python-$PYTHON_VERSION.tgz \
&& cd Python-$PYTHON_VERSION \
&& ./configure --enable-shared --prefix=/usr/local \
Expand All @@ -107,19 +107,19 @@
&& rm -rf ~/.cache/pip/*

# Install AWS CLI
RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \

Check failure on line 110 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL4001 warning: Either use Wget or Curl but not both
&& 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.
# This image installs pip packages into the base interpreter (no venv). Ubuntu 26 — like 24 —
# ships its system Python as an externally-managed environment (PEP 668), which blocks that, so
# set the flag to allow it. (Interpreter here is the source-built Python 3.14 in /usr/local.)
# https://peps.python.org/pep-0668/
RUN ${PIP} config set global.break-system-packages true

# Install EFA
RUN apt-get update \

Check failure on line 122 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

DL3047 info: Avoid use of wget without progress bar. Use `wget --progress=dot:giga <url>`. Or consider using `-q` or `-nv` (shorthands for `--quiet` or `--no-verbose`).

Check failure on line 122 in jax/training/0.10.0/Dockerfile.neuronx

View workflow job for this annotation

GitHub Actions / dockerfile-linter 

SC2086 info: Double quote to prevent globbing and word splitting.
&& 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 \
Expand Down Expand Up @@ -160,37 +160,50 @@

# 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
RUN --mount=type=secret,id=neuron_repo \
mkdir -p /etc/apt/keyrings \
&& echo "deb [signed-by=/etc/apt/keyrings/neuron.gpg] https://${NEURON_APT_REPO} jammy main" > /etc/apt/sources.list.d/neuron.list \
&& NEURON_REPO_CRED="$([ -s /run/secrets/neuron_repo ] && cat /run/secrets/neuron_repo || true)" \
&& curl $([ -n "${NEURON_REPO_CRED}" ] && echo "-u ${NEURON_REPO_CRED}") --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.34.10.0-ac18d186d
ARG NEURONX_COLLECTIVES_LIB_VERSION=2.34.10.0-74eaafac6
ARG NEURONX_TOOLS_VERSION=2.32.28.0-526c2b7f6
ARG NEURONX_RUNTIME_LIB_VERSION=2.35.30.0-d2725ccb8
ARG NEURONX_COLLECTIVES_LIB_VERSION=2.35.43.0-354ad3064
ARG NEURONX_TOOLS_VERSION=2.33.38.0-f874d5618

ARG NEURONX_CC_VERSION=2.27.5334.0+f702b353
ARG NEURONX_JAX_TRAINING_VERSION=0.10.0.1.0.11459+1e1fc1db
ARG NKI_VERSION=0.6.0+31049202112.g85070674
ARG NEURON_AGENTIC_DEVELOPMENT_VERSION=1.3
ARG NEURONX_CC_VERSION=2.28.7631.0+8f7b0cf5
ARG NEURONX_JAX_TRAINING_VERSION=0.10.0.1.0.14198+40d442c2
ARG NKI_VERSION=0.7.0+37713544100.gfede04bb
ARG NEURON_AGENTIC_DEVELOPMENT_VERSION=1.4

FROM base AS repo

# Install Neuron components from the apt and pip repos (latest versions)
RUN apt-get update \
RUN --mount=type=secret,id=neuron_repo \
mkdir -p /etc/apt/auth.conf.d \
&& if [ -s /run/secrets/neuron_repo ]; then \
CRED="$(cat /run/secrets/neuron_repo)"; \
printf 'machine %s login %s password %s\n' "${NEURON_APT_REPO}" "${CRED%%:*}" "${CRED#*:}" > /etc/apt/auth.conf.d/neuron.conf; \
chmod 600 /etc/apt/auth.conf.d/neuron.conf; \
fi \
&& 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
&& apt-get clean \
&& rm -f /etc/apt/auth.conf.d/neuron.conf

RUN PIP_REPO_URL=$([ -n "${NEURON_PIP_REPO_KEY}" ] && echo "https://${NEURON_PIP_REPO_KEY}@${NEURON_PIP_REPO}" || echo "https://${NEURON_PIP_REPO}") \
RUN --mount=type=secret,id=neuron_repo \
PIP_REPO_URL="https://${NEURON_PIP_REPO}" \
&& if [ -s /run/secrets/neuron_repo ]; then \
CRED="$(cat /run/secrets/neuron_repo)"; \
ENC="$(printf '%s' "$CRED" | python3 -c 'import urllib.parse,sys; u,p=sys.stdin.read().split(":",1); print(urllib.parse.quote(u,safe="")+":"+urllib.parse.quote(p,safe=""))')"; \
PIP_REPO_URL="https://${ENC}@${NEURON_PIP_REPO}"; \
fi \
&& ${PIP} install --no-cache-dir --force-reinstall \
--index-url ${PIP_REPO_URL} \
--extra-index-url ${PYPI_SIMPLE_URL} \
Expand All @@ -203,19 +216,31 @@

FROM base AS prod

# Install Neuron components
# Install Neuron Driver, Runtime and Tools
RUN apt-get update \
# Install Neuron components from the apt and pip repos
RUN --mount=type=secret,id=neuron_repo \
mkdir -p /etc/apt/auth.conf.d \
&& if [ -s /run/secrets/neuron_repo ]; then \
CRED="$(cat /run/secrets/neuron_repo)"; \
printf 'machine %s login %s password %s\n' "${NEURON_APT_REPO}" "${CRED%%:*}" "${CRED#*:}" > /etc/apt/auth.conf.d/neuron.conf; \
chmod 600 /etc/apt/auth.conf.d/neuron.conf; \
fi \
&& 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
&& apt-get clean \
&& rm -f /etc/apt/auth.conf.d/neuron.conf

# 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}") \
RUN --mount=type=secret,id=neuron_repo \
PIP_REPO_URL="https://${NEURON_PIP_REPO}" \
&& if [ -s /run/secrets/neuron_repo ]; then \
CRED="$(cat /run/secrets/neuron_repo)"; \
ENC="$(printf '%s' "$CRED" | python3 -c 'import urllib.parse,sys; u,p=sys.stdin.read().split(":",1); print(urllib.parse.quote(u,safe="")+":"+urllib.parse.quote(p,safe=""))')"; \
PIP_REPO_URL="https://${ENC}@${NEURON_PIP_REPO}"; \
fi \
&& ${PIP} install --no-cache-dir --force-reinstall \
--index-url ${PIP_REPO_URL} \
--trusted-host ${NEURON_PIP_REPO} \
Expand All @@ -229,7 +254,8 @@
FROM ${BUILD_STAGE} AS final

# Upgrade OS packages to latest versions
RUN apt-get update \
RUN rm -f /etc/apt/sources.list.d/neuron.list /etc/apt/keyrings/neuron.gpg \
&& apt-get update \
&& apt-get upgrade -y \
&& rm -rf /var/lib/apt/lists/* \
&& apt-get clean
Expand Down
Loading
Loading