Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions .github/assets/ngc/containers/nmp-customizer-tasks.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
---
labels:
- Fine Tuning
- NeMo
---
## NeMo Customizer Tasks Container
Comment thread
coderabbitai[bot] marked this conversation as resolved.

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).
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ labels:
- Fine Tuning
- NeMo
---
## NeMo Automodel Tasks Container
## NeMo-RL Training Container
Comment thread
coderabbitai[bot] marked this conversation as resolved.

This container image provides job support for NeMo Automodel workflows. It is designed to run as part of NeMo Platform.
This container image provides a NeMo-RL-based DPO training runtime on Ray. It is designed to run as part of NeMo Platform.
Comment thread
soluwalana marked this conversation as resolved.
Outdated

### Resources

Expand Down
4 changes: 2 additions & 2 deletions .github/workflows/release.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ on:
container-ids:
description: >-
Custom only. Comma-separated container IDs.
Allowed: nmp-api, nmp-cpu-tasks, nmp-automodel-tasks,
Allowed: nmp-api, nmp-cpu-tasks, nmp-customizer-tasks,
nmp-automodel-training, nmp-unsloth-training, auditor-tasks,
safe-synthesizer-tasks.
required: false
Expand Down Expand Up @@ -112,7 +112,7 @@ 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":"auditor-tasks","target":"auditor-tasks-docker"},
Expand Down
4 changes: 2 additions & 2 deletions RELEASING.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`, `auditor-tasks`, `safe-synthesizer-tasks` |
| Helm chart | `nemo-platform` |

For every selected wheel, the workflow checks that its package configuration
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
102 changes: 53 additions & 49 deletions docker-bake.hcl
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]
}

Expand All @@ -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" {
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 = "."
Expand All @@ -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 = "."
Expand Down
69 changes: 69 additions & 0 deletions docker/Dockerfile.nmp-customizer-tasks
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
# 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

# Model-spec needs accelerate alongside base transformers; Nemotron/Mamba configs
# need the prebuilt CUDA wheels (installed --no-deps after sync).
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
"accelerate>=1.0.0"

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

RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --package nmp-customization-common --package nmp-models --no-dev --inexact

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
73 changes: 0 additions & 73 deletions docker/Dockerfile.nmp-rl-tasks

This file was deleted.

Loading
Loading