fix(customizer): reduce complexity of the tasks images used by customizer backends

.github/assets/ngc/containers/nmp-customizer-tasks.md

@@ -0,0 +1,16 @@
+---
+labels:
+  - Fine Tuning
+  - NeMo
+---
+## NeMo Customizer Tasks Container
+
+This container image provides shared CPU task support for NeMo Platform customization workflows (automodel, unsloth, and RL). It hosts file I/O, model entity registration, model spec analysis, and the LoRA adapter sidecar.
+
+### Resources
+
+[Documentation](https://docs.nvidia.com/nemo-platform)
+
+### License
+
+This container is licensed under the [Apache License 2.0](https://github.com/NVIDIA-NeMo/nemo-platform/blob/main/LICENSE).

.github/assets/ngc/containers/nmp-automodel-tasks.md → .github/assets/ngc/containers/nmp-rl-training.md

@@ -3,9 +3,9 @@ labels:
   - Fine Tuning
   - NeMo
 ---
-## NeMo Automodel Tasks Container
+## NeMo-RL Training Container
 
-This container image provides job support for NeMo Automodel workflows. It is designed to run as part of NeMo Platform.
+This container image provides the NeMo-RL training runtime on Ray. It is designed to run as part of NeMo Platform.
 
 ### Resources
 

.github/workflows/release.yaml

@@ -49,9 +49,9 @@ on:
       container-ids:
         description: >-
           Custom only. Comma-separated container IDs.
-          Allowed: nmp-api, nmp-cpu-tasks, nmp-automodel-tasks,
-          nmp-automodel-training, nmp-unsloth-training, auditor-tasks,
-          safe-synthesizer-tasks.
+          Allowed: nmp-api, nmp-cpu-tasks, nmp-customizer-tasks,
+          nmp-automodel-training, nmp-unsloth-training, nmp-rl-training,
+          auditor-tasks, safe-synthesizer-tasks.
         required: false
         type: string
         default: ""
@@ -112,9 +112,10 @@ env:
     [
       {"id":"nmp-api","target":"nmp-api-docker"},
       {"id":"nmp-cpu-tasks","target":"nmp-cpu-tasks-docker"},
-      {"id":"nmp-automodel-tasks","target":"nmp-automodel-tasks-docker"},
+      {"id":"nmp-customizer-tasks","target":"nmp-customizer-tasks"},
       {"id":"nmp-automodel-training","target":"nmp-automodel-training-docker"},
       {"id":"nmp-unsloth-training","target":"nmp-unsloth-training"},
+      {"id":"nmp-rl-training","target":"nmp-rl-training"},
       {"id":"auditor-tasks","target":"auditor-tasks-docker"},
       {"id":"safe-synthesizer-tasks","target":"safe-synthesizer-tasks-docker"}
     ]

RELEASING.md

@@ -34,7 +34,7 @@ input description in that workflow.
 | Type | IDs |
 | --- | --- |
 | Wheels | `nemo-platform`, `nemo-platform-plugin` |
-| Containers | `nmp-api`, `nmp-cpu-tasks`, `nmp-automodel-tasks`, `nmp-automodel-training`, `nmp-unsloth-training`, `auditor-tasks`, `safe-synthesizer-tasks` |
+| Containers | `nmp-api`, `nmp-cpu-tasks`, `nmp-customizer-tasks`, `nmp-automodel-training`, `nmp-unsloth-training`, `nmp-rl-training`, `auditor-tasks`, `safe-synthesizer-tasks` |
 | Helm chart | `nemo-platform` |
 
 For every selected wheel, the workflow checks that its package configuration
@@ -66,7 +66,7 @@ Examples:
 | --- | --- |
 | Scheduled-style nightly | Leave `release-type` as `nightly` and use the default `all` scope. |
 | Stable full release | `release-type: stable`, `source-sha: <40-character SHA>`, `version: <MAJOR.MINOR.PATCH>`, `release-scope: all`. |
-| One container | `release-scope: custom`, `container-ids: nmp-automodel-tasks`. |
+| One container | `release-scope: custom`, `container-ids: nmp-customizer-tasks`. |
 | Helm-only validation | `release-scope: helm`, `dry-run: true`. |
 
 Nightlies also run automatically Monday through Friday at 8:00 PM

conftest.py

@@ -214,7 +214,7 @@ def pytest_collection_modifyitems(config, items):
         "e2e",
         "auth_idp",
         "smoke_gpu_tasks",
-        "smoke_nmp_automodel_tasks",
+        "smoke_nmp_customizer_tasks",
         "smoke_nmp_automodel_training",
         "integration",
         "regression",

docker-bake.hcl

@@ -313,10 +313,10 @@ group "nmp-automodel-gpu-wheels" {
 group "nmp-automodel" {
   targets = [
     "nmp-automodel-base-builder",
-    "nmp-automodel-tasks-docker",
     "nmp-automodel-training-docker",
-    "nmp-automodel-tasks-smoke-test",
     "nmp-automodel-training-smoke-test",
+    "nmp-customizer-tasks",
+    "nmp-customizer-tasks-smoke-test",
   ]
 }
 
@@ -329,11 +329,61 @@ group "nmp-unsloth" {
 group "nmp-rl" {
   targets = [
     "nmp-rl-base-builder",
-    "nmp-rl-tasks",
     "nmp-rl-training",
+    "nmp-customizer-tasks",
   ]
 }
 
+group "nmp-customizer" {
+  targets = [
+    "nmp-customizer-tasks",
+    "nmp-customizer-tasks-smoke-test",
+  ]
+}
+
+# Pruned workspace slice for nmp-customizer-tasks (keep in sync with
+# docker/customizer/pyproject.workspace.toml + Dockerfile.platform-workspace members).
+target "customizer-platform-workspace" {
+  target     = "platform-workspace"
+  context    = "."
+  dockerfile = "docker/customizer/Dockerfile.platform-workspace"
+  output     = ["type=cacheonly"]
+  platforms  = get_platforms()
+}
+
+target "nmp-customizer-tasks" {
+  target     = "runtime"
+  context    = "."
+  dockerfile = "docker/Dockerfile.nmp-customizer-tasks"
+  contexts = {
+    platform-workspace       = "target:customizer-platform-workspace"
+    causal-conv1d-wheel-src  = causal_conv1d_wheel_context()
+    mamba-ssm-wheel-src      = mamba_ssm_wheel_context()
+  }
+  cache-to   = maybe_registry_cache_to("nmp-customizer-tasks")
+  cache-from = maybe_registry_cache_from("nmp-customizer-tasks")
+  tags       = sha_and_maybe_latest_tags("nmp-customizer-tasks")
+  output     = image_output()
+  platforms  = get_platforms()
+}
+
+target "nmp-customizer-tasks-smoke-test" {
+  target     = "smoke-test"
+  context    = "."
+  dockerfile = "docker/Dockerfile.nmp-customizer-tasks"
+  contexts = {
+    platform-workspace       = "target:customizer-platform-workspace"
+    causal-conv1d-wheel-src  = causal_conv1d_wheel_context()
+    mamba-ssm-wheel-src      = mamba_ssm_wheel_context()
+  }
+  args = {
+    SMOKE_MARKER = "smoke_nmp_customizer_tasks"
+  }
+  cache-from = maybe_registry_cache_from("nmp-customizer-tasks")
+  output     = ["type=cacheonly"]
+  platforms  = get_platforms()
+}
+
 # Pruned workspace slice for nmp-rl images (keep in sync with
 # docker/rl/pyproject.workspace.toml + Dockerfile.platform-workspace members).
 target "rl-platform-workspace" {
@@ -372,21 +422,6 @@ target "nmp-rl-training" {
   platforms  = get_platforms()
 }
 
-# Lighter CPU image for the file_io / model_entity steps (no NeMo-RL/Ray).
-target "nmp-rl-tasks" {
-  target     = "runtime"
-  context    = "."
-  dockerfile = "docker/Dockerfile.nmp-rl-tasks"
-  contexts = {
-    platform-workspace = "target:rl-platform-workspace"
-  }
-  cache-to   = maybe_registry_cache_to("nmp-rl-tasks")
-  cache-from = maybe_registry_cache_from("nmp-rl-tasks")
-  tags       = sha_and_maybe_latest_tags("nmp-rl-tasks")
-  output     = image_output()
-  platforms  = get_platforms()
-}
-
 # Base images for consolidated containers
 target "nmp-python-base" {
   target     = python_base_target()
@@ -797,21 +832,6 @@ target "nmp-automodel-base-builder" {
   platforms = get_platforms()
 }
 
-target "nmp-automodel-tasks-docker" {
-  target     = "runtime"
-  context    = "."
-  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
-  contexts = {
-    platform-workspace = "target:automodel-platform-workspace"
-    nmp-automodel-base = automodel_base_context()
-  }
-  cache-to   = maybe_registry_cache_to("nmp-automodel-tasks")
-  cache-from = maybe_registry_cache_from("nmp-automodel-tasks")
-  tags       = sha_and_maybe_latest_tags("nmp-automodel-tasks")
-  output     = image_output()
-  platforms = get_platforms()
-}
-
 target "nmp-automodel-training-docker" {
   target     = "runtime"
   context    = "."
@@ -827,22 +847,6 @@ target "nmp-automodel-training-docker" {
   platforms = get_platforms()
 }
 
-target "nmp-automodel-tasks-smoke-test" {
-  target     = "smoke-test"
-  context    = "."
-  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
-  contexts = {
-    platform-workspace = "target:automodel-platform-workspace"
-    nmp-automodel-base = automodel_base_context()
-  }
-  args = {
-    SMOKE_MARKER       = "smoke_nmp_automodel_tasks"
-  }
-  cache-from = maybe_registry_cache_from("nmp-automodel-tasks")
-  output     = ["type=cacheonly"]
-  platforms  = get_platforms()
-}
-
 target "nmp-automodel-training-smoke-test" {
   target     = "smoke-test"
   context    = "."

docker/Dockerfile.nmp-customizer-tasks

@@ -0,0 +1,83 @@
+# syntax=docker/dockerfile:1
+# nmp-customizer-tasks - shared CPU task image for customization backends.
+#
+# Hosts file_io / model_entity (nmp.customization_common.tasks.*), model_spec
+# analysis (nmp.core.models.tasks.model_spec), and the LoRA adapter sidecar.
+# Built on NGC PyTorch 26.05 + transformers/accelerate + Nemotron wheels
+# (mamba-ssm, causal-conv1d) + platform glue + nmp-models.
+
+ARG SMOKE_MARKER=smoke_nmp_customizer_tasks
+ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
+
+FROM ${PYTORCH_BASE} AS base
+
+WORKDIR /app
+
+COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv
+ENV PATH="/bin:${PATH}"
+
+ENV VIRTUAL_ENV=/opt/venv \
+    UV_PROJECT_ENVIRONMENT=/opt/venv \
+    UV_LINK_MODE=copy \
+    UV_COMPILE_BYTECODE=1 \
+    HF_HUB_ENABLE_HF_TRANSFER=1 \
+    OTEL_PYTHON_EXCLUDED_URLS="health"
+ENV PATH="/opt/venv/bin:/root/.local/bin:${PATH}"
+
+RUN uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+
+FROM base AS runtime
+
+ARG USERNAME=ubuntu
+ARG USER_UID=1000
+ARG USER_GID=1000
+
+COPY --from=platform-workspace / /app
+WORKDIR /app
+
+RUN mkdir -p /home/${USERNAME}/.cache && \
+    chown -R ${USER_UID}:${USER_GID} /home/${USERNAME} /app/services/core/models
+
+# HF stack for model_spec + Nemotron configs. preserve_base_torch.txt blocks uv from
+# pulling PyPI torch (same idea as Automodel update_pyproject_pytorch.sh /
+# docker/common/uv-pytorch.toml). TRANSFORMERS_VERSION matches Automodel r0.5.0 pin.
+ARG TRANSFORMERS_VERSION=5.8.1
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
+    --overrides /app/docker/customizer/preserve_base_torch.txt \
+    "transformers==${TRANSFORMERS_VERSION}" \
+    "accelerate>=1.0.0"
+
+# Nemotron/Mamba CUDA wheels — compiled against NGC 26.05 torch; install --no-deps
+# after the HF stack so a PyPI torch never lands in the venv first.
+RUN --mount=from=causal-conv1d-wheel-src,target=/tmp/causal-conv1d-wheel-src,readonly \
+    --mount=from=mamba-ssm-wheel-src,target=/tmp/mamba-ssm-wheel-src,readonly \
+    uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache-dir --no-deps \
+    /tmp/causal-conv1d-wheel-src/wheels/cu13.2/causal_conv1d-*cp312*.whl \
+    /tmp/mamba-ssm-wheel-src/wheels/cu13.2/mamba_ssm-2.3.0-cp312*.whl
+
+# Platform glue — no_override_requirements.txt keeps uv from clobbering NGC torch
+# or re-resolving the HF stack (same pattern as nmp-automodel-training).
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
+    --overrides /app/docker/customizer/no_override_requirements.txt \
+    -e /app/sdk/python/nemo-platform \
+    -e /app/packages/nemo_platform_plugin \
+    -e /app/packages/nmp_common \
+    -e /app/packages/nmp_customization_common \
+    -e /app/services/core/models
+
+ENV PATH="${VIRTUAL_ENV}/bin:${PATH}"
+ENTRYPOINT ["/opt/venv/bin/python"]
+CMD ["-m", "nmp.customization_common.tasks.file_io", "--help"]
+
+USER ${USER_UID}:${USER_GID}
+
+FROM runtime AS smoke-test
+ARG SMOKE_MARKER
+USER 0
+COPY tests/smoke_gpu/ /smoke_test/
+RUN uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --reinstall pytest && \
+    ${VIRTUAL_ENV}/bin/pytest /smoke_test/ -m ${SMOKE_MARKER} -v
+
+FROM runtime