Skip to content

Commit 5204747

Browse files
JyothirmaikottuJyothirmaikottu
committed
add uv and pin sparkmagic to 0.23
1 parent 8198cd3 commit 5204747

File tree

4 files changed

+50
-47
lines changed

4 files changed

+50
-47
lines changed

‎dlc_developer_config.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ dlc-pr-base = ""
125125

126126
# Standard Framework Training
127127
dlc-pr-pytorch-training = ""
128-
dlc-pr-tensorflow-2-training = "tensorflow/training/buildspec-2-19-sm.yml"
128+
dlc-pr-tensorflow-2-training = ""
129129
dlc-pr-autogluon-training = ""
130130

131131
# ARM64 Training

‎tensorflow/training/buildspec-2-19-sm.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ images:
3636
BuildTensorflowSageMakerCpuPy310TrainingDockerImage:
3737
<<: *TRAINING_REPOSITORY
3838
build: &TENSORFLOW_CPU_TRAINING_PY3 false
39-
image_size_baseline: &IMAGE_SIZE_BASELINE 14000
39+
image_size_baseline: &IMAGE_SIZE_BASELINE 7500
4040
device_type: &DEVICE_TYPE cpu
4141
python_version: &DOCKER_PYTHON_VERSION py3
4242
tag_python_version: &TAG_PYTHON_VERSION py312
@@ -52,7 +52,7 @@ images:
5252
BuildTensorflowSageMakerGpuPy310Cu125TrainingDockerImage:
5353
<<: *TRAINING_REPOSITORY
5454
build: &TENSORFLOW_GPU_TRAINING_PY3 false
55-
image_size_baseline: &IMAGE_SIZE_BASELINE 21000
55+
image_size_baseline: &IMAGE_SIZE_BASELINE 11998
5656
device_type: &DEVICE_TYPE gpu
5757
python_version: &DOCKER_PYTHON_VERSION py3
5858
tag_python_version: &TAG_PYTHON_VERSION py312

‎tensorflow/training/docker/2.19/py3/Dockerfile.cpu‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -130,15 +130,20 @@ RUN wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSIO
130130
&& make install \
131131
&& rm -rf ../Python-$PYTHON_VERSION*
132132

133-
RUN ${PIP} --no-cache-dir install --upgrade \
133+
# Install uv for faster package management
134+
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh
135+
136+
# Upgrade pip and setuptools (setuptools>=81.0.0 fixes CVE-2026-23949)
137+
RUN uv pip install --system --no-cache \
134138
pip \
135139
"setuptools>=81.0.0"
136140

137141
# Some TF tools expect a "python" binary
138142
RUN ln -s $(which ${PYTHON}) /usr/local/bin/python \
139143
&& ln -s $(which ${PIP}) /usr/bin/pip
140144

141-
RUN ${PIP} install --no-cache-dir -U \
145+
# Install common dependencies using uv
146+
RUN uv pip install --system --no-cache \
142147
pybind11 \
143148
cmake \
144149
scipy \
@@ -148,9 +153,6 @@ RUN ${PIP} install --no-cache-dir -U \
148153
"awscli<2" \
149154
urllib3 \
150155
mpi4py \
151-
# Let's install TensorFlow separately in the end to avoid
152-
# the library version to be overwritten
153-
&& ${PIP} install --no-cache-dir -U \
154156
h5py \
155157
absl-py \
156158
opencv-python \
@@ -189,7 +191,7 @@ RUN curl https://aws-dlc-licenses.s3.amazonaws.com/tensorflow-${TF_VERSION}/lice
189191
FROM common AS ec2
190192
ARG TF_URL
191193

192-
RUN ${PIP} install --no-cache-dir -U \
194+
RUN uv pip install --system --no-cache \
193195
${TF_URL} \
194196
"tensorflow-io==0.37.*" \
195197
tensorflow-datasets
@@ -254,32 +256,28 @@ RUN apt-get update \
254256
# the sagemaker package will revert pyyaml back to 6 for its requirement
255257
# and this is fine since sagemaker is more important than the models and
256258
# the models still work on pyyaml 6 in this context.
257-
# Need to install wheel before we can fix the pyyaml issue below
258-
RUN pip install --no-cache-dir -U \
259+
RUN uv pip install --system --no-cache \
259260
wheel \
260261
"cython<3" \
261-
"pyyaml<6" \
262-
--no-build-isolation
262+
"pyyaml<6"
263263

