fix(customizer): address CVEs in unsloth and automodel

docker-bake.hcl

@@ -52,6 +52,10 @@ variable "MAMBA_SSM_WHEEL_CONTEXT" {
   default = ""
 }
 
+variable "FFMPEG_VLM_WHEEL_CONTEXT" {
+  default = ""
+}
+
 variable "DISTROLESS_BASE" {
   default = "nvcr.io/nvidia/distroless/python:3.11-v4.0.8"
 }
@@ -75,12 +79,12 @@ variable "BASE_TAG_PYTHON" {
 
 # Pin for nmp-automodel-base.
 variable "BASE_TAG_AUTOMODEL" {
-  default = "f8239353044d71cbd53e209b60c0600ead484b58"
+  default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
 }
 
 # The tag for base images if needed
 variable "WHEELS_TAG" {
-  default = "65632527b258367fb8855ef431f596abb8538577"
+  default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
 }
 
 variable "BAKE_CACHE_SOURCE_BRANCH" {
@@ -138,6 +142,11 @@ function "get_mamba_ssm_wheel_image" {
   result = "${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}"
 }
 
+function "get_ffmpeg_vlm_wheel_image" {
+  params = []
+  result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}"
+}
+
 function "get_arch_tag" {
   params = []
   result = BUILD_ARCH == "linux/arm64" ? "linux-arm64" : "linux-amd64"
@@ -175,6 +184,11 @@ function "mamba_ssm_wheel_context" {
   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()}"
 }
 
+function "ffmpeg_vlm_wheel_context" {
+  params = []
+  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()}"
+}
+
 function "wheel_tags" {
   params = [name]
   result = [
@@ -293,6 +307,7 @@ group "nmp-automodel-gpu-wheels" {
   targets = [
     "causal-conv1d-wheel",
     "mamba-ssm-wheel",
+    "ffmpeg-vlm-wheel",
   ]
 }
 
@@ -500,15 +515,14 @@ target "nmp-cpu-tasks-docker" {
   platforms  = get_platforms()
 }
 
-# Mamba Wheel Builders
-# Builds Python wheels for mamba-ssm and causal-conv1d in parallel.
-# Both only ship source distributions on PyPI; this pre-builds them for
-# amd64 and arm64. The wheels live at /wheels/*.whl inside the image.
+# Python wheel builders (causal-conv1d, mamba-ssm, av, opencv-python-headless).
+# CUDA extensions only ship source on PyPI; av/opencv bundle FFmpeg. Pre-built for
+# amd64 and arm64. Wheels live at /wheels/*.whl inside each image.
 
 target "causal-conv1d-wheel" {
   target     = "causal-conv1d-wheel"
   context    = "."
-  dockerfile = "docker/base/Dockerfile.mamba-wheel"
+  dockerfile = "docker/base/Dockerfile.python-wheels"
   cache-to   = maybe_registry_cache_to("causal-conv1d-wheel")
   cache-from = maybe_registry_cache_from("causal-conv1d-wheel")
   tags       = wheel_tags("causal-conv1d-wheel")
@@ -523,7 +537,7 @@ target "causal-conv1d-wheel" {
 target "mamba-ssm-wheel" {
   target     = "mamba-ssm-wheel"
   context    = "."
-  dockerfile = "docker/base/Dockerfile.mamba-wheel"
+  dockerfile = "docker/base/Dockerfile.python-wheels"
   cache-to   = maybe_registry_cache_to("mamba-ssm-wheel")
   cache-from = maybe_registry_cache_from("mamba-ssm-wheel")
   tags       = wheel_tags("mamba-ssm-wheel")
@@ -536,6 +550,17 @@ target "mamba-ssm-wheel" {
   platforms = get_platforms()
 }
 
+target "ffmpeg-vlm-wheel" {
+  target     = "ffmpeg-vlm-wheel"
+  context    = "."
+  dockerfile = "docker/base/Dockerfile.python-wheels"
+  cache-to   = maybe_registry_cache_to("ffmpeg-vlm-wheel")
+  cache-from = maybe_registry_cache_from("ffmpeg-vlm-wheel")
+  tags       = wheel_tags("ffmpeg-vlm-wheel")
+  output     = image_output()
+  platforms  = get_platforms()
+}
+
 
 target "safe-synthesizer-tasks-docker" {
   target     = "runtime"
@@ -698,7 +723,7 @@ target "automodel-platform-workspace" {
 target "nmp-automodel-base-builder" {
   target          = "nmp-automodel-base"
   context         = "."
-  dockerfile      = "docker/Dockerfile.nmp-automodel-base"
+  dockerfile      = "docker/automodel/Dockerfile.nmp-automodel-base"
   no-cache-filter = ["automodel-clone"]
   cache-to        = maybe_registry_cache_to("nmp-automodel-base")
   cache-from      = maybe_registry_cache_from("nmp-automodel-base")
@@ -707,14 +732,15 @@ target "nmp-automodel-base-builder" {
   contexts = {
     causal-conv1d-wheel-image = causal_conv1d_wheel_context()
     mamba-ssm-wheel-image     = mamba_ssm_wheel_context()
+    ffmpeg-vlm-wheel-image    = ffmpeg_vlm_wheel_context()
   }
   platforms = get_platforms()
 }
 
 target "nmp-automodel-tasks-docker" {
   target     = "runtime"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -729,7 +755,7 @@ target "nmp-automodel-tasks-docker" {
 target "nmp-automodel-training-docker" {
   target     = "runtime"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-training"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -744,7 +770,7 @@ target "nmp-automodel-training-docker" {
 target "nmp-automodel-tasks-smoke-test" {
   target     = "smoke-test"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -760,7 +786,7 @@ target "nmp-automodel-tasks-smoke-test" {
 target "nmp-automodel-training-smoke-test" {
   target     = "smoke-test"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-training"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()

docker/Dockerfile.nmp-unsloth-training

@@ -13,10 +13,10 @@
 #      blocks uv from installing/upgrading torch into the venv so the NGC
 #      base's PyTorch + CUDA remain the runtime stack.
 #   1b. bitsandbytes — compiled from source against NGC CUDA 13.2 (same pattern
-#       as docker/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
+#       as docker/automodel/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
 #       cuda130; source build replaces the wheel from step 1.
 #   1c. mamba-ssm + causal-conv1d — prebuilt cu13.2 / cp312 wheels (shared with
-#       docker/Dockerfile.nmp-automodel-base).
+#       docker/automodel/Dockerfile.nmp-automodel-base).
 #   1d. flash-attn — optional for unsloth and currently NOT installed (see the
 #       commented TODO below). Without it Unsloth falls back when xformers is
 #       also missing (common on newer CUDA stacks), logging "FA2 = False /
@@ -28,14 +28,14 @@
 # Default tag:    `local` (override via BAKE_TAG at build time).
 
 # NGC PyTorch base. 26.05-py3 ships PyTorch 2.12 + CUDA 13.2 + Python 3.12 (aligned with
-# docker/Dockerfile.nmp-automodel-base). Override at build time:
+# docker/automodel/Dockerfile.nmp-automodel-base). Override at build time:
 # --set nmp-unsloth-training.args.PYTORCH_BASE=...
 ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
 
 # Prebuilt CUDA-extension wheels (mamba-ssm + causal-conv1d), shared with
-# docker/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
+# docker/automodel/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
 # resolves these contexts to the causal-conv1d-wheel / mamba-ssm-wheel images
-# (see docker-bake.hcl + docker/base/Dockerfile.mamba-wheel). Local builds need
+# (see docker-bake.hcl + docker/base/Dockerfile.python-wheels). Local builds need
 # `USE_LOCAL_WHEELS=1` (build the wheel targets) or a reachable WHEELS_REGISTRY.
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
@@ -153,7 +153,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     -e /app/packages/nemo_platform_plugin \
     -e /app/packages/nmp_common \
     -e /app/packages/nmp_customization_common \
-    -e "/app/services/unsloth[integrations]"
+    -e "/app/services/unsloth"
 
 # hf-transfer can pull huggingface-hub 1.x — install then re-pin hub + transformers.
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -179,6 +179,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "mlflow-skinny>=3.11.1,<3.12.0"
 
 # Stale NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
+RUN apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --only-upgrade openssl libssl3t64 && \
+    rm -rf /var/lib/apt/lists/*
+
 RUN rm -rf \
     /usr/local/lib/python3.12/dist-packages/aiohttp \
     /usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
@@ -199,7 +203,13 @@ RUN rm -rf \
     /usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
     /usr/local/lib/python3.12/dist-packages/mlflow \
     /usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
-    /usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info
+    /usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info \
+    /usr/local/lib/python3.12/dist-packages/jupyter_server \
+    /usr/local/lib/python3.12/dist-packages/jupyter_server-*.dist-info \
+    /usr/local/cuda/NsightSystems-cli-*  \
+    /usr/local/bin/nsys \
+    /usr/local/cuda/bin/nsys
+
 
 ENTRYPOINT ["/opt/venv/bin/python"]
 CMD ["-m", "nmp.unsloth.tasks.training", "--help"]

docker/Dockerfile.nmp-automodel-base → docker/automodel/Dockerfile.nmp-automodel-base

@@ -11,6 +11,7 @@ ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
 
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
+FROM ffmpeg-vlm-wheel-image AS ffmpeg-vlm-wheel-src
 
 FROM alpine/git AS automodel-clone
 ARG AUTOMODEL_COMMIT
@@ -95,18 +96,36 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "black>=26.3.1,<27" \
     "GitPython>=3.1.50,<4" \
     "jaraco-context>=6.1.0,<7" \
-    "jupyter-server>=2.18.0,<3" \
+    "jupyter-server>=2.20.0,<3" \
     "jupyterlab>=4.5.7,<5" \
     "mistune>=3.2.1,<4" \
     "notebook>=7.5.6,<8" \
     "onnx>=1.21.0,<2" \
     "pillow>=12.2.0,<13" \
     "pyasn1>=0.6.3,<1" \
+    "pyarrow>=23.0.1,<24" \
+    "msgpack>=1.2.1,<2" \
+    "cryptography>=48.0.1,<49" \
     "tornado>=6.5.5,<7" \
     "urllib3>=2.7.0,<3" \
     "mlflow-skinny>=3.11.1,<3.12.0" \
     "grpcio>=1.81.1,<2" \
-    "wandb==0.27.2"
+    "wandb>=0.28.0,<1"
+
+# Replace PyPI av/opencv (bundled FFmpeg < 8.1.2) with custom wheels from ffmpeg-vlm-wheel.
+# vlm sync pulls PyPI copies transitively; drop them before installing our manylinux wheels.
+# PyAV 17.x uses cp311-abi3 (stable ABI) — compatible with cp312, not cp312-cp312 tagged.
+RUN --mount=from=ffmpeg-vlm-wheel-src,target=/tmp/ffmpeg-vlm-wheel-src,readonly \
+    rm -rf \
+    /opt/venv/lib/python*/site-packages/av \
+    /opt/venv/lib/python*/site-packages/av-*.dist-info \
+    /opt/venv/lib/python*/site-packages/av.libs \
+    /opt/venv/lib/python*/site-packages/cv2 \
+    /opt/venv/lib/python*/site-packages/opencv_python_headless.libs \
+    /opt/venv/lib/python*/site-packages/opencv_python_headless-*.dist-info && \
+    uv pip install --no-cache-dir --no-deps \
+    /tmp/ffmpeg-vlm-wheel-src/wheels/av-*-abi3-manylinux_2_28_*.whl \
+    /tmp/ffmpeg-vlm-wheel-src/wheels/opencv_python_headless-*cp312*.whl
 
 # Published base image (same filesystem as builder).
 FROM ${PYTORCH_BASE} AS nmp-automodel-base
@@ -119,6 +138,12 @@ COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv
 # Stale copies under NGC system site-packages (verified on nvcr.io/nvidia/pytorch:26.05-py3).
 # CVE scanners read dist-packages even when --system-site-packages resolves imports from /opt/venv.
 # rm -rf (not pip uninstall) avoids corrupted dist-info RECORDs on the NGC base image.
+RUN apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --only-upgrade openssl libssl3t64 && \
+    rm -rf /var/lib/apt/lists/*
+
+# NGC ships Nsight Systems (~585 MB); not used by nmp-automodel training.
+# Cleanup stale NGC image packages.
 RUN rm -rf \
     /usr/local/lib/python3.12/dist-packages/aiohttp \
     /usr/local/lib/python3.12/dist-packages/aiohttp-*.dist-info \
@@ -141,7 +166,12 @@ RUN rm -rf \
     /usr/local/lib/python3.12/dist-packages/wandb-*.dist-info \
     /usr/local/lib/python3.12/dist-packages/mlflow \
     /usr/local/lib/python3.12/dist-packages/mlflow-*.dist-info \
-    /usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info
+    /usr/local/lib/python3.12/dist-packages/mlflow_skinny-*.dist-info \
+    /usr/local/lib/python3.12/dist-packages/jupyter_server \
+    /usr/local/lib/python3.12/dist-packages/jupyter_server-*.dist-info \
+    /usr/local/cuda/NsightSystems-cli-* \
+    /usr/local/bin/nsys \
+    /usr/local/cuda/bin/nsys
 
 ENV VIRTUAL_ENV=/opt/venv \
     UV_PROJECT_ENVIRONMENT=/opt/venv \

docker/Dockerfile.nmp-automodel-training → docker/automodel/Dockerfile.nmp-automodel-training

@@ -40,7 +40,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps \
     -e /opt/Automodel
 
-
 ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
 ENTRYPOINT ["/opt/venv/bin/python"]
 CMD ["-m", "nmp.automodel.tasks.training", "--help"]

docker/automodel/README.md

@@ -1,12 +1,14 @@
 # nmp-automodel container images
 
+All Automodel Docker build files live under **`docker/automodel/`** (wheel vendor script: `docker/base/build-ffmpeg-vendor.sh`).
+
 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).
 
 | Image | Dockerfile | Role |
 |-------|------------|------|
-| `nmp-automodel-base` | `Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
-| `nmp-automodel-tasks` | `Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
-| `nmp-automodel-training` | `Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
+| `nmp-automodel-base` | `docker/automodel/Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
+| `nmp-automodel-tasks` | `docker/automodel/Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
+| `nmp-automodel-training` | `docker/automodel/Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
 
 Full references (default tag `local`):
 

docker/automodel/no_override_requirements.txt

@@ -10,4 +10,4 @@ safetensors; sys_platform == 'never'
 numpy; sys_platform == 'never'
 
 # Keep W&B's bundled wandb-core binary on a build with patched go-git/go-billy.
-wandb==0.27.2
+wandb==0.28.0

docker/automodel/pyproject.workspace.toml

@@ -8,6 +8,8 @@ requires-python = ">=3.11,<3.14"
 
 [tool.uv]
 required-version = ">=0.9.14,<0.10.0"
+# greenlet 3.5.x ships x86_64-only wheels today; cap until aarch64 wheels land.
+constraint-dependencies = ["greenlet>=3.0.0,<3.5"]
 
 [tool.uv.workspace]
 members = [