Skip to content

Commit 331f755

Browse files
neuron-containers-cineuron-containers-ci
authored
Updated Dockerfiles for 2.24.1 release (#111)
Updated Dockerfiles
1 parent d82d02c commit 331f755

4 files changed

Lines changed: 71 additions & 251 deletions

File tree

‎docker/jax/training/0.6/Dockerfile.neuronx‎

Lines changed: 4 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -171,83 +171,23 @@ ARG NEURONX_TOOLS_VERSION=2.24.54.0
171171
ARG NEURONX_CC_VERSION=2.19.8089.0
172172
ARG NEURONX_JAX_TRAINING_VERSION=0.6.0.1.0.1296+1f770067
173173

174-
FROM base AS dev
175-
176-
RUN --mount=type=bind,source=apt,target=${NEURON_ARTIFACT_PATH}/apt \
177-
install_apt_package() { \
178-
pkg_name=$1; \
179-
version_arg=$2; \
180-
if [ -f "${NEURON_ARTIFACT_PATH}/apt/${version_arg}" ]; then \
181-
apt-get install -y ${NEURON_ARTIFACT_PATH}/apt/${version_arg}; \
182-
elif [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
183-
apt-get install -y ${pkg_name}=${version_arg}; \
184-
else \
185-
echo "Ignoring package ${pkg_name}"; \
186-
fi; \
187-
} \
188-
&& apt-get update \
189-
&& install_apt_package "aws-neuronx-collectives" "${NEURONX_COLLECTIVES_LIB_VERSION}" \
190-
&& install_apt_package "aws-neuronx-runtime-lib" "${NEURONX_RUNTIME_LIB_VERSION}" \
191-
&& install_apt_package "aws-neuronx-tools" "${NEURONX_TOOLS_VERSION}" \
192-
&& rm -rf /var/lib/apt/lists/* \
193-
&& rm -rf /tmp/tmp* \
194-
&& apt-get clean
195-
196-
RUN --mount=type=bind,source=pip,target=${NEURON_ARTIFACT_PATH}/pip \
197-
install_pip_package() { \
198-
packages=""; \
199-
flags=""; \
200-
while [ "$#" -gt 0 ]; do \
201-
pkg_name=$(echo $1 | cut -d: -f1); \
202-
version_arg=$(echo $1 | cut -d: -f2); \
203-
extra_flags=$(echo $1 | cut -d: -f3); \
204-
if [ -f "${NEURON_ARTIFACT_PATH}/pip/${version_arg}" ]; then \
205-
packages="${packages} ${NEURON_ARTIFACT_PATH}/pip/${version_arg}"; \
206-
else \
207-
if [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
208-
packages="${packages} ${pkg_name}==${version_arg}"; \
209-
else \
210-
echo "Ignoring package ${pkg_name}"; \
211-
fi; \
212-
fi; \
213-
# Store unique flags
214-
if [ ! -z "${extra_flags}" ]; then \
215-
for flag in $(echo "${extra_flags}" | tr ' ' '\n'); do \
216-
case " ${flags} " in \
217-
*" ${flag} "*) ;; \
218-
*) flags="${flags} ${flag}" ;; \
219-
esac \
220-
done; \
221-
fi; \
222-
shift; \
223-
done; \
224-
if [ ! -z "${packages}" ]; then \
225-
echo "Installing packages: ${packages} with flags ${flags}"; \
226-
${PIP} install --no-cache-dir --force-reinstall \
227-
--extra-index-url="file:///${NEURON_ARTIFACT_PATH}/pip" \
228-
${packages} ${flags}; \
229-
fi; \
230-
} \
231-
&& install_pip_package "neuronx-cc:${NEURONX_CC_VERSION}:" "jax-neuronx:${NEURONX_JAX_TRAINING_VERSION}:" \
232-
&& rm -rf ~/.cache/pip/*
233-
234174
FROM base AS repo
235175

236-
# Install Neuron components from the apt and pip repos
176+
# Install Neuron components from the apt and pip repos (latest versions)
237177
RUN apt-get update \
238178
&& apt-get install -y \
239179
aws-neuronx-tools \
240180
aws-neuronx-collectives \
241181
aws-neuronx-runtime-lib \
242182
&& rm -rf /var/lib/apt/lists/* \
243183
&& rm -rf /tmp/tmp* \
244-
&& apt-get clean \
245-
&& ${PIP} install --no-cache-dir --force-reinstall \
184+
&& apt-get clean
185+
186+
RUN ${PIP} install --no-cache-dir --force-reinstall \
246187
neuronx-cc \
247188
jax-neuronx \
248189
&& rm -rf ~/.cache/pip/*
249190

250-
251191
FROM base AS prod
252192

253193
# Install Neuron components

‎docker/pytorch/inference/2.7.0/Dockerfile.neuronx‎

Lines changed: 22 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -159,107 +159,46 @@ ARG NEURONX_TRANSFORMERS_VERSION=0.13.985
159159
ARG NEURONX_DISTRIBUTED_VERSION=0.13.14393+b8569585
160160
ARG NEURONX_DISTRIBUTED_INFERENCE_VERSION=0.4.7422+9483d307
161161

162-
FROM base AS dev
162+
FROM base AS repo
163163

164-
RUN --mount=type=bind,source=apt,target=${NEURON_ARTIFACT_PATH}/apt \
165-
install_apt_package() { \
166-
pkg_name=$1; \
167-
version_arg=$2; \
168-
if [ -f "${NEURON_ARTIFACT_PATH}/apt/${version_arg}" ]; then \
169-
apt-get install -y ${NEURON_ARTIFACT_PATH}/apt/${version_arg}; \
170-
elif [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
171-
apt-get install -y ${pkg_name}=${version_arg}; \
172-
else \
173-
echo "Ignoring package ${pkg_name}"; \
174-
fi; \
175-
} \
176-
&& apt-get update \
177-
&& install_apt_package "aws-neuronx-collectives" "${NEURONX_COLLECTIVES_LIB_VERSION}" \
178-
&& install_apt_package "aws-neuronx-runtime-lib" "${NEURONX_RUNTIME_LIB_VERSION}" \
179-
&& install_apt_package "aws-neuronx-tools" "${NEURONX_TOOLS_VERSION}" \
164+
# Install Neuron components from the apt and pip repos (latest versions)
165+
RUN apt-get update \
166+
&& apt-get install -y \
167+
aws-neuronx-tools \
168+
aws-neuronx-collectives \
169+
aws-neuronx-runtime-lib \
180170
&& rm -rf /var/lib/apt/lists/* \
181171
&& rm -rf /tmp/tmp* \
182172
&& apt-get clean
183173

184-
RUN --mount=type=bind,source=pip,target=${NEURON_ARTIFACT_PATH}/pip \
185-
install_pip_package() { \
186-
packages=""; \
187-
flags=""; \
188-
while [ "$#" -gt 0 ]; do \
189-
pkg_name=$(echo $1 | cut -d: -f1); \
190-
version_arg=$(echo $1 | cut -d: -f2); \
191-
extra_flags=$(echo $1 | cut -d: -f3); \
192-
if [ -f "${NEURON_ARTIFACT_PATH}/pip/${version_arg}" ]; then \
193-
packages="${packages} ${NEURON_ARTIFACT_PATH}/pip/${version_arg}"; \
194-
else \
195-
if [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
196-
packages="${packages} ${pkg_name}==${version_arg}"; \
197-
else \
198-
echo "Ignoring package ${pkg_name}"; \
199-
fi; \
200-
fi; \
201-
# Store unique flags
202-
if [ ! -z "${extra_flags}" ]; then \
203-
for flag in $(echo "${extra_flags}" | tr ' ' '\n'); do \
204-
case " ${flags} " in \
205-
*" ${flag} "*) ;; \
206-
*) flags="${flags} ${flag}" ;; \
207-
esac \
208-
done; \
209-
fi; \
210-
shift; \
211-
done; \
212-
if [ ! -z "${packages}" ]; then \
213-
echo "Installing packages: ${packages} with flags ${flags}"; \
214-
${PIP} install --no-cache-dir \
215-
--extra-index-url="file:///${NEURON_ARTIFACT_PATH}/pip" \
216-
${packages} ${flags}; \
217-
fi; \
218-
} \
219-
&& install_pip_package "neuronx-cc:${NEURONX_CC_VERSION}:" "torch-neuronx:${NEURONX_FRAMEWORK_VERSION}:" \
220-
&& install_pip_package "transformers-neuronx:${NEURONX_TRANSFORMERS_VERSION}:" \
221-
&& install_pip_package "neuronx_distributed:${NEURONX_DISTRIBUTED_VERSION}:--no-deps" \
222-
&& install_pip_package "neuronx_distributed_inference:${NEURONX_DISTRIBUTED_INFERENCE_VERSION}:--no-deps" \
174+
RUN ${PIP} install --no-cache-dir \
175+
neuronx-cc \
176+
torch-neuronx \
177+
transformers-neuronx \
178+
&& ${PIP} install --no-cache-dir --no-deps \
179+
neuronx_distributed \
180+
neuronx_distributed_inference \
223181
&& rm -rf ~/.cache/pip/*
224182

225-
FROM base AS repo
226-
227-
# Install Neuron components from the apt and pip repos
228-
RUN apt-get update \
229-
&& apt-get install -y \
230-
aws-neuronx-tools \
231-
aws-neuronx-collectives \
232-
aws-neuronx-runtime-lib \
233-
&& pip install --no-cache-dir \
234-
neuronx-cc \
235-
torch-neuronx \
236-
transformers-neuronx \
237-
&& pip install --no-cache-dir --no-deps \
238-
neuronx_distributed \
239-
neuronx_distributed_inference \
240-
&& rm -rf /var/lib/apt/lists/* \
241-
&& rm -rf /tmp/tmp* \
242-
&& rm -rf ~/.cache/pip/* \
243-
&& apt-get clean
244-
245183
FROM base AS prod
246184

185+
# Install Neuron components with specific versions
247186
RUN apt-get update \
248187
&& apt-get install -y \
249-
aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
250-
aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
251-
aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
188+
aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
189+
aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
190+
aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
252191
&& rm -rf /var/lib/apt/lists/* \
253192
&& rm -rf /tmp/tmp* \
254193
&& apt-get clean
255194

256-
RUN pip install --no-cache-dir \
195+
RUN ${PIP} install --no-cache-dir \
257196
neuronx-cc==$NEURONX_CC_VERSION \
258197
torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
259198
transformers-neuronx==$NEURONX_TRANSFORMERS_VERSION \
260-
&& pip install --no-cache-dir --no-deps \
261-
neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
262-
neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION \
199+
&& ${PIP} install --no-cache-dir --no-deps \
200+
neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
201+
neuronx_distributed_inference==$NEURONX_DISTRIBUTED_INFERENCE_VERSION \
263202
&& rm -rf ~/.cache/pip/*
264203

265204
FROM ${BUILD_STAGE} AS final

‎docker/pytorch/training/2.7.0/Dockerfile.neuronx‎

Lines changed: 16 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -231,106 +231,46 @@ ARG NEURONX_TOOLS_VERSION=2.24.54.0
231231
ARG NEURONX_FRAMEWORK_VERSION=2.7.0.2.8.6734+ac864f72
232232
ARG NEURONX_CC_VERSION=2.19.8089.0
233233
ARG NEURONX_DISTRIBUTED_VERSION=0.13.14393+b8569585
234-
ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.4.0
235-
236-
FROM base AS dev
237-
238-
RUN --mount=type=bind,source=apt,target=${NEURON_ARTIFACT_PATH}/apt \
239-
install_apt_package() { \
240-
pkg_name=$1; \
241-
version_arg=$2; \
242-
if [ -f "${NEURON_ARTIFACT_PATH}/apt/${version_arg}" ]; then \
243-
apt-get install -y ${NEURON_ARTIFACT_PATH}/apt/${version_arg}; \
244-
elif [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
245-
apt-get install -y ${pkg_name}=${version_arg}; \
246-
else \
247-
echo "Ignoring package ${pkg_name}"; \
248-
fi; \
249-
} \
250-
&& apt-get update \
251-
&& install_apt_package "aws-neuronx-collectives" "${NEURONX_COLLECTIVES_LIB_VERSION}" \
252-
&& install_apt_package "aws-neuronx-runtime-lib" "${NEURONX_RUNTIME_LIB_VERSION}" \
253-
&& install_apt_package "aws-neuronx-tools" "${NEURONX_TOOLS_VERSION}" \
254-
&& rm -rf /var/lib/apt/lists/* \
255-
&& rm -rf /tmp/tmp* \
256-
&& apt-get clean
257-
258-
RUN --mount=type=bind,source=pip,target=${NEURON_ARTIFACT_PATH}/pip \
259-
install_pip_package() { \
260-
packages=""; \
261-
flags=""; \
262-
while [ "$#" -gt 0 ]; do \
263-
pkg_name=$(echo $1 | cut -d: -f1); \
264-
version_arg=$(echo $1 | cut -d: -f2); \
265-
extra_flags=$(echo $1 | cut -d: -f3); \
266-
if [ -f "${NEURON_ARTIFACT_PATH}/pip/${version_arg}" ]; then \
267-
packages="${packages} ${NEURON_ARTIFACT_PATH}/pip/${version_arg}"; \
268-
else \
269-
if [ "${IGNORE_MISSING_NEURON_COMPONENTS}" = "false" ]; then \
270-
packages="${packages} ${pkg_name}==${version_arg}"; \
271-
else \
272-
echo "Ignoring package ${pkg_name}"; \
273-
fi; \
274-
fi; \
275-
# Store unique flags
276-
if [ ! -z "${extra_flags}" ]; then \
277-
for flag in $(echo "${extra_flags}" | tr ' ' '\n'); do \
278-
case " ${flags} " in \
279-
*" ${flag} "*) ;; \
280-
*) flags="${flags} ${flag}" ;; \
281-
esac \
282-
done; \
283-
fi; \
284-
shift; \
285-
done; \
286-
if [ ! -z "${packages}" ]; then \
287-
echo "Installing packages: ${packages} with flags ${flags}"; \
288-
${PIP} install --no-cache-dir --force-reinstall \
289-
--extra-index-url="file:///${NEURON_ARTIFACT_PATH}/pip" \
290-
${packages} ${flags}; \
291-
fi; \
292-
} \
293-
&& install_pip_package "neuronx-cc:${NEURONX_CC_VERSION}:" "torch-neuronx:${NEURONX_FRAMEWORK_VERSION}:" \
294-
&& install_pip_package "neuronx_distributed:${NEURONX_DISTRIBUTED_VERSION}:--no-deps" \
295-
&& install_pip_package "neuronx_distributed_training:${NEURONX_DISTRIBUTED_TRAINING_VERSION}:--no-deps" \
296-
&& rm -rf ~/.cache/pip/*
234+
ARG NEURONX_DISTRIBUTED_TRAINING_VERSION=1.4.1
297235

298236
FROM base AS repo
299237

300-
# Install Neuron components from the apt and pip repos
238+
# Install Neuron components from the apt and pip repos (latest versions)
301239
RUN apt-get update \
302240
&& apt-get install -y \
303241
aws-neuronx-tools \
304242
aws-neuronx-collectives \
305243
aws-neuronx-runtime-lib \
306-
&& ${PIP} install --no-cache-dir --force-reinstall \
244+
&& rm -rf /var/lib/apt/lists/* \
245+
&& rm -rf /tmp/tmp* \
246+
&& apt-get clean
247+
248+
RUN ${PIP} install --no-cache-dir --force-reinstall \
307249
torch-neuronx \
308250
neuronx-cc \
309251
&& ${PIP} install --no-cache-dir --no-deps \
310252
neuronx_distributed \
311253
neuronx_distributed_training \
312-
&& rm -rf /var/lib/apt/lists/* \
313-
&& rm -rf /tmp/tmp* \
314-
&& rm -rf ~/.cache/pip/* \
315-
&& apt-get clean
254+
&& rm -rf ~/.cache/pip/*
316255

317256
FROM base AS prod
318257

258+
# Install Neuron components with specific versions
319259
RUN apt-get update \
320260
&& apt-get install -y \
321-
aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
322-
aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
323-
aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
261+
aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
262+
aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
263+
aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
324264
&& rm -rf /var/lib/apt/lists/* \
325265
&& rm -rf /tmp/tmp* \
326266
&& apt-get clean
327267

328268
RUN ${PIP} install --force-reinstall \
329-
torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
330-
neuronx-cc==$NEURONX_CC_VERSION \
269+
torch-neuronx==$NEURONX_FRAMEWORK_VERSION \
270+
neuronx-cc==$NEURONX_CC_VERSION \
331271
&& ${PIP} install --force-reinstall --no-deps \
332-
neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
333-
neuronx_distributed_training==$NEURONX_DISTRIBUTED_TRAINING_VERSION \
272+
neuronx_distributed==$NEURONX_DISTRIBUTED_VERSION \
273+
neuronx_distributed_training==$NEURONX_DISTRIBUTED_TRAINING_VERSION \
334274
&& rm -rf ~/.cache/pip/*
335275

336276
FROM ${BUILD_STAGE} AS final

0 commit comments

Comments
 (0)