Skip to content

Commit f664ec2

Browse files
soluwalanasoluwalana
authored
fix(customizer): address CVEs in unsloth and automodel (#474)
* fix(customizer): address CVEs in unsloth and automodel Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * fix(customizer): duplicate pin Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * rebuild and retest Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Remove the ffmpeg bundled in, and comment how it breaks things Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Rollback file not used Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * lint fix Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Update liscences Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * fix(customizer): fix dockerfile organization for customizer, add AV opencv build with ffmpeg 8.1.2+ Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * wire up ffmpeg Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Fixes Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * fix ref to av Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * update sha for automodel build Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> * Cap greenlet to allow arm64 to build Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com> --------- Signed-off-by: Sam Oluwalana <soluwalana@nvidia.com>
1 parent 7b058c6 commit f664ec2

21 files changed

Lines changed: 494 additions & 69 deletions

‎docker-bake.hcl‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ variable "MAMBA_SSM_WHEEL_CONTEXT" {
5252
default = ""
5353
}
5454

55+
variable "FFMPEG_VLM_WHEEL_CONTEXT" {
56+
default = ""
57+
}
58+
5559
variable "DISTROLESS_BASE" {
5660
default = "nvcr.io/nvidia/distroless/python:3.11-v4.0.8"
5761
}
@@ -75,12 +79,12 @@ variable "BASE_TAG_PYTHON" {
7579

7680
# Pin for nmp-automodel-base.
7781
variable "BASE_TAG_AUTOMODEL" {
78-
default = "f8239353044d71cbd53e209b60c0600ead484b58"
82+
default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
7983
}
8084

8185
# The tag for base images if needed
8286
variable "WHEELS_TAG" {
83-
default = "65632527b258367fb8855ef431f596abb8538577"
87+
default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
8488
}
8589

8690
variable "BAKE_CACHE_SOURCE_BRANCH" {
@@ -138,6 +142,11 @@ function "get_mamba_ssm_wheel_image" {
138142
result = "${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}"
139143
}
140144

145+
function "get_ffmpeg_vlm_wheel_image" {
146+
params = []
147+
result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}"
148+
}
149+
141150
function "get_arch_tag" {
142151
params = []
143152
result = BUILD_ARCH == "linux/arm64" ? "linux-arm64" : "linux-amd64"
@@ -175,6 +184,11 @@ function "mamba_ssm_wheel_context" {
175184
result = notequal(MAMBA_SSM_WHEEL_CONTEXT, "") ? MAMBA_SSM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:mamba-ssm-wheel" : "docker-image://${get_mamba_ssm_wheel_image()}"
176185
}
177186

187+
function "ffmpeg_vlm_wheel_context" {
188+
params = []
189+
result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}"
190+
}
191+
178192
function "wheel_tags" {
179193
params = [name]
180194
result = [
@@ -293,6 +307,7 @@ group "nmp-automodel-gpu-wheels" {
293307
targets = [
294308
"causal-conv1d-wheel",
295309
"mamba-ssm-wheel",
310+
"ffmpeg-vlm-wheel",
296311
]
297312
}
298313

@@ -500,15 +515,14 @@ target "nmp-cpu-tasks-docker" {
500515
platforms = get_platforms()
501516
}
502517

503-
# Mamba Wheel Builders
504-
# Builds Python wheels for mamba-ssm and causal-conv1d in parallel.
505-
# Both only ship source distributions on PyPI; this pre-builds them for
506-
# amd64 and arm64. The wheels live at /wheels/*.whl inside the image.
518+
# Python wheel builders (causal-conv1d, mamba-ssm, av, opencv-python-headless).
519+
# CUDA extensions only ship source on PyPI; av/opencv bundle FFmpeg. Pre-built for
520+
# amd64 and arm64. Wheels live at /wheels/*.whl inside each image.
507521

508522
target "causal-conv1d-wheel" {
509523
target = "causal-conv1d-wheel"
510524
context = "."
511-
dockerfile = "docker/base/Dockerfile.mamba-wheel"
525+
dockerfile = "docker/base/Dockerfile.python-wheels"
512526
cache-to = maybe_registry_cache_to("causal-conv1d-wheel")
513527
cache-from = maybe_registry_cache_from("causal-conv1d-wheel")
514528
tags = wheel_tags("causal-conv1d-wheel")
@@ -523,7 +537,7 @@ target "causal-conv1d-wheel" {
523537
target "mamba-ssm-wheel" {
524538
target = "mamba-ssm-wheel"
525539
context = "."
526-
dockerfile = "docker/base/Dockerfile.mamba-wheel"
540+
dockerfile = "docker/base/Dockerfile.python-wheels"
527541
cache-to = maybe_registry_cache_to("mamba-ssm-wheel")
528542
cache-from = maybe_registry_cache_from("mamba-ssm-wheel")
529543
tags = wheel_tags("mamba-ssm-wheel")
@@ -536,6 +550,17 @@ target "mamba-ssm-wheel" {
536550
platforms = get_platforms()
537551
}
538552

553+
target "ffmpeg-vlm-wheel" {
554+
target = "ffmpeg-vlm-wheel"
555+
context = "."
556+
dockerfile = "docker/base/Dockerfile.python-wheels"
557+
cache-to = maybe_registry_cache_to("ffmpeg-vlm-wheel")
558+
cache-from = maybe_registry_cache_from("ffmpeg-vlm-wheel")
559+
tags = wheel_tags("ffmpeg-vlm-wheel")
560+
output = image_output()
561+
platforms = get_platforms()
562+
}
563+
539564

540565
target "safe-synthesizer-tasks-docker" {
541566
target = "runtime"
@@ -698,7 +723,7 @@ target "automodel-platform-workspace" {
698723
target "nmp-automodel-base-builder" {
699724
target = "nmp-automodel-base"
700725
context = "."
701-
dockerfile = "docker/Dockerfile.nmp-automodel-base"
726+
dockerfile = "docker/automodel/Dockerfile.nmp-automodel-base"
702727
no-cache-filter = ["automodel-clone"]
703728
cache-to = maybe_registry_cache_to("nmp-automodel-base")
704729
cache-from = maybe_registry_cache_from("nmp-automodel-base")
@@ -707,14 +732,15 @@ target "nmp-automodel-base-builder" {
707732
contexts = {
708733
causal-conv1d-wheel-image = causal_conv1d_wheel_context()
709734
mamba-ssm-wheel-image = mamba_ssm_wheel_context()
735+
ffmpeg-vlm-wheel-image = ffmpeg_vlm_wheel_context()
710736
}
711737
platforms = get_platforms()
712738
}
713739

714740
target "nmp-automodel-tasks-docker" {
715741
target = "runtime"
716742
context = "."
717-
dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
743+
dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
718744
contexts = {
719745
platform-workspace = "target:automodel-platform-workspace"
720746
nmp-automodel-base = automodel_base_context()
@@ -729,7 +755,7 @@ target "nmp-automodel-tasks-docker" {
729755
target "nmp-automodel-training-docker" {
730756
target = "runtime"
731757
context = "."
732-
dockerfile = "docker/Dockerfile.nmp-automodel-training"
758+
dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
733759
contexts = {
734760
platform-workspace = "target:automodel-platform-workspace"
735761
nmp-automodel-base = automodel_base_context()
@@ -744,7 +770,7 @@ target "nmp-automodel-training-docker" {
744770
target "nmp-automodel-tasks-smoke-test" {
745771
target = "smoke-test"
746772
context = "."
747-
dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
773+
dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
748774
contexts = {
749775
platform-workspace = "target:automodel-platform-workspace"
750776
nmp-automodel-base = automodel_base_context()
@@ -760,7 +786,7 @@ target "nmp-automodel-tasks-smoke-test" {
760786
target "nmp-automodel-training-smoke-test" {
761787
target = "smoke-test"
762788
context = "."
763-
dockerfile = "docker/Dockerfile.nmp-automodel-training"
789+
dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
764790
contexts = {
765791
platform-workspace = "target:automodel-platform-workspace"
766792
nmp-automodel-base = automodel_base_context()

‎docker/Dockerfile.nmp-unsloth-training‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@
1313
# blocks uv from installing/upgrading torch into the venv so the NGC
1414
# base's PyTorch + CUDA remain the runtime stack.
1515
# 1b. bitsandbytes — compiled from source against NGC CUDA 13.2 (same pattern
16-
# as docker/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
16+
# as docker/automodel/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
1717
# cuda130; source build replaces the wheel from step 1.
1818
# 1c. mamba-ssm + causal-conv1d — prebuilt cu13.2 / cp312 wheels (shared with
19-
# docker/Dockerfile.nmp-automodel-base).
19+
# docker/automodel/Dockerfile.nmp-automodel-base).
2020
# 1d. flash-attn — optional for unsloth and currently NOT installed (see the
2121
# commented TODO below). Without it Unsloth falls back when xformers is
2222
# also missing (common on newer CUDA stacks), logging "FA2 = False /
@@ -28,14 +28,14 @@
2828
# Default tag: `local` (override via BAKE_TAG at build time).
2929

3030
# NGC PyTorch base. 26.05-py3 ships PyTorch 2.12 + CUDA 13.2 + Python 3.12 (aligned with
31-
# docker/Dockerfile.nmp-automodel-base). Override at build time:
31+
# docker/automodel/Dockerfile.nmp-automodel-base). Override at build time:
3232
# --set nmp-unsloth-training.args.PYTORCH_BASE=...
3333
ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
3434

3535
# Prebuilt CUDA-extension wheels (mamba-ssm + causal-conv1d), shared with
36-
# docker/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
36+
# docker/automodel/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
3737
# resolves these contexts to the causal-conv1d-wheel / mamba-ssm-wheel images
38-
# (see docker-bake.hcl + docker/base/Dockerfile.mamba-wheel). Local builds need
38+
# (see docker-bake.hcl + docker/base/Dockerfile.python-wheels). Local builds need
3939
# `USE_LOCAL_WHEELS=1` (build the wheel targets) or a reachable WHEELS_REGISTRY.
4040
FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
4141
FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
@@ -153,7 +153,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
153153
-e /app/packages/nemo_platform_plugin \
154154
-e /app/packages/nmp_common \
155155
-e /app/packages/nmp_customization_common \
156-
-e "/app/services/unsloth[integrations]"
156+
-e "/app/services/unsloth"
157157

158158
# hf-transfer can pull huggingface-hub 1.x — install then re-pin hub + transformers.
159159
RUN --mount=type=cache,target=/root/.cache/uv \
@@ -179,6 +179,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
179179
"mlflow-skinny>=3.11.1,<3.12.0"
180180

181181
# Stale NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
182+
RUN apt-get update && \
183+
DEBIAN_FRONTEND=noninteractive apt-get install -y --only-upgrade openssl libssl3t64 && \
184+
rm -rf /var/lib/apt/lists/*
185+
182186
RUN rm -rf \
183187
/usr/local/lib/python3.12/dist-packages/aiohttp \
184188
/usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
@@ -199,7 +203,13 @@ RUN rm -rf \
199203
/usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
200204
/usr/local/lib/python3.12/dist-packages/mlflow \
201205
/usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
202-
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info
206+
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info \
207+
/usr/local/lib/python3.12/dist-packages/jupyter_server \
208+
/usr/local/lib/python3.12/dist-packages/jupyter_server-*.dist-info \
209+
/usr/local/cuda/NsightSystems-cli-* \
210+
/usr/local/bin/nsys \
211+
/usr/local/cuda/bin/nsys
212+
203213

204214
ENTRYPOINT ["/opt/venv/bin/python"]
205215
CMD ["-m", "nmp.unsloth.tasks.training", "--help"]

docker/Dockerfile.nmp-automodel-base renamed to docker/automodel/Dockerfile.nmp-automodel-base

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
1111

1212
FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
1313
FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
14+
FROM ffmpeg-vlm-wheel-image AS ffmpeg-vlm-wheel-src
1415

1516
FROM alpine/git AS automodel-clone
1617
ARG AUTOMODEL_COMMIT
@@ -95,18 +96,36 @@ RUN --mount=type=cache,target=/root/.cache/uv \
9596
"black>=26.3.1,<27" \
9697
"GitPython>=3.1.50,<4" \
9798
"jaraco-context>=6.1.0,<7" \
98-
"jupyter-server>=2.18.0,<3" \
99+
"jupyter-server>=2.20.0,<3" \
99100
"jupyterlab>=4.5.7,<5" \
100101
"mistune>=3.2.1,<4" \
101102
"notebook>=7.5.6,<8" \
102103
"onnx>=1.21.0,<2" \
103104
"pillow>=12.2.0,<13" \
104105
"pyasn1>=0.6.3,<1" \
106+
"pyarrow>=23.0.1,<24" \
107+
"msgpack>=1.2.1,<2" \
108+
"cryptography>=48.0.1,<49" \
105109
"tornado>=6.5.5,<7" \
106110
"urllib3>=2.7.0,<3" \
107111
"mlflow-skinny>=3.11.1,<3.12.0" \
108112
"grpcio>=1.81.1,<2" \
109-
"wandb==0.27.2"
113+
"wandb>=0.28.0,<1"
114+
115+
# Replace PyPI av/opencv (bundled FFmpeg < 8.1.2) with custom wheels from ffmpeg-vlm-wheel.
116+
# vlm sync pulls PyPI copies transitively; drop them before installing our manylinux wheels.
117+
# PyAV 17.x uses cp311-abi3 (stable ABI) — compatible with cp312, not cp312-cp312 tagged.
118+
RUN --mount=from=ffmpeg-vlm-wheel-src,target=/tmp/ffmpeg-vlm-wheel-src,readonly \
119+
rm -rf \
120+
/opt/venv/lib/python*/site-packages/av \
121+
/opt/venv/lib/python*/site-packages/av-*.dist-info \
122+
/opt/venv/lib/python*/site-packages/av.libs \
123+
/opt/venv/lib/python*/site-packages/cv2 \
124+
/opt/venv/lib/python*/site-packages/opencv_python_headless.libs \
125+
/opt/venv/lib/python*/site-packages/opencv_python_headless-*.dist-info && \
126+
uv pip install --no-cache-dir --no-deps \
127+
/tmp/ffmpeg-vlm-wheel-src/wheels/av-*-abi3-manylinux_2_28_*.whl \
128+
/tmp/ffmpeg-vlm-wheel-src/wheels/opencv_python_headless-*cp312*.whl
110129

111130
# Published base image (same filesystem as builder).
112131
FROM ${PYTORCH_BASE} AS nmp-automodel-base
@@ -119,6 +138,12 @@ COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv
119138
# Stale copies under NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
120139
# CVE scanners read dist-packages even when --system-site-packages resolves imports from /opt/venv.
121140
# rm -rf (not pip uninstall) avoids corrupted dist-info RECORDs on the NGC base image.
141+
RUN apt-get update && \
142+
DEBIAN_FRONTEND=noninteractive apt-get install -y --only-upgrade openssl libssl3t64 && \
143+
rm -rf /var/lib/apt/lists/*
144+
145+
# NGC ships Nsight Systems (~585 MB); not used by nmp-automodel training.
146+
# Cleanup stale NGC image packages.
122147
RUN rm -rf \
123148
/usr/local/lib/python3.12/dist-packages/aiohttp \
124149
/usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
@@ -141,7 +166,12 @@ RUN rm -rf \
141166
/usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
142167
/usr/local/lib/python3.12/dist-packages/mlflow \
143168
/usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
144-
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info
169+
/usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info \
170+
/usr/local/lib/python3.12/dist-packages/jupyter_server \
171+
/usr/local/lib/python3.12/dist-packages/jupyter_server-*.dist-info \
172+
/usr/local/cuda/NsightSystems-cli-* \
173+
/usr/local/bin/nsys \
174+
/usr/local/cuda/bin/nsys
145175

146176
ENV VIRTUAL_ENV=/opt/venv \
147177
UV_PROJECT_ENVIRONMENT=/opt/venv \

docker/Dockerfile.nmp-automodel-tasks renamed to docker/automodel/Dockerfile.nmp-automodel-tasks

File renamed without changes.

docker/Dockerfile.nmp-automodel-training renamed to docker/automodel/Dockerfile.nmp-automodel-training

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
4040
uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps \
4141
-e /opt/Automodel
4242

43-
4443
ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
4544
ENTRYPOINT ["/opt/venv/bin/python"]
4645
CMD ["-m", "nmp.automodel.tasks.training", "--help"]

‎docker/automodel/README.md‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
# nmp-automodel container images
22

3+
All Automodel Docker build files live under **`docker/automodel/`** (wheel vendor script: `docker/base/build-ffmpeg-vendor.sh`).
4+
35
Three images for the **nmp-automodel** customization backend. Published as flat repo names under **`my-registry/nemo-platform-dev/nmp-automodel-*`** (no nested `nmp/...` path — some registries reject that on push).
46

57
| Image | Dockerfile | Role |
68
|-------|------------|------|
7-
| `nmp-automodel-base` | `Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
8-
| `nmp-automodel-tasks` | `Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
9-
| `nmp-automodel-training` | `Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
9+
| `nmp-automodel-base` | `docker/automodel/Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
10+
| `nmp-automodel-tasks` | `docker/automodel/Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
11+
| `nmp-automodel-training` | `docker/automodel/Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
1012

1113
Full references (default tag `local`):
1214

‎docker/automodel/no_override_requirements.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,4 @@ safetensors; sys_platform == 'never'
1010
numpy; sys_platform == 'never'
1111

1212
# Keep W&B's bundled wandb-core binary on a build with patched go-git/go-billy.
13-
wandb==0.27.2
13+
wandb==0.28.0

‎docker/automodel/pyproject.workspace.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ requires-python = ">=3.11,<3.14"
88

99
[tool.uv]
1010
required-version = ">=0.9.14,<0.10.0"
11+
# greenlet 3.5.x ships x86_64-only wheels today; cap until aarch64 wheels land.
12+
constraint-dependencies = ["greenlet>=3.0.0,<3.5"]
1113

1214
[tool.uv.workspace]
1315
members = [

0 commit comments

Comments
 (0)