264264
# https://github.com/tensorflow/models/issues/9267
265265
# tf-models does not respect existing installations of TF and always installs open source TF
266-
RUN ${PIP} install \
267-
--default-timeout=300 \
268-
--retries 5 \
269-
--no-cache-dir -U \
266+
RUN uv pip install --system --no-cache \
270267
"tf-models-official==${TF_VERSION}.1" \
271268
"tensorflow-text==${TF_VERSION}.0" \
272-
&& ${PIP} uninstall -y tensorflow tensorflow-gpu \
273-
&& ${PIP} install --no-cache-dir -U \
269+
&& uv pip uninstall --system tensorflow tensorflow-gpu \
270+
&& uv pip install --system --no-cache \
274271
${TF_URL} \
275272
"tensorflow-io==0.37.*" \
276273
tensorflow-datasets
277274

278-
# Install rust and cargo
275+
# Install rust and cargo (required for some packages)
279276
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
280277
ENV PATH="/root/.cargo/bin:${PATH}"
281278

282-
RUN $PYTHON -m pip install --no-cache-dir -U \
279+
# Install data science packages
280+
RUN uv pip install --system --no-cache \
283281
numba \
284282
bokeh \
285283
imageio \
@@ -288,14 +286,15 @@ RUN $PYTHON -m pip install --no-cache-dir -U \
288286
seaborn \
289287
shap
290288

291-
RUN $PYTHON -m pip install --no-cache-dir -U \
289+
# Install SageMaker packages
290+
RUN uv pip install --system --no-cache \
292291
sagemaker \
293292
sagemaker-experiments \
294293
sagemaker-tensorflow-training \
295294
sagemaker-training \
296295
y-py \
297296
sagemaker-studio-analytics-extension \
298-
"sparkmagic<1" \
297+
"sparkmagic==0.23.0" \
299298
sagemaker-studio-sparkmagic-lib \
300299
smclarify
301300

@@ -305,6 +304,10 @@ RUN /usr/local/bin/jupyter-kernelspec remove -f python3
305304
# Remove ipykernel
306305
RUN rm -rf /usr/local/share/jupyter/kernels/python3/kernel.json
307306

307+
# Clean up JupyterLab staging and test directories to reduce image size
308+
RUN rm -rf /usr/local/lib/python*/site-packages/jupyterlab/staging \
309+
&& rm -rf /usr/local/lib/python*/site-packages/jupyterlab/tests
310+
308311
# remove tmp files
309312
RUN rm -rf /tmp/*
310313

‎tensorflow/training/docker/2.19/py3/cu125/Dockerfile.gpu‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,6 @@ RUN apt-get update \
179179
ffmpeg \
180180
libsm6 \
181181
libxext6 \
182-
# remove libsasl2-2 after Nvidia docker upstream fix CVE-2022-24407
183182
&& rm -rf /var/lib/apt/lists/* \
184183
&& apt-get clean
185184

@@ -191,7 +190,11 @@ RUN wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSIO
191190
&& make install \
192191
&& rm -rf ../Python-$PYTHON_VERSION*
193192

194-
RUN ${PIP} --no-cache-dir install --upgrade \
193+
# Install uv for faster package management
194+
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh
195+
196+
# Upgrade pip and setuptools (setuptools>=81.0.0 fixes CVE-2026-23949)
197+
RUN uv pip install --system --no-cache \
195198
pip \
196199
"setuptools>=81.0.0" \
197200
wheel
@@ -200,7 +203,8 @@ RUN ${PIP} --no-cache-dir install --upgrade \
200203
RUN ln -s $(which ${PYTHON}) /usr/local/bin/python \
201204
&& ln -s $(which ${PIP}) /usr/bin/pip
202205

203-
RUN ${PIP} install --no-cache-dir -U \
206+
# Install common dependencies using uv
207+
RUN uv pip install --system --no-cache \
204208
pybind11 \
205209
cmake \
206210
scipy \
@@ -269,7 +273,7 @@ FROM common AS ec2
269273

270274
ARG TF_URL
271275

272-
RUN ${PIP} install --no-cache-dir -U \
276+
RUN uv pip install --system --no-cache \
273277
${TF_URL} \
274278
"tensorflow-io==0.37.*" \
275279
tensorflow-datasets
@@ -318,40 +322,31 @@ ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main
318322
# https://github.com/yaml/pyyaml/issues/601
319323
# PyYaml less than 6.0.1 failes to build with cython v3 and above.
320324
# tf-models-official uses older versions, breaking the install.
321-
# going to install the older pyyaml and cython to get tfd-models-official
322-
# the sagemaker package will revert pyyaml back to 6 for its requirement
323-
# and this is fine since sagemaker is more important than the models and
324-
# the models still work on pyyaml 6 in this context.
325-
# Need to install wheel before we can fix the pyyaml issue below
326-
RUN pip install --no-cache-dir -U \
325+
RUN uv pip install --system --no-cache \
327326
wheel \
328327
"cython<3" \
329-
"pyyaml<6" \
330-
--no-build-isolation
331-
328+
"pyyaml<6"
332329

333330
# Pin numpy to 1.26.4
334-
RUN ${PIP} install --no-cache-dir -U "numpy==1.26.4"
331+
RUN uv pip install --system --no-cache "numpy==1.26.4"
335332

336333
# https://github.com/tensorflow/models/issues/9267
337334
# tf-models does not respect existing installations of TF and always installs open source TF
338-
RUN ${PIP} install \
339-
--default-timeout=300 \
340-
--retries 5 \
341-
--no-cache-dir -U \
335+
RUN uv pip install --system --no-cache \
342336
"tf-models-official==${TF_VERSION}.1" \
343337
"tensorflow-text==${TF_VERSION}.0" \
344-
&& ${PIP} uninstall -y tensorflow tensorflow-gpu \
345-
&& ${PIP} install --no-cache-dir -U \
338+
&& uv pip uninstall --system tensorflow tensorflow-gpu \
339+
&& uv pip install --system --no-cache \
346340
${TF_URL} \
347341
"tensorflow-io==0.37.*" \
348342
tensorflow-datasets
349343

350-
# Install rust and cargo
344+
# Install rust and cargo (required for some packages)
351345
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
352346
ENV PATH="/root/.cargo/bin:${PATH}"
353347

354-
RUN $PYTHON -m pip install --no-cache-dir -U \
348+
# Install data science packages
349+
RUN uv pip install --system --no-cache \
355350
numba \
356351
bokeh \
357352
imageio \
@@ -360,14 +355,15 @@ RUN $PYTHON -m pip install --no-cache-dir -U \
360355
seaborn \
361356
shap
362357

363-
RUN $PYTHON -m pip install --no-cache-dir -U \
358+
# Install SageMaker packages
359+
RUN uv pip install --system --no-cache \
364360
sagemaker \
365361
sagemaker-experiments==0.* \
366362
sagemaker-tensorflow-training \
367363
sagemaker-training \
368364
y-py \
369365
sagemaker-studio-analytics-extension \
370-
"sparkmagic<1" \
366+
"sparkmagic==0.23.0" \
371367
sagemaker-studio-sparkmagic-lib \
372368
smclarify
373369

@@ -392,6 +388,10 @@ RUN /usr/local/bin/jupyter-kernelspec remove -f python3
392388
# Remove ipykernel
393389
RUN rm -rf /usr/local/share/jupyter/kernels/python3/kernel.json
394390

391+
# Clean up JupyterLab staging and test directories to reduce image size
392+
RUN rm -rf /usr/local/lib/python*/site-packages/jupyterlab/staging \
393+
&& rm -rf /usr/local/lib/python*/site-packages/jupyterlab/tests
394+
395395
# remove tmp files
396396
RUN rm -rf /tmp/*
397397

0 commit comments

Comments
 (0)