diff --git a/.github/assets/ngc/containers/nmp-customizer-tasks.md b/.github/assets/ngc/containers/nmp-customizer-tasks.md new file mode 100644 index 0000000000..5dcffcd20a --- /dev/null +++ b/.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). diff --git a/.github/assets/ngc/containers/nmp-automodel-tasks.md b/.github/assets/ngc/containers/nmp-rl-training.md similarity index 61% rename from .github/assets/ngc/containers/nmp-automodel-tasks.md rename to .github/assets/ngc/containers/nmp-rl-training.md index 7cadd76408..3e0cc87f5b 100644 --- a/.github/assets/ngc/containers/nmp-automodel-tasks.md +++ b/.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 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index d9a6ea19c9..5dbd2774cb 100644 --- a/.github/workflows/release.yaml +++ b/.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"} ] diff --git a/RELEASING.md b/RELEASING.md index 060d959335..cd9f57cae6 100644 --- a/RELEASING.md +++ b/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: `, `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 diff --git a/conftest.py b/conftest.py index 64b1ec532f..c90bfa5af7 100644 --- a/conftest.py +++ b/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", diff --git a/docker-bake.hcl b/docker-bake.hcl index cafad318de..de391eb987 100644 --- a/docker-bake.hcl +++ b/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 = "." diff --git a/docker/Dockerfile.nmp-customizer-tasks b/docker/Dockerfile.nmp-customizer-tasks new file mode 100644 index 0000000000..6b5114c529 --- /dev/null +++ b/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 diff --git a/docker/Dockerfile.nmp-rl-tasks b/docker/Dockerfile.nmp-rl-tasks deleted file mode 100644 index a9e2270c86..0000000000 --- a/docker/Dockerfile.nmp-rl-tasks +++ /dev/null @@ -1,73 +0,0 @@ -# syntax=docker/dockerfile:1 -# nmp-rl tasks - CPU file_io / model_entity steps for an nmp-rl DPO job. -# -# Deliberately does NOT build on nmp-rl-base: the download/upload/model-entity -# steps only need the platform glue (SDK + customization-common), not NeMo-RL / -# Ray / vLLM. Basing on the NGC image keeps it consistent with the platform's -# CUDA userspace while staying far lighter than the training image. - -ARG SMOKE_MARKER=smoke_nmp_rl_tasks -ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-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/rl - -# Only the glue + nmp-rl (compile/tasks side). No NeMo-RL extra → no Ray/vLLM. -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \ - -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/rl - -ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" -# Default CMD is a harmless help invocation; the platform overrides `command` -# per step (e.g. -m nmp.rl.tasks.file_io / nmp.rl.tasks.model_entity). -ENTRYPOINT ["/opt/venv/bin/python"] -CMD ["-m", "nmp.rl.tasks.file_io", "--help"] - -USER ${USER_UID}:${USER_GID} - -# NOTE: this smoke-test stage is intentionally NOT wired into docker-bake.hcl -# (unlike nmp-automodel-{tasks,training}-smoke-test). There are no RL smoke tests -# yet — tests/smoke_gpu/ carries no tests marked `smoke_nmp_rl_tasks`, and that -# marker isn't registered in its conftest.py. Wiring a bake target now would run -# `pytest -m smoke_nmp_rl_tasks` against zero collected tests, which exits 5 -# ("no tests ran") and fails the bake. Add RL import smoke tests + register the -# marker first, then add the bake target (see nmp-automodel-tasks-smoke-test for -# the pattern). The stage is kept so it's ready to wire up once tests exist. -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 diff --git a/docker/automodel/Dockerfile.nmp-automodel-tasks b/docker/automodel/Dockerfile.nmp-automodel-tasks deleted file mode 100644 index 4e70b8a99c..0000000000 --- a/docker/automodel/Dockerfile.nmp-automodel-tasks +++ /dev/null @@ -1,52 +0,0 @@ -# syntax=docker/dockerfile:1 -# nmp-automodel tasks - file_io, model_entity, and other platform task steps. -# Built on nmp-automodel-base (GPU-capable; runs on CPU or GPU nodes). - -ARG SMOKE_MARKER=smoke_nmp_automodel_tasks - -# Supplied by bake. Local builds use the base builder target; CI can override -# the context to a pinned docker-image:// reference. -FROM nmp-automodel-base AS automodel-base - -FROM automodel-base AS runtime - -# Pin uv for platform workspace installs (base may lack /bin/uv or PATH may prefer 0.10.x). -COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv -ENV PATH="/bin:${PATH}" - -ARG USERNAME=ubuntu -ARG USER_UID=1000 -ARG USER_GID=1000 - -ENV HF_HUB_ENABLE_HF_TRANSFER=1 \ - OTEL_PYTHON_EXCLUDED_URLS="health" - -COPY --from=platform-workspace / /app -WORKDIR /app - -RUN mkdir -p /home/${USERNAME}/.cache && \ - chown -R ${USER_UID}:${USER_GID} /home/${USERNAME} /app/services/automodel - -# /app/pyproject.toml is pyproject.workspace.toml (see Dockerfile.platform-workspace). -# --inexact: keep PyTorch / Automodel packages already in the base venv; add nmp-automodel + deps. -ENV UV_LINK_MODE=copy \ - UV_PROJECT_ENVIRONMENT=${VIRTUAL_ENV} - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --package nmp-automodel --package nmp-models --no-dev --inexact - -ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" -ENTRYPOINT ["/opt/venv/bin/python"] -CMD ["-m", "nmp.automodel.tasks", "--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 - -# Default stage for untargeted builds (bake uses --target runtime; smoke-test is opt-in). -FROM runtime diff --git a/docker/automodel/Dockerfile.platform-workspace b/docker/automodel/Dockerfile.platform-workspace index 9d4854a15c..bef830d140 100644 --- a/docker/automodel/Dockerfile.platform-workspace +++ b/docker/automodel/Dockerfile.platform-workspace @@ -19,6 +19,7 @@ COPY packages/models packages/models COPY packages/nmp_common packages/nmp_common COPY packages/nmp_customization_common packages/nmp_customization_common COPY packages/nemo_platform_plugin packages/nemo_platform_plugin +COPY plugins/nemo-deployments plugins/nemo-deployments COPY sdk/python/nemo-platform sdk/python/nemo-platform COPY services/automodel services/automodel COPY services/core/models services/core/models diff --git a/docker/automodel/README.md b/docker/automodel/README.md index a31e370aa3..1ba462ea2a 100644 --- a/docker/automodel/README.md +++ b/docker/automodel/README.md @@ -2,18 +2,18 @@ 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). +Two images for the **nmp-automodel** customization backend, plus the shared **`nmp-customizer-tasks`** CPU image used by all customization backends. Published as flat repo names under **`my-registry/nemo-platform-dev/nmp-*`** (no nested `nmp/...` path — some registries reject that on push). | Image | Dockerfile | Role | |-------|------------|------| | `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-customizer-tasks` | `docker/Dockerfile.nmp-customizer-tasks` | Shared CPU tasks (`file_io`, `model_entity`, `model_spec`, LoRA sidecar) | | `nmp-automodel-training` | `docker/automodel/Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) | Full references (default tag `local`): - `my-registry/nemo-platform-dev/nmp-automodel-base:local` -- `my-registry/nemo-platform-dev/nmp-automodel-tasks:local` +- `my-registry/nemo-platform-dev/nmp-customizer-tasks:local` - `my-registry/nemo-platform-dev/nmp-automodel-training:local` Bake file: **`docker-bake.hcl`** at the Platform repo root (`context = "."`). Run all commands from the Platform repo root. @@ -87,20 +87,24 @@ Override registry: `export WHEELS_REGISTRY=...` and `export IMAGE_REGISTRY=...` |-------|---------| | `3d98f6e3.diff` | Drop `decord` + `imageio-ffmpeg` (old bundled ffmpeg); use `torchcodec` for VLM video (`FORCE_QWENVL_VIDEO_READER=torchcodec`) | -**Tasks image:** `uv sync --package nmp-automodel --no-dev --inexact` from the minimal workspace. CPU steps only need platform SDK glue; upgrading ancillary packages here does not affect training. +**Customizer tasks image (`nmp-customizer-tasks`):** `uv sync --package nmp-customization-common --package nmp-models --no-dev --inexact` from the customizer workspace slice (`docker/customizer/`). Hosts shared CPU steps (`file_io`, `model_entity`, `model_spec`, LoRA sidecar) for all customization backends. **Training image:** Do **not** use `uv sync` — it upgrades `transformers` and breaks `PreTrainedModel`. Use **`uv pip install -e`** with **`--overrides no_override_requirements.txt`**, then `uv pip install --no-deps -e /opt/Automodel` to re-pin `nemo_automodel` from the base clone (not PyPI). ## Runtime -Entrypoint is `/opt/venv/bin/python`. Job steps pass `-m nmp.automodel.tasks.` (see `nmp.automodel.app.jobs.compiler`). Local smoke: +Entrypoint is `/opt/venv/bin/python` on both images. The compiler routes CPU steps to `nmp-customizer-tasks` and the GPU training step to `nmp-automodel-training` (see `nmp.automodel.app.jobs.compiler` and `nmp.automodel.images`). Local smoke: ```bash -# No extra args → uses image CMD (python -m nmp.automodel.tasks --help). -docker run --rm $NMP_AUTOMODEL_TASKS_IMAGE +# Customizer tasks image — default CMD prints file_io help. +docker run --rm my-registry/nemo-platform-dev/nmp-customizer-tasks:local -# Extra args replace CMD; include -m nmp.automodel.tasks or you get plain `python --help`. -docker run --rm $NMP_AUTOMODEL_TASKS_IMAGE -m nmp.automodel.tasks --list +# Automodel CPU step (file_io). +docker run --rm my-registry/nemo-platform-dev/nmp-customizer-tasks:local \ + -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer --help + +# Training image — default CMD prints training help. +docker run --rm my-registry/nemo-platform-dev/nmp-automodel-training:local ``` -The job compiler resolves `nmp-automodel-tasks` and `nmp-automodel-training` under `NMP_AUTOMODEL_IMAGE_REGISTRY` (default `my-registry/nemo-platform-dev`). See `nmp.automodel.images`. +The job compiler resolves `nmp-customizer-tasks` and `nmp-automodel-training` under `NMP_AUTOMODEL_IMAGE_REGISTRY` (default `my-registry/nemo-platform-dev`). See `nmp.automodel.images`. diff --git a/docker/automodel/pyproject.workspace.toml b/docker/automodel/pyproject.workspace.toml index f576d9da0b..483d205861 100644 --- a/docker/automodel/pyproject.workspace.toml +++ b/docker/automodel/pyproject.workspace.toml @@ -19,6 +19,7 @@ members = [ "packages/nemo_platform_plugin", "packages/nmp_common", "packages/nmp_customization_common", + "plugins/nemo-deployments", "services/automodel", "services/core/models", ] @@ -28,6 +29,7 @@ nmp-build-tools = { workspace = true } models = { workspace = true } nemo-platform-sdk = { workspace = true } nemo-platform-plugin = { workspace = true } +nemo-deployments-plugin = { workspace = true } nmp-common = { workspace = true } nmp-customization-common = { workspace = true } nmp-automodel = { workspace = true } diff --git a/docker/customizer/Dockerfile.platform-workspace b/docker/customizer/Dockerfile.platform-workspace new file mode 100644 index 0000000000..e1f882db7c --- /dev/null +++ b/docker/customizer/Dockerfile.platform-workspace @@ -0,0 +1,20 @@ +# syntax=docker/dockerfile:1 +# Minimal Platform workspace slice for nmp-customizer-tasks container installs. +# Used as a named build context (platform-workspace). +# Keep in sync with docker/customizer/pyproject.workspace.toml members. + +FROM scratch AS platform-workspace +COPY docker/customizer/pyproject.workspace.toml pyproject.toml +# uv --overrides files the tasks Dockerfile reads at /app/docker/customizer/. +COPY docker/customizer/preserve_base_torch.txt docker/customizer/preserve_base_torch.txt +COPY docker/customizer/no_override_requirements.txt docker/customizer/no_override_requirements.txt +COPY docs docs +COPY openapi openapi +COPY packages/nmp_build_tools packages/nmp_build_tools +COPY packages/models packages/models +COPY packages/nmp_common packages/nmp_common +COPY packages/nmp_customization_common packages/nmp_customization_common +COPY packages/nemo_platform_plugin packages/nemo_platform_plugin +COPY plugins/nemo-deployments plugins/nemo-deployments +COPY sdk/python/nemo-platform sdk/python/nemo-platform +COPY services/core/models services/core/models diff --git a/docker/customizer/no_override_requirements.txt b/docker/customizer/no_override_requirements.txt new file mode 100644 index 0000000000..44e87189d2 --- /dev/null +++ b/docker/customizer/no_override_requirements.txt @@ -0,0 +1,10 @@ +# Preserve NGC PyTorch + HF stack installed before platform glue. +# Same pattern as docker/automodel/no_override_requirements.txt: impossible +# markers block uv from re-resolving these when adding editable platform packages. +transformers; sys_platform == 'never' +torch; sys_platform == 'never' +torchvision; sys_platform == 'never' +tokenizers; sys_platform == 'never' +accelerate; sys_platform == 'never' +safetensors; sys_platform == 'never' +numpy; sys_platform == 'never' diff --git a/docker/customizer/preserve_base_torch.txt b/docker/customizer/preserve_base_torch.txt new file mode 100644 index 0000000000..8e667f08c5 --- /dev/null +++ b/docker/customizer/preserve_base_torch.txt @@ -0,0 +1,8 @@ +# Block uv from installing or upgrading PyTorch wheels into the venv. +# The NGC base image ships torch + CUDA; the venv uses --system-site-packages +# to inherit that stack. Impossible markers (sys_platform == 'never') are the +# same pattern as Automodel docker/common/uv-pytorch.toml and +# docker/automodel/no_override_requirements.txt. +torch; sys_platform == 'never' +torchvision; sys_platform == 'never' +torchaudio; sys_platform == 'never' diff --git a/docker/customizer/pyproject.workspace.toml b/docker/customizer/pyproject.workspace.toml new file mode 100644 index 0000000000..7afb6e67fc --- /dev/null +++ b/docker/customizer/pyproject.workspace.toml @@ -0,0 +1,33 @@ +# Minimal uv workspace for nmp-customizer-tasks container image builds only. +# Keep in sync with docker/customizer/Dockerfile.platform-workspace members. + +[project] +name = "nemo-platform-customizer-tasks-image" +version = "0.0.0" +requires-python = ">=3.11,<3.14" + +[tool.uv] +required-version = ">=0.9.14,<0.10.0" +constraint-dependencies = ["greenlet>=3.0.0,<3.5"] + +[tool.uv.workspace] +members = [ + "packages/nmp_build_tools", + "packages/models", + "sdk/python/nemo-platform", + "packages/nemo_platform_plugin", + "packages/nmp_common", + "packages/nmp_customization_common", + "plugins/nemo-deployments", + "services/core/models", +] + +[tool.uv.sources] +nmp-build-tools = { workspace = true } +models = { workspace = true } +nemo-platform-sdk = { workspace = true } +nemo-platform-plugin = { workspace = true } +nemo-deployments-plugin = { workspace = true } +nmp-common = { workspace = true } +nmp-customization-common = { workspace = true } +nmp-models = { workspace = true } diff --git a/docker/unsloth/README.md b/docker/unsloth/README.md index 6f1ce93d21..0cad732bff 100644 --- a/docker/unsloth/README.md +++ b/docker/unsloth/README.md @@ -1,23 +1,19 @@ -# nmp-unsloth container image +# nmp-unsloth container images -Single image — `nmp-unsloth-training` — used for all four steps of an -Unsloth customization job (file_io download, training, file_io upload, -model_entity). +Unsloth customization jobs use the shared CPU image for file_io / model_entity +steps and a dedicated GPU image for training. | Image | Dockerfile | Role | |-------|------------|------| -| `nmp-unsloth-training` | `Dockerfile.nmp-unsloth-training` | NGC PyTorch base + Unsloth ML stack + platform glue. ENTRYPOINT is `/opt/venv/bin/python`. | +| `nmp-customizer-tasks` | `docker/Dockerfile.nmp-customizer-tasks` | Shared CPU steps (`file_io`, `model_entity`) | +| `nmp-unsloth-training` | `docker/Dockerfile.nmp-unsloth-training` | NGC PyTorch base + Unsloth ML stack + platform glue. ENTRYPOINT is `/opt/venv/bin/python`. | Bake file: **`docker-bake.hcl`** at the Platform repo root (`context = "."`). Run all commands from the Platform repo root. Tags use the same bake variables as automodel (`IMAGE_REGISTRY`, `BAKE_TAG`; defaults in `docker-bake.hcl`): - Default: `${IMAGE_REGISTRY}/nmp-unsloth-training:${BAKE_TAG}` (e.g. `…/nmp-unsloth-training:local` with `--load`) - -Future-proofing: a leaner CPU image (`nmp-unsloth-tasks`) can be added -later for the file_io / model_entity steps. The compiler already routes -those steps through `get_tasks_image()`, which falls back to the training -image when `NMP_UNSLOTH_TASKS_IMAGE` is not set. +- CPU steps: `${IMAGE_REGISTRY}/nmp-customizer-tasks:${BAKE_TAG}` via `get_tasks_image()` (override with `NMP_CUSTOMIZER_TASKS_IMAGE` or `NMP_UNSLOTH_TASKS_IMAGE`) --- @@ -57,11 +53,11 @@ docker buildx bake \ The build pulls the NGC PyTorch base, then: -1. `uv pip install unsloth --torch-backend=auto transformers==4.57.6 huggingface-hub==0.36.2` with +1. `uv pip install unsloth --torch-backend=auto transformers==5.5.0 huggingface-hub==1.5.0` with `preserve_base_torch.txt` overrides so the NGC base's PyTorch + CUDA are not replaced. Unsloth's resolver still pulls `unsloth_zoo`, trl, peft, accelerate, datasets, bitsandbytes, and xformers. **transformers is pinned - explicitly** to `4.57.6` (override at build time via + explicitly** to `5.5.0` (override at build time via `--build-arg TRANSFORMERS_VERSION=...`). 1b. bitsandbytes — compiled from source against the NGC CUDA 13.1 toolkit (PyPI wheels only ship through cuda130), replacing the wheel from step 1. @@ -318,9 +314,9 @@ nemo files filesets delete qwen-unsloth-smoke-out -w default separate ML stack. If you need both backends on the same cluster, run both images side by side; jobs from each backend route to their own `nmp-{backend}-training` image via env-var overrides. -- **transformers + huggingface-hub pins** — the training image pins `transformers==4.57.6` - and `huggingface-hub==0.36.2` in - `Dockerfile.nmp-unsloth-training` (compatible with unsloth's upstream +- **transformers + huggingface-hub pins** — the training image pins `transformers==5.5.0` + and `huggingface-hub==1.5.0` in + `docker/Dockerfile.nmp-unsloth-training` (compatible with unsloth's upstream blocklists). Other HF deps (trl, peft, bitsandbytes, etc.) still come from unsloth's resolver. **PyTorch + CUDA** stay on the NGC base stack via `--system-site-packages` and `preserve_base_torch.txt` / `no_override_requirements.txt` diff --git a/docs/customizer/index.mdx b/docs/customizer/index.mdx index 6f91734b29..2f2a507b03 100644 --- a/docs/customizer/index.mdx +++ b/docs/customizer/index.mdx @@ -28,14 +28,14 @@ Fine-tuning jobs run in container images published to NVIDIA NGC. For NeMo Platf | Image | Purpose | |-------|---------| -| `nvcr.io/nvidia/nemo-platform/nmp-automodel-tasks:0.2.0` | Automodel file transfer and model entity task steps | +| `nvcr.io/nvidia/nemo-platform/nmp-customizer-tasks:0.2.0` | Shared CPU task steps (file I/O, model entity, model spec) for all customization backends | | `nvcr.io/nvidia/nemo-platform/nmp-automodel-training:0.2.0` | Automodel GPU training step | -| `nvcr.io/nvidia/nemo-platform/nmp-unsloth-training:0.2.0` | Unsloth training and task steps | +| `nvcr.io/nvidia/nemo-platform/nmp-unsloth-training:0.2.0` | Unsloth GPU training step | These public images can be pulled directly from `nvcr.io`: ```bash -docker pull nvcr.io/nvidia/nemo-platform/nmp-automodel-tasks:0.2.0 +docker pull nvcr.io/nvidia/nemo-platform/nmp-customizer-tasks:0.2.0 docker pull nvcr.io/nvidia/nemo-platform/nmp-automodel-training:0.2.0 docker pull nvcr.io/nvidia/nemo-platform/nmp-unsloth-training:0.2.0 ``` diff --git a/docs/set-up/config-reference.mdx b/docs/set-up/config-reference.mdx index 78f64fa2d5..6b44244d63 100644 --- a/docs/set-up/config-reference.mdx +++ b/docs/set-up/config-reference.mdx @@ -528,16 +528,16 @@ models: peft_source: http://nemo-entity-store:8000 # PEFT refresh interval in seconds (only used when lora_enabled is true) | default: 30 peft_refresh_interval: 30 - # Image name (without registry/tag) for the LoRA adapters sidecar container. Registry and tag are taken from the platform config (NMP_IMAGE_REGISTRY / NMP_IMAGE_TAG). Override to 'nmp-automodel-tasks' for local dev when that image is already available but nmp-api is not. | default: 'nmp-api' + # Image name (without registry/tag) for the LoRA adapters sidecar container. Registry and tag are taken from the platform config (NMP_IMAGE_REGISTRY / NMP_IMAGE_TAG). Override to 'nmp-customizer-tasks' for local dev when that image is already available but nmp-api is not. | default: 'nmp-api' lora_sidecar_image_name: nmp-api - # Kubernetes container command (entrypoint) for the LoRA sidecar. Default uses the nmp-platform-runner entrypoint present in nmp-api. When using nmp-automodel-tasks set to ['python'] and set lora_sidecar_args to ['-m', 'nmp.core.models.sidecars.adapters.main']. + # Kubernetes container command (entrypoint) for the LoRA sidecar. Default uses the nmp-platform-runner entrypoint present in nmp-api. When using nmp-customizer-tasks set to ['python'] and set lora_sidecar_args to ['-m', 'nmp.core.models.sidecars.adapters.main']. lora_sidecar_command: - nemo - services - run - --sidecars - adapters - # Kubernetes container args for the LoRA sidecar (appended after lora_sidecar_command). Leave empty for nmp-api. Set to ['-m', 'nmp.core.models.sidecars.adapters.main'] when using nmp-automodel-tasks. | default: [] + # Kubernetes container args for the LoRA sidecar (appended after lora_sidecar_command). Leave empty for nmp-api. Set to ['-m', 'nmp.core.models.sidecars.adapters.main'] when using nmp-customizer-tasks. | default: [] lora_sidecar_args: [] # Default user ID for NIM containers (security context) default_user_id: @@ -876,7 +876,7 @@ Environment variables use the NMP_AUTOMODEL_ prefix. ```yaml automodel: - # Registry host/path prefix for nmp-automodel-tasks and nmp-automodel-training. Override via NMP_AUTOMODEL_IMAGE_REGISTRY for other environments, defaults to the platform's image registry. + # Registry host/path prefix for nmp-customizer-tasks and nmp-automodel-training. Override via NMP_AUTOMODEL_IMAGE_REGISTRY for other environments, defaults to the platform's image registry. image_registry: # Override entire GPU training image (registry/name:tag). training_image: diff --git a/packages/nmp_customization_common/src/nmp/customization_common/config.py b/packages/nmp_customization_common/src/nmp/customization_common/config.py new file mode 100644 index 0000000000..7293c77718 --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/config.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared customization configuration (CPU tasks image override).""" + +from nmp.common.config import create_service_config_class, get_service_config +from pydantic import Field + + +class CustomizationCommonConfig(create_service_config_class("customizer")): # type: ignore[misc] + """Environment variables use the ``NMP_CUSTOMIZER_`` prefix.""" + + tasks_image: str | None = Field( + default=None, + description=( + "Override entire CPU tasks image (registry/name:tag) for all customization backends. " + "Set via NMP_CUSTOMIZER_TASKS_IMAGE." + ), + ) + + +config = get_service_config(CustomizationCommonConfig) diff --git a/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py b/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py index 28343c4e95..83da9ead9b 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py @@ -5,7 +5,7 @@ Both backends submit a 4-step ``PlatformJobSpec`` (download → train → upload → model-entity) executed on the platform GPU cluster. ``to_spec`` and the -Docker-runtime guard are shared here; ``compile`` genuinely diverges (compiler +runtime guards are shared here; ``compile`` genuinely diverges (compiler call convention, schema validation, profile resolution) and stays per-backend. """ @@ -20,31 +20,60 @@ from pydantic import BaseModel -def require_docker_runtime(backend_label: str) -> None: - """Refuse to compile when the platform isn't configured for Docker. +def require_container_runtime(backend_label: str, *, num_nodes: int = 1) -> None: + """Refuse to compile unless the platform can run the requested container job. - The compile step builds Docker container specs; surface the misconfiguration - before the Jobs API rejects the spec. + SFT backends (automodel / unsloth) build a container ``PlatformJobSpec`` the + platform runs on either supported target: + + - **Kubernetes** — the platform schedules GPU pods, including multi-node + ``gpu_distributed`` jobs via Volcano. + - **Docker** — the platform's local Docker GPU executor (single host). + + Single-node jobs accept either runtime. **Multi-node jobs** (``num_nodes > + 1``) compile to a ``gpu_distributed`` executor that only the Volcano + (Kubernetes) backend can place — Docker has no multi-node/``gpu_distributed`` + backend — so they require ``platform.runtime: kubernetes``. Failing here + surfaces the misconfiguration at compile time instead of as an opaque + "no backend found" scheduling error (or, for ``runtime: none``, before the + Jobs API rejects the spec). """ platform_config = NemoPlatformConfig.get() - if platform_config.runtime != Runtime.DOCKER: - raise PlatformJobCompilationError( - f"{backend_label} training requires platform.runtime: docker with GPU-backed container execution.", - ) - from nemo_platform_plugin.config import validate_docker_available + runtime = platform_config.runtime - if not validate_docker_available(): + if num_nodes > 1 and runtime != Runtime.KUBERNETES: raise PlatformJobCompilationError( - f"{backend_label} training requires a reachable Docker daemon (platform.runtime: docker).", + f"{backend_label} multi-node training (num_nodes={num_nodes}) requires " + "platform.runtime: kubernetes — multi-node jobs run on the Volcano " + "(gpu_distributed) backend, which has no Docker equivalent. " + f"Current runtime: {runtime.value}.", ) + if runtime == Runtime.KUBERNETES: + return + + if runtime == Runtime.DOCKER: + from nemo_platform_plugin.config import validate_docker_available + + if not validate_docker_available(): + raise PlatformJobCompilationError( + f"{backend_label} training requires a reachable Docker daemon (platform.runtime: docker).", + ) + return + + raise PlatformJobCompilationError( + f"{backend_label} training requires a container runtime: set platform.runtime to " + "'kubernetes' (schedules GPU pods) or 'docker' (local GPU executor). " + f"Current runtime: {runtime.value}.", + ) + def require_distributed_runtime(backend_label: str) -> None: """Refuse to compile when the platform isn't a remote Kubernetes cluster. - Sibling to :func:`require_docker_runtime` for backends that provision a Ray - cluster (e.g. NeMo-RL DPO). Unlike the SFT backends, these have no local - single-node Docker fallback: they need the platform's Kubernetes/Volcano + Sibling to :func:`require_container_runtime` for backends that provision a Ray + cluster (e.g. NeMo-RL DPO). Unlike the SFT backends, these accept **only** + Kubernetes (no Docker fallback): they need the platform's Kubernetes/Volcano scheduler to place GPU pods and inject the distributed env (``RANK``/``WORLD_SIZE``/``MASTER_ADDR``). Surface the misconfiguration before the Jobs API rejects the spec. @@ -62,12 +91,12 @@ class BaseSubmitJob(NemoJob): Subclasses set the ``NemoJob`` ClassVars (``name``, ``description``, ``job_collection_path``, ``input_spec_schema``, ``spec_schema``), implement - :meth:`_transform` and :meth:`compile`, and may set :attr:`docker_runtime_label`. + :meth:`_transform` and :meth:`compile`, and may set :attr:`runtime_label`. """ dependencies: ClassVar[list[str]] = ["entities", "auth", "jobs", "secrets", "files", "models"] - #: Human-readable backend name used in the Docker-runtime guard messages. - docker_runtime_label: ClassVar[str] = "Training" + #: Human-readable backend name used in the runtime guard messages. + runtime_label: ClassVar[str] = "Training" @classmethod async def _transform(cls, job_input: BaseModel, workspace: str, async_sdk: AsyncNeMoPlatform) -> BaseModel: diff --git a/packages/nmp_customization_common/src/nmp/customization_common/service/images.py b/packages/nmp_customization_common/src/nmp/customization_common/service/images.py index b4d4b53b74..9ebd546ca6 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/service/images.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/service/images.py @@ -3,15 +3,20 @@ """Shared Docker image resolution for customization job steps. -Each backend keeps its own image-name constants and ``get_tasks_image`` / -``get_training_image`` (their fallback behavior differs); this module holds the -common registry-resolution logic. +Each backend keeps its own ``get_training_image``; CPU task steps share +``nmp-customizer-tasks`` via :func:`get_customizer_tasks_image`. """ from __future__ import annotations from nemo_platform_plugin.config import get_platform_config from nemo_platform_plugin.jobs.image import get_qualified_image +from nmp.customization_common.config import config as customization_common_config + +CUSTOMIZER_TASKS_IMAGE_NAME = "nmp-customizer-tasks" + +# Must match ENTRYPOINT in docker/Dockerfile.nmp-customizer-tasks. +CUSTOMIZER_PYTHON_ENTRYPOINT = ["/opt/venv/bin/python"] def resolve_qualified_image(name: str, override: str | None, image_registry: str | None) -> str: @@ -31,3 +36,18 @@ def resolve_qualified_image(name: str, override: str | None, image_registry: str platform_config = get_platform_config() registry = image_registry or platform_config.image_registry return get_qualified_image(name, registry=registry) + + +def get_customizer_tasks_image( + *, + backend_override: str | None = None, + image_registry: str | None = None, +) -> str: + """Resolve the shared CPU tasks image for customization job steps. + + Precedence: ``NMP_CUSTOMIZER_TASKS_IMAGE`` (global), then the + per-backend ``tasks_image`` override (e.g. ``NMP_AUTOMODEL_TASKS_IMAGE``), + then ``{registry}/nmp-customizer-tasks:{tag}``. + """ + override = customization_common_config.tasks_image or backend_override + return resolve_qualified_image(CUSTOMIZER_TASKS_IMAGE_NAME, override, image_registry) diff --git a/services/rl/src/nmp/rl/tasks/file_io/__init__.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__init__.py similarity index 51% rename from services/rl/src/nmp/rl/tasks/file_io/__init__.py rename to packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__init__.py index 2d2c9ad611..afeaee25ae 100644 --- a/services/rl/src/nmp/rl/tasks/file_io/__init__.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__init__.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""File I/O task for rl customization jobs.""" +"""File I/O task package.""" -from nmp.rl.tasks.file_io.run import run +from nmp.customization_common.tasks.file_io.run import FileIORunner, run -__all__ = ["run"] +__all__ = ["FileIORunner", "run"] diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__main__.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__main__.py new file mode 100644 index 0000000000..8584e6f720 --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/__main__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Entry point for the shared file_io container task.""" + +import argparse +import sys + +from nmp.customization_common.tasks.file_io.run import run + + +def main() -> int: + parser = argparse.ArgumentParser(description="NeMo customization file I/O task") + parser.add_argument( + "--service-source", + required=True, + help="Value stamped on upload-created filesets (custom_fields.service_source)", + ) + parser.add_argument( + "--service-name", + required=True, + help="Platform service identity for SDK auth/telemetry (e.g. customizer, unsloth, rl)", + ) + args = parser.parse_args() + return run(service_source=args.service_source, service_name=args.service_name) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/services/unsloth/src/nmp/unsloth/tasks/file_io/callbacks.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/callbacks.py similarity index 91% rename from services/unsloth/src/nmp/unsloth/tasks/file_io/callbacks.py rename to packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/callbacks.py index e72777198b..3da1747811 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/callbacks.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/callbacks.py @@ -32,8 +32,6 @@ def get_percentage(current: int, total: int) -> int: if total <= 0: return 0 if current > total or current < 0: - # Benign (see docstring) but worth a breadcrumb now that we no longer - # raise — the old hard error is what previously surfaced count drift. logger.debug("get_percentage clamping out-of-range progress: current=%s total=%s", current, total) current = max(0, min(current, total)) return int((current / total) * 100) @@ -94,8 +92,6 @@ def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> Tq except ValueError: relative_path = dest_full_path.name - # full_src_path looks like "workspace/fileset/relative/path/file.txt". - # Strip the prefix to look up the size by relative path. relative_file_path = full_src_path if full_src_path.startswith(self.fileset_path): relative_file_path = full_src_path[len(self.fileset_path) :].lstrip("/") @@ -112,8 +108,6 @@ def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> Tq }, ) - # set_size() rather than tqdm_kwargs["total"] so the SDK can also - # call set_size() from a Content-Length header without conflict. if file_size is not None: callback.set_size(file_size) @@ -121,15 +115,7 @@ def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> Tq class BaseProgressCallback(Callback): - """Base class for file upload/download progress callbacks. - - Tracks file transfer progress and reports to the Jobs service. - Subclasses implement upload-vs-download behavior. - - Thread Safety: - Uses ``threading.Lock`` to protect stats updates because - FilesetFileSystem transfers files concurrently. - """ + """Base class for file upload/download progress callbacks.""" progress_reporter: ProgressReporter fileset_name: str @@ -167,13 +153,7 @@ def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "BaseSing class BaseSingleFileCallback(Callback): - """Base class for per-file callbacks within a batch operation. - - Uses the template-method pattern: ``close()`` runs the shared - state-update + progress-report sequence, while subclasses customize - via ``_get_phase``, ``_get_file_display_path``, ``_update_stats``, - ``_get_files_count``, and ``_build_status_details``. - """ + """Base class for per-file callbacks within a batch operation.""" parent: BaseProgressCallback source_path: str @@ -234,7 +214,6 @@ def close(self) -> None: logger.debug(f"File transferred: {current_file} ({files_count}/{parent.total_files})") - # Report outside the lock — don't block other threads on the network call. parent.progress_reporter.update_progress( status=PlatformJobStatus.ACTIVE, status_details=self._build_status_details(files_count, total_bytes, current_file), @@ -409,12 +388,7 @@ def _build_status_details(self, files_count: int, total_bytes: int, current_file class CompositeCallback(Callback): - """A callback that delegates to multiple child callbacks. - - Lets us combine console-side ``TqdmCallback`` and Jobs-service - ``File{Upload,Download}ProgressCallback`` into one callback object - passed to fsspec operations. - """ + """A callback that delegates to multiple child callbacks.""" def __init__(self, *callbacks: Callback, **kwargs: Any): super().__init__(**kwargs) diff --git a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py similarity index 78% rename from services/unsloth/src/nmp/unsloth/tasks/file_io/run.py rename to packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py index b26bc12842..02b5c7051b 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/run.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io/run.py @@ -5,21 +5,24 @@ Handles file operations between NeMo Platform Files Service and the job's shared PVC. -The task reads configuration and performs: -- Downloads: If config.download is non-empty, download files from FileSets to local paths -- Uploads: If config.upload is non-empty, upload files from local paths to FileSets +Each backend's compiler sets ``--service-source`` and ``--service-name`` via +``FILE_IO_TASK_COMMAND`` in that backend's ``images.py`` (the job step passes +these flags; do not hardcode one backend when documenting or invoking locally): + +- automodel: ``--service-source automodel --service-name customizer`` +- unsloth: ``--service-source unsloth --service-name unsloth`` +- rl: ``--service-source rl --service-name rl`` + +Usage (example — match the backend you are exercising):: -Usage: export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.unsloth.tasks.file_io + python -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer """ import logging from pathlib import Path import httpx - -# https://docs.nvidia.com/nemo/microservices/latest/pysdk/index.html#handling-errors from nemo_platform import ( APIConnectionError, APITimeoutError, @@ -52,6 +55,13 @@ UploadStats, ) from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.tasks.file_io.callbacks import ( + CompositeCallback, + FileDownloadProgressCallback, + FileUploadProgressCallback, + TqdmPerFileDownloadCallback, + TqdmPerFileUploadCallback, +) from nmp.customization_common.tasks.file_io_progress_reporter import JobsServiceProgressReporter, ProgressReporter from nmp.customization_common.tasks.file_io_utils import ( filesystem_sdk_error_handler, @@ -60,46 +70,24 @@ validate_safe_path, validate_storage_path, ) -from nmp.unsloth.app.constants import SERVICE_NAME -from nmp.unsloth.tasks.file_io.callbacks import ( - CompositeCallback, - FileDownloadProgressCallback, - FileUploadProgressCallback, - TqdmPerFileDownloadCallback, - TqdmPerFileUploadCallback, -) from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential logger = logging.getLogger(__name__) -# Service-source tag stamped onto every upload-created fileset. Lets operators -# filter filesets by training backend. -SERVICE_SOURCE = "unsloth" - CREATE_FILESET_TIMEOUT = 10.0 LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) -# Timeout configurations for FilesetFileSystem operations. Passed via -# sdk.with_options(timeout=...). httpx.Timeout(read=...) is per-chunk -# (the SDK chunks at 16MB), NOT total transfer time — it's a socket-level -# timeout. SDK defaults are httpx.Timeout(timeout=60, connect=5.0). DOWNLOAD_TIMEOUT = httpx.Timeout(30.0, read=5 * 60) UPLOAD_TIMEOUT = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) -# Retry configuration. MAX_RETRIES = 3 INITIAL_BACKOFF_SECONDS = 1.0 MAX_BACKOFF_SECONDS = 30.0 -# Transient exceptions that should trigger retries for filesystem operations. -# FilesetFileSystem uses httpx under the hood, so we retry httpx transients -# in addition to SDK-level transients. TRANSIENT_FILESYSTEM_EXCEPTIONS = ( httpx.TimeoutException, httpx.ConnectError, httpx.ReadTimeout, - # Connection dropped mid-transfer (CDN/proxy closed the socket before the - # full body arrived). Common on large multi-GB model shards; safe to retry. httpx.RemoteProtocolError, httpx.ReadError, ) @@ -113,10 +101,13 @@ def __init__( sdk: NeMoPlatform, progress_reporter: ProgressReporter, job_ctx: NMPJobContext, + *, + service_source: str, ): self.sdk = sdk self.progress_reporter = progress_reporter self.job_ctx = job_ctx + self.service_source = service_source def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFile]: """List files in a FileSet. Returns a list of ``FilesetFile`` objects.""" @@ -134,58 +125,35 @@ def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFile]: ) from e def download_fileset(self, fileset: FileSetRef, dest_dir: Path) -> DownloadStats: - """Download all files from a FileSet to a destination directory. - - Uses ``FilesetFileSystem.get()`` with ``recursive=True`` for efficient batch - downloads. Progress is tracked via two callbacks combined in a - ``CompositeCallback``: - - - ``TqdmPerFileDownloadCallback`` — separate console progress bar per file - - ``FileDownloadProgressCallback`` — reports to Jobs service after each file - - Raises: - FileDownloadError: If the download fails. - """ - stats = DownloadStats() + """Download all files from a FileSet to a destination directory.""" fileset_name = str(fileset) files = self.list_fileset_files(fileset) if not files: logger.warning(f"FileSet {fileset_name} contains no files") - return stats + return DownloadStats() total_files = len(files) total_size = sum(f.size for f in files) dest_dir.mkdir(parents=True, exist_ok=True) - # Maps relative file paths to byte sizes for tqdm percent display. file_sizes = {f.path.lstrip("/"): f.size for f in files} - tqdm_callback = TqdmPerFileDownloadCallback( - dest_path=dest_dir, - fileset_path=fileset_name, - file_sizes=file_sizes, - ) - jobs_callback = FileDownloadProgressCallback( - progress_reporter=self.progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - ) - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - with filesystem_sdk_error_handler( FileDownloadError, f"download from '{fileset_name}' to '{dest_dir}'", ): - self._download_with_retry( + stats = self._download_with_retry( fileset_name=fileset.name, fileset_workspace=fileset.workspace, dest_dir=str(dest_dir), - callback=composite_callback, + fileset_display_name=fileset_name, + dest_path=dest_dir, + file_sizes=file_sizes, + total_files=total_files, + total_size=total_size, ) logger.info(f"Download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") @@ -203,42 +171,40 @@ def _download_with_retry( fileset_name: str, fileset_workspace: str | None, dest_dir: str, - callback: CompositeCallback, - ) -> None: + fileset_display_name: str, + dest_path: Path, + file_sizes: dict[str, int], + total_files: int, + total_size: int, + ) -> DownloadStats: """Internal method with retry logic for downloading from FilesetFileSystem.""" + stats = DownloadStats() + tqdm_callback = TqdmPerFileDownloadCallback( + dest_path=dest_path, + fileset_path=fileset_display_name, + file_sizes=file_sizes, + ) + jobs_callback = FileDownloadProgressCallback( + progress_reporter=self.progress_reporter, + fileset_name=fileset_display_name, + total_files=total_files, + total_size=total_size, + stats=stats, + ) + composite_callback = CompositeCallback(tqdm_callback, jobs_callback) + self.sdk.with_options(timeout=DOWNLOAD_TIMEOUT).files.download( fileset=fileset_name, workspace=fileset_workspace, local_path=dest_dir, - callback=callback, + callback=composite_callback, ) + return stats def upload_fileset(self, fileset: FileSetRef, src_path: Path) -> UploadStats: - """Upload all files from a source path (file or directory) to a FileSet. - - Uses ``FilesetFileSystem.put()`` with ``recursive=True`` for efficient batch - uploads. Progress is tracked via the same composite-callback pattern as - downloads. - - Raises: - FileUploadError: If the upload fails. - """ - stats = UploadStats() + """Upload all files from a source path (file or directory) to a FileSet.""" fileset_name = str(fileset) - tqdm_callback = TqdmPerFileUploadCallback(src_path=src_path) - jobs_callback = FileUploadProgressCallback( - progress_reporter=self.progress_reporter, - src_path=src_path, - fileset_name=fileset_name, - stats=stats, - ) - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - - # Build local and remote paths for upload. ``remote_path`` is relative within - # the fileset ("" for root, "filename" for single file). Trailing slash on - # ``local_path`` follows rsync/scp convention: "dir/" copies contents, - # "dir" copies the directory itself. if src_path.is_dir(): local_path = f"{src_path}/" remote_path = "" @@ -250,12 +216,13 @@ def upload_fileset(self, fileset: FileSetRef, src_path: Path) -> UploadStats: FileUploadError, f"upload from '{src_path}' to '{fileset_name}'", ): - self._upload_with_retry( + stats = self._upload_with_retry( local_path=local_path, remote_path=remote_path, fileset_name=fileset.name, fileset_workspace=fileset.workspace, - callback=composite_callback, + fileset_display_name=fileset_name, + src_path=src_path, ) logger.info(f"Upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") @@ -274,28 +241,34 @@ def _upload_with_retry( remote_path: str, fileset_name: str, fileset_workspace: str | None, - callback: CompositeCallback, - ) -> None: + fileset_display_name: str, + src_path: Path, + ) -> UploadStats: """Internal method with retry logic for uploading to FilesetFileSystem.""" + stats = UploadStats() + tqdm_callback = TqdmPerFileUploadCallback(src_path=src_path) + jobs_callback = FileUploadProgressCallback( + progress_reporter=self.progress_reporter, + src_path=src_path, + fileset_name=fileset_display_name, + stats=stats, + ) + composite_callback = CompositeCallback(tqdm_callback, jobs_callback) + self.sdk.with_options(timeout=UPLOAD_TIMEOUT).files.upload( local_path=local_path, remote_path=remote_path, fileset=fileset_name, workspace=fileset_workspace, - callback=callback, + callback=composite_callback, ) + return stats def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> None: - """Create a FileSet. Skip if it already exists. - - Wraps the retry with ``sdk_error_handler`` to convert exceptions after - all retries exhaust. - """ + """Create a FileSet. Skip if it already exists.""" with sdk_error_handler(FileUploadError, f"create fileset {fileset}", passthrough=(ConflictError,)): self._create_fileset_with_retry(fileset, metadata) - # We don't use the SDK's built-in retry: it would retry on ConflictError, - # which is expected here and would just waste calls. @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), @@ -319,7 +292,7 @@ def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None try: body_kwargs: dict = { "name": fileset.name, - "custom_fields": {"service_source": SERVICE_SOURCE}, + "custom_fields": {"service_source": self.service_source}, } if metadata is not None: body_kwargs["metadata"] = metadata @@ -445,27 +418,28 @@ def run_upload(self, uploads: list[UploadItem]) -> None: logger.info(f"All uploads complete: {total_stats.files_uploaded} files, {total_stats.total_bytes} bytes total") -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the file I/O task. - - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - - Returns: - Exit code (0 for success, non-zero for failure). - """ +def run( + sdk: NeMoPlatform | None = None, + job_ctx: NMPJobContext | None = None, + *, + service_source: str, + service_name: str, +) -> int: + """Execute the file I/O task.""" job_ctx = job_ctx or NMPJobContext.from_env() validate_storage_path(job_ctx.storage_path) sdk_owned = sdk is None progress_reporter: ProgressReporter | None = None try: - sdk = sdk or get_task_sdk(SERVICE_NAME) + sdk = sdk or get_task_sdk(service_name) progress_reporter = JobsServiceProgressReporter.create_progress_reporter(sdk, job_ctx) - runner = FileIORunner(sdk=sdk, progress_reporter=progress_reporter, job_ctx=job_ctx) + runner = FileIORunner( + sdk=sdk, + progress_reporter=progress_reporter, + job_ctx=job_ctx, + service_source=service_source, + ) config = get_config(job_ctx.config_path) @@ -509,18 +483,3 @@ def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) - finally: if sdk_owned and sdk is not None: sdk.close() - - -def build_output_metadata(spec) -> dict: - """Build the metadata dict stamped onto the output fileset. - - Captures the bits a downstream consumer (model-entity creation, - deployment) needs about this artefact without re-deriving them - from the training spec. - """ - return { - "model": spec.model.name, - "finetuning_type": spec.training.finetuning_type, - "save_method": spec.output.save_method, - "output_type": spec.output.type, - } diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py new file mode 100644 index 0000000000..f48c50c569 --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_metadata.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Helpers for stamping output fileset metadata during customization uploads.""" + + +def build_output_metadata( + *, + model: str, + finetuning_type: str, + output_type: str, + save_method: str | None = None, +) -> dict: + """Build the metadata dict stamped onto the output fileset. + + Captures the bits a downstream consumer (model-entity creation, + deployment) needs about this artefact without re-deriving them + from the training spec. + """ + metadata: dict[str, str] = { + "model": model, + "finetuning_type": finetuning_type, + "output_type": output_type, + } + if save_method is not None: + metadata["save_method"] = save_method + return metadata diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__init__.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__init__.py new file mode 100644 index 0000000000..5a5c8632d4 --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model entity task package.""" + +from nmp.customization_common.tasks.model_entity.run import ModelEntityRunner, run + +__all__ = ["ModelEntityRunner", "run"] diff --git a/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__main__.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__main__.py new file mode 100644 index 0000000000..4eb79c5fed --- /dev/null +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/__main__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Entry point for the shared model_entity container task.""" + +import argparse +import sys + +from nmp.customization_common.tasks.model_entity.run import run + + +def main() -> int: + parser = argparse.ArgumentParser(description="NeMo customization model entity task") + parser.add_argument( + "--service-name", + required=True, + help="Platform service identity for SDK auth/telemetry (e.g. customizer, unsloth, rl)", + ) + args = parser.parse_args() + return run(service_name=args.service_name) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/run.py similarity index 81% rename from services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py rename to packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/run.py index 96874ea20c..1da52d3c1a 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/model_entity/run.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/tasks/model_entity/run.py @@ -3,15 +3,9 @@ """Model entity task entry point. -Handles creating model entities in the Models service after customization completes. - -The task reads configuration and creates a Model Entity that references the -uploaded model artifacts in the Files service. When ``deployment_config`` is set -on the task config, the task also launches an inference deployment. - Usage: export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.unsloth.tasks.model_entity + python -m nmp.customization_common.tasks.model_entity --service-name customizer """ import json @@ -20,6 +14,7 @@ import time from pathlib import Path +import httpx from nemo_platform import ( APIConnectionError, APITimeoutError, @@ -38,6 +33,7 @@ from nemo_platform.types.models import LoraParam, ModelEntity from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import InternalServerError as ClientInternalServerError from nemo_platform_plugin.files.client import FilesClient from nmp.common.sdk_factory import get_task_sdk from nmp.customization_common.schemas.model_entity import ( @@ -45,14 +41,12 @@ ModelEntityCreationError, ModelEntityTaskConfig, ) +from nmp.customization_common.schemas.values import FinetuningType from nmp.customization_common.service.context import NMPJobContext -from nmp.unsloth.app.constants import SERVICE_NAME -from nmp.unsloth.entities.values import FinetuningType from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential logger = logging.getLogger(__name__) -# Retry configuration. MAX_RETRIES = 3 INITIAL_BACKOFF_SECONDS = 1.0 MAX_BACKOFF_SECONDS = 30.0 @@ -62,6 +56,15 @@ SPEC_POLL_INTERVAL_SECONDS = 10 SPEC_POLL_TIMEOUT_SECONDS = 600 +TRANSIENT_RETRYABLE_EXCEPTIONS = ( + InternalServerError, + APITimeoutError, + APIConnectionError, + ClientInternalServerError, + httpx.TimeoutException, + httpx.ConnectError, +) + def get_config(config_path: Path) -> ModelEntityTaskConfig: """Load and validate the model_entity step config from disk.""" @@ -70,11 +73,7 @@ def get_config(config_path: Path) -> ModelEntityTaskConfig: def sanitize_name(prefix: str, name: str) -> str: - """Build a deployment-safe name from a free-form model name. - - Must match the API's ``{'pattern': '^[a-z](?!.*--)[a-z0-9\\-@.+_]{1,62}(? ModelEntity: - """Poll until the model_spec task has populated the model's spec. - - The spec must be populated before creating a deployment because the - inference service relies on ``spec.family`` and ``spec.base_num_parameters`` - to select the correct NIM profile. - - Raises: - ModelEntityCreationError: If the spec is not populated within the timeout. - """ + """Poll until the model_spec task has populated the model's spec.""" logger.info(f"Waiting for model_spec to populate spec on {workspace}/{name}") start = time.monotonic() while time.monotonic() - start < SPEC_POLL_TIMEOUT_SECONDS: try: target = self.sdk.models.retrieve(name=name, workspace=workspace) - if target.spec: - logger.info(f"Spec populated on {workspace}/{name}") - return target - except (APIConnectionError, APITimeoutError, InternalServerError) as e: + spec = target.spec + if spec is not None: + family = getattr(spec, "family", None) + base_num_parameters = getattr(spec, "base_num_parameters", None) + if family and base_num_parameters is not None: + logger.info(f"Spec populated on {workspace}/{name}") + return target + raise ModelEntityCreationError( + f"Model spec on {workspace}/{name} is missing required fields: " + "family and base_num_parameters must be set (typically by the " + "platform model_spec task). Verify the model checkpoint is valid " + "and in a supported format." + ) + except ModelEntityCreationError: + raise + except TRANSIENT_RETRYABLE_EXCEPTIONS as e: logger.warning(f"Transient error polling spec for {workspace}/{name}: {e}") time.sleep(SPEC_POLL_INTERVAL_SECONDS) @@ -119,10 +122,14 @@ def _wait_for_spec(self, workspace: str, name: str) -> ModelEntity: def get_model_entity(self, model_entity: str, fileset_workspace: str) -> ModelEntity: """Resolve ``"workspace/name"`` (or bare ``"name"``) to a ``ModelEntity``.""" parts = model_entity.split("/") - if len(parts) == 1: + if len(parts) == 1 and parts[0]: me_workspace, me_name = fileset_workspace, parts[0] - else: + elif len(parts) == 2 and all(parts): me_workspace, me_name = parts[0], parts[1] + else: + raise ModelEntityCreationError( + f"Invalid model entity reference '{model_entity}': expected 'name' or 'workspace/name'." + ) try: me: ModelEntity = self.sdk.models.retrieve(name=me_name, workspace=me_workspace) @@ -134,23 +141,15 @@ def get_model_entity(self, model_entity: str, fileset_workspace: str) -> ModelEn @retry( stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), + retry=retry_if_exception_type(TRANSIENT_RETRYABLE_EXCEPTIONS), reraise=True, ) def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, ModelEntity]: - """Create a model entity in the Models service. + """Create a model entity in the Models service.""" + output_workspace = config.workspace + logger.info(f"Creating model entity: {output_workspace}/{config.name}") - Returns: - Tuple of (result dict, deploy target). For LoRA the deploy target is the - *base* model entity; for SFT it is the newly created output model entity. - - Raises: - ModelEntityCreationError: If creation fails. - """ - workspace = self.job_ctx.workspace - logger.info(f"Creating model entity: {workspace}/{config.name}") - - fileset_workspace = config.fileset.workspace or workspace + fileset_workspace = config.fileset.workspace or self.job_ctx.workspace fileset_ref = f"{fileset_workspace}/{config.fileset.name}" logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") @@ -159,6 +158,8 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode workspace=fileset_workspace, name=config.fileset.name ) logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") + except TRANSIENT_RETRYABLE_EXCEPTIONS: + raise except Exception as e: logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") raise ModelEntityCreationError( @@ -170,7 +171,7 @@ def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, Mode if config.peft is not None and config.peft.type == FinetuningType.LORA: return self._create_or_update_adapter(config, base_me, fileset_ref) - return self._create_or_update_full_entity(config, base_me, fileset_ref, workspace) + return self._create_or_update_full_entity(config, fileset_ref, output_workspace) def _create_or_update_adapter( self, @@ -179,7 +180,7 @@ def _create_or_update_adapter( fileset_ref: str, ) -> tuple[dict, ModelEntity]: """Create or update a LoRA adapter on ``base_me``. Returns (result, base_me).""" - assert config.peft is not None # type narrowing — caller already checked + assert config.peft is not None try: output_me = self.sdk.models.adapters.create( model_name=base_me.name, @@ -214,7 +215,7 @@ def _create_or_update_adapter( f"for base model {base_me.workspace}/{base_me.name}" ) return output_me.model_dump(), base_me - except (InternalServerError, APITimeoutError, APIConnectionError): + except TRANSIENT_RETRYABLE_EXCEPTIONS: raise except Exception as update_error: logger.exception( @@ -230,7 +231,6 @@ def _create_or_update_adapter( def _create_or_update_full_entity( self, config: ModelEntityTaskConfig, - base_me: ModelEntity, fileset_ref: str, workspace: str, ) -> tuple[dict, ModelEntity]: @@ -242,7 +242,7 @@ def _create_or_update_full_entity( "description": config.description, "fileset": fileset_ref, "finetuning_type": ft_type, - "trust_remote_code": base_me.trust_remote_code, + "trust_remote_code": config.trust_remote_code, } if config.base_model: request_body["base_model"] = config.base_model @@ -262,7 +262,7 @@ def _create_or_update_full_entity( ) logger.info(f"Successfully updated model entity: {output_me.workspace}/{output_me.name}") return output_me.model_dump(), output_me - except (InternalServerError, APITimeoutError, APIConnectionError): + except TRANSIENT_RETRYABLE_EXCEPTIONS: raise except Exception as update_error: logger.exception(f"Failed to update existing model entity: {update_error}") @@ -274,17 +274,11 @@ def _create_or_update_full_entity( raise ModelEntityCreationError(f"Failed to create model entity: {e}") from e def launch_model(self, config: ModelEntityTaskConfig, me: ModelEntity) -> None: - """Deploy a model entity after creation. - - For LoRA jobs, ``me`` should be the base model entity. - For SFT jobs, ``me`` should be the output model entity. - """ + """Deploy a model entity after creation.""" dc = config.deployment_config if dc is None: return - # LORA_MERGED produces a full-weight model, so it's deployed like SFT and - # is intentionally excluded from the LoRA-only checks below. is_lora = config.peft is not None and config.peft.type == FinetuningType.LORA if is_lora and self._has_active_deployment(me): return @@ -379,7 +373,7 @@ def _create_deployment_config(self, deploy_params: DeploymentParameters, me: Mod def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: ModelEntity) -> None: """Create a deployment from the given ``ModelDeploymentConfig``.""" - logger.info(f"Deployment config: {deployment_config}") + logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") if not me.spec: _ = self._wait_for_spec(me.workspace, me.name) @@ -391,7 +385,7 @@ def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: Model name=deployment_name, config=deployment_config.name, ) - logger.info(f"Deployment created: {deployment}") + logger.info(f"Deployment created: {deployment.workspace}/{deployment.name}") except ConflictError: logger.info(f"Deployment {deployment_config.workspace}/{deployment_name} already exists") deployment = self.sdk.inference.deployments.retrieve( @@ -403,32 +397,36 @@ def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: Model workspace=deployment.workspace, name=deployment.name, ) - logger.info(f"Deployment status: {deployment_status}") - - -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the model entity creation task. + logger.info( + f"Deployment {deployment_status.workspace}/{deployment_status.name} status: {deployment_status.status}" + ) - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - Returns: - Exit code (0 for success, non-zero for failure). - """ +def run( + sdk: NeMoPlatform | None = None, + job_ctx: NMPJobContext | None = None, + *, + service_name: str, +) -> int: + """Execute the model entity creation task.""" job_ctx = job_ctx or NMPJobContext.from_env() sdk_owned = sdk is None try: - sdk = sdk or get_task_sdk(SERVICE_NAME).with_options(workspace=job_ctx.workspace) + sdk = sdk or get_task_sdk(service_name).with_options(workspace=job_ctx.workspace) runner = ModelEntityRunner(sdk=sdk, job_ctx=job_ctx) config = get_config(job_ctx.config_path) - logger.info(f"Starting model entity task with job context: {job_ctx}") - logger.info(f"Config: {config.model_dump_json(indent=2)}") + logger.info( + "Starting model entity task: job_id=%s, name=%s, workspace=%s, fileset=%s/%s, deployment_configured=%s", + job_ctx.job_id, + config.name, + config.workspace, + config.fileset.workspace or job_ctx.workspace, + config.fileset.name, + config.deployment_config is not None, + ) logger.info(f"NeMo Platform service URL: {sdk.base_url}") result, deploy_target = runner.create_model_entity(config) diff --git a/packages/nmp_customization_common/tests/contributor/test_jobs.py b/packages/nmp_customization_common/tests/contributor/test_jobs.py new file mode 100644 index 0000000000..b2bce9f5ed --- /dev/null +++ b/packages/nmp_customization_common/tests/contributor/test_jobs.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime-gate tests for the shared customization submit scaffold. + +Covers :func:`require_container_runtime` (automodel / unsloth) and +:func:`require_distributed_runtime` (rl). We patch the platform config and the +Docker-availability probe so the checks are exercised without a real runtime. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from nemo_platform_plugin.config import Runtime +from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError +from nmp.customization_common.contributor import jobs as jobs_mod +from nmp.customization_common.contributor.jobs import ( + require_container_runtime, + require_distributed_runtime, +) + + +@pytest.fixture +def _patch_runtime(monkeypatch: pytest.MonkeyPatch): + """Return a helper that pins the platform runtime and Docker availability.""" + + def _apply(runtime: Runtime, *, docker_available: bool = True) -> None: + monkeypatch.setattr( + jobs_mod.NemoPlatformConfig, + "get", + classmethod(lambda cls: SimpleNamespace(runtime=runtime)), + ) + monkeypatch.setattr( + "nemo_platform_plugin.config.validate_docker_available", + lambda: docker_available, + ) + + return _apply + + +class TestRequireContainerRuntime: + def test_kubernetes_single_node_ok(self, _patch_runtime) -> None: + _patch_runtime(Runtime.KUBERNETES) + require_container_runtime("Automodel") # no raise + + def test_kubernetes_multi_node_ok(self, _patch_runtime) -> None: + _patch_runtime(Runtime.KUBERNETES) + require_container_runtime("Automodel", num_nodes=2) # no raise + + def test_docker_single_node_ok(self, _patch_runtime) -> None: + _patch_runtime(Runtime.DOCKER, docker_available=True) + require_container_runtime("Automodel") # no raise + + def test_docker_daemon_unavailable_raises(self, _patch_runtime) -> None: + _patch_runtime(Runtime.DOCKER, docker_available=False) + with pytest.raises(PlatformJobCompilationError, match="reachable Docker daemon"): + require_container_runtime("Automodel") + + def test_docker_multi_node_requires_kubernetes(self, _patch_runtime) -> None: + _patch_runtime(Runtime.DOCKER, docker_available=True) + with pytest.raises(PlatformJobCompilationError, match="multi-node training .* requires"): + require_container_runtime("Automodel", num_nodes=2) + + def test_none_runtime_raises(self, _patch_runtime) -> None: + _patch_runtime(Runtime.NONE) + with pytest.raises(PlatformJobCompilationError, match="requires a container runtime"): + require_container_runtime("Automodel") + + def test_none_runtime_multi_node_raises_kubernetes_message(self, _patch_runtime) -> None: + _patch_runtime(Runtime.NONE) + with pytest.raises(PlatformJobCompilationError, match="multi-node training .* requires"): + require_container_runtime("Automodel", num_nodes=4) + + +class TestRequireDistributedRuntime: + def test_kubernetes_ok(self, _patch_runtime) -> None: + _patch_runtime(Runtime.KUBERNETES) + require_distributed_runtime("NeMo-RL") # no raise + + def test_docker_raises(self, _patch_runtime) -> None: + _patch_runtime(Runtime.DOCKER) + with pytest.raises(PlatformJobCompilationError, match="requires platform.runtime: kubernetes"): + require_distributed_runtime("NeMo-RL") diff --git a/services/unsloth/tests/test_file_io.py b/packages/nmp_customization_common/tests/tasks/test_file_io.py similarity index 51% rename from services/unsloth/tests/test_file_io.py rename to packages/nmp_customization_common/tests/tasks/test_file_io.py index 9832038e03..81960e94e7 100644 --- a/services/unsloth/tests/test_file_io.py +++ b/packages/nmp_customization_common/tests/tasks/test_file_io.py @@ -1,20 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the Unsloth file_io runner. - -Uses a hybrid testing strategy: - -- ``MagicMock`` for the SDK because the runner now uses chained calls like - ``self.sdk.with_options(timeout=...).files.upload(callback=...)``. Mock - fluents handle that cleanly. -- A ``NoOpProgressReporter`` to suppress Jobs-service reporting during tests. -- Light record-keeping wrappers for fileset / file objects to keep - assertions readable. - -The Files service contract is exercised in the SDK's own integration tests; -these focus on "what shape of arguments did this call site emit?". -""" +"""Tests for the shared customization file_io runner.""" from __future__ import annotations @@ -25,10 +12,10 @@ import pytest -def _make_runner(sdk, workspace: str = "default", storage_path: Path | None = None): +def _make_runner(sdk, *, service_source: str = "unsloth", workspace: str = "default", storage_path: Path | None = None): from nmp.customization_common.service.context import NMPJobContext + from nmp.customization_common.tasks.file_io.run import FileIORunner from nmp.customization_common.tasks.file_io_progress_reporter import NoOpProgressReporter - from nmp.unsloth.tasks.file_io.run import FileIORunner job_ctx = NMPJobContext( workspace=workspace, @@ -41,32 +28,24 @@ def _make_runner(sdk, workspace: str = "default", storage_path: Path | None = No storage_path=storage_path or Path("/tmp"), config_path=Path("/tmp/cfg.json"), ) - return FileIORunner(sdk=sdk, progress_reporter=NoOpProgressReporter(), job_ctx=job_ctx) + return FileIORunner( + sdk=sdk, + progress_reporter=NoOpProgressReporter(), + job_ctx=job_ctx, + service_source=service_source, + ) def _make_sdk() -> MagicMock: - """Build a MagicMock SDK with sensible defaults for fluent chaining. - - ``with_options`` returns the same SDK so chained timeouts don't break - attribute access in tests. - """ sdk = MagicMock() sdk.with_options.return_value = sdk return sdk def _raise_runner_conflict() -> None: - """Raise the exact ``ConflictError`` class bound in tasks.file_io.run. - - The real ``nemo_platform.ConflictError`` is an ``APIStatusError`` that - needs ``response`` + ``body`` kwargs to instantiate. We grab the class - via ``sys.modules`` (the package ``__init__.py`` re-exports ``run`` as - a function, which shadows the submodule for plain attribute access) - and construct via ``__new__`` to dodge the constructor signature. - """ import sys - run_mod = sys.modules["nmp.unsloth.tasks.file_io.run"] + run_mod = sys.modules["nmp.customization_common.tasks.file_io.run"] raise run_mod.ConflictError.__new__(run_mod.ConflictError, "already exists") @@ -78,21 +57,17 @@ def _make_dir(tmp_path: Path) -> Path: return src -# --------------------------------------------------------------------------- -# FileIORunner.create_fileset -# --------------------------------------------------------------------------- - - class TestCreateFileset: - @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + @patch("nmp.customization_common.tasks.file_io.run.client_from_platform") def test_creates_fileset_with_service_source_and_metadata(self, mock_cfp) -> None: from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.customization_common.schemas.file_io import FileSetRef mock_fc = MagicMock() + mock_fc.with_options.return_value = mock_fc mock_cfp.return_value = mock_fc sdk = _make_sdk() - runner = _make_runner(sdk) + runner = _make_runner(sdk, service_source="automodel") metadata = {"model": {"tool_calling": {"tool_call_parser": "llama3_json"}}} dest = FileSetRef(workspace="default", name="qwen-test") @@ -104,21 +79,22 @@ def test_creates_fileset_with_service_source_and_metadata(self, mock_cfp) -> Non body = call.kwargs["body"] assert isinstance(body, CreateFilesetRequest) assert body.name == "qwen-test" - assert body.custom_fields == {"service_source": "unsloth"} + assert body.custom_fields == {"service_source": "automodel"} assert body.metadata is not None assert body.metadata.model is not None assert body.metadata.model.tool_calling.tool_call_parser == "llama3_json" - @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + @patch("nmp.customization_common.tasks.file_io.run.client_from_platform") def test_conflict_patches_metadata_on_existing(self, mock_cfp) -> None: from nemo_platform_plugin.files.types import UpdateFilesetRequest from nmp.customization_common.schemas.file_io import FileSetRef mock_fc = MagicMock() + mock_fc.with_options.return_value = mock_fc mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() mock_cfp.return_value = mock_fc sdk = _make_sdk() - runner = _make_runner(sdk) + runner = _make_runner(sdk, service_source="rl") dest = FileSetRef(workspace="default", name="exists") metadata = {"model": {"tool_calling": {"tool_call_parser": "hermes"}}} @@ -134,11 +110,12 @@ def test_conflict_patches_metadata_on_existing(self, mock_cfp) -> None: assert body.metadata.model is not None assert body.metadata.model.tool_calling.tool_call_parser == "hermes" - @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") + @patch("nmp.customization_common.tasks.file_io.run.client_from_platform") def test_conflict_no_metadata_skips_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef mock_fc = MagicMock() + mock_fc.with_options.return_value = mock_fc mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() mock_cfp.return_value = mock_fc sdk = _make_sdk() @@ -149,33 +126,6 @@ def test_conflict_no_metadata_skips_update(self, mock_cfp) -> None: mock_fc.update_fileset.assert_not_called() - @patch("nmp.unsloth.tasks.file_io.run.client_from_platform") - def test_update_failure_is_warning_not_fatal( - self, - mock_cfp, - caplog: pytest.LogCaptureFixture, - ) -> None: - from nmp.customization_common.schemas.file_io import FileSetRef - - mock_fc = MagicMock() - mock_fc.create_fileset.side_effect = lambda **_: _raise_runner_conflict() - mock_fc.update_fileset.side_effect = RuntimeError("backend down") - mock_cfp.return_value = mock_fc - sdk = _make_sdk() - runner = _make_runner(sdk) - dest = FileSetRef(workspace="default", name="exists") - metadata = {"model": {"tool_calling": {"tool_call_parser": "hermes"}}} - - with caplog.at_level("WARNING"): - runner.create_fileset(dest, metadata=metadata) - - assert any("Could not patch metadata" in r.getMessage() for r in caplog.records) - - -# --------------------------------------------------------------------------- -# FileIORunner.upload_fileset (uses FilesetFileSystem with callbacks) -# --------------------------------------------------------------------------- - class TestUploadFileset: def test_directory_uploads_with_trailing_slash(self, tmp_path: Path) -> None: @@ -190,27 +140,11 @@ def test_directory_uploads_with_trailing_slash(self, tmp_path: Path) -> None: sdk.files.upload.assert_called_once() call = sdk.files.upload.call_args - # Trailing slash → directory contents go to fileset root. assert call.kwargs["local_path"] == f"{src.resolve()}/" assert call.kwargs["remote_path"] == "" assert call.kwargs["fileset"] == "qwen-test" assert call.kwargs["workspace"] == "default" - def test_single_file_uploads_to_basename(self, tmp_path: Path) -> None: - from nmp.customization_common.schemas.file_io import FileSetRef - - sdk = _make_sdk() - runner = _make_runner(sdk) - src = tmp_path / "result.json" - src.write_text("{}") - dest = FileSetRef(workspace="default", name="single-file") - - runner.upload_fileset(dest, src.resolve()) - - call = sdk.files.upload.call_args - assert call.kwargs["local_path"] == str(src.resolve()) - assert call.kwargs["remote_path"] == src.name - def test_upload_failure_propagates_as_file_upload_error(self, tmp_path: Path) -> None: from nmp.customization_common.schemas.file_io import FileSetRef, FileUploadError @@ -224,17 +158,11 @@ def test_upload_failure_propagates_as_file_upload_error(self, tmp_path: Path) -> runner.upload_fileset(dest, src.resolve()) -# --------------------------------------------------------------------------- -# FileIORunner.download_fileset -# --------------------------------------------------------------------------- - - class TestDownloadFileset: def test_lists_then_downloads(self, tmp_path: Path) -> None: from nmp.customization_common.schemas.file_io import FileSetRef sdk = _make_sdk() - # files.list returns an object with .data (a list of FilesetFile-ish objects). sdk.files.list.return_value = types.SimpleNamespace( data=[ types.SimpleNamespace(path="model.safetensors", size=100), @@ -269,64 +197,3 @@ def test_empty_fileset_returns_zero_stats_without_downloading(self, tmp_path: Pa assert stats.files_downloaded == 0 assert stats.total_bytes == 0 sdk.files.download.assert_not_called() - - -# --------------------------------------------------------------------------- -# build_output_metadata -# --------------------------------------------------------------------------- - - -class TestBuildOutputMetadata: - def test_extracts_canonical_fields(self) -> None: - from nmp.unsloth.schemas import ( - DatasetSpec, - ModelLoadSpec, - OutputResponse, - TrainingSpec, - UnslothJobOutput, - ) - from nmp.unsloth.tasks.file_io.run import build_output_metadata - - spec = UnslothJobOutput( - model=ModelLoadSpec( - name="Qwen/Qwen3-0.6B", - load_in_4bit=False, - load_in_8bit=False, - ), - dataset=DatasetSpec(path="/data/sample.jsonl"), - training=TrainingSpec(finetuning_type="all_weights", lora=None), - output=OutputResponse( - name="qwen-out", - type="model", - save_method="lora", - fileset="qwen-out", - ), - ) - - meta = build_output_metadata(spec) - assert meta == { - "model": "Qwen/Qwen3-0.6B", - "finetuning_type": "all_weights", - "save_method": "lora", - "output_type": "model", - } - - -# --------------------------------------------------------------------------- -# validate_safe_path (path traversal protection) -# --------------------------------------------------------------------------- - - -class TestValidateSafePath: - def test_safe_path_resolves(self, tmp_path: Path) -> None: - from nmp.customization_common.tasks.file_io_utils import validate_safe_path - - result = validate_safe_path(tmp_path, "subdir/file.txt") - assert result == (tmp_path / "subdir/file.txt").resolve() - - def test_traversal_raises(self, tmp_path: Path) -> None: - from nmp.customization_common.schemas.file_io import PathTraversalError - from nmp.customization_common.tasks.file_io_utils import validate_safe_path - - with pytest.raises(PathTraversalError): - validate_safe_path(tmp_path, "../../etc/passwd") diff --git a/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py new file mode 100644 index 0000000000..966c31f758 --- /dev/null +++ b/packages/nmp_customization_common/tests/tasks/test_file_io_metadata.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for output fileset metadata helpers.""" + +from nmp.customization_common.tasks.file_io_metadata import build_output_metadata + + +class TestBuildOutputMetadata: + def test_extracts_canonical_fields(self) -> None: + meta = build_output_metadata( + model="Qwen/Qwen3-0.6B", + finetuning_type="all_weights", + save_method="lora", + output_type="model", + ) + assert meta == { + "model": "Qwen/Qwen3-0.6B", + "finetuning_type": "all_weights", + "save_method": "lora", + "output_type": "model", + } + + def test_omits_save_method_when_not_provided(self) -> None: + meta = build_output_metadata( + model="default/base-model", + finetuning_type="all_weights", + output_type="model", + ) + assert meta == { + "model": "default/base-model", + "finetuning_type": "all_weights", + "output_type": "model", + } diff --git a/packages/nmp_platform/README.md b/packages/nmp_platform/README.md index 0fb9083927..6e124264ca 100644 --- a/packages/nmp_platform/README.md +++ b/packages/nmp_platform/README.md @@ -24,9 +24,9 @@ that points callers at `nemo services run`. A handful of task container images and seed jobs invoke `nemo-platform run task` as their entrypoint: -- `nmp-automodel-tasks` — used by the automodel file_io task - (`services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml` - runs `nmp.automodel.tasks.file_io`). +- `nmp-customizer-tasks` — shared CPU task image (`nmp.customization_common.tasks.file_io` / + `model_entity`, plus `model_spec` and the LoRA sidecar). Local compose example: + `services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml`. - `services/platform-seed` — recommended invocation in its README is `nemo-platform run task --task nmp.platform_seed`. diff --git a/packages/nmp_platform/config/local.yaml b/packages/nmp_platform/config/local.yaml index 051a82399e..337d577803 100644 --- a/packages/nmp_platform/config/local.yaml +++ b/packages/nmp_platform/config/local.yaml @@ -123,9 +123,9 @@ models: enabled: true # LoRA adapter-download sidecar for deployments with lora_enabled. The # docker backend has no separate args field (unlike k8s) and keeps the - # image ENTRYPOINT — nmp-automodel-tasks is `/opt/venv/bin/python`, so + # image ENTRYPOINT — nmp-customizer-tasks is `/opt/venv/bin/python`, so # the module goes in lora_sidecar_command (runs `python -m `). - lora_sidecar_image_name: nmp-automodel-tasks + lora_sidecar_image_name: nmp-customizer-tasks lora_sidecar_command: ["-m", "nmp.core.models.sidecars.adapters.main"] # Inference gateway configuration - uses default values diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/config/local.yaml b/packages/nmp_platform_runner/src/nmp/platform_runner/config/local.yaml index 3383800d06..637e52fa85 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/config/local.yaml +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/config/local.yaml @@ -65,9 +65,9 @@ models: enabled: true # LoRA adapter-download sidecar for deployments with lora_enabled. The # docker backend has no separate args field (unlike k8s) and keeps the - # image ENTRYPOINT — nmp-automodel-tasks is `/opt/venv/bin/python`, so + # image ENTRYPOINT — nmp-customizer-tasks is `/opt/venv/bin/python`, so # the module goes in lora_sidecar_command (runs `python -m `). - lora_sidecar_image_name: nmp-automodel-tasks + lora_sidecar_image_name: nmp-customizer-tasks lora_sidecar_command: ["-m", "nmp.core.models.sidecars.adapters.main"] inference_gateway: {} diff --git a/plugins/nemo-automodel/src/nemo_automodel_plugin/jobs/jobs.py b/plugins/nemo-automodel/src/nemo_automodel_plugin/jobs/jobs.py index 926622cf41..0ca2045bb1 100644 --- a/plugins/nemo-automodel/src/nemo_automodel_plugin/jobs/jobs.py +++ b/plugins/nemo-automodel/src/nemo_automodel_plugin/jobs/jobs.py @@ -20,7 +20,7 @@ from nemo_platform_plugin.jobs.api_factory import PlatformJobSpec from nemo_platform_plugin.jobs.docker import validate_gpu_available_for_docker from nmp.automodel.compile import platform_job_config_compiler -from nmp.customization_common.contributor.jobs import BaseSubmitJob, require_docker_runtime +from nmp.customization_common.contributor.jobs import BaseSubmitJob, require_container_runtime from pydantic import BaseModel @@ -32,7 +32,7 @@ class AutomodelJob(BaseSubmitJob): job_collection_path: ClassVar[str | None] = "/automodel/jobs" input_spec_schema: ClassVar[type[BaseModel] | None] = AutomodelJobInput spec_schema: ClassVar[type[BaseModel] | None] = AutomodelJobOutput - docker_runtime_label: ClassVar[str] = "Automodel" + runtime_label: ClassVar[str] = "Automodel" @classmethod async def _transform(cls, job_input: BaseModel, workspace: str, async_sdk: AsyncNeMoPlatform) -> AutomodelJobOutput: @@ -50,10 +50,12 @@ async def compile( options: dict | None = None, ) -> PlatformJobSpec: del entity_client, options - require_docker_runtime(cls.docker_runtime_label) canonical = ( spec if isinstance(spec, AutomodelJobOutput) else AutomodelJobOutput.model_validate(spec.model_dump()) ) + # Multi-node jobs compile to a gpu_distributed (Volcano) executor, which + # only exists on Kubernetes; gate here so docker platforms fail fast. + require_container_runtime(cls.runtime_label, num_nodes=canonical.parallelism.num_nodes) canonical.validate_for_training() plugin_config = get_config() diff --git a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/SKILL.md b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/SKILL.md index 053349bdb2..991bfc369b 100644 --- a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/SKILL.md +++ b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/SKILL.md @@ -153,7 +153,7 @@ For **`automodel`/`unsloth`**, training never runs inside the `nemo` CLI process - **Unsloth validation defaults** — when `dataset.validation_path` is set and `schedule.eval_steps` is omitted, the trainer runs validation once per effective epoch automatically. Report final `metrics.val_loss` from job status (see `references/reporting.md`). Set `eval_steps` explicitly to override cadence. - **Do not use local `docker info`** to pick automodel vs unsloth. Run `nemo jobs list-execution-profiles -f json` against the user's platform (login first only if auth is enabled — see **Authentication**; see `references/troubleshooting.md`). Default output is a table — **`-f json` is required** for scripting; parse **stdout only** (do not pipe `2>&1` into `json.load`). - **Do not merge stderr into stdout when parsing JSON** — `submit`, `explain`, and `-f json` commands write **JSON on stdout**; harmless warnings like `Configuration file not found, using defaults` go to **stderr**. Piping with **`2>&1`** before `json.load` raises `JSONDecodeError` even when submit **succeeded** — a common cause of **duplicate jobs** when the agent re-submits after a parse error. Parse stdout only; redirect stderr if needed (`2>/dev/null`). See `references/troubleshooting.md` § **Parsing CLI JSON**. -- For submit/image/plugin errors (all backends), read `references/troubleshooting.md`. Unsloth needs the `nmp-unsloth-training` container image on the **platform host's** Docker daemon (see `docker/unsloth/README.md`); rl needs the `nmp-rl-tasks` / `nmp-rl-training` images on the Kubernetes cluster (see **rl (DPO) gotchas** and `references/rl-kubernetes-runtime.md`). +- For submit/image/plugin errors (all backends), read `references/troubleshooting.md`. Unsloth needs the `nmp-unsloth-training` container image on the **platform host's** Docker daemon (see `docker/unsloth/README.md`); rl needs the `nmp-customizer-tasks` / `nmp-rl-training` images on the Kubernetes cluster (see **rl (DPO) gotchas** and `references/rl-kubernetes-runtime.md`). - **Missing training image on a remote platform** — if the user gave a non-localhost `NMP_BASE_URL` and the job errors with `Failed to pull image`, `manifest unknown`, or missing `nmp-unsloth-training` / automodel training image: **do not** run `docker build`, `docker pull`, or `docker buildx bake` on the agent machine. Report with the template in `references/reporting.md` (use **Output adapter fileset (planned):** on error), then append on-target build steps from `references/troubleshooting.md` § **Missing training images**. - **Gated HuggingFace models** (Llama, Gemma, …) — confirm `hf-token` + fileset `token_secret` before submit; download fails with `Failed to access upstream storage` / 502 when missing. See **HuggingFace token (gated models)** and `references/troubleshooting.md` § **Gated HuggingFace models**. - **Post-training eval format** — use the same CHAT `messages` JSONL as training. **Do not** flatten rows to `prompt`/`expected` for the evaluator. Send `messages[:-1]` at inference (exclude final assistant label); score against `messages[-1].content`. See `references/post-training-eval.md` and `references/eval_helpers.py`. @@ -164,7 +164,7 @@ For **`automodel`/`unsloth`**, training never runs inside the `nemo` CLI process - **rl is DPO, not SFT** — it trains on **preference pairs** `{prompt, chosen, rejected}`, full-weight (no LoRA/adapter; `finetuning_type` is not user-set). Don't route SFT/LoRA work here, and don't route DPO to automodel/unsloth. - **One preference fileset, two files** — `dataset` is a **single string** ref to a fileset that holds **both** `training.jsonl` and `validation.jsonl` (uploaded with `--remote-path`). Unlike automodel (`dataset.training`/`dataset.validation`) and unsloth (`dataset.path`/`validation_path`), there is no separate validation ref. See `references/dataset-formats.md` § NeMo-RL. - **String refs** — `model` and `dataset` are plain strings (`"workspace/name"`), not objects. The training method goes under `training` with `type: "dpo"`. -- **Kubernetes job backend, not Docker** — rl steps run as Kubernetes pods via the `kubernetes_job` backend; the docker job backend cannot run rl. `rl submit` fails fast on a docker-runtime platform. The target cluster must have the **job-step images** (`nmp-rl-tasks`, `nmp-rl-training`), the **jobs-launcher** image (the per-step init container), and a **job-storage PVC**. Verify the platform with `nemo jobs list-execution-profiles -f json` (expect `backend: kubernetes_job`); to configure one, see `references/rl-kubernetes-runtime.md`. Multi-node (`parallelism.num_nodes > 1`) also needs the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH` (shared FS for Ray coordination) or compile fails fast. +- **Kubernetes job backend, not Docker** — rl steps run as Kubernetes pods via the `kubernetes_job` backend; the docker job backend cannot run rl. `rl submit` fails fast on a docker-runtime platform. The target cluster must have the **job-step images** (`nmp-customizer-tasks`, `nmp-rl-training`), the **jobs-launcher** image (the per-step init container), and a **job-storage PVC**. Verify the platform with `nemo jobs list-execution-profiles -f json` (expect `backend: kubernetes_job`); to configure one, see `references/rl-kubernetes-runtime.md`. Multi-node (`parallelism.num_nodes > 1`) also needs the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH` (shared FS for Ray coordination) or compile fails fast. - **Job id prefix is `rl-`** and the platform auto-generates it — `rl submit` has **no `--name` flag** (the job JSON `name` is the *output* name, not the job id). Derive the job id from `nemo jobs list` (newest `rl-*`) for polling; `poll_customization_job.sh rl-` works. - **DPO main knob is `ref_policy_kl_penalty`** (β). For OOM, enable `activation_checkpointing: true` first. Full DPO field reference: `references/hyperparameters-rl.md`. - **`max_steps` + `epochs`** — same caveat as the other backends: `max_steps` caps mid-epoch; it's in the smoke fixture (`plugins/nemo-rl/tests/fixtures/minimal_dpo.json`) — omit for real runs. diff --git a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/rl-kubernetes-runtime.md b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/rl-kubernetes-runtime.md index 68807ac106..8b328020e8 100644 --- a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/rl-kubernetes-runtime.md +++ b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/rl-kubernetes-runtime.md @@ -43,7 +43,7 @@ be configured with all of: (the platform rewrites the `NMP_*_URL` it injects into pods to this, so the download/upload steps can call the files/jobs APIs). 4. The target GPU cluster has, available as pullable/loaded images: the job-step - images (`nmp-rl-tasks`, `nmp-rl-training`), the **jobs-launcher** image (each + images (`nmp-customizer-tasks`, `nmp-rl-training`), the **jobs-launcher** image (each step runs a launcher init container), and a **job-storage PVC** the steps share. 5. Multi-node only (`parallelism.num_nodes > 1`): `NMP_RL_MULTINODE_SHARED_STORAGE_PATH` (a shared filesystem for Ray's cross-node coordination). diff --git a/plugins/nemo-deployments/pyproject.toml b/plugins/nemo-deployments/pyproject.toml index e1a0bb0d6a..ab3acfab1c 100644 --- a/plugins/nemo-deployments/pyproject.toml +++ b/plugins/nemo-deployments/pyproject.toml @@ -6,7 +6,7 @@ requires-python = ">=3.11,<3.14" dependencies = [ "fastapi>=0.115", "httpx>=0.27", - "nemo-platform", + "nemo-platform-sdk", "nemo-platform-plugin", "pydantic>=2.10.6", ] @@ -31,7 +31,7 @@ packages = ["src/nemo_deployments_plugin"] [tool.uv.sources] -nemo-platform = { workspace = true } +nemo-platform-sdk = { workspace = true } nemo-platform-plugin = { workspace = true } diff --git a/plugins/nemo-rl/README.md b/plugins/nemo-rl/README.md index 701ed26d85..6171c70ada 100644 --- a/plugins/nemo-rl/README.md +++ b/plugins/nemo-rl/README.md @@ -59,7 +59,7 @@ authoritative input shape; `nemo customization rl explain` prints it live. `submit` → `RlJobInput` → transform → `RlJobOutput` → compiled `PlatformJobSpec`: -1. **download** — model fileset + preference dataset → PVC (CPU, `nmp-rl-tasks`) +1. **download** — model fileset + preference dataset → PVC (CPU, `nmp-customizer-tasks`) 2. **dpo-training** — Ray DPO step (GPU, `nmp-rl-training`); single-node `gpu` or multi-node `gpu_distributed` executor, selected by `parallelism.num_nodes` 3. **upload** — trained checkpoint → output fileset (CPU) @@ -72,4 +72,4 @@ authoritative input shape; `nemo customization rl explain` prints it live. - **Design:** [`docs/customizer/nemo-rl-dpo-plugin-design.md`](../../docs/customizer/nemo-rl-dpo-plugin-design.md). - **GPU e2e smoke test:** [`scripts/gpu-dpo-smoke/`](../../scripts/gpu-dpo-smoke). - **Images:** [`docker/Dockerfile.nmp-rl-base`](../../docker/Dockerfile.nmp-rl-base), - `Dockerfile.nmp-rl-training`, `Dockerfile.nmp-rl-tasks`. + `Dockerfile.nmp-rl-training`, `Dockerfile.nmp-customizer-tasks`. diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py b/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py index bfb5cd11e7..022c6add5f 100644 --- a/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py +++ b/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py @@ -35,7 +35,7 @@ class RlJob(BaseSubmitJob): job_collection_path: ClassVar[str | None] = "/rl/jobs" input_spec_schema: ClassVar[type[BaseModel] | None] = RlJobInput spec_schema: ClassVar[type[BaseModel] | None] = RlJobOutput - docker_runtime_label: ClassVar[str] = "NeMo-RL" + runtime_label: ClassVar[str] = "NeMo-RL" @classmethod async def _transform(cls, job_input: BaseModel, workspace: str, async_sdk: AsyncNeMoPlatform) -> RlJobOutput: @@ -61,7 +61,7 @@ async def compile( (single-node ``gpu`` vs multi-node ``gpu_distributed``). """ del entity_client, options - require_distributed_runtime(cls.docker_runtime_label) + require_distributed_runtime(cls.runtime_label) canonical = spec if isinstance(spec, RlJobOutput) else RlJobOutput.model_validate(spec.model_dump()) canonical.validate_for_training() diff --git a/plugins/nemo-unsloth/README.md b/plugins/nemo-unsloth/README.md index 235da26535..013bc3127d 100644 --- a/plugins/nemo-unsloth/README.md +++ b/plugins/nemo-unsloth/README.md @@ -32,11 +32,16 @@ Job JSON uses the `UnslothJobInput` schema (see `nemo_unsloth_plugin/schema.py`) What happens after submit: 1. The plugin's `to_spec` validates the model entity + dataset fileset against the live platform. -2. `UnslothJob.compile` produces a 4-step `PlatformJobSpec`: - 1. **`model-and-dataset-download`** — CPU step, `nmp.unsloth.tasks.file_io` pulls the model entity's fileset + the dataset fileset to the shared PVC. - 2. **`training`** — GPU step, `nmp.unsloth.tasks.training` runs `train_sft` against the local paths. - 3. **`model-upload`** — CPU step, `nmp.unsloth.tasks.file_io` uploads the saved checkpoint to a new fileset (named after `output.fileset`). - 4. **`model-entity-creation`** — CPU step, `nmp.unsloth.tasks.model_entity` registers the output entity (adapter for LoRA, full model entity otherwise). +2. `UnslothJob.compile` produces a 4-step `PlatformJobSpec` (see `nmp.unsloth.images` for the exact container commands): + + | Step | Image | Container command | + |------|-------|-------------------| + | **`model-and-dataset-download`** | `nmp-customizer-tasks` | `python -m nmp.customization_common.tasks.file_io --service-source unsloth --service-name unsloth` | + | **`training`** | `nmp-unsloth-training` | `python -m nmp.unsloth.tasks.training` | + | **`model-upload`** | `nmp-customizer-tasks` | `python -m nmp.customization_common.tasks.file_io --service-source unsloth --service-name unsloth` | + | **`model-entity-creation`** | `nmp-customizer-tasks` | `python -m nmp.customization_common.tasks.model_entity --service-name unsloth` | + + Shared CPU tasks require explicit identity flags on the module entrypoints (`nmp.customization_common.tasks.file_io.__main__` requires both `--service-source` and `--service-name`; `model_entity.__main__` requires `--service-name`). The compiler sets these via `FILE_IO_TASK_COMMAND` and `MODEL_ENTITY_TASK_COMMAND` in `services/unsloth/src/nmp/unsloth/images.py`. `--service-source` is stamped on upload-created filesets; `--service-name` drives SDK auth/telemetry. 3. The platform Jobs runner schedules each step; tail logs with the standard jobs API. ## CLI surface @@ -75,7 +80,7 @@ The container image targets the same compute capabilities NVIDIA's stock `pytorc This plugin is the **thin contributor wrapper**. The heavy code lives in `services/unsloth/` (`nmp-unsloth`): - **Plugin** (`plugins/nemo-unsloth/`, `nemo_unsloth_plugin`) — `UnslothContributor`, `UnslothJob` (lifecycle + `compile()`), submitter-facing schema (`UnslothJobInput`), CLI overrides, SDK shapes, contributor wiring. -- **Service** (`services/unsloth/`, `nmp.unsloth`) — canonical schemas (`UnslothJobOutput` and shared sub-shapes), the `train_sft` training driver, the three container task entrypoints (`tasks/file_io`, `tasks/model_entity`, `tasks/training`), and the `platform_job_config_compiler`. +- **Service** (`services/unsloth/`, `nmp.unsloth`) — canonical schemas (`UnslothJobOutput` and shared sub-shapes), the `train_sft` training driver, CPU task commands wired through `nmp.customization_common.tasks.*` (`nmp-customizer-tasks` image), the `nmp.unsloth.tasks.training` GPU entrypoint, and the `platform_job_config_compiler`. The plugin imports two things from the service: diff --git a/plugins/nemo-unsloth/src/nemo_unsloth_plugin/jobs/jobs.py b/plugins/nemo-unsloth/src/nemo_unsloth_plugin/jobs/jobs.py index 2ccc6f1e75..196d651a5f 100644 --- a/plugins/nemo-unsloth/src/nemo_unsloth_plugin/jobs/jobs.py +++ b/plugins/nemo-unsloth/src/nemo_unsloth_plugin/jobs/jobs.py @@ -20,7 +20,7 @@ from nemo_platform_plugin.jobs.docker import validate_gpu_available_for_docker from nemo_unsloth_plugin.schema import UnslothJobInput from nemo_unsloth_plugin.transform import transform_input_to_output -from nmp.customization_common.contributor.jobs import BaseSubmitJob, require_docker_runtime +from nmp.customization_common.contributor.jobs import BaseSubmitJob, require_container_runtime from nmp.unsloth.compile import platform_job_config_compiler from nmp.unsloth.config import config as unsloth_config from nmp.unsloth.schemas import UnslothJobOutput @@ -35,7 +35,7 @@ class UnslothJob(BaseSubmitJob): job_collection_path: ClassVar[str | None] = "/unsloth/jobs" input_spec_schema: ClassVar[type[BaseModel] | None] = UnslothJobInput spec_schema: ClassVar[type[BaseModel] | None] = UnslothJobOutput - docker_runtime_label: ClassVar[str] = "Unsloth" + runtime_label: ClassVar[str] = "Unsloth" @classmethod async def _transform(cls, job_input: BaseModel, workspace: str, async_sdk: AsyncNeMoPlatform) -> UnslothJobOutput: @@ -59,7 +59,7 @@ async def compile( ``unsloth_config.default_training_execution_profile``. """ del entity_client, options - require_docker_runtime(cls.docker_runtime_label) + require_container_runtime(cls.runtime_label) canonical = spec if isinstance(spec, UnslothJobOutput) else UnslothJobOutput.model_validate(spec.model_dump()) execution_profile = profile or unsloth_config.default_training_execution_profile diff --git a/plugins/nemo-unsloth/tests/test_jobs.py b/plugins/nemo-unsloth/tests/test_jobs.py index 8a1fa5c00b..121963df0f 100644 --- a/plugins/nemo-unsloth/tests/test_jobs.py +++ b/plugins/nemo-unsloth/tests/test_jobs.py @@ -9,10 +9,10 @@ - ``to_spec`` resolves output naming + fileset against a stub SDK. - ``compile`` delegates to the service-side compiler (we patch it out) - and returns the resulting ``PlatformJobSpec`` after the Docker + and returns the resulting ``PlatformJobSpec`` after the container runtime check. -- The Docker runtime check fires when the platform isn't configured for - Docker. +- The container runtime check fires when the platform has no usable + container runtime (neither Kubernetes nor Docker). """ from __future__ import annotations @@ -99,7 +99,7 @@ def test_compile_delegates_to_service_compiler(self) -> None: ) with ( - patch("nemo_unsloth_plugin.jobs.jobs.require_docker_runtime"), + patch("nemo_unsloth_plugin.jobs.jobs.require_container_runtime"), patch( "nemo_unsloth_plugin.jobs.jobs.platform_job_config_compiler", new=AsyncMock(return_value=fake_spec), @@ -132,7 +132,7 @@ def test_compile_delegates_to_service_compiler(self) -> None: def test_compile_passes_caller_profile_override(self) -> None: canonical = _make_canonical() with ( - patch("nemo_unsloth_plugin.jobs.jobs.require_docker_runtime"), + patch("nemo_unsloth_plugin.jobs.jobs.require_container_runtime"), patch( "nemo_unsloth_plugin.jobs.jobs.platform_job_config_compiler", new=AsyncMock(return_value=SimpleNamespace(steps=[])), @@ -152,15 +152,15 @@ def test_compile_passes_caller_profile_override(self) -> None: assert compile_mock.await_args.kwargs["profile"] == "gpu_distributed" - def test_compile_rejects_non_docker_runtime(self) -> None: + def test_compile_rejects_runtime_without_container_support(self) -> None: canonical = _make_canonical() - # Force the runtime check to raise so we don't need a Docker daemon + # Force the runtime check to raise so we don't need a real runtime # in CI. The check is what runs first; the rest never executes. with patch( - "nemo_unsloth_plugin.jobs.jobs.require_docker_runtime", - side_effect=PlatformJobCompilationError("not docker"), + "nemo_unsloth_plugin.jobs.jobs.require_container_runtime", + side_effect=PlatformJobCompilationError("no container runtime"), ): - with pytest.raises(PlatformJobCompilationError, match="not docker"): + with pytest.raises(PlatformJobCompilationError, match="no container runtime"): asyncio.run( UnslothJob.compile( workspace="default", diff --git a/pytest.ini b/pytest.ini index 29033c341a..0980e96dfc 100644 --- a/pytest.ini +++ b/pytest.ini @@ -61,7 +61,7 @@ markers = integration: Service integration tests - test individual service interfaces and interactions (uses ASGI, mocks external services via SDK) gpu_integration: GPU integration tests - test individual services that utilize AI dependencies/ a GPU smoke_gpu_tasks: Import smoke tests for the nmp-gpu-tasks image - smoke_nmp_automodel_tasks: Import smoke tests for the nmp-automodel-tasks image + smoke_nmp_customizer_tasks: Import smoke tests for the nmp-customizer-tasks image smoke_nmp_automodel_training: Import smoke tests for the nmp-automodel-training image e2e: End-to-end tests - test complete customer workflows on deployed infrastructure (Helm/Docker Compose) auth_idp: Auth IdP e2e tests - provider-backed auth compose coverage through the gateway diff --git a/services/automodel/README.md b/services/automodel/README.md index 10a2edc2f0..9961fed16d 100644 --- a/services/automodel/README.md +++ b/services/automodel/README.md @@ -1,5 +1,5 @@ # nmp-automodel -Compiler and task entrypoints for NeMo Automodel training jobs on the platform. **No HTTP server** — consumed by `nemo-automodel-plugin` and Jobs task images (`my-registry/nemo-platform-dev/nmp-automodel-tasks`, `.../nmp-automodel-training`). +Compiler and task entrypoints for NeMo Automodel training jobs on the platform. **No HTTP server** — consumed by `nemo-automodel-plugin` and Jobs task images (`my-registry/nemo-platform-dev/nmp-customizer-tasks`, `.../nmp-automodel-training`). Runtime exceptions from `nemo_automodel` are mapped to user-facing error types via `src/nmp/automodel/tasks/training/errors/error_rules.yaml`. See [docs/automodel_errors.md](docs/automodel_errors.md) for the full catalog and validation status of each Automodel error. diff --git a/services/automodel/pyproject.toml b/services/automodel/pyproject.toml index cc1c575a67..9f044480de 100644 --- a/services/automodel/pyproject.toml +++ b/services/automodel/pyproject.toml @@ -21,9 +21,7 @@ version = "0.0.0" dev = ["pytest>=8.3.4", "pytest-asyncio>=0.25.3", "pytest-mock>=3.14.0"] [project.scripts] -nmp-automodel-file-io = "nmp.automodel.tasks.file_io:run" nmp-automodel-training = "nmp.automodel.tasks.training.__main__:run" -nmp-automodel-model-entity = "nmp.automodel.tasks.model_entity.__main__:run" [build-system] requires = ["hatchling"] diff --git a/services/automodel/src/nmp/automodel/app/jobs/compiler.py b/services/automodel/src/nmp/automodel/app/jobs/compiler.py index 047b73798d..b086482fc3 100644 --- a/services/automodel/src/nmp/automodel/app/jobs/compiler.py +++ b/services/automodel/src/nmp/automodel/app/jobs/compiler.py @@ -37,7 +37,12 @@ ) from nmp.automodel.config import config from nmp.automodel.entities.values import FinetuningType -from nmp.automodel.images import AUTOMODEL_PYTHON_ENTRYPOINT, get_tasks_image +from nmp.automodel.images import ( + AUTOMODEL_PYTHON_ENTRYPOINT, + FILE_IO_TASK_COMMAND, + MODEL_ENTITY_TASK_COMMAND, + get_tasks_image, +) from nmp.common.auth import AuthClient, auth_client_context from nmp.common.entities.utils import parse_entity_ref from nmp.common.jobs.constants import DEFAULT_JOB_STORAGE_PATH, PERSISTENT_JOB_STORAGE_PATH_ENVVAR @@ -452,7 +457,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=AUTOMODEL_PYTHON_ENTRYPOINT, - command=["-m", "nmp.automodel.tasks.file_io"], + command=FILE_IO_TASK_COMMAND, ), resources=cpu_resources, ), @@ -474,7 +479,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=AUTOMODEL_PYTHON_ENTRYPOINT, - command=["-m", "nmp.automodel.tasks.file_io"], + command=FILE_IO_TASK_COMMAND, ), resources=cpu_resources, ), @@ -489,7 +494,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=AUTOMODEL_PYTHON_ENTRYPOINT, - command=["-m", "nmp.automodel.tasks.model_entity"], + command=MODEL_ENTITY_TASK_COMMAND, ), resources=cpu_resources, ), diff --git a/services/automodel/src/nmp/automodel/config.py b/services/automodel/src/nmp/automodel/config.py index 2c189ab9c2..1141467a6e 100644 --- a/services/automodel/src/nmp/automodel/config.py +++ b/services/automodel/src/nmp/automodel/config.py @@ -13,7 +13,7 @@ class AutomodelConfig(create_service_config_class("automodel")): # type: ignore image_registry: str | None = Field( default=None, description=( - "Registry host/path prefix for nmp-automodel-tasks and nmp-automodel-training. " + "Registry host/path prefix for nmp-customizer-tasks and nmp-automodel-training. " "Override via NMP_AUTOMODEL_IMAGE_REGISTRY for other environments, defaults to the platform's image registry." ), ) diff --git a/services/automodel/src/nmp/automodel/images.py b/services/automodel/src/nmp/automodel/images.py index 9889745483..b8867e0b50 100644 --- a/services/automodel/src/nmp/automodel/images.py +++ b/services/automodel/src/nmp/automodel/images.py @@ -6,16 +6,32 @@ from __future__ import annotations from nmp.automodel.config import config -from nmp.customization_common.service.images import resolve_qualified_image +from nmp.customization_common.service.images import ( + CUSTOMIZER_PYTHON_ENTRYPOINT, + get_customizer_tasks_image, + resolve_qualified_image, +) BASE_IMAGE_NAME = "nmp-automodel-base" -TASKS_IMAGE_NAME = "nmp-automodel-tasks" TRAINING_IMAGE_NAME = "nmp-automodel-training" -# Must match ENTRYPOINT in docker/automodel/Dockerfile.nmp-automodel-{tasks,training}. -# Job specs must set this explicitly: Docker API create() replaces the image -# entrypoint when the platform passes entrypoint=[]. -AUTOMODEL_PYTHON_ENTRYPOINT = ["/opt/venv/bin/python"] +# Alias for backward compatibility in compiler imports. +AUTOMODEL_PYTHON_ENTRYPOINT = CUSTOMIZER_PYTHON_ENTRYPOINT + +FILE_IO_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.file_io", + "--service-source", + "automodel", + "--service-name", + "customizer", +] +MODEL_ENTITY_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.model_entity", + "--service-name", + "customizer", +] def get_automodel_qualified_image(name: str, override: str | None = None) -> str: @@ -24,8 +40,8 @@ def get_automodel_qualified_image(name: str, override: str | None = None) -> str def get_tasks_image() -> str: - """CPU task steps (file_io, model_entity).""" - return get_automodel_qualified_image(TASKS_IMAGE_NAME, config.tasks_image) + """CPU task steps (file_io, model_entity) — shared ``nmp-customizer-tasks`` image.""" + return get_customizer_tasks_image(backend_override=config.tasks_image, image_registry=config.image_registry) def get_training_image() -> str: diff --git a/services/automodel/src/nmp/automodel/tasks/__main__.py b/services/automodel/src/nmp/automodel/tasks/__main__.py index 6e6e865482..c411f3615d 100644 --- a/services/automodel/src/nmp/automodel/tasks/__main__.py +++ b/services/automodel/src/nmp/automodel/tasks/__main__.py @@ -1,10 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Default entrypoint for the nmp-automodel-tasks image (help / task listing). +"""Default entrypoint for the nmp-customizer-tasks image (help / task listing). Production job steps invoke a specific module directly, e.g. -``python -m nmp.automodel.tasks.file_io``. +``python -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer``. """ from __future__ import annotations @@ -13,21 +13,39 @@ import sys _TASK_MODULES = ( - ("file_io", "nmp.automodel.tasks.file_io", "Download/upload model and dataset files"), - ("model_entity", "nmp.automodel.tasks.model_entity", "Create output model entity"), + ( + "file_io", + "nmp.customization_common.tasks.file_io", + "Download/upload model and dataset files (pass --service-source / --service-name)", + ), + ( + "model_entity", + "nmp.customization_common.tasks.model_entity", + "Create output model entity (pass --service-name)", + ), + ( + "model_spec", + "nmp.core.models.tasks.model_spec", + "Populate model entity spec from checkpoint metadata", + ), + ( + "lora_sidecar", + "nmp.core.models.sidecars.adapters.main", + "LoRA adapter sidecar for NIM/vLLM deployments", + ), ) def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( prog="python -m nmp.automodel.tasks", - description="NeMo Automodel CPU task image. The jobs compiler runs one module per step.", + description="NeMo customization CPU task image. The jobs compiler runs one module per step.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="Examples:\n" " python -m nmp.automodel.tasks --help\n" - " python -m nmp.automodel.tasks.file_io\n" - " python -m nmp.automodel.tasks.model_entity\n\n" - "GPU training uses the nmp-automodel-training image:\n" + " python -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer\n" + " python -m nmp.customization_common.tasks.model_entity --service-name customizer\n\n" + "GPU training uses backend-specific training images:\n" " python -m nmp.automodel.tasks.training\n", ) parser.add_argument( diff --git a/services/automodel/src/nmp/automodel/tasks/docker/README.md b/services/automodel/src/nmp/automodel/tasks/docker/README.md index bf41efc344..d2c8a55e74 100644 --- a/services/automodel/src/nmp/automodel/tasks/docker/README.md +++ b/services/automodel/src/nmp/automodel/tasks/docker/README.md @@ -8,7 +8,7 @@ Scripts for running the file_io task container locally. ```bash cd /path/to/nemo-platform - docker buildx bake -f docker-bake.hcl nmp-automodel-tasks-docker + docker buildx bake -f docker-bake.hcl nmp-customizer-tasks ``` 2. **Have NeMo Platform running** (files service) at `http://localhost:8080` @@ -24,10 +24,10 @@ cd services/automodel/src/nmp/automodel/tasks/docker docker compose up # Run with custom image -FILE_IO_IMAGE=my-registry/nemo-platform-dev/nmp-automodel-tasks:dev docker compose up +FILE_IO_IMAGE=my-registry/nemo-platform-dev/nmp-customizer-tasks:dev docker compose up # Run interactively -docker compose run --rm file-io run task --task nmp.automodel.tasks.file_io +docker compose run --rm file-io -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer ``` ## Configuration @@ -44,7 +44,7 @@ docker compose run --rm file-io run task --task nmp.automodel.tasks.file_io | `NEMO_JOB_TASK` | Task identifier | `file-io-task` | | `NEMO_JOB_WORKSPACE` | Workspace name | `default` | | `LOG_LEVEL` | Logging level | `INFO` | -| `FILE_IO_IMAGE` | Docker image to use | `my-registry/nemo-platform-dev/nmp-automodel-tasks:local` | +| `FILE_IO_IMAGE` | Docker image to use | `my-registry/nemo-platform-dev/nmp-customizer-tasks:local` | ### Config File Format @@ -71,3 +71,12 @@ The `sample_config.json` defines what files to upload/download: - `upload[].dest`: Target FileSet in format `workspace/fileset-name` - `download[].src`: Source FileSet in format `workspace/fileset-name` - `download[].dest`: Path relative to job storage defined by NEMO_JOB_PERSISTENT_JOB_STORAGE_PATH + +## Next Steps + +- **[`nmp-customizer-tasks` image build & runtime](../../../../../../../docker/automodel/README.md)** — bake targets, workspace slice, and smoke commands for the shared CPU image. +- **[Automodel job compiler](../../app/jobs/compiler.py)** — how download / upload / model-entity steps are compiled onto `nmp-customizer-tasks` with `--service-source automodel --service-name customizer`. +- **[Shared customization task runners](../../../../../../../packages/nmp_customization_common/README.md)** — `nmp.customization_common.tasks.file_io` and `model_entity` (used by automodel, unsloth, and rl). +- **[`nmp-automodel` service overview](../../../../../README.md)** — package layout, training image, and plugin integration. +- **[Customizer docs](../../../../../../../docs/customizer/index.mdx)** — published container images and the end-to-end fine-tuning workflow on the platform. + diff --git a/services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml b/services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml index 75b6b8b31b..f094a091c5 100644 --- a/services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml +++ b/services/automodel/src/nmp/automodel/tasks/docker/docker-compose.yaml @@ -5,17 +5,17 @@ # docker compose up # # # Run with custom command -# docker compose run --rm file-io run task --task nmp.automodel.tasks.file_io +# docker compose run --rm file-io -m nmp.customization_common.tasks.file_io --service-source automodel --service-name customizer # # Prerequisites: # - Build the image first (from Platform repo root): -# docker buildx bake -f docker-bake.hcl nmp-automodel-tasks-docker +# docker buildx bake -f docker-bake.hcl nmp-customizer-tasks # - Have NeMo Platform running at http://localhost:8080 # - Create sample_config.json (or use the one provided) services: file-io: - image: ${FILE_IO_IMAGE:-my-registry/nemo-platform-dev/nmp-automodel-tasks:local} + image: ${FILE_IO_IMAGE:-my-registry/nemo-platform-dev/nmp-customizer-tasks:local} container_name: file-io-task # Mount config file and storage directory @@ -43,9 +43,15 @@ services: # Logging LOG_LEVEL: ${LOG_LEVEL:-INFO} - # Default command - run the file_io task - # Note: ENTRYPOINT is already "nemo-platform", so command should NOT include it - command: ["run", "task", "--task", "nmp.automodel.tasks.file_io"] + # Default command - run the shared file_io task for automodel + # Note: ENTRYPOINT is already "/opt/venv/bin/python" + command: + - "-m" + - "nmp.customization_common.tasks.file_io" + - "--service-source" + - "automodel" + - "--service-name" + - "customizer" # For macOS/Windows Docker Desktop - allows reaching host services extra_hosts: diff --git a/services/automodel/src/nmp/automodel/tasks/file_io/__init__.py b/services/automodel/src/nmp/automodel/tasks/file_io/__init__.py deleted file mode 100644 index 8df0acb9ed..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/file_io/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""File I/O task for Automodel customization jobs.""" - -from nmp.automodel.tasks.file_io.run import run - -__all__ = ["run"] diff --git a/services/automodel/src/nmp/automodel/tasks/file_io/__main__.py b/services/automodel/src/nmp/automodel/tasks/file_io/__main__.py deleted file mode 100644 index 68981865e6..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/file_io/__main__.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import sys - -from nmp.automodel.tasks.file_io.run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/automodel/src/nmp/automodel/tasks/file_io/callbacks.py b/services/automodel/src/nmp/automodel/tasks/file_io/callbacks.py deleted file mode 100644 index 0668b5fd97..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/file_io/callbacks.py +++ /dev/null @@ -1,745 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Custom fsspec callbacks for progress reporting during file I/O operations.""" - -import logging -import threading -from abc import abstractmethod -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -from fsspec.callbacks import Callback, TqdmCallback -from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.customization_common.schemas.file_io import DownloadStats, TaskPhase, UploadStats -from nmp.customization_common.tasks.file_io_progress_reporter import ProgressReporter -from nmp.customization_common.tasks.file_io_utils import list_local_files as _list_local_files - -logger = logging.getLogger(__name__) - - -def get_percentage(current: int, total: int) -> int: - """Get percentage of current / total. - - Args: - current: The current value (numerator). - total: The total value (denominator). - - Returns: - Integer percentage from 0-100. - - Raises: - ValueError: If current > total, or if either value is negative. - - """ - if current > total: - raise ValueError( - f"Unexpected value of the current and total values: current={current} cannot be greater than total={total}", - ) - if total < 0: - raise ValueError(f"Unexpected negative value of the total value: total={total}, current={current}") - if current < 0: - raise ValueError(f"Unexpected negative value of the current value: current={current}, total={total}") - - if total == 0: - return 0 - return int((current / total) * 100) - - -@dataclass -class FileInfo: - """A dataclass for file information.""" - - path: str - size: int - - -class TqdmPerFileUploadCallback(Callback): - """A callback that creates a separate tqdm progress bar for each file. - - Unlike TqdmCallback which shows overall progress, this callback creates a new - tqdm progress bar when branched() is called for each file. Each file's progress - bar shows byte-level progress for that individual file. - - Usage: - callback = TqdmPerFileUploadCallback() - filesystem_sdk.put(src, dest, recursive=True, callback=callback) - # Creates a separate progress bar for each file being uploaded - """ - - def __init__(self, src_path: Path, **kwargs: Any): - """Initialize the per-file tqdm callback. - - Args: - **kwargs: Additional arguments passed to the base Callback. - - """ - self.src_path = src_path - super().__init__(**kwargs) - - def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: - """Create a TqdmCallback for this specific file transfer. - - Args: - full_src_path: Source file path. - full_dest_path: Destination file path. - **kwargs: Additional keyword arguments. - - Returns: - A TqdmCallback configured for byte-level progress of this file. - - """ - # Extract just the filename for the progress bar description - if self.src_path.is_file(): - relative_path_upload_dir = self.src_path.name - else: - relative_path_upload_dir = Path(full_src_path).relative_to(self.src_path) - return TqdmCallback( - # https://tqdm.github.io/docs/tqdm - tqdm_kwargs={ - "desc": f"Uploading {relative_path_upload_dir!s}", - # use bytes as the unit - "unit": "B", - # scale the unit to be more readable (e.g. 1024 bytes = 1 KB) - "unit_scale": True, - # divide the unit by 1024 to get the next unit - "unit_divisor": 1024, - # The minimum number of iterations (bytes processed) that must occur before the progress bar refreshes - "miniters": 1, - }, - ) - - -class TqdmPerFileDownloadCallback(Callback): - """A callback that creates a separate tqdm progress bar for each file download. - - Similar to TqdmPerFileUploadCallback but for download operations. Creates a new - tqdm progress bar when branched() is called for each file being downloaded. - - The callback accepts a file_sizes dict to set the total size for each file's - progress bar. This is necessary because the SDK may not receive Content-Length - headers for streaming downloads (e.g., when chunked transfer encoding is used). - - Usage: - # Build file_sizes from listing - files = list_fileset_files(fileset) - file_sizes = {f.path.lstrip("/"): f.size for f in files} - - callback = TqdmPerFileDownloadCallback( - dest_path=dest_dir, - fileset_path="workspace/fileset", - file_sizes=file_sizes, - ) - filesystem_sdk.get(src, dest, recursive=True, callback=callback) - # Creates a separate progress bar for each file being downloaded - """ - - def __init__(self, dest_path: Path, fileset_path: str, file_sizes: dict[str, int] | None = None, **kwargs: Any): - """Initialize the per-file tqdm download callback. - - Args: - dest_path: The local destination directory path. - fileset_path: The fileset path (e.g., "workspace/fileset") used to extract - relative file paths from full source paths. - file_sizes: Optional dict mapping relative file paths to their sizes in bytes. - Used to set the progress bar's total for percentage display. - **kwargs: Additional arguments passed to the base Callback. - - """ - self.dest_path = dest_path - self.fileset_path = fileset_path.rstrip("/") - self.file_sizes = file_sizes or {} - super().__init__(**kwargs) - - def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: - """Create a TqdmCallback for this specific file download. - - Args: - full_src_path: Source file path in the fileset (e.g., "workspace/fileset/dir/file.txt"). - full_dest_path: Destination local file path. - **kwargs: Additional keyword arguments. - - Returns: - A TqdmCallback configured for byte-level progress of this file. - - """ - # Extract relative path for the progress bar description - # full_dest_path is the full local path, we want to show just the filename or relative path - dest_full_path = Path(full_dest_path) - if self.dest_path.is_file(): - relative_path = dest_full_path.name - else: - try: - relative_path = dest_full_path.relative_to(self.dest_path) - except ValueError: - # If can't compute relative path, use filename - relative_path = dest_full_path.name - - # Extract relative file path from full source path to look up size - # full_src_path format: "workspace/fileset/relative/path/to/file.txt" - # We need to extract "relative/path/to/file.txt" - relative_file_path = full_src_path - if full_src_path.startswith(self.fileset_path): - relative_file_path = full_src_path[len(self.fileset_path) :].lstrip("/") - - # Look up file size from pre-computed mapping - file_size = self.file_sizes.get(relative_file_path) - - callback = TqdmCallback( - tqdm_kwargs={ - "desc": f"Downloading {relative_path!s}", - "unit": "B", - "unit_scale": True, - "unit_divisor": 1024, - "miniters": 1, - }, - ) - - # Set size if we know it - this enables percentage display in tqdm - # Must be called via set_size() rather than tqdm_kwargs["total"] because - # the SDK may also call set_size() from Content-Length header - if file_size is not None: - callback.set_size(file_size) - - return callback - - -class BaseProgressCallback(Callback): - """Base class for file upload/download progress callbacks. - - This abstract base class provides common functionality for tracking file transfer - progress and reporting to the Jobs service. Subclasses implement operation-specific - behavior (upload vs download). - - Thread Safety: - This callback uses threading.Lock for synchronization. FilesetFileSystem is - async-first and transfers files concurrently. The lock protects against - concurrent access when multiple files complete simultaneously. - - Attributes: - progress_reporter: The progress reporter for sending updates to Jobs service. - fileset_name: The name of the fileset (workspace/name format). - total_files: Total number of files to transfer. - total_size: Total size of all files in bytes. - stats: Mutable stats object to track progress (UploadStats or DownloadStats). - _lock: Threading lock for thread-safe stats updates. - - """ - - progress_reporter: ProgressReporter - fileset_name: str - total_files: int - total_size: int - stats: UploadStats | DownloadStats - _lock: threading.Lock - - def __init__( - self, - progress_reporter: ProgressReporter, - fileset_name: str, - total_files: int, - total_size: int, - stats: UploadStats | DownloadStats, - **kwargs: Any, - ): - """Initialize the progress callback. - - Args: - progress_reporter: The progress reporter for sending updates to Jobs service. - fileset_name: The name of the fileset (workspace/name format). - total_files: Total number of files to transfer. - total_size: Total size of all files in bytes. - stats: Mutable stats object to track progress. - **kwargs: Additional arguments passed to the base Callback. - - """ - super().__init__(**kwargs) - self.progress_reporter = progress_reporter - self.fileset_name = str(fileset_name) - self.total_files = total_files - self.total_size = total_size - self.stats = stats - self._lock = threading.Lock() - - @staticmethod - def list_local_files(src_path: Path) -> list[FileInfo]: - """List all files from a local path (see shared ``list_local_files``).""" - return [FileInfo(path=f.path, size=f.size) for f in _list_local_files(src_path)] - - @abstractmethod - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "BaseSingleFileCallback": - """Create a child callback for a single file transfer. - - Args: - source_path: Source file path. - dest_path: Destination file path. - **kwargs: Additional keyword arguments. - - Returns: - A BaseSingleFileCallback subclass for tracking this file's transfer. - - """ - ... - - -class BaseSingleFileCallback(Callback): - """Base class for single file upload/download callbacks. - - This abstract base class provides common functionality for tracking individual - file transfers within a batch operation. Subclasses implement operation-specific - behavior via the template method pattern. - - The close() method uses the template method pattern, calling abstract methods - that subclasses override to provide operation-specific behavior: - - _get_phase(): Returns the TaskPhase for this operation - - _get_file_display_path(): Returns the path to display for logging - - _update_stats(): Updates the parent's stats for this operation - - _build_status_details(): Builds the status_details dict for progress reporting - """ - - parent: BaseProgressCallback - source_path: str - dest_path: str - _completed: bool - - def __init__( - self, - parent: BaseProgressCallback, - source_path: str, - dest_path: str, - **kwargs: Any, - ): - """Initialize the single file callback. - - Args: - parent: The parent progress callback. - source_path: Path to the source file. - dest_path: Destination path for the file. - **kwargs: Additional arguments passed to the base Callback. - - """ - super().__init__(**kwargs) - self.parent = parent - self.source_path = source_path - self.dest_path = dest_path - self._completed = False - - @abstractmethod - def _get_phase(self) -> str: - """Return the TaskPhase for this operation.""" - ... - - @abstractmethod - def _get_file_display_path(self) -> str: - """Return the path to use for display/logging.""" - ... - - @abstractmethod - def _update_stats(self) -> None: - """Update the parent's stats for this operation (called within lock).""" - ... - - @abstractmethod - def _get_files_count(self) -> int: - """Return the current files count from stats (called within lock).""" - ... - - @abstractmethod - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - """Build the status_details dict for progress reporting. - - Args: - files_count: Number of files transferred so far. - total_bytes: Total bytes transferred so far. - current_file: Name of the current file for display. - - Returns: - Dictionary with status details for the progress report. - - """ - ... - - def close(self) -> None: - """Called when the file transfer completes. - - Updates the parent's statistics and reports progress to the Jobs service. - Thread-safe: uses parent's lock to protect stats updates. - """ - if self._completed: - return - - self._completed = True - parent = self.parent - - # Extract the filename for logging/display - current_file = self._get_file_display_path() - - # Thread-safe stats update - with parent._lock: - # Update stats (operation-specific) - self._update_stats() - - # Capture current values while holding the lock - files_count = self._get_files_count() - total_bytes = parent.stats.total_bytes - - logger.debug(f"File transferred: {current_file} ({files_count}/{parent.total_files})") - - # Report progress to Jobs service (outside lock to avoid holding it during I/O) - parent.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details=self._build_status_details(files_count, total_bytes, current_file), - ) - - def __enter__(self) -> "BaseSingleFileCallback": - return self - - def __exit__(self, *exc_args: object) -> None: - self.close() - - -class FileUploadProgressCallback(BaseProgressCallback): - """Callback for tracking file upload progress and reporting to the Jobs service. - - This callback integrates with fsspec's callback mechanism to report progress - after each file is uploaded. It uses the branched callback pattern where: - - The parent callback tracks overall upload statistics - - Child callbacks are created for each file via `branched()` - - When a child callback closes, it signals file completion to the parent - - Usage: - callback = FileUploadProgressCallback( - progress_reporter=reporter, - src_path=src_path, - fileset_name="workspace/fileset", - stats=upload_stats, - ) - filesystem_sdk.put(src, dest, recursive=True, callback=callback) - """ - - stats: UploadStats - - def __init__( - self, - progress_reporter: ProgressReporter, - src_path: Path, - fileset_name: str, - stats: UploadStats, - **kwargs: Any, - ): - """Initialize the upload progress callback. - - Args: - progress_reporter: The progress reporter for sending updates to Jobs service. - src_path: The source path (file or directory) to upload. - fileset_name: The name of the target fileset (workspace/name format). - stats: Mutable UploadStats object to track progress. - **kwargs: Additional arguments passed to the base Callback. - - """ - # List files to get stats before upload - files = self.list_local_files(src_path) - - if not files: - logger.warning(f"Source path {src_path} contains no files") - - total_files = len(files) - total_size = sum(f.size for f in files) - - # Initialize base class with computed values - super().__init__( - progress_reporter=progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - **kwargs, - ) - - logger.info(f"Uploading {total_files} files ({total_size} bytes) to {self.fileset_name}") - - # Report initial progress - progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "fileset": self.fileset_name, - "total_files": total_files, - "total_size": total_size, - "uploaded_files": 0, - "uploaded_bytes": 0, - }, - ) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileUploadCallback": - """Create a child callback for a single file upload. - - This method is called by fsspec when starting a file transfer within - a recursive put operation. It returns a child callback that tracks - the individual file's progress and reports completion to the parent. - - Args: - source_path: Source file path. - path_2: Destination file path. - **kwargs: Additional keyword arguments. - - Returns: - A SingleFileUploadCallback for tracking this file's upload. - - """ - return SingleFileUploadCallback( - parent=self, - source_path=source_path, - dest_path=dest_path, - **kwargs, - ) - - -class SingleFileUploadCallback(BaseSingleFileCallback): - """Callback for tracking a single file upload within a batch operation. - - This child callback is created by FileUploadProgressCallback.branched() - for each file being uploaded. When the upload completes and this callback - is closed, it notifies the parent to update overall progress. - """ - - parent: FileUploadProgressCallback - - def _get_phase(self) -> str: - """Return the TaskPhase for upload operations.""" - return TaskPhase.UPLOADING - - def _get_file_display_path(self) -> str: - """Return the destination filename for display.""" - return self.dest_path.split("/")[-1] if "/" in self.dest_path else self.dest_path - - def _update_stats(self) -> None: - """Update the parent's upload stats.""" - self.parent.stats.files_uploaded += 1 - if self.size is not None: - self.parent.stats.total_bytes += self.size - - def _get_files_count(self) -> int: - """Return the current uploaded files count.""" - return self.parent.stats.files_uploaded - - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - """Build the status_details dict for upload progress reporting.""" - return { - "phase": TaskPhase.UPLOADING, - "fileset": self.parent.fileset_name, - "total_files": self.parent.total_files, - "total_size": self.parent.total_size, - "uploaded_files": files_count, - "uploaded_bytes": total_bytes, - "current_file": current_file, - "progress_pct": get_percentage(files_count, self.parent.total_files), - } - - -class FileDownloadProgressCallback(BaseProgressCallback): - """Callback for tracking file download progress and reporting to the Jobs service. - - Similar to FileUploadProgressCallback but for download operations. - - Usage: - callback = FileDownloadProgressCallback( - progress_reporter=reporter, - fileset_name="workspace/fileset", - total_files=10, - total_size=1024000, - stats=download_stats, - ) - filesystem_sdk.get(src, dest, recursive=True, callback=callback) - """ - - stats: DownloadStats - - def __init__( - self, - progress_reporter: ProgressReporter, - fileset_name: str, - total_files: int, - total_size: int, - stats: DownloadStats, - **kwargs: Any, - ): - """Initialize the download progress callback. - - Args: - progress_reporter: The progress reporter for sending updates to Jobs service. - fileset_name: The name of the source fileset (workspace/name format). - total_files: Total number of files to download. - total_size: Total size of all files in bytes. - stats: Mutable DownloadStats object to track progress. - **kwargs: Additional arguments passed to the base Callback. - - """ - super().__init__( - progress_reporter=progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - **kwargs, - ) - - logger.info(f"Downloading {total_files} files ({total_size} bytes) from {self.fileset_name}") - - # Report initial progress - progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "fileset": self.fileset_name, - "total_files": total_files, - "total_size": total_size, - "downloaded_files": 0, - "downloaded_bytes": 0, - }, - ) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileDownloadCallback": - """Create a child callback for a single file download. - - Args: - source_path: Source file path in the fileset. - dest_path: Destination local file path. - **kwargs: Additional keyword arguments. - - Returns: - A SingleFileDownloadCallback for tracking this file's download. - - """ - return SingleFileDownloadCallback( - parent=self, - source_path=source_path, - dest_path=dest_path, - **kwargs, - ) - - -class SingleFileDownloadCallback(BaseSingleFileCallback): - """Callback for tracking a single file download within a batch operation. - - This child callback is created by FileDownloadProgressCallback.branched() - for each file being downloaded. When the download completes and this callback - is closed, it notifies the parent to update overall progress. - """ - - parent: FileDownloadProgressCallback - - def _get_phase(self) -> str: - """Return the TaskPhase for download operations.""" - return TaskPhase.DOWNLOADING - - def _get_file_display_path(self) -> str: - """Return the source filename for display.""" - return self.source_path.split("/")[-1] if "/" in self.source_path else self.source_path - - def _update_stats(self) -> None: - """Update the parent's download stats.""" - self.parent.stats.files_downloaded += 1 - if self.size is not None: - self.parent.stats.total_bytes += self.size - - def _get_files_count(self) -> int: - """Return the current downloaded files count.""" - return self.parent.stats.files_downloaded - - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - """Build the status_details dict for download progress reporting.""" - return { - "phase": TaskPhase.DOWNLOADING, - "fileset": self.parent.fileset_name, - "total_files": self.parent.total_files, - "total_size": self.parent.total_size, - "downloaded_files": files_count, - "downloaded_bytes": total_bytes, - "current_file": current_file, - "progress_pct": get_percentage(files_count, self.parent.total_files), - } - - -class CompositeCallback(Callback): - """A callback that delegates to multiple child callbacks. - - This allows combining multiple callbacks (e.g., TqdmCallback for console progress - and FileUploadProgressCallback for Jobs service reporting) into a single callback - that can be passed to fsspec operations. - - All callback methods are forwarded to each child callback in order. - - Usage: - tqdm_cb = TqdmCallback(tqdm_kwargs={"desc": "Uploading"}) - progress_cb = FileUploadProgressCallback(...) - composite = CompositeCallback(tqdm_cb, progress_cb) - filesystem_sdk.put(src, dest, recursive=True, callback=composite) - """ - - def __init__(self, *callbacks: Callback, **kwargs: Any): - """Initialize with multiple callbacks. - - Args: - *callbacks: Variable number of Callback instances to delegate to. - **kwargs: Additional arguments passed to the base Callback. - - """ - super().__init__(**kwargs) - self.callbacks = list(callbacks) - - def set_size(self, size: int) -> None: - """Set size on all child callbacks.""" - self.size = size - for cb in self.callbacks: - cb.set_size(size) - - def absolute_update(self, value: int) -> None: - """Update absolute value on all child callbacks.""" - self.value = value - for cb in self.callbacks: - cb.absolute_update(value) - - def relative_update(self, inc: int = 1) -> None: - """Update relative value on all child callbacks.""" - self.value += inc - for cb in self.callbacks: - cb.relative_update(inc) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "CompositeCallback": - """Create a composite child callback from all child callbacks' branched results. - - Each child callback's branched() method is called, and the results are - wrapped in a new CompositeCallback. - - Args: - source_path: Source path. - dest_path: Destination path. - **kwargs: Additional keyword arguments. - - Returns: - A new CompositeCallback wrapping all child callbacks' branched results. - - """ - child_callbacks = [cb.branched(source_path, dest_path, **kwargs) for cb in self.callbacks] - return CompositeCallback(*child_callbacks) - - def call(self, hook_name: str | None = None, **kwargs: Any) -> None: - """Call hooks on all child callbacks.""" - for cb in self.callbacks: - cb.call(hook_name, **kwargs) - - def close(self) -> None: - """Close all child callbacks.""" - for cb in self.callbacks: - cb.close() - - def __enter__(self) -> "CompositeCallback": - for cb in self.callbacks: - cb.__enter__() - return self - - def __exit__(self, *exc_args: object) -> None: - for cb in self.callbacks: - cb.__exit__(*exc_args) diff --git a/services/automodel/src/nmp/automodel/tasks/file_io/run.py b/services/automodel/src/nmp/automodel/tasks/file_io/run.py deleted file mode 100644 index a91b27e8e8..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/file_io/run.py +++ /dev/null @@ -1,580 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""File I/O task entry point. - -Handles file operations between NeMo Platform Files Service and the job's shared PVC. - -The task reads configuration and performs: -- Downloads: If config.download is non-empty, download files from FileSets to local paths -- Uploads: If config.upload is non-empty, upload files from local paths to FileSets - -Usage: - export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.automodel.tasks.file_io -""" - -import logging -from pathlib import Path - -import httpx - -# https://docs.nvidia.com/nemo/microservices/latest/pysdk/index.html#handling-errors -from nemo_platform import ( - APIConnectionError, - APITimeoutError, - InternalServerError, - NeMoPlatform, - NotFoundError, -) -from nemo_platform.types.files.fileset_file import FilesetFile -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import ( - ConflictError, -) -from nemo_platform_plugin.client.errors import ( - InternalServerError as ClientInternalServerError, -) -from nemo_platform_plugin.client.types import RetryPolicy -from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest -from nmp.automodel.app.constants import SERVICE_NAME -from nmp.automodel.tasks.file_io.callbacks import ( - CompositeCallback, - FileDownloadProgressCallback, - FileUploadProgressCallback, - TqdmPerFileDownloadCallback, - TqdmPerFileUploadCallback, -) -from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.common.sdk_factory import get_task_sdk -from nmp.customization_common.schemas.file_io import ( - DownloadItem, - DownloadStats, - FileDownloadError, - FileSetRef, - FileUploadError, - PathTraversalError, - TaskPhase, - UploadItem, - UploadStats, -) -from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.tasks.file_io_progress_reporter import JobsServiceProgressReporter, ProgressReporter -from nmp.customization_common.tasks.file_io_utils import ( - filesystem_sdk_error_handler, - get_config, - sdk_error_handler, - validate_safe_path, - validate_storage_path, -) -from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential - -logger = logging.getLogger(__name__) - -CREATE_FILESET_TIMEOUT = 10.0 -LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) - -# Timeout configurations for FilesetFileSystem operations. -# These are passed via sdk.with_options(timeout=...) and control the httpx client. -# httpx.Timeout(read=...) is the max wait for a single chunk (16MB by default), NOT total transfer time. -# nemo-platform/src/nemo_platform/filesets/filesystem/filesystem.py > blocksize = 16 * 1024 * 1024 # 16MB -# It's a socket-level timeout. Each individual socket read has its own timeout window. -# SDK defaults httpx.Timeout(timeout=60, connect=5.0) nemo-platform/src/nemo_platform/_constants.py -DOWNLOAD_TIMEOUT = httpx.Timeout(30.0, read=5 * 60) # 30s connect/pool, 5min per-chunk read -UPLOAD_TIMEOUT = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) # 30s connect/pool, 10min write, 5min read - -# Retry configuration -MAX_RETRIES = 3 -INITIAL_BACKOFF_SECONDS = 1.0 -MAX_BACKOFF_SECONDS = 30.0 - -# Transient exceptions that should trigger retries for filesystem operations. -# FilesetFileSystem uses httpx under the hood, so we retry on httpx transient errors -# in addition to SDK-level transient errors. -TRANSIENT_FILESYSTEM_EXCEPTIONS = ( - httpx.TimeoutException, - httpx.ConnectError, - httpx.ReadTimeout, - # Connection dropped mid-transfer (CDN/proxy closed the socket before the - # full body arrived). Common on large multi-GB model shards; safe to retry. - httpx.RemoteProtocolError, - httpx.ReadError, -) - - -class FileIORunner: - def __init__( - self, - sdk: NeMoPlatform, - progress_reporter: ProgressReporter, - job_ctx: NMPJobContext, - ): - self.sdk = sdk - self.progress_reporter = progress_reporter - self.job_ctx = job_ctx - - def list_fileset_files( - self, - fileset: FileSetRef, - ) -> list[FilesetFile]: - """List files in a FileSet. - - Returns list of file info dicts with 'path' and 'size' keys. - """ - try: - with sdk_error_handler(FileDownloadError, f"list files in fileset {fileset}", passthrough=(NotFoundError,)): - response = self.sdk.with_options(timeout=LIST_FILES_TIMEOUT).files.list( - fileset=fileset.name, - workspace=fileset.workspace, - ) - logger.info(f"Found {len(response.data)} files in FileSet {fileset!s}") - return response.data - except NotFoundError as e: - raise FileDownloadError( - f"FileSet {fileset!s} not found. Please ensure the FileSet exists and contains the expected files.", - ) from e - - def download_fileset( - self, - fileset: FileSetRef, - dest_dir: Path, - ) -> DownloadStats: - """Download all files from a FileSet to a destination directory. - - Uses FilesetFileSystem.get() with recursive=True for efficient batch downloads. - Progress is tracked via two callbacks combined in a CompositeCallback: - - TqdmPerFileDownloadCallback: Creates a separate console progress bar per file (shows bytes) - - FileDownloadProgressCallback: Reports progress to Jobs service after each file - - Args: - fileset: The source FileSet reference. - dest_dir: The destination directory path. - - Returns: - DownloadStats with files_downloaded, total_bytes, and failed_files counts. - - Raises: - FileDownloadError: If the download fails. - - """ - stats = DownloadStats() - fileset_name = str(fileset) - - # List files in the fileset to get total count and size - files = self.list_fileset_files(fileset) - - if not files: - logger.warning(f"FileSet {fileset_name} contains no files") - return stats - - total_files = len(files) - total_size = sum(f.size for f in files) - - # Ensure destination directory exists - dest_dir.mkdir(parents=True, exist_ok=True) - - # Build file sizes mapping for progress bar display - # Maps relative file paths to their sizes in bytes - file_sizes = {f.path.lstrip("/"): f.size for f in files} - - # Create callbacks: - # 1. TqdmPerFileDownloadCallback for console progress - creates a separate progress bar per file - tqdm_callback = TqdmPerFileDownloadCallback( - dest_path=dest_dir, - fileset_path=fileset_name, - file_sizes=file_sizes, - ) - - # 2. FileDownloadProgressCallback for Jobs service reporting - jobs_callback = FileDownloadProgressCallback( - progress_reporter=self.progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - ) - - # Combine both callbacks into a composite that delegates to both - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - - with filesystem_sdk_error_handler( - FileDownloadError, - f"download from '{fileset_name}' to '{dest_dir}'", - ): - self._download_with_retry( - fileset_name=fileset.name, - fileset_workspace=fileset.workspace, - dest_dir=str(dest_dir), - callback=composite_callback, - ) - - logger.info(f"Download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") - return stats - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), - reraise=True, - before_sleep=before_sleep_log(logger, logging.WARNING), - ) - def _download_with_retry( - self, - fileset_name: str, - fileset_workspace: str | None, - dest_dir: str, - callback: CompositeCallback, - ) -> None: - """Internal method with retry logic for downloading from FilesetFileSystem.""" - self.sdk.with_options(timeout=DOWNLOAD_TIMEOUT).files.download( - fileset=fileset_name, - workspace=fileset_workspace, - local_path=dest_dir, - callback=callback, - ) - - def upload_fileset( - self, - fileset: FileSetRef, - src_path: Path, - ) -> UploadStats: - """Upload all files from a source path (file or directory) to a FileSet. - - Uses FilesetFileSystem.put() with recursive=True for efficient batch uploads. - Progress is tracked via two callbacks combined in a CompositeCallback: - - TqdmPerFileCallback: Creates a separate console progress bar per file (shows bytes) - - FileUploadProgressCallback: Reports progress to Jobs service after each file - - Args: - fileset: The target FileSet reference. - src_path: The source path, can be a single file or a directory. - progress_reporter: Progress reporter for status updates. - - Returns: - UploadStats with files_uploaded, total_bytes, and failed_files counts. - - Raises: - FileUploadError: If the upload fails. - - """ - stats = UploadStats() - fileset_name = str(fileset) - - # Create callbacks: - # 1. TqdmPerFileCallback for console progress - creates a separate progress bar per file - tqdm_callback = TqdmPerFileUploadCallback(src_path=src_path) - - # 2. FileUploadProgressCallback for Jobs service reporting - jobs_callback = FileUploadProgressCallback( - progress_reporter=self.progress_reporter, - src_path=src_path, - fileset_name=fileset_name, - stats=stats, - ) - - # Combine both callbacks into a composite that delegates to both - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - - # Build local and remote paths for upload - # remote_path is relative within the fileset (e.g., "" for root, "filename" for single file) - if src_path.is_dir(): - # Add trailing slash to source to copy directory CONTENTS (not the directory itself) - # This follows rsync/scp convention: "dir/" copies contents, "dir" copies the directory - local_path = f"{src_path}/" - remote_path = "" # Upload to fileset root - else: - # Single file: upload to fileset root with same filename - local_path = str(src_path) - remote_path = src_path.name - - with filesystem_sdk_error_handler( - FileUploadError, - f"upload from '{src_path}' to '{fileset_name}'", - ): - self._upload_with_retry( - local_path=local_path, - remote_path=remote_path, - fileset_name=fileset.name, - fileset_workspace=fileset.workspace, - callback=composite_callback, - ) - - logger.info(f"Upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") - return stats - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), - reraise=True, - before_sleep=before_sleep_log(logger, logging.WARNING, exc_info=True), - ) - def _upload_with_retry( - self, - local_path: str, - remote_path: str, - fileset_name: str, - fileset_workspace: str | None, - callback: CompositeCallback, - ) -> None: - """Internal method with retry logic for uploading to FilesetFileSystem.""" - self.sdk.with_options(timeout=UPLOAD_TIMEOUT).files.upload( - local_path=local_path, - remote_path=remote_path, - fileset=fileset_name, - workspace=fileset_workspace, - callback=callback, - ) - - def run_download(self, downloads: list[DownloadItem]) -> None: - """Execute download operations. - - Downloads files from FileSets to job storage based on downloads list. - """ - if not downloads: - logger.info("No downloads configured, skipping download operation") - return - - storage_path = validate_storage_path(self.job_ctx.storage_path) - - logger.info(f"Starting download operation: {len(downloads)} fileset(s) to download") - - # Report task started - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "total_filesets": len(downloads), - "completed_filesets": 0, - }, - ) - - total_stats = DownloadStats() - - for idx, item in enumerate(downloads): - fileset = item.src - # Validate destination path to prevent path traversal attacks - dest_dir = validate_safe_path(storage_path, item.dest) - - logger.info(f"[{idx + 1}/{len(downloads)}] Downloading from {fileset!s} to {dest_dir}") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "total_filesets": len(downloads), - "completed_filesets": idx, - "current_fileset": f"{fileset!s}", - }, - ) - - stats = self.download_fileset( - fileset, - dest_dir, - ) - total_stats.files_downloaded += stats.files_downloaded - total_stats.total_bytes += stats.total_bytes - - logger.info(f"FileSet download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") - - logger.info( - f"All downloads complete: {total_stats.files_downloaded} files, {total_stats.total_bytes} bytes total", - ) - - def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> None: - """Create a FileSet. Skip if it already exists. - - Uses retry logic for transient errors and converts SDK exceptions to FileUploadError. - """ - # sdk_error_handler wraps the retry to convert exceptions after all retries exhaust - with sdk_error_handler(FileUploadError, f"create fileset {fileset}", passthrough=(ConflictError,)): - self._create_fileset_with_retry(fileset, metadata) - - # we don't use sdk retry because it would retry on ConflictError which is expected and would be wasteful - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type( - ( - InternalServerError, - APITimeoutError, - APIConnectionError, - ClientInternalServerError, - httpx.TimeoutException, - httpx.ConnectError, - ) - ), - reraise=True, - ) - def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: - """Internal method with retry logic for creating a FileSet.""" - files = client_from_platform(self.sdk, FilesClient).with_options( - timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) - ) - try: - body_kwargs: dict = { - "name": fileset.name, - "custom_fields": {"service_source": "automodel"}, - } - if metadata is not None: - body_kwargs["metadata"] = metadata - result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() - logger.info(f"Created FileSet: {result.workspace}/{result.name}") - except ConflictError: - workspace = fileset.workspace or self.job_ctx.workspace - if metadata is not None: - try: - files.update_fileset( - workspace=workspace, - name=fileset.name, - body=UpdateFilesetRequest(metadata=metadata), - ) - logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") - except Exception as e: - logger.warning( - f"Could not patch metadata on existing fileset {workspace}/{fileset.name}: {e}. " - "Upload will continue; model-spec may lack tool_calling/chat_template from source." - ) - - def run_upload(self, uploads: list[UploadItem]) -> None: - """Execute upload operations. - - Uploads files from job storage to FileSets based on uploads list. - - Args: - uploads: List of upload items to process. - """ - if not uploads: - logger.info("No uploads configured, skipping upload operation") - return - - storage_path = validate_storage_path(self.job_ctx.storage_path) - - logger.info(f"Starting upload operation: {len(uploads)} fileset(s) to upload") - - # Report task started - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "total_filesets": len(uploads), - "completed_filesets": 0, - }, - ) - - total_stats = UploadStats() - - for idx, item in enumerate(uploads): - if item.dest.workspace is None: - item.dest.workspace = self.job_ctx.workspace - fileset = item.dest - # Validate source path to prevent path traversal attacks - src_path = validate_safe_path(storage_path, item.src) - if not src_path.exists(): - raise FileUploadError(f"Source path does not exist: {src_path}. Ensure the source path exists.") - if not src_path.is_dir() and not src_path.is_file(): - raise FileUploadError( - f"Source path is not a file or directory: {src_path}. Ensure the source path is a file or directory.", - ) - - logger.info(f"[{idx + 1}/{len(uploads)}] Uploading from {src_path} to {fileset!s}") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "total_filesets": len(uploads), - "completed_filesets": idx, - "current_fileset": str(fileset), - }, - ) - - self.create_fileset(fileset, metadata=item.metadata) - - stats = self.upload_fileset( - fileset, - src_path, - ) - total_stats.files_uploaded += stats.files_uploaded - total_stats.total_bytes += stats.total_bytes - - logger.info(f"FileSet upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") - - logger.info(f"All uploads complete: {total_stats.files_uploaded} files, {total_stats.total_bytes} bytes total") - - -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the file I/O task. - - Processes downloads and uploads based on the configuration. - - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - - Returns: - Exit code (0 for success, non-zero for failure). - - """ - job_ctx = job_ctx or NMPJobContext.from_env() - validate_storage_path(job_ctx.storage_path) - - sdk_owned = sdk is None - progress_reporter: ProgressReporter | None = None - try: - sdk = sdk or get_task_sdk(SERVICE_NAME) - # Initialize progress reporter (no-op if Jobs URL not configured) - progress_reporter = JobsServiceProgressReporter.create_progress_reporter(sdk, job_ctx) - runner = FileIORunner(sdk=sdk, progress_reporter=progress_reporter, job_ctx=job_ctx) - - config = get_config(job_ctx.config_path) - - logger.info(f"Starting file I/O task with job context: {job_ctx}") - logger.info(f"Config: {config.model_dump_json(indent=2)}") - logger.info(f"NeMo Platform service URL: {sdk.base_url}") - - # Execute uploads if configured - runner.run_upload(config.upload) - - # Execute downloads if configured - runner.run_download(config.download) - - # Report overall completion - progress_reporter.update_progress( - status=PlatformJobStatus.COMPLETED, - status_details={"phase": TaskPhase.COMPLETED, "message": "File I/O task completed successfully"}, - ) - - return 0 - except PathTraversalError as e: - logger.error(f"Security error - path traversal detected: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - except (FileDownloadError, FileUploadError) as e: - logger.exception(f"File operation failed: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - - except Exception as e: - logger.exception(f"File I/O task failed: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - - finally: - if sdk_owned and sdk is not None: - sdk.close() diff --git a/services/automodel/src/nmp/automodel/tasks/model_entity/__init__.py b/services/automodel/src/nmp/automodel/tasks/model_entity/__init__.py deleted file mode 100644 index 49784f0f55..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/model_entity/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Model entity task for creating model entities after customization.""" - -from nmp.automodel.tasks.model_entity.run import run - -__all__ = ["run"] diff --git a/services/automodel/src/nmp/automodel/tasks/model_entity/__main__.py b/services/automodel/src/nmp/automodel/tasks/model_entity/__main__.py deleted file mode 100644 index 90a4ffe62f..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/model_entity/__main__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Entry point for model_entity task. - -Usage: - python -m nmp.automodel.tasks.model_entity -""" - -import sys - -from .run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/automodel/src/nmp/automodel/tasks/model_entity/run.py b/services/automodel/src/nmp/automodel/tasks/model_entity/run.py deleted file mode 100644 index 246992ff88..0000000000 --- a/services/automodel/src/nmp/automodel/tasks/model_entity/run.py +++ /dev/null @@ -1,447 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Model entity task entry point. - -Handles creating model entities in the Models service after customization completes. - -The task reads configuration and creates a Model Entity that references the -uploaded model artifacts in the Files service. - -Usage: - export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.automodel.tasks.model_entity -""" - -import json -import logging -import re -import time -from pathlib import Path - -from nemo_platform import ( - APIConnectionError, - APITimeoutError, - ConflictError, - InternalServerError, - NeMoPlatform, - NotFoundError, -) -from nemo_platform.types.inference import ( - ContainerExecutorConfigParam, - ModelDeploymentConfig, - ModelDeploymentConfigFilterParam, - ModelDeploymentConfigModelSpecParam, - ModelDeploymentFilterParam, -) -from nemo_platform.types.models import LoraParam, ModelEntity -from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.files.client import FilesClient -from nmp.automodel.app.constants import SERVICE_NAME -from nmp.automodel.entities.values import FinetuningType -from nmp.common.sdk_factory import get_task_sdk -from nmp.customization_common.schemas.model_entity import ( - DeploymentParameters, - ModelEntityCreationError, - ModelEntityTaskConfig, -) -from nmp.customization_common.service.context import NMPJobContext -from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential - -logger = logging.getLogger(__name__) - -# Retry configuration -MAX_RETRIES = 3 -INITIAL_BACKOFF_SECONDS = 1.0 -MAX_BACKOFF_SECONDS = 30.0 - -ACTIVE_DEPLOYMENT_STATUSES = frozenset({"CREATED", "PENDING", "READY"}) - -SPEC_POLL_INTERVAL_SECONDS = 10 -SPEC_POLL_TIMEOUT_SECONDS = 600 - - -def get_config(config_path: Path) -> ModelEntityTaskConfig: - """Get typed task configuration from a config file. - - Loads the JSON config file and validates it against the ModelEntityTaskConfig schema. - - Args: - config_path: Path to the JSON configuration file. - - Returns: - Validated ModelEntityTaskConfig. - """ - with open(config_path) as f: - data = json.load(f) - return ModelEntityTaskConfig.model_validate(data) - - -def sanitize_name(prefix: str, name: str) -> str: - """Sanitize model_name: keep only allowed chars, replace invalid with hyphen, avoid consecutive/trailing hyphens - - Must be compatible with - {'pattern': '^[a-z](?!.*--)[a-z0-9\\-@.+_]{1,62}(? ModelEntity: - """Poll until the model_spec task has populated the model's spec. - - The spec must be populated before creating a deployment because the - inference service relies on ``spec.family`` and ``spec.base_num_parameters`` - to select the correct NIM profile. - - Raises: - ModelEntityCreationError: If the spec is not populated within the timeout. - """ - logger.info(f"Waiting for model_spec to populate spec on {workspace}/{name}") - start = time.monotonic() - - while time.monotonic() - start < SPEC_POLL_TIMEOUT_SECONDS: - try: - target = self.sdk.models.retrieve(name=name, workspace=workspace) - if target.spec: - logger.info(f"Spec populated on {workspace}/{name}") - return target - except (APIConnectionError, APITimeoutError, InternalServerError) as e: - logger.warning(f"Transient error polling spec for {workspace}/{name}: {e}") - time.sleep(SPEC_POLL_INTERVAL_SECONDS) - - raise ModelEntityCreationError( - f"Timed out waiting for model spec on {workspace}/{name} " - f"after {SPEC_POLL_TIMEOUT_SECONDS}s. The platform could not auto-detect the " - f"model's specifications. Verify the model checkpoint is valid and in a supported format." - ) - - def get_model_entity(self, model_entity: str, fileset_workspace: str) -> ModelEntity: - parts = model_entity.split("/") - if len(parts) == 1: - me_workspace = fileset_workspace - me_name = parts[0] - else: - me_workspace = parts[0] - me_name = parts[1] - - try: - me: ModelEntity = self.sdk.models.retrieve(name=me_name, workspace=me_workspace) - except NotFoundError as e: - raise ModelEntityCreationError(f"Model entity {me_workspace}/{me_name} not found") from e - - return me - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), - reraise=True, - ) - def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, ModelEntity]: - """Create a model entity in the Models service. - - Args: - config: Configuration for the model entity to create. - - Returns: - Tuple of (result dict, deploy target). For LoRA the deploy target is the - base model entity; for SFT it is the newly created output model entity. - - Raises: - ModelEntityCreationError: If creation fails. - """ - workspace = self.job_ctx.workspace - logger.info(f"Creating model entity: {workspace}/{config.name}") - - fileset_workspace = config.fileset.workspace or workspace - fileset_ref = f"{fileset_workspace}/{config.fileset.name}" - - logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") - try: - client_from_platform(self.sdk, FilesClient).get_fileset( - workspace=fileset_workspace, name=config.fileset.name - ) - logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") - except Exception as e: - logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") - raise ModelEntityCreationError( - f"Cannot create model entity: fileset '{fileset_workspace}/{config.fileset.name}' does not exist or is not accessible" - ) from e - - base_me: ModelEntity = self.get_model_entity(config.model_entity, fileset_workspace) - - if config.peft is not None and config.peft.type == FinetuningType.LORA: - try: - output_me = self.sdk.models.adapters.create( - model_name=base_me.name, - workspace=base_me.workspace, - name=config.name, - description=config.description, - fileset=fileset_ref, - finetuning_type=config.peft.type.value, - lora_config=LoraParam( - alpha=config.peft.alpha, - rank=config.peft.rank, - ), - enabled=True, - ) - return output_me.model_dump(), base_me - except ConflictError: - logger.warning( - f"adapter {base_me.workspace}/{config.name} already exists for model {base_me.workspace}/{base_me.name}, updating with new fileset" - ) - try: - output_me = self.sdk.models.adapters.update( - adapter=config.name, - model_name=base_me.name, - workspace=base_me.workspace, - fileset=fileset_ref, - description=config.description, - enabled=True, - ) - logger.info( - f"Successfully updated adapter: {base_me.workspace}/{config.name} for base model {base_me.workspace}/{base_me.name}" - ) - return output_me.model_dump(), base_me - except (InternalServerError, APITimeoutError, APIConnectionError): - raise - except Exception as update_error: - logger.exception( - f"Failed to update existing adapter, {base_me.workspace}/{config.name}: {update_error}" - ) - raise ModelEntityCreationError( - f"Adapter '{config.name}' already exists but update failed: {update_error}" - ) from update_error - except Exception as e: - logger.exception(f"Failed to create model adapter: {e}") - raise ModelEntityCreationError(f"Failed to create model adapter: {e}") from e - else: - ft_type = config.peft.type.value if config.peft else FinetuningType.ALL_WEIGHTS.value - - request_body = { - "name": config.name, - "description": config.description, - "fileset": fileset_ref, - "finetuning_type": ft_type, - "trust_remote_code": base_me.trust_remote_code, - } - - if config.base_model: - request_body["base_model"] = config.base_model - - try: - output_me = self.sdk.models.create( - workspace=workspace, - **request_body, - ) - logger.info(f"Successfully created model entity: {output_me.workspace}/{output_me.name}") - return output_me.model_dump(), output_me - - except ConflictError: - logger.warning(f"Model entity already exists: {workspace}/{config.name}, updating existing model") - try: - update_body = {k: v for k, v in request_body.items() if k != "name"} - output_me = self.sdk.models.update( - name=config.name, - workspace=workspace, - **update_body, - ) - logger.info(f"Successfully updated model entity: {output_me.workspace}/{output_me.name}") - return output_me.model_dump(), output_me - except (InternalServerError, APITimeoutError, APIConnectionError): - raise - except Exception as update_error: - logger.exception(f"Failed to update existing model entity: {update_error}") - raise ModelEntityCreationError( - f"Model entity '{config.name}' already exists and update failed: {update_error}" - ) from update_error - - except Exception as e: - logger.exception(f"Failed to create model entity: {e}") - raise ModelEntityCreationError(f"Failed to create model entity: {e}") from e - - def launch_model(self, config: ModelEntityTaskConfig, me: ModelEntity): - """Deploy a model entity after creation. - - For LoRA jobs, ``me`` should be the base model entity. - For SFT jobs, ``me`` should be the output model entity. - """ - dc = config.deployment_config - if dc is None: - return - - # LORA_MERGED produces a full-weight model, so it is deployed like SFT - # and intentionally excluded from LoRA-specific checks below. - is_lora = config.peft is not None and config.peft.type == FinetuningType.LORA - if is_lora and self._has_active_deployment(me): - return - - if is_lora and isinstance(dc, DeploymentParameters) and not dc.lora_enabled: - logger.warning(f"Deployment requested but lora_enabled is false for a LoRA job: {dc}") - return - - # Resolve an existing config or create a new one from inline params. - if isinstance(dc, str): - logger.info(f"Resolving deployment config reference: {dc}") - deployment_config = self._resolve_config_ref(dc, me.workspace) - logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") - else: - deployment_config = self._create_deployment_config(dc, me) - - self._create_deployment(deployment_config, me) - - def _has_active_deployment(self, me: ModelEntity) -> bool: - """Check if the model entity already has an active deployment.""" - deployment_configs = self.sdk.inference.deployment_configs.list( - workspace=me.workspace, - filter=ModelDeploymentConfigFilterParam(model_entity_id=f"{me.workspace}/{me.name}"), - ).data - - for c in deployment_configs: - deployments = self.sdk.inference.deployments.list( - filter=ModelDeploymentFilterParam(config=c.name, workspace=me.workspace) - ).data - for d in deployments: - if d.status in ACTIVE_DEPLOYMENT_STATUSES: - logger.info(f"Active deployment (status={d.status}) exists for config {c.name}, skipping") - return True - - return False - - def _resolve_config_ref(self, config_ref: str, me_workspace: str) -> ModelDeploymentConfig: - """Resolve a ``name`` or ``workspace/name`` reference to a ModelDeploymentConfig.""" - parts = config_ref.split("/") - if len(parts) == 2: - workspace = parts[0] - name = parts[1] - elif len(parts) == 1: - workspace = me_workspace - name = parts[0] - else: - raise ModelEntityCreationError( - f"Invalid deployment config reference '{config_ref}': expected 'name' or 'workspace/name'" - ) - - try: - return self.sdk.inference.deployment_configs.retrieve(workspace=workspace, name=name) - except Exception as e: - raise ModelEntityCreationError( - f"Failed to resolve deployment config '{config_ref}' in workspace '{workspace}': {e}" - ) from e - - def _create_deployment_config(self, deploy_params: DeploymentParameters, me: ModelEntity) -> ModelDeploymentConfig: - """Create (or update) a ModelDeploymentConfig from inline parameters.""" - model_spec = ModelDeploymentConfigModelSpecParam( - model_name=me.name, - model_namespace=me.workspace, - lora_enabled=deploy_params.lora_enabled, - ) - executor_config = ContainerExecutorConfigParam( - image_name=deploy_params.image_name, - image_tag=deploy_params.image_tag, - gpu=deploy_params.gpu, - additional_envs=deploy_params.additional_envs, - ) - - if deploy_params.tool_call_config: - model_spec["tool_call_config"] = ToolCallConfigParam( - **deploy_params.tool_call_config.model_dump(exclude_none=True) - ) - - deployment_cfg_name = sanitize_name("sft-cfg", me.name) - try: - return self.sdk.inference.deployment_configs.create( - workspace=me.workspace, - name=deployment_cfg_name, - engine="nim", - model_spec=model_spec, - executor_config=executor_config, - ) - except ConflictError: - logger.info(f"Deployment config {me.workspace}/{deployment_cfg_name} already exists, updating") - return self.sdk.inference.deployment_configs.update( - workspace=me.workspace, - name=deployment_cfg_name, - engine="nim", - model_spec=model_spec, - executor_config=executor_config, - ) - - def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: ModelEntity) -> None: - """Create a deployment from the given ModelDeploymentConfig.""" - logger.info(f"Deployment config: {deployment_config}") - - if not me.spec: - _ = self._wait_for_spec(me.workspace, me.name) - - deployment_name = sanitize_name("sft-deploy", me.name) - try: - deployment = self.sdk.inference.deployments.create( - workspace=deployment_config.workspace, - name=deployment_name, - config=deployment_config.name, - ) - logger.info(f"Deployment created: {deployment}") - except ConflictError: - logger.info(f"Deployment {deployment_config.workspace}/{deployment_name} already exists") - deployment = self.sdk.inference.deployments.retrieve( - workspace=deployment_config.workspace, - name=deployment_name, - ) - - deployment_status = self.sdk.inference.deployments.retrieve( - workspace=deployment.workspace, - name=deployment.name, - ) - logger.info(f"Deployment status: {deployment_status}") - - -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the model entity creation task. - - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - - Returns: - Exit code (0 for success, non-zero for failure). - """ - job_ctx = job_ctx or NMPJobContext.from_env() - - sdk_owned = sdk is None - try: - sdk = sdk or get_task_sdk(SERVICE_NAME).with_options(workspace=job_ctx.workspace) - runner = ModelEntityRunner(sdk=sdk, job_ctx=job_ctx) - - config = get_config(job_ctx.config_path) - - logger.info(f"Starting model entity task with job context: {job_ctx}") - logger.info(f"Config: {config.model_dump_json(indent=2)}") - logger.info(f"NeMo Platform service URL: {sdk.base_url}") - - result, deploy_target = runner.create_model_entity(config) - logger.info(f"Model entity creation complete: {result}") - - runner.launch_model(config, deploy_target) - return 0 - - except ModelEntityCreationError as e: - logger.exception(f"Model entity creation failed: {e}") - return 1 - except Exception as e: - logger.exception(f"Model entity task failed: {e}") - return 1 - finally: - if sdk_owned and sdk is not None: - sdk.close() diff --git a/services/automodel/tests/test_compiler.py b/services/automodel/tests/test_compiler.py index 022da249c0..60464c8841 100644 --- a/services/automodel/tests/test_compiler.py +++ b/services/automodel/tests/test_compiler.py @@ -162,7 +162,14 @@ async def test_platform_job_config_compiler_sft_lora(mock_sdk, monkeypatch): if hasattr(steps[0], "executor") else steps[0]["executor"]["container"]["command"] ) - assert download_cmd[-1] == "nmp.automodel.tasks.file_io" + assert download_cmd == [ + "-m", + "nmp.customization_common.tasks.file_io", + "--service-source", + "automodel", + "--service-name", + "customizer", + ] download_entrypoint = ( steps[0].executor.container.entrypoint if hasattr(steps[0], "executor") diff --git a/services/automodel/tests/test_images.py b/services/automodel/tests/test_images.py index 8782e0d079..e875e7caa5 100644 --- a/services/automodel/tests/test_images.py +++ b/services/automodel/tests/test_images.py @@ -9,12 +9,12 @@ import pytest from nmp.automodel.config import AutomodelConfig from nmp.automodel.images import ( - TASKS_IMAGE_NAME, TRAINING_IMAGE_NAME, get_automodel_qualified_image, get_tasks_image, get_training_image, ) +from nmp.customization_common.service.images import CUSTOMIZER_TASKS_IMAGE_NAME @pytest.fixture @@ -31,9 +31,9 @@ def test_default_automodel_images_use_platform_registry(monkeypatch, platform_co tasks = get_tasks_image() training = get_training_image() - assert tasks == f"{platform_config.image_registry}/{TASKS_IMAGE_NAME}:{platform_config.image_tag}" + assert tasks == f"{platform_config.image_registry}/{CUSTOMIZER_TASKS_IMAGE_NAME}:{platform_config.image_tag}" assert training == f"{platform_config.image_registry}/{TRAINING_IMAGE_NAME}:{platform_config.image_tag}" - assert TASKS_IMAGE_NAME.count("/") == 0 # single repo segment, no nested paths + assert CUSTOMIZER_TASKS_IMAGE_NAME.count("/") == 0 # single repo segment, no nested paths def test_automodel_image_registry_override(monkeypatch, platform_config): @@ -44,8 +44,8 @@ def test_automodel_image_registry_override(monkeypatch, platform_config): ) assert ( - get_automodel_qualified_image(TASKS_IMAGE_NAME) - == f"my-registry/other-registry/{TASKS_IMAGE_NAME}:{platform_config.image_tag}" + get_automodel_qualified_image(TRAINING_IMAGE_NAME) + == f"my-registry/other-registry/{TRAINING_IMAGE_NAME}:{platform_config.image_tag}" ) @@ -54,8 +54,8 @@ def test_automodel_full_image_override(monkeypatch, platform_config): automodel_images, "config", AutomodelConfig( - tasks_image="my-registry/nemo-platform-dev/nmp-automodel-tasks:dev", + tasks_image="my-registry/nemo-platform-dev/nmp-customizer-tasks:dev", ), ) - assert get_tasks_image() == "my-registry/nemo-platform-dev/nmp-automodel-tasks:dev" + assert get_tasks_image() == "my-registry/nemo-platform-dev/nmp-customizer-tasks:dev" diff --git a/services/core/models/pyproject.toml b/services/core/models/pyproject.toml index e4a98b4cb7..4b5cfe6cd7 100644 --- a/services/core/models/pyproject.toml +++ b/services/core/models/pyproject.toml @@ -40,7 +40,7 @@ packages = ["src/nmp"] [dependency-groups] # No task deps for models — `nmp.core.models.parallelism` is the only consumer of # torch/transformers/accelerate, and it's only invoked from the `model_spec` batch -# task. That task runs in the `nmp-automodel-tasks` image (PyTorch from +# task. That task runs in the `nmp-customizer-tasks` image (PyTorch from # nmp-automodel-base); parallelism tests in this repo guard with `pytest.importorskip("torch")`. dev = [ diff --git a/services/core/models/src/nmp/core/models/api/v2/models.py b/services/core/models/src/nmp/core/models/api/v2/models.py index 4b7a449b7d..74fc7a31b5 100644 --- a/services/core/models/src/nmp/core/models/api/v2/models.py +++ b/services/core/models/src/nmp/core/models/api/v2/models.py @@ -279,7 +279,7 @@ async def start_update_model_spec_job(model_entity: ModelEntity): # host subprocess backend when subprocess/default is registered. profile="gpu", container=ContainerSpec( - image=get_qualified_image("nmp-automodel-tasks"), + image=get_qualified_image("nmp-customizer-tasks"), entrypoint=["/opt/venv/bin/python"], command=["-m", "nmp.core.models.tasks.model_spec"], ), diff --git a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py index d20cb4b4e8..75e8a3e938 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/k8s_nim_operator/config.py @@ -40,7 +40,7 @@ class K8sNimOperatorConfig(BaseModel): description=( "Image name (without registry/tag) for the LoRA adapters sidecar container. " "Registry and tag are taken from the platform config (NMP_IMAGE_REGISTRY / NMP_IMAGE_TAG). " - "Override to 'nmp-automodel-tasks' for local dev when that image is already available " + "Override to 'nmp-customizer-tasks' for local dev when that image is already available " "but nmp-api is not." ), ) @@ -49,7 +49,7 @@ class K8sNimOperatorConfig(BaseModel): description=( "Kubernetes container command (entrypoint) for the LoRA sidecar. " "Default uses the nmp-platform-runner entrypoint present in nmp-api. " - "When using nmp-automodel-tasks set to ['python'] and set lora_sidecar_args to " + "When using nmp-customizer-tasks set to ['python'] and set lora_sidecar_args to " "['-m', 'nmp.core.models.sidecars.adapters.main']." ), ) @@ -58,7 +58,7 @@ class K8sNimOperatorConfig(BaseModel): description=( "Kubernetes container args for the LoRA sidecar (appended after lora_sidecar_command). " "Leave empty for nmp-api. " - "Set to ['-m', 'nmp.core.models.sidecars.adapters.main'] when using nmp-automodel-tasks." + "Set to ['-m', 'nmp.core.models.sidecars.adapters.main'] when using nmp-customizer-tasks." ), ) diff --git a/services/rl/README.md b/services/rl/README.md index 4c5d85569e..4a3b08feb1 100644 --- a/services/rl/README.md +++ b/services/rl/README.md @@ -14,9 +14,9 @@ No HTTP server. The thin contributor layer lives in (download → DPO train → upload → model-entity). The training step's executor is chosen by `parallelism.num_nodes` (single-node `gpu` vs multi-node `gpu_distributed`). -- `nmp.rl.tasks.*` — container entrypoints (`file_io`, `model_entity`, - `training`). The training task bootstraps a Ray cluster and runs the DPO - driver against the NeMo-RL library. +- `nmp.rl.tasks.training` — GPU training entrypoint (bootstraps Ray and runs the DPO + driver). CPU `file_io` / `model_entity` steps run from the shared + `nmp-customizer-tasks` image (`nmp.customization_common.tasks.*`). ## Scope diff --git a/services/rl/pyproject.toml b/services/rl/pyproject.toml index e426ae506f..42cb3af303 100644 --- a/services/rl/pyproject.toml +++ b/services/rl/pyproject.toml @@ -30,8 +30,6 @@ integrations = [ [project.scripts] # Container entrypoints. Names match the automodel/unsloth pattern for parity. nmp-rl-training = "nmp.rl.tasks.training.__main__:main" -nmp-rl-file-io = "nmp.rl.tasks.file_io.run:run" -nmp-rl-model-entity = "nmp.rl.tasks.model_entity.__main__:run" [build-system] requires = ["hatchling"] diff --git a/services/rl/src/nmp/rl/app/jobs/compiler.py b/services/rl/src/nmp/rl/app/jobs/compiler.py index feda73b43d..cc63bdc352 100644 --- a/services/rl/src/nmp/rl/app/jobs/compiler.py +++ b/services/rl/src/nmp/rl/app/jobs/compiler.py @@ -50,6 +50,7 @@ ) from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig from nmp.customization_common.service.platform_client import fetch_model_entity +from nmp.customization_common.tasks.file_io_metadata import build_output_metadata from nmp.rl.app.constants import ( BASE_LOG_DIR_ENVVAR, DEFAULT_DATASET_PATH, @@ -66,7 +67,13 @@ ) from nmp.rl.config import config from nmp.rl.entities.values import FinetuningType, TrainingType -from nmp.rl.images import RL_PYTHON_ENTRYPOINT, get_tasks_image, get_training_image +from nmp.rl.images import ( + FILE_IO_TASK_COMMAND, + MODEL_ENTITY_TASK_COMMAND, + RL_PYTHON_ENTRYPOINT, + get_tasks_image, + get_training_image, +) from nmp.rl.schemas import DPOTraining, RlJobOutput logger = logging.getLogger(__name__) @@ -114,9 +121,19 @@ def _build_download_config(job_spec: RlJobOutput, me: ModelEntity, *, workspace: ) -def _build_upload_config(output_fileset_name: str) -> FileIOTaskConfig: +def _build_upload_config(job_spec: RlJobOutput) -> FileIOTaskConfig: return FileIOTaskConfig( - upload=[UploadItem(src=DEFAULT_OUTPUT_MODEL_PATH, dest=FileSetRef(workspace=None, name=output_fileset_name))], + upload=[ + UploadItem( + src=DEFAULT_OUTPUT_MODEL_PATH, + dest=FileSetRef(workspace=None, name=job_spec.output.fileset), + metadata=build_output_metadata( + model=job_spec.model, + finetuning_type=FinetuningType.ALL_WEIGHTS.value, + output_type=str(job_spec.output.type), + ), + ), + ], ) @@ -332,7 +349,9 @@ async def platform_job_config_compiler( base_env = _base_environment() def _cpu_task_step( - name: str, command: str, task_config: FileIOTaskConfig | ModelEntityTaskConfig + name: str, + command: list[str], + task_config: FileIOTaskConfig | ModelEntityTaskConfig, ) -> PlatformJobStep: return PlatformJobStep( name=name, @@ -341,7 +360,7 @@ def _cpu_task_step( container=ContainerSpec( image=get_tasks_image(), entrypoint=RL_PYTHON_ENTRYPOINT, - command=["-m", command], + command=command, ), resources=cpu_resources, ), @@ -352,14 +371,14 @@ def _cpu_task_step( steps: list[PlatformJobStep] = [ _cpu_task_step( "model-and-dataset-download", - "nmp.rl.tasks.file_io", + FILE_IO_TASK_COMMAND, _build_download_config(job_spec, me, workspace=workspace), ), _build_training_step(job_spec, base_env, trust_remote_code=trust_remote_code, profile=profile), - _cpu_task_step("model-upload", "nmp.rl.tasks.file_io", _build_upload_config(job_spec.output.fileset)), + _cpu_task_step("model-upload", FILE_IO_TASK_COMMAND, _build_upload_config(job_spec)), _cpu_task_step( "model-entity-creation", - "nmp.rl.tasks.model_entity", + MODEL_ENTITY_TASK_COMMAND, _build_model_entity_config(workspace, job_spec, trust_remote_code=trust_remote_code), ), ] diff --git a/services/rl/src/nmp/rl/config.py b/services/rl/src/nmp/rl/config.py index 89837a1694..22f007fdc0 100644 --- a/services/rl/src/nmp/rl/config.py +++ b/services/rl/src/nmp/rl/config.py @@ -18,7 +18,7 @@ class RlConfig(create_service_config_class("rl")): # type: ignore[misc] image_registry: str | None = Field( default=None, description=( - "Registry host/path prefix for nmp-rl-tasks and nmp-rl-training. " + "Registry host/path prefix for nmp-customizer-tasks and nmp-rl-training. " "Override via NMP_RL_IMAGE_REGISTRY; defaults to the platform's image registry." ), ) diff --git a/services/rl/src/nmp/rl/images.py b/services/rl/src/nmp/rl/images.py index 7aea0caeb7..e3d179cf70 100644 --- a/services/rl/src/nmp/rl/images.py +++ b/services/rl/src/nmp/rl/images.py @@ -4,24 +4,38 @@ """Docker image resolution for nmp-rl job steps. Unlike unsloth (single image), nmp-rl follows the automodel split: a heavy -``nmp-rl-training`` image (NGC + NeMo-RL + Ray) for the GPU training step and a -lighter ``nmp-rl-tasks`` image for the CPU file_io / model_entity steps. Both -build on ``nmp-rl-base``. Override via ``NMP_RL_TRAINING_IMAGE`` / -``NMP_RL_TASKS_IMAGE``. +``nmp-rl-training`` image for the GPU training step and the shared +``nmp-customizer-tasks`` image for CPU file_io / model_entity steps. """ from __future__ import annotations -from nmp.customization_common.service.images import resolve_qualified_image +from nmp.customization_common.service.images import ( + CUSTOMIZER_PYTHON_ENTRYPOINT, + get_customizer_tasks_image, + resolve_qualified_image, +) from nmp.rl.config import config BASE_IMAGE_NAME = "nmp-rl-base" -TASKS_IMAGE_NAME = "nmp-rl-tasks" TRAINING_IMAGE_NAME = "nmp-rl-training" -# Must match ENTRYPOINT in Dockerfile.nmp-rl-{tasks,training}. Job specs set this -# explicitly: Docker API create() replaces the image entrypoint when passed []. -RL_PYTHON_ENTRYPOINT = ["/opt/venv/bin/python"] +RL_PYTHON_ENTRYPOINT = CUSTOMIZER_PYTHON_ENTRYPOINT + +FILE_IO_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.file_io", + "--service-source", + "rl", + "--service-name", + "rl", +] +MODEL_ENTITY_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.model_entity", + "--service-name", + "rl", +] def get_rl_qualified_image(name: str, override: str | None = None) -> str: @@ -30,8 +44,8 @@ def get_rl_qualified_image(name: str, override: str | None = None) -> str: def get_tasks_image() -> str: - """CPU task steps (file_io, model_entity) — lighter image, no NeMo-RL/vLLM.""" - return get_rl_qualified_image(TASKS_IMAGE_NAME, config.tasks_image) + """CPU task steps (file_io, model_entity) — shared ``nmp-customizer-tasks`` image.""" + return get_customizer_tasks_image(backend_override=config.tasks_image, image_registry=config.image_registry) def get_training_image() -> str: diff --git a/services/rl/src/nmp/rl/tasks/file_io/__main__.py b/services/rl/src/nmp/rl/tasks/file_io/__main__.py deleted file mode 100644 index c34e051201..0000000000 --- a/services/rl/src/nmp/rl/tasks/file_io/__main__.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import sys - -from nmp.rl.tasks.file_io.run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/rl/src/nmp/rl/tasks/file_io/callbacks.py b/services/rl/src/nmp/rl/tasks/file_io/callbacks.py deleted file mode 100644 index e72777198b..0000000000 --- a/services/rl/src/nmp/rl/tasks/file_io/callbacks.py +++ /dev/null @@ -1,457 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Custom fsspec callbacks for progress reporting during file I/O operations.""" - -import logging -import threading -from abc import abstractmethod -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -from fsspec.callbacks import Callback, TqdmCallback -from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.customization_common.schemas.file_io import DownloadStats, TaskPhase, UploadStats -from nmp.customization_common.tasks.file_io_progress_reporter import ProgressReporter -from nmp.customization_common.tasks.file_io_utils import list_local_files as _list_local_files - -logger = logging.getLogger(__name__) - - -def get_percentage(current: int, total: int) -> int: - """Get integer percentage 0-100, clamped to the valid range. - - Progress accounting must never abort the underlying transfer. Inputs - can fall outside ``[0, total]`` for benign reasons — most commonly when - the pre-transfer file listing under-counts a source that contains nested - directories, so the live ``current`` count exceeds ``total`` by the - number of nested files. Clamp rather than raise so a cosmetic progress - number can't fail a multi-GB download. - """ - if total <= 0: - return 0 - if current > total or current < 0: - # Benign (see docstring) but worth a breadcrumb now that we no longer - # raise — the old hard error is what previously surfaced count drift. - logger.debug("get_percentage clamping out-of-range progress: current=%s total=%s", current, total) - current = max(0, min(current, total)) - return int((current / total) * 100) - - -@dataclass -class FileInfo: - """A dataclass for file information.""" - - path: str - size: int - - -class TqdmPerFileUploadCallback(Callback): - """A callback that creates a separate tqdm progress bar for each file upload.""" - - def __init__(self, src_path: Path, **kwargs: Any): - self.src_path = src_path - super().__init__(**kwargs) - - def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: - if self.src_path.is_file(): - relative_path = self.src_path.name - else: - relative_path = Path(full_src_path).relative_to(self.src_path) - return TqdmCallback( - tqdm_kwargs={ - "desc": f"Uploading {relative_path!s}", - "unit": "B", - "unit_scale": True, - "unit_divisor": 1024, - "miniters": 1, - }, - ) - - -class TqdmPerFileDownloadCallback(Callback): - """A callback that creates a separate tqdm progress bar for each file download. - - Accepts a ``file_sizes`` dict (relative path -> byte size) so each - progress bar can show percent-complete even when the SDK streams the - file without a Content-Length header. - """ - - def __init__(self, dest_path: Path, fileset_path: str, file_sizes: dict[str, int] | None = None, **kwargs: Any): - self.dest_path = dest_path - self.fileset_path = fileset_path.rstrip("/") - self.file_sizes = file_sizes or {} - super().__init__(**kwargs) - - def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: - dest_full_path = Path(full_dest_path) - if self.dest_path.is_file(): - relative_path = dest_full_path.name - else: - try: - relative_path = dest_full_path.relative_to(self.dest_path) - except ValueError: - relative_path = dest_full_path.name - - # full_src_path looks like "workspace/fileset/relative/path/file.txt". - # Strip the prefix to look up the size by relative path. - relative_file_path = full_src_path - if full_src_path.startswith(self.fileset_path): - relative_file_path = full_src_path[len(self.fileset_path) :].lstrip("/") - - file_size = self.file_sizes.get(relative_file_path) - - callback = TqdmCallback( - tqdm_kwargs={ - "desc": f"Downloading {relative_path!s}", - "unit": "B", - "unit_scale": True, - "unit_divisor": 1024, - "miniters": 1, - }, - ) - - # set_size() rather than tqdm_kwargs["total"] so the SDK can also - # call set_size() from a Content-Length header without conflict. - if file_size is not None: - callback.set_size(file_size) - - return callback - - -class BaseProgressCallback(Callback): - """Base class for file upload/download progress callbacks. - - Tracks file transfer progress and reports to the Jobs service. - Subclasses implement upload-vs-download behavior. - - Thread Safety: - Uses ``threading.Lock`` to protect stats updates because - FilesetFileSystem transfers files concurrently. - """ - - progress_reporter: ProgressReporter - fileset_name: str - total_files: int - total_size: int - stats: UploadStats | DownloadStats - _lock: threading.Lock - - def __init__( - self, - progress_reporter: ProgressReporter, - fileset_name: str, - total_files: int, - total_size: int, - stats: UploadStats | DownloadStats, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.progress_reporter = progress_reporter - self.fileset_name = str(fileset_name) - self.total_files = total_files - self.total_size = total_size - self.stats = stats - self._lock = threading.Lock() - - @staticmethod - def list_local_files(src_path: Path) -> list[FileInfo]: - """List all files under *src_path* (see shared ``list_local_files``).""" - return [FileInfo(path=f.path, size=f.size) for f in _list_local_files(src_path)] - - @abstractmethod - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "BaseSingleFileCallback": - """Create a child callback for a single file transfer.""" - ... - - -class BaseSingleFileCallback(Callback): - """Base class for per-file callbacks within a batch operation. - - Uses the template-method pattern: ``close()`` runs the shared - state-update + progress-report sequence, while subclasses customize - via ``_get_phase``, ``_get_file_display_path``, ``_update_stats``, - ``_get_files_count``, and ``_build_status_details``. - """ - - parent: BaseProgressCallback - source_path: str - dest_path: str - _completed: bool - - def __init__( - self, - parent: BaseProgressCallback, - source_path: str, - dest_path: str, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.parent = parent - self.source_path = source_path - self.dest_path = dest_path - self._completed = False - - @abstractmethod - def _get_phase(self) -> str: - """Return the TaskPhase for this operation.""" - ... - - @abstractmethod - def _get_file_display_path(self) -> str: - """Return the path to use for display/logging.""" - ... - - @abstractmethod - def _update_stats(self) -> None: - """Update the parent's stats for this operation (called within lock).""" - ... - - @abstractmethod - def _get_files_count(self) -> int: - """Return the current files count from stats (called within lock).""" - ... - - @abstractmethod - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - """Build the status_details dict for progress reporting.""" - ... - - def close(self) -> None: - """Called when the file transfer completes.""" - if self._completed: - return - - self._completed = True - parent = self.parent - current_file = self._get_file_display_path() - - with parent._lock: - self._update_stats() - files_count = self._get_files_count() - total_bytes = parent.stats.total_bytes - - logger.debug(f"File transferred: {current_file} ({files_count}/{parent.total_files})") - - # Report outside the lock — don't block other threads on the network call. - parent.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details=self._build_status_details(files_count, total_bytes, current_file), - ) - - def __enter__(self) -> "BaseSingleFileCallback": - return self - - def __exit__(self, *exc_args: object) -> None: - self.close() - - -class FileUploadProgressCallback(BaseProgressCallback): - """Callback for tracking file upload progress and reporting to the Jobs service.""" - - stats: UploadStats - - def __init__( - self, - progress_reporter: ProgressReporter, - src_path: Path, - fileset_name: str, - stats: UploadStats, - **kwargs: Any, - ): - files = self.list_local_files(src_path) - if not files: - logger.warning(f"Source path {src_path} contains no files") - total_files = len(files) - total_size = sum(f.size for f in files) - - super().__init__( - progress_reporter=progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - **kwargs, - ) - - logger.info(f"Uploading {total_files} files ({total_size} bytes) to {self.fileset_name}") - - progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "fileset": self.fileset_name, - "total_files": total_files, - "total_size": total_size, - "uploaded_files": 0, - "uploaded_bytes": 0, - }, - ) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileUploadCallback": - return SingleFileUploadCallback( - parent=self, - source_path=source_path, - dest_path=dest_path, - **kwargs, - ) - - -class SingleFileUploadCallback(BaseSingleFileCallback): - """Per-file upload callback. Notifies parent on completion.""" - - parent: FileUploadProgressCallback - - def _get_phase(self) -> str: - return TaskPhase.UPLOADING - - def _get_file_display_path(self) -> str: - return self.dest_path.split("/")[-1] if "/" in self.dest_path else self.dest_path - - def _update_stats(self) -> None: - self.parent.stats.files_uploaded += 1 - if self.size is not None: - self.parent.stats.total_bytes += self.size - - def _get_files_count(self) -> int: - return self.parent.stats.files_uploaded - - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - return { - "phase": TaskPhase.UPLOADING, - "fileset": self.parent.fileset_name, - "total_files": self.parent.total_files, - "total_size": self.parent.total_size, - "uploaded_files": files_count, - "uploaded_bytes": total_bytes, - "current_file": current_file, - "progress_pct": get_percentage(files_count, self.parent.total_files), - } - - -class FileDownloadProgressCallback(BaseProgressCallback): - """Callback for tracking file download progress and reporting to the Jobs service.""" - - stats: DownloadStats - - def __init__( - self, - progress_reporter: ProgressReporter, - fileset_name: str, - total_files: int, - total_size: int, - stats: DownloadStats, - **kwargs: Any, - ): - super().__init__( - progress_reporter=progress_reporter, - fileset_name=fileset_name, - total_files=total_files, - total_size=total_size, - stats=stats, - **kwargs, - ) - - logger.info(f"Downloading {total_files} files ({total_size} bytes) from {self.fileset_name}") - - progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "fileset": self.fileset_name, - "total_files": total_files, - "total_size": total_size, - "downloaded_files": 0, - "downloaded_bytes": 0, - }, - ) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileDownloadCallback": - return SingleFileDownloadCallback( - parent=self, - source_path=source_path, - dest_path=dest_path, - **kwargs, - ) - - -class SingleFileDownloadCallback(BaseSingleFileCallback): - """Per-file download callback. Notifies parent on completion.""" - - parent: FileDownloadProgressCallback - - def _get_phase(self) -> str: - return TaskPhase.DOWNLOADING - - def _get_file_display_path(self) -> str: - return self.source_path.split("/")[-1] if "/" in self.source_path else self.source_path - - def _update_stats(self) -> None: - self.parent.stats.files_downloaded += 1 - if self.size is not None: - self.parent.stats.total_bytes += self.size - - def _get_files_count(self) -> int: - return self.parent.stats.files_downloaded - - def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: - return { - "phase": TaskPhase.DOWNLOADING, - "fileset": self.parent.fileset_name, - "total_files": self.parent.total_files, - "total_size": self.parent.total_size, - "downloaded_files": files_count, - "downloaded_bytes": total_bytes, - "current_file": current_file, - "progress_pct": get_percentage(files_count, self.parent.total_files), - } - - -class CompositeCallback(Callback): - """A callback that delegates to multiple child callbacks. - - Lets us combine console-side ``TqdmCallback`` and Jobs-service - ``File{Upload,Download}ProgressCallback`` into one callback object - passed to fsspec operations. - """ - - def __init__(self, *callbacks: Callback, **kwargs: Any): - super().__init__(**kwargs) - self.callbacks = list(callbacks) - - def set_size(self, size: int) -> None: - self.size = size - for cb in self.callbacks: - cb.set_size(size) - - def absolute_update(self, value: int) -> None: - self.value = value - for cb in self.callbacks: - cb.absolute_update(value) - - def relative_update(self, inc: int = 1) -> None: - self.value += inc - for cb in self.callbacks: - cb.relative_update(inc) - - def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "CompositeCallback": - child_callbacks = [cb.branched(source_path, dest_path, **kwargs) for cb in self.callbacks] - return CompositeCallback(*child_callbacks) - - def call(self, hook_name: str | None = None, **kwargs: Any) -> None: - for cb in self.callbacks: - cb.call(hook_name, **kwargs) - - def close(self) -> None: - for cb in self.callbacks: - cb.close() - - def __enter__(self) -> "CompositeCallback": - for cb in self.callbacks: - cb.__enter__() - return self - - def __exit__(self, *exc_args: object) -> None: - for cb in self.callbacks: - cb.__exit__(*exc_args) diff --git a/services/rl/src/nmp/rl/tasks/file_io/run.py b/services/rl/src/nmp/rl/tasks/file_io/run.py deleted file mode 100644 index fa86215864..0000000000 --- a/services/rl/src/nmp/rl/tasks/file_io/run.py +++ /dev/null @@ -1,552 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""File I/O task entry point. - -Handles file operations between NeMo Platform Files Service and the job's shared PVC. - -The task reads configuration and performs: -- Downloads: If config.download is non-empty, download files from FileSets to local paths -- Uploads: If config.upload is non-empty, upload files from local paths to FileSets - -Usage: - export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.rl.tasks.file_io -""" - -import logging -from pathlib import Path - -import httpx - -# https://docs.nvidia.com/nemo/microservices/latest/pysdk/index.html#handling-errors -from nemo_platform import ( - APIConnectionError, - APITimeoutError, - InternalServerError, - NeMoPlatform, - NotFoundError, -) -from nemo_platform.types.files.fileset_file import FilesetFile -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import ( - ConflictError, -) -from nemo_platform_plugin.client.errors import ( - InternalServerError as ClientInternalServerError, -) -from nemo_platform_plugin.client.types import RetryPolicy -from nemo_platform_plugin.files.client import FilesClient -from nemo_platform_plugin.files.types import CreateFilesetRequest, UpdateFilesetRequest -from nmp.common.jobs.schemas import PlatformJobStatus -from nmp.common.sdk_factory import get_task_sdk -from nmp.customization_common.schemas.file_io import ( - DownloadItem, - DownloadStats, - FileDownloadError, - FileSetRef, - FileUploadError, - PathTraversalError, - TaskPhase, - UploadItem, - UploadStats, -) -from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.tasks.file_io_progress_reporter import JobsServiceProgressReporter, ProgressReporter -from nmp.customization_common.tasks.file_io_utils import ( - filesystem_sdk_error_handler, - get_config, - sdk_error_handler, - validate_safe_path, - validate_storage_path, -) -from nmp.rl.app.constants import SERVICE_NAME -from nmp.rl.tasks.file_io.callbacks import ( - CompositeCallback, - FileDownloadProgressCallback, - FileUploadProgressCallback, - TqdmPerFileDownloadCallback, - TqdmPerFileUploadCallback, -) -from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential - -logger = logging.getLogger(__name__) - -# Service-source tag stamped onto every upload-created fileset. Lets operators -# filter filesets by training backend. -SERVICE_SOURCE = "rl" - -CREATE_FILESET_TIMEOUT = 10.0 -LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) - -# Timeout configurations for FilesetFileSystem operations. Passed via -# sdk.with_options(timeout=...). httpx.Timeout(read=...) is per-chunk -# (the SDK chunks at 16MB), NOT total transfer time — it's a socket-level -# timeout. SDK defaults are httpx.Timeout(timeout=60, connect=5.0). -DOWNLOAD_TIMEOUT = httpx.Timeout(30.0, read=5 * 60) -UPLOAD_TIMEOUT = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) - -# Retry configuration. -MAX_RETRIES = 3 -INITIAL_BACKOFF_SECONDS = 1.0 -MAX_BACKOFF_SECONDS = 30.0 - -# Transient exceptions that should trigger retries for filesystem operations. -# FilesetFileSystem uses httpx under the hood, so we retry httpx transients -# in addition to SDK-level transients. -TRANSIENT_FILESYSTEM_EXCEPTIONS = ( - httpx.TimeoutException, - httpx.ConnectError, - httpx.ReadTimeout, - # Connection dropped mid-transfer (CDN/proxy closed the socket before the - # full body arrived). Common on large multi-GB model shards; safe to retry. - httpx.RemoteProtocolError, - httpx.ReadError, -) - - -class FileIORunner: - """Runner for file I/O operations against the Files service.""" - - def __init__( - self, - sdk: NeMoPlatform, - progress_reporter: ProgressReporter, - job_ctx: NMPJobContext, - ): - self.sdk = sdk - self.progress_reporter = progress_reporter - self.job_ctx = job_ctx - - def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFile]: - """List files in a FileSet. Returns a list of ``FilesetFile`` objects.""" - try: - with sdk_error_handler(FileDownloadError, f"list files in fileset {fileset}", passthrough=(NotFoundError,)): - response = self.sdk.with_options(timeout=LIST_FILES_TIMEOUT).files.list( - fileset=fileset.name, - workspace=fileset.workspace, - ) - logger.info(f"Found {len(response.data)} files in FileSet {fileset!s}") - return response.data - except NotFoundError as e: - raise FileDownloadError( - f"FileSet {fileset!s} not found. Please ensure the FileSet exists and contains the expected files.", - ) from e - - def download_fileset(self, fileset: FileSetRef, dest_dir: Path) -> DownloadStats: - """Download all files from a FileSet to a destination directory. - - Uses ``FilesetFileSystem.get()`` with ``recursive=True`` for efficient batch - downloads. Progress is tracked via two callbacks combined in a - ``CompositeCallback``: - - - ``TqdmPerFileDownloadCallback`` — separate console progress bar per file - - ``FileDownloadProgressCallback`` — reports to Jobs service after each file - - Raises: - FileDownloadError: If the download fails. - """ - fileset_name = str(fileset) - - files = self.list_fileset_files(fileset) - - if not files: - logger.warning(f"FileSet {fileset_name} contains no files") - return DownloadStats() - - total_files = len(files) - total_size = sum(f.size for f in files) - - dest_dir.mkdir(parents=True, exist_ok=True) - - # Maps relative file paths to byte sizes for tqdm percent display. - file_sizes = {f.path.lstrip("/"): f.size for f in files} - - with filesystem_sdk_error_handler( - FileDownloadError, - f"download from '{fileset_name}' to '{dest_dir}'", - ): - # Progress state (stats + callbacks) is built inside the retried call so - # each attempt starts from zero. A retry restarts the transfer from - # scratch, so reusing a single stats/callback object across attempts - # would double-count earlier progress and inflate the reported totals. - stats = self._download_with_retry( - fileset_name=fileset.name, - fileset_workspace=fileset.workspace, - dest_dir=str(dest_dir), - fileset_display_name=fileset_name, - dest_path=dest_dir, - file_sizes=file_sizes, - total_files=total_files, - total_size=total_size, - ) - - logger.info(f"Download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") - return stats - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), - reraise=True, - before_sleep=before_sleep_log(logger, logging.WARNING), - ) - def _download_with_retry( - self, - fileset_name: str, - fileset_workspace: str | None, - dest_dir: str, - fileset_display_name: str, - dest_path: Path, - file_sizes: dict[str, int], - total_files: int, - total_size: int, - ) -> DownloadStats: - """Internal method with retry logic for downloading from FilesetFileSystem. - - Builds fresh ``DownloadStats`` and callbacks on every attempt so retried - transfers report progress from zero rather than accumulating across attempts. - """ - stats = DownloadStats() - tqdm_callback = TqdmPerFileDownloadCallback( - dest_path=dest_path, - fileset_path=fileset_display_name, - file_sizes=file_sizes, - ) - jobs_callback = FileDownloadProgressCallback( - progress_reporter=self.progress_reporter, - fileset_name=fileset_display_name, - total_files=total_files, - total_size=total_size, - stats=stats, - ) - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - - self.sdk.with_options(timeout=DOWNLOAD_TIMEOUT).files.download( - fileset=fileset_name, - workspace=fileset_workspace, - local_path=dest_dir, - callback=composite_callback, - ) - return stats - - def upload_fileset(self, fileset: FileSetRef, src_path: Path) -> UploadStats: - """Upload all files from a source path (file or directory) to a FileSet. - - Uses ``FilesetFileSystem.put()`` with ``recursive=True`` for efficient batch - uploads. Progress is tracked via the same composite-callback pattern as - downloads. - - Raises: - FileUploadError: If the upload fails. - """ - fileset_name = str(fileset) - - # Build local and remote paths for upload. ``remote_path`` is relative within - # the fileset ("" for root, "filename" for single file). Trailing slash on - # ``local_path`` follows rsync/scp convention: "dir/" copies contents, - # "dir" copies the directory itself. - if src_path.is_dir(): - local_path = f"{src_path}/" - remote_path = "" - else: - local_path = str(src_path) - remote_path = src_path.name - - with filesystem_sdk_error_handler( - FileUploadError, - f"upload from '{src_path}' to '{fileset_name}'", - ): - # Fresh progress state per attempt — see _download_with_retry for why - # reusing stats/callbacks across retries double-counts progress. - stats = self._upload_with_retry( - local_path=local_path, - remote_path=remote_path, - fileset_name=fileset.name, - fileset_workspace=fileset.workspace, - fileset_display_name=fileset_name, - src_path=src_path, - ) - - logger.info(f"Upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") - return stats - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), - reraise=True, - before_sleep=before_sleep_log(logger, logging.WARNING, exc_info=True), - ) - def _upload_with_retry( - self, - local_path: str, - remote_path: str, - fileset_name: str, - fileset_workspace: str | None, - fileset_display_name: str, - src_path: Path, - ) -> UploadStats: - """Internal method with retry logic for uploading to FilesetFileSystem. - - Builds fresh ``UploadStats`` and callbacks on every attempt so retried - transfers report progress from zero rather than accumulating across attempts. - """ - stats = UploadStats() - tqdm_callback = TqdmPerFileUploadCallback(src_path=src_path) - jobs_callback = FileUploadProgressCallback( - progress_reporter=self.progress_reporter, - src_path=src_path, - fileset_name=fileset_display_name, - stats=stats, - ) - composite_callback = CompositeCallback(tqdm_callback, jobs_callback) - - self.sdk.with_options(timeout=UPLOAD_TIMEOUT).files.upload( - local_path=local_path, - remote_path=remote_path, - fileset=fileset_name, - workspace=fileset_workspace, - callback=composite_callback, - ) - return stats - - def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> None: - """Create a FileSet. Skip if it already exists. - - Wraps the retry with ``sdk_error_handler`` to convert exceptions after - all retries exhaust. - """ - with sdk_error_handler(FileUploadError, f"create fileset {fileset}", passthrough=(ConflictError,)): - self._create_fileset_with_retry(fileset, metadata) - - # We don't use the SDK's built-in retry: it would retry on ConflictError, - # which is expected here and would just waste calls. - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type( - ( - InternalServerError, - APITimeoutError, - APIConnectionError, - ClientInternalServerError, - httpx.TimeoutException, - httpx.ConnectError, - ) - ), - reraise=True, - ) - def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: - """Internal method with retry logic for creating a FileSet.""" - files = client_from_platform(self.sdk, FilesClient).with_options( - timeout=CREATE_FILESET_TIMEOUT, retry=RetryPolicy(max_retries=0) - ) - try: - body_kwargs: dict = { - "name": fileset.name, - "custom_fields": {"service_source": SERVICE_SOURCE}, - } - if metadata is not None: - body_kwargs["metadata"] = metadata - result = files.create_fileset(workspace=fileset.workspace, body=CreateFilesetRequest(**body_kwargs)).data() - logger.info(f"Created FileSet: {result.workspace}/{result.name}") - except ConflictError: - workspace = fileset.workspace or self.job_ctx.workspace - if metadata is not None: - try: - files.update_fileset( - workspace=workspace, - name=fileset.name, - body=UpdateFilesetRequest(metadata=metadata), - ) - logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") - except Exception as e: - logger.warning( - f"Could not patch metadata on existing fileset {workspace}/{fileset.name}: {e}. " - "Upload will continue; downstream consumers may lack the latest metadata.", - ) - - def run_download(self, downloads: list[DownloadItem]) -> None: - """Execute download operations.""" - if not downloads: - logger.info("No downloads configured, skipping download operation") - return - - storage_path = validate_storage_path(self.job_ctx.storage_path) - - logger.info(f"Starting download operation: {len(downloads)} fileset(s) to download") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "total_filesets": len(downloads), - "completed_filesets": 0, - }, - ) - - total_stats = DownloadStats() - - for idx, item in enumerate(downloads): - fileset = item.src - dest_dir = validate_safe_path(storage_path, item.dest) - - logger.info(f"[{idx + 1}/{len(downloads)}] Downloading from {fileset!s} to {dest_dir}") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.DOWNLOADING, - "total_filesets": len(downloads), - "completed_filesets": idx, - "current_fileset": f"{fileset!s}", - }, - ) - - stats = self.download_fileset(fileset, dest_dir) - total_stats.files_downloaded += stats.files_downloaded - total_stats.total_bytes += stats.total_bytes - - logger.info(f"FileSet download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") - - logger.info( - f"All downloads complete: {total_stats.files_downloaded} files, {total_stats.total_bytes} bytes total", - ) - - def run_upload(self, uploads: list[UploadItem]) -> None: - """Execute upload operations.""" - if not uploads: - logger.info("No uploads configured, skipping upload operation") - return - - storage_path = validate_storage_path(self.job_ctx.storage_path) - - logger.info(f"Starting upload operation: {len(uploads)} fileset(s) to upload") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "total_filesets": len(uploads), - "completed_filesets": 0, - }, - ) - - total_stats = UploadStats() - - for idx, item in enumerate(uploads): - if item.dest.workspace is None: - item.dest.workspace = self.job_ctx.workspace - fileset = item.dest - src_path = validate_safe_path(storage_path, item.src) - if not src_path.exists(): - raise FileUploadError(f"Source path does not exist: {src_path}. Ensure the source path exists.") - if not src_path.is_dir() and not src_path.is_file(): - raise FileUploadError( - f"Source path is not a file or directory: {src_path}. " - "Ensure the source path is a file or directory.", - ) - - logger.info(f"[{idx + 1}/{len(uploads)}] Uploading from {src_path} to {fileset!s}") - - self.progress_reporter.update_progress( - status=PlatformJobStatus.ACTIVE, - status_details={ - "phase": TaskPhase.UPLOADING, - "total_filesets": len(uploads), - "completed_filesets": idx, - "current_fileset": str(fileset), - }, - ) - - self.create_fileset(fileset, metadata=item.metadata) - - stats = self.upload_fileset(fileset, src_path) - total_stats.files_uploaded += stats.files_uploaded - total_stats.total_bytes += stats.total_bytes - - logger.info(f"FileSet upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") - - logger.info(f"All uploads complete: {total_stats.files_uploaded} files, {total_stats.total_bytes} bytes total") - - -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the file I/O task. - - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - - Returns: - Exit code (0 for success, non-zero for failure). - """ - job_ctx = job_ctx or NMPJobContext.from_env() - validate_storage_path(job_ctx.storage_path) - - sdk_owned = sdk is None - progress_reporter: ProgressReporter | None = None - try: - sdk = sdk or get_task_sdk(SERVICE_NAME) - progress_reporter = JobsServiceProgressReporter.create_progress_reporter(sdk, job_ctx) - runner = FileIORunner(sdk=sdk, progress_reporter=progress_reporter, job_ctx=job_ctx) - - config = get_config(job_ctx.config_path) - - logger.info(f"Starting file I/O task with job context: {job_ctx}") - logger.info(f"Config: {config.model_dump_json(indent=2)}") - logger.info(f"NeMo Platform service URL: {sdk.base_url}") - - runner.run_upload(config.upload) - runner.run_download(config.download) - - progress_reporter.update_progress( - status=PlatformJobStatus.COMPLETED, - status_details={"phase": TaskPhase.COMPLETED, "message": "File I/O task completed successfully"}, - ) - - return 0 - except PathTraversalError as e: - logger.error(f"Security error - path traversal detected: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - except (FileDownloadError, FileUploadError) as e: - logger.exception(f"File operation failed: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - except Exception as e: - logger.exception(f"File I/O task failed: {e}") - if progress_reporter: - progress_reporter.update_progress( - status=PlatformJobStatus.ERROR, - error_details={"message": str(e), "type": type(e).__name__}, - ) - return 1 - finally: - if sdk_owned and sdk is not None: - sdk.close() - - -def build_output_metadata(spec) -> dict: - """Build the metadata dict stamped onto the output fileset. - - Captures the bits a downstream consumer (model-entity creation, - deployment) needs about this artefact without re-deriving them - from the training spec. - """ - return { - "model": spec.model.name, - "finetuning_type": spec.training.finetuning_type, - "save_method": spec.output.save_method, - "output_type": spec.output.type, - } diff --git a/services/rl/src/nmp/rl/tasks/model_entity/__init__.py b/services/rl/src/nmp/rl/tasks/model_entity/__init__.py deleted file mode 100644 index ebe5a7cf8f..0000000000 --- a/services/rl/src/nmp/rl/tasks/model_entity/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Model entity task for creating model entities after rl customization.""" - -from nmp.rl.tasks.model_entity.run import run - -__all__ = ["run"] diff --git a/services/rl/src/nmp/rl/tasks/model_entity/__main__.py b/services/rl/src/nmp/rl/tasks/model_entity/__main__.py deleted file mode 100644 index 3a2739e1b9..0000000000 --- a/services/rl/src/nmp/rl/tasks/model_entity/__main__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Entry point for model_entity task. - -Usage: - python -m nmp.rl.tasks.model_entity -""" - -import sys - -from .run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/rl/src/nmp/rl/tasks/model_entity/run.py b/services/rl/src/nmp/rl/tasks/model_entity/run.py deleted file mode 100644 index 192fc36f58..0000000000 --- a/services/rl/src/nmp/rl/tasks/model_entity/run.py +++ /dev/null @@ -1,487 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Model entity task entry point. - -Handles creating model entities in the Models service after customization completes. - -The task reads configuration and creates a Model Entity that references the -uploaded model artifacts in the Files service. When ``deployment_config`` is set -on the task config, the task also launches an inference deployment. - -Usage: - export NEMO_JOB_STEP_CONFIG_FILE_PATH= - python -m nmp.rl.tasks.model_entity -""" - -import json -import logging -import re -import time -from pathlib import Path - -import httpx -from nemo_platform import ( - APIConnectionError, - APITimeoutError, - ConflictError, - InternalServerError, - NeMoPlatform, - NotFoundError, -) -from nemo_platform.types.inference import ( - ContainerExecutorConfigParam, - ModelDeploymentConfig, - ModelDeploymentConfigFilterParam, - ModelDeploymentConfigModelSpecParam, - ModelDeploymentFilterParam, -) -from nemo_platform.types.models import LoraParam, ModelEntity -from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import InternalServerError as ClientInternalServerError -from nemo_platform_plugin.files.client import FilesClient -from nmp.common.sdk_factory import get_task_sdk -from nmp.customization_common.schemas.model_entity import ( - DeploymentParameters, - ModelEntityCreationError, - ModelEntityTaskConfig, -) -from nmp.customization_common.service.context import NMPJobContext -from nmp.rl.app.constants import SERVICE_NAME -from nmp.rl.entities.values import FinetuningType -from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential - -logger = logging.getLogger(__name__) - -# Retry configuration. -MAX_RETRIES = 3 -INITIAL_BACKOFF_SECONDS = 1.0 -MAX_BACKOFF_SECONDS = 30.0 - -ACTIVE_DEPLOYMENT_STATUSES = frozenset({"CREATED", "PENDING", "READY"}) - -SPEC_POLL_INTERVAL_SECONDS = 10 -SPEC_POLL_TIMEOUT_SECONDS = 600 - - -def get_config(config_path: Path) -> ModelEntityTaskConfig: - """Load and validate the model_entity step config from disk.""" - with open(config_path) as f: - return ModelEntityTaskConfig.model_validate(json.load(f)) - - -def sanitize_name(prefix: str, name: str) -> str: - """Build a deployment-safe name from a free-form model name. - - Must match the API's ``{'pattern': '^[a-z](?!.*--)[a-z0-9\\-@.+_]{1,62}(? ModelEntity: - """Poll until the model_spec task has populated the model's spec. - - The spec must be populated before creating a deployment because the - inference service relies on ``spec.family`` and ``spec.base_num_parameters`` - to select the correct NIM profile. - - Raises: - ModelEntityCreationError: If the spec is not populated within the timeout. - """ - logger.info(f"Waiting for model_spec to populate spec on {workspace}/{name}") - start = time.monotonic() - - while time.monotonic() - start < SPEC_POLL_TIMEOUT_SECONDS: - try: - target = self.sdk.models.retrieve(name=name, workspace=workspace) - spec = target.spec - # Deployment/NIM profile selection needs both spec.family and - # spec.base_num_parameters; a partially-populated spec isn't enough. - if spec and getattr(spec, "family", None) and getattr(spec, "base_num_parameters", None) is not None: - logger.info(f"Spec populated on {workspace}/{name}") - return target - except (APIConnectionError, APITimeoutError, InternalServerError) as e: - logger.warning(f"Transient error polling spec for {workspace}/{name}: {e}") - time.sleep(SPEC_POLL_INTERVAL_SECONDS) - - raise ModelEntityCreationError( - f"Timed out waiting for model spec on {workspace}/{name} " - f"after {SPEC_POLL_TIMEOUT_SECONDS}s. The platform could not auto-detect the " - f"model's specifications. Verify the model checkpoint is valid and in a supported format." - ) - - def get_model_entity(self, model_entity: str, fileset_workspace: str) -> ModelEntity: - """Resolve ``"workspace/name"`` (or bare ``"name"``) to a ``ModelEntity``.""" - parts = model_entity.split("/") - if len(parts) == 1 and parts[0]: - me_workspace, me_name = fileset_workspace, parts[0] - elif len(parts) == 2 and all(parts): - me_workspace, me_name = parts[0], parts[1] - else: - # Reject anything that isn't exactly 'name' or 'workspace/name' (e.g. - # 'a/b/c', '/b', 'a/') instead of silently dropping extra segments. - raise ModelEntityCreationError( - f"Invalid model entity reference '{model_entity}': expected 'name' or 'workspace/name'." - ) - - try: - me: ModelEntity = self.sdk.models.retrieve(name=me_name, workspace=me_workspace) - except NotFoundError as e: - raise ModelEntityCreationError(f"Model entity {me_workspace}/{me_name} not found") from e - - return me - - @retry( - stop=stop_after_attempt(MAX_RETRIES), - wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), - retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), - reraise=True, - ) - def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, ModelEntity]: - """Create a model entity in the Models service. - - Returns: - Tuple of (result dict, deploy target). For LoRA the deploy target is the - *base* model entity; for SFT it is the newly created output model entity. - - Raises: - ModelEntityCreationError: If creation fails. - """ - # The output entity is created in the workspace declared on the config - # (the "workspace of the model entity to create" contract), not the - # ambient job workspace — the two can differ for cross-workspace jobs. - output_workspace = config.workspace - logger.info(f"Creating model entity: {output_workspace}/{config.name}") - - fileset_workspace = config.fileset.workspace or self.job_ctx.workspace - fileset_ref = f"{fileset_workspace}/{config.fileset.name}" - - logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") - try: - client_from_platform(self.sdk, FilesClient).get_fileset( - workspace=fileset_workspace, name=config.fileset.name - ) - logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") - except ( - InternalServerError, - APITimeoutError, - APIConnectionError, - ClientInternalServerError, - httpx.TimeoutException, - httpx.ConnectError, - ): - raise - except Exception as e: - logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") - raise ModelEntityCreationError( - f"Cannot create model entity: fileset '{fileset_workspace}/{config.fileset.name}' " - "does not exist or is not accessible" - ) from e - - base_me: ModelEntity = self.get_model_entity(config.model_entity, fileset_workspace) - - if config.peft is not None and config.peft.type == FinetuningType.LORA: - return self._create_or_update_adapter(config, base_me, fileset_ref) - return self._create_or_update_full_entity(config, fileset_ref, output_workspace) - - def _create_or_update_adapter( - self, - config: ModelEntityTaskConfig, - base_me: ModelEntity, - fileset_ref: str, - ) -> tuple[dict, ModelEntity]: - """Create or update a LoRA adapter on ``base_me``. Returns (result, base_me).""" - assert config.peft is not None # type narrowing — caller already checked - try: - output_me = self.sdk.models.adapters.create( - model_name=base_me.name, - workspace=base_me.workspace, - name=config.name, - description=config.description, - fileset=fileset_ref, - finetuning_type=config.peft.type.value, - lora_config=LoraParam( - alpha=config.peft.alpha, - rank=config.peft.rank, - ), - enabled=True, - ) - return output_me.model_dump(), base_me - except ConflictError: - logger.warning( - f"Adapter {base_me.workspace}/{config.name} already exists for model " - f"{base_me.workspace}/{base_me.name}, updating with new fileset" - ) - try: - output_me = self.sdk.models.adapters.update( - adapter=config.name, - model_name=base_me.name, - workspace=base_me.workspace, - fileset=fileset_ref, - description=config.description, - enabled=True, - ) - logger.info( - f"Successfully updated adapter: {base_me.workspace}/{config.name} " - f"for base model {base_me.workspace}/{base_me.name}" - ) - return output_me.model_dump(), base_me - except (InternalServerError, APITimeoutError, APIConnectionError): - raise - except Exception as update_error: - logger.exception( - f"Failed to update existing adapter, {base_me.workspace}/{config.name}: {update_error}" - ) - raise ModelEntityCreationError( - f"Adapter '{config.name}' already exists but update failed: {update_error}" - ) from update_error - except Exception as e: - logger.exception(f"Failed to create model adapter: {e}") - raise ModelEntityCreationError(f"Failed to create model adapter: {e}") from e - - def _create_or_update_full_entity( - self, - config: ModelEntityTaskConfig, - fileset_ref: str, - workspace: str, - ) -> tuple[dict, ModelEntity]: - """Create or update a full / merged model entity. Returns (result, output_me).""" - ft_type = config.peft.type.value if config.peft else FinetuningType.ALL_WEIGHTS.value - - request_body: dict = { - "name": config.name, - "description": config.description, - "fileset": fileset_ref, - "finetuning_type": ft_type, - # Honor the task config's flag (resolved by the compiler from the base - # model entity) rather than re-reading it off a freshly fetched entity. - "trust_remote_code": config.trust_remote_code, - } - if config.base_model: - request_body["base_model"] = config.base_model - - try: - output_me = self.sdk.models.create(workspace=workspace, **request_body) - logger.info(f"Successfully created model entity: {output_me.workspace}/{output_me.name}") - return output_me.model_dump(), output_me - except ConflictError: - logger.warning(f"Model entity already exists: {workspace}/{config.name}, updating existing model") - try: - update_body = {k: v for k, v in request_body.items() if k != "name"} - output_me = self.sdk.models.update( - name=config.name, - workspace=workspace, - **update_body, - ) - logger.info(f"Successfully updated model entity: {output_me.workspace}/{output_me.name}") - return output_me.model_dump(), output_me - except (InternalServerError, APITimeoutError, APIConnectionError): - raise - except Exception as update_error: - logger.exception(f"Failed to update existing model entity: {update_error}") - raise ModelEntityCreationError( - f"Model entity '{config.name}' already exists and update failed: {update_error}" - ) from update_error - except Exception as e: - logger.exception(f"Failed to create model entity: {e}") - raise ModelEntityCreationError(f"Failed to create model entity: {e}") from e - - def launch_model(self, config: ModelEntityTaskConfig, me: ModelEntity) -> None: - """Deploy a model entity after creation. - - For LoRA jobs, ``me`` should be the base model entity. - For SFT jobs, ``me`` should be the output model entity. - """ - dc = config.deployment_config - if dc is None: - return - - # LORA_MERGED produces a full-weight model, so it's deployed like SFT and - # is intentionally excluded from the LoRA-only checks below. - is_lora = config.peft is not None and config.peft.type == FinetuningType.LORA - if is_lora and self._has_active_deployment(me): - return - - if is_lora and isinstance(dc, DeploymentParameters) and not dc.lora_enabled: - logger.warning(f"Deployment requested but lora_enabled is false for a LoRA job: {dc}") - return - - if isinstance(dc, str): - logger.info(f"Resolving deployment config reference: {dc}") - deployment_config = self._resolve_config_ref(dc, me.workspace) - logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") - else: - deployment_config = self._create_deployment_config(dc, me) - - self._create_deployment(deployment_config, me) - - def _has_active_deployment(self, me: ModelEntity) -> bool: - """Check if the model entity already has an active deployment.""" - deployment_configs = self.sdk.inference.deployment_configs.list( - workspace=me.workspace, - filter=ModelDeploymentConfigFilterParam(model_entity_id=f"{me.workspace}/{me.name}"), - ).data - - for c in deployment_configs: - deployments = self.sdk.inference.deployments.list( - filter=ModelDeploymentFilterParam(config=c.name, workspace=me.workspace) - ).data - for d in deployments: - if d.status in ACTIVE_DEPLOYMENT_STATUSES: - logger.info(f"Active deployment (status={d.status}) exists for config {c.name}, skipping") - return True - - return False - - def _resolve_config_ref(self, config_ref: str, me_workspace: str) -> ModelDeploymentConfig: - """Resolve a ``name`` or ``workspace/name`` reference to a ``ModelDeploymentConfig``.""" - parts = config_ref.split("/") - if len(parts) == 2: - workspace, name = parts[0], parts[1] - elif len(parts) == 1: - workspace, name = me_workspace, parts[0] - else: - raise ModelEntityCreationError( - f"Invalid deployment config reference '{config_ref}': expected 'name' or 'workspace/name'" - ) - - try: - return self.sdk.inference.deployment_configs.retrieve(workspace=workspace, name=name) - except Exception as e: - raise ModelEntityCreationError( - f"Failed to resolve deployment config '{config_ref}' in workspace '{workspace}': {e}" - ) from e - - def _create_deployment_config(self, deploy_params: DeploymentParameters, me: ModelEntity) -> ModelDeploymentConfig: - """Create (or update) a ``ModelDeploymentConfig`` from inline parameters.""" - model_spec = ModelDeploymentConfigModelSpecParam( - model_name=me.name, - model_namespace=me.workspace, - lora_enabled=deploy_params.lora_enabled, - ) - executor_config = ContainerExecutorConfigParam( - image_name=deploy_params.image_name, - image_tag=deploy_params.image_tag, - gpu=deploy_params.gpu, - additional_envs=deploy_params.additional_envs, - ) - - if deploy_params.tool_call_config: - model_spec["tool_call_config"] = ToolCallConfigParam( - **deploy_params.tool_call_config.model_dump(exclude_none=True) - ) - - deployment_cfg_name = sanitize_name("sft-cfg", me.name) - try: - return self.sdk.inference.deployment_configs.create( - workspace=me.workspace, - name=deployment_cfg_name, - engine="nim", - model_spec=model_spec, - executor_config=executor_config, - ) - except ConflictError: - logger.info(f"Deployment config {me.workspace}/{deployment_cfg_name} already exists, updating") - return self.sdk.inference.deployment_configs.update( - workspace=me.workspace, - name=deployment_cfg_name, - engine="nim", - model_spec=model_spec, - executor_config=executor_config, - ) - - def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: ModelEntity) -> None: - """Create a deployment from the given ``ModelDeploymentConfig``.""" - # Log identifiers only: the full ModelDeploymentConfig embeds - # executor_config.additional_envs (deployment secrets), which would - # otherwise become durable in the job logs. - logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") - - if not me.spec: - _ = self._wait_for_spec(me.workspace, me.name) - - deployment_name = sanitize_name("sft-deploy", me.name) - try: - deployment = self.sdk.inference.deployments.create( - workspace=deployment_config.workspace, - name=deployment_name, - config=deployment_config.name, - ) - logger.info(f"Deployment created: {deployment.workspace}/{deployment.name}") - except ConflictError: - logger.info(f"Deployment {deployment_config.workspace}/{deployment_name} already exists") - deployment = self.sdk.inference.deployments.retrieve( - workspace=deployment_config.workspace, - name=deployment_name, - ) - - deployment_status = self.sdk.inference.deployments.retrieve( - workspace=deployment.workspace, - name=deployment.name, - ) - logger.info( - f"Deployment {deployment_status.workspace}/{deployment_status.name} status: {deployment_status.status}" - ) - - -def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: - """Execute the model entity creation task. - - Args: - sdk: Optional SDK instance for dependency injection (for testing). - If None, creates one via get_task_sdk(). - job_ctx: Optional job context for dependency injection (for testing). - If None, creates one via NMPJobContext.from_env(). - - Returns: - Exit code (0 for success, non-zero for failure). - """ - job_ctx = job_ctx or NMPJobContext.from_env() - - sdk_owned = sdk is None - try: - sdk = sdk or get_task_sdk(SERVICE_NAME).with_options(workspace=job_ctx.workspace) - runner = ModelEntityRunner(sdk=sdk, job_ctx=job_ctx) - - config = get_config(job_ctx.config_path) - - # Log only a non-sensitive summary. The full job context carries service - # URLs/identifiers and the config's deployment_config.additional_envs can - # carry deployment secrets, so neither is dumped wholesale. - logger.info( - "Starting model entity task: job_id=%s, name=%s, workspace=%s, fileset=%s/%s, deployment_configured=%s", - job_ctx.job_id, - config.name, - config.workspace, - config.fileset.workspace or job_ctx.workspace, - config.fileset.name, - config.deployment_config is not None, - ) - logger.info(f"NeMo Platform service URL: {sdk.base_url}") - - result, deploy_target = runner.create_model_entity(config) - logger.info(f"Model entity creation complete: {result}") - - runner.launch_model(config, deploy_target) - return 0 - - except ModelEntityCreationError as e: - logger.exception(f"Model entity creation failed: {e}") - return 1 - except Exception as e: - logger.exception(f"Model entity task failed: {e}") - return 1 - finally: - if sdk_owned and sdk is not None: - sdk.close() diff --git a/services/rl/tests/test_compiler.py b/services/rl/tests/test_compiler.py index ad5fa4c9c9..41f0b46d18 100644 --- a/services/rl/tests/test_compiler.py +++ b/services/rl/tests/test_compiler.py @@ -198,12 +198,31 @@ async def test_compiler_emits_four_steps(monkeypatch: pytest.MonkeyPatch, mock_s names = [s["name"] for s in steps] assert names == ["model-and-dataset-download", "dpo-training", "model-upload", "model-entity-creation"] - # CPU task steps share the lighter tasks image; the GPU step uses the training image. - assert "nmp-rl-tasks" in _container(steps[0])["image"] + # CPU task steps share the lighter customizer-tasks image; the GPU step uses the training image. + assert "nmp-customizer-tasks" in _container(steps[0])["image"] assert "nmp-rl-training" in _container(steps[1])["image"] - assert "nmp-rl-tasks" in _container(steps[2])["image"] - assert _container(steps[0])["command"] == ["-m", "nmp.rl.tasks.file_io"] - assert _container(steps[3])["command"] == ["-m", "nmp.rl.tasks.model_entity"] + assert "nmp-customizer-tasks" in _container(steps[2])["image"] + assert _container(steps[0])["command"] == [ + "-m", + "nmp.customization_common.tasks.file_io", + "--service-source", + "rl", + "--service-name", + "rl", + ] + assert _container(steps[3])["command"] == [ + "-m", + "nmp.customization_common.tasks.model_entity", + "--service-name", + "rl", + ] + + upload_meta = steps[2]["config"]["upload"][0]["metadata"] + assert upload_meta == { + "model": "default/base-model", + "finetuning_type": "all_weights", + "output_type": "model", + } @pytest.mark.asyncio diff --git a/services/unsloth/README.md b/services/unsloth/README.md index f2e94ba1cf..07aa5bf687 100644 --- a/services/unsloth/README.md +++ b/services/unsloth/README.md @@ -7,9 +7,8 @@ This package owns the heavy code that runs *inside* the platform's GPU container - **Canonical schemas** (`nmp.unsloth.schemas`) — `UnslothJobOutput` and shared sub-shapes consumed by both compile-time and runtime code. - **Training driver** (`nmp.unsloth.tasks.training.backends.unsloth_sft.train_sft`) — runs SFT inside the training container's baked venv (`unsloth` + `torch` + `transformers` + `trl` + `peft` + `bitsandbytes`). Heavy imports are localized to the function body so the parent process can import this module without dragging in the ML stack. - **Container entrypoints**: - - `nmp.unsloth.tasks.file_io` — handles model + dataset download (pre-train) and checkpoint upload (post-train). - - `nmp.unsloth.tasks.training` — runs `train_sft` against the paths the file_io step populated. - - `nmp.unsloth.tasks.model_entity` — registers the output model entity / adapter. + - `nmp.customization_common.tasks.file_io` / `model_entity` — shared CPU steps (run from `nmp-customizer-tasks`; compiler passes `--service-source unsloth --service-name unsloth`). + - `nmp.unsloth.tasks.training` — runs `train_sft` against the paths the file_io step populated (`nmp-unsloth-training` image). - **Compile glue** (`nmp.unsloth.compile.platform_job_config_compiler`) — turns a canonical `UnslothJobOutput` into a 4-step `PlatformJobSpec`. Invoked by the plugin's `UnslothJob.compile`. The thin contributor wrapper that registers Unsloth with the customization hub lives in `plugins/nemo-unsloth/`. That plugin owns submitter-facing schema (`UnslothJobInput`), the `UnslothContributor`, the `UnslothJob` lifecycle (`to_spec` + `compile`), the SDK shapes, and CLI overrides (`submit` reshaped, `run` disabled). @@ -41,20 +40,22 @@ services/unsloth/ │ ├── compiler.py # GPU training PlatformJobStep │ └── schemas.py # TrainingStepConfig └── tasks/ - ├── file_io/__main__.py + run.py - ├── model_entity/__main__.py + run.py └── training/ ├── __main__.py (entrypoint: reads step config, calls train_sft) └── backends/ └── unsloth_sft.py (train_sft) ``` +CPU `file_io` / `model_entity` runners live in `packages/nmp_customization_common` +(`nmp.customization_common.tasks.*`) and ship in the shared `nmp-customizer-tasks` +image — not under `services/unsloth/src/nmp/unsloth/tasks/`. + ## Why a service package, not just a plugin module? Two reasons: -1. **Container-process boundary.** The training driver and the file_io/model_entity tasks run inside containers built from this package. The plugin (compile-time) and the containers (runtime) need to share schemas (`UnslothJobOutput`, `FileIOTaskConfig`, `ModelEntityTaskConfig`, `TrainingStepConfig`) — co-locating them with the runtime code avoids a circular dep where the plugin owns canonical schemas the container needs to import. -2. **Image isolation.** The plugin process stays lightweight (no `unsloth` / `torch`). Heavy ML deps are installed only inside `nmp-unsloth-training` (see `docker/`). +1. **Container-process boundary.** The training driver runs inside `nmp-unsloth-training`. CPU tasks run from `nmp-customizer-tasks`. The plugin (compile-time) and containers (runtime) share schemas (`UnslothJobOutput`, `FileIOTaskConfig`, `ModelEntityTaskConfig`, `TrainingStepConfig`) via `nmp-customization-common` and this package. +2. **Image isolation.** The plugin process stays lightweight (no `unsloth` / `torch`). Heavy ML deps are installed only inside `nmp-unsloth-training` (see `docker/unsloth/README.md`). CPU steps use the lighter `nmp-customizer-tasks` image. ## Status diff --git a/services/unsloth/pyproject.toml b/services/unsloth/pyproject.toml index c208973100..5eeb307fd0 100644 --- a/services/unsloth/pyproject.toml +++ b/services/unsloth/pyproject.toml @@ -38,8 +38,6 @@ unsloth = ["unsloth[huggingface]"] [project.scripts] # Container entrypoints. Each script matches the automodel name pattern for parity. nmp-unsloth-training = "nmp.unsloth.tasks.training.__main__:main" -nmp-unsloth-file-io = "nmp.unsloth.tasks.file_io.run:run" -nmp-unsloth-model-entity = "nmp.unsloth.tasks.model_entity.__main__:run" [build-system] requires = ["hatchling"] diff --git a/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py b/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py index 5ce588af7b..2743974f4a 100644 --- a/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py +++ b/services/unsloth/src/nmp/unsloth/app/jobs/compiler.py @@ -41,6 +41,7 @@ ) from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig from nmp.customization_common.service.platform_client import fetch_model_entity +from nmp.customization_common.tasks.file_io_metadata import build_output_metadata from nmp.unsloth.app.constants import ( DEFAULT_DATASET_PATH, DEFAULT_MODEL_PATH, @@ -50,7 +51,12 @@ from nmp.unsloth.app.jobs.training.compiler import compile_training_step from nmp.unsloth.config import config from nmp.unsloth.entities.values import FinetuningType -from nmp.unsloth.images import UNSLOTH_PYTHON_ENTRYPOINT, get_tasks_image +from nmp.unsloth.images import ( + FILE_IO_TASK_COMMAND, + MODEL_ENTITY_TASK_COMMAND, + UNSLOTH_PYTHON_ENTRYPOINT, + get_tasks_image, +) from nmp.unsloth.schemas import UnslothJobOutput logger = logging.getLogger(__name__) @@ -165,7 +171,7 @@ def _build_file_download_config( return FileIOTaskConfig(download=downloads) -def _build_file_upload_config(output_fileset_name: str) -> FileIOTaskConfig: +def _build_file_upload_config(job_spec: UnslothJobOutput) -> FileIOTaskConfig: """Compile the upload step. ``workspace=None`` tells the file_io task to use the job's workspace @@ -175,7 +181,13 @@ def _build_file_upload_config(output_fileset_name: str) -> FileIOTaskConfig: upload=[ UploadItem( src=DEFAULT_OUTPUT_MODEL_PATH, - dest=FileSetRef(workspace=None, name=output_fileset_name), + dest=FileSetRef(workspace=None, name=job_spec.output.fileset), + metadata=build_output_metadata( + model=job_spec.model.name, + finetuning_type=job_spec.training.finetuning_type, + save_method=job_spec.output.save_method, + output_type=job_spec.output.type, + ), ), ], ) @@ -230,7 +242,7 @@ async def platform_job_config_compiler( validation_dataset_path = _resolve_validation_dataset_path(job_spec, workspace=workspace) download_config = _build_file_download_config(job_spec, me, workspace=workspace) - upload_config = _build_file_upload_config(job_spec.output.fileset) + upload_config = _build_file_upload_config(job_spec) model_entity_config = _build_model_entity_config( workspace, job_spec, @@ -245,7 +257,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=UNSLOTH_PYTHON_ENTRYPOINT, - command=["-m", "nmp.unsloth.tasks.file_io"], + command=FILE_IO_TASK_COMMAND, ), resources=cpu_resources, ), @@ -265,7 +277,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=UNSLOTH_PYTHON_ENTRYPOINT, - command=["-m", "nmp.unsloth.tasks.file_io"], + command=FILE_IO_TASK_COMMAND, ), resources=cpu_resources, ), @@ -279,7 +291,7 @@ async def platform_job_config_compiler( container=ContainerSpec( image=get_tasks_image(), entrypoint=UNSLOTH_PYTHON_ENTRYPOINT, - command=["-m", "nmp.unsloth.tasks.model_entity"], + command=MODEL_ENTITY_TASK_COMMAND, ), resources=cpu_resources, ), diff --git a/services/unsloth/src/nmp/unsloth/images.py b/services/unsloth/src/nmp/unsloth/images.py index b96c189b11..7ca62b60cd 100644 --- a/services/unsloth/src/nmp/unsloth/images.py +++ b/services/unsloth/src/nmp/unsloth/images.py @@ -1,31 +1,36 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Docker image resolution for nmp-unsloth job steps. - -Mirrors :mod:`nmp.automodel.images`. Consumed by the compiler in -:mod:`nmp.unsloth.app.jobs.compiler` (and its training sub-compiler) to -stamp the right image refs onto each container step. - -Unsloth ships a **single** image, ``nmp-unsloth-training``, used by all four -steps (file_io, model_entity, training) — the CPU task steps reuse the training -image rather than a separate ``nmp-unsloth-tasks`` build. Override the whole -image via ``NMP_UNSLOTH_TRAINING_IMAGE``. -""" +"""Docker image resolution for nmp-unsloth job steps.""" from __future__ import annotations -from nmp.customization_common.service.images import resolve_qualified_image +from nmp.customization_common.service.images import ( + CUSTOMIZER_PYTHON_ENTRYPOINT, + get_customizer_tasks_image, + resolve_qualified_image, +) from nmp.unsloth.config import config BASE_IMAGE_NAME = "nmp-unsloth-base" -TASKS_IMAGE_NAME = "nmp-unsloth-tasks" TRAINING_IMAGE_NAME = "nmp-unsloth-training" -# Must match ENTRYPOINT in Dockerfile.nmp-unsloth-{tasks,training}. -# Job specs must set this explicitly: Docker API ``create()`` replaces the -# image entrypoint when the platform passes ``entrypoint=[]``. -UNSLOTH_PYTHON_ENTRYPOINT = ["/opt/venv/bin/python"] +UNSLOTH_PYTHON_ENTRYPOINT = CUSTOMIZER_PYTHON_ENTRYPOINT + +FILE_IO_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.file_io", + "--service-source", + "unsloth", + "--service-name", + "unsloth", +] +MODEL_ENTITY_TASK_COMMAND = [ + "-m", + "nmp.customization_common.tasks.model_entity", + "--service-name", + "unsloth", +] def get_unsloth_qualified_image(name: str, override: str | None = None) -> str: @@ -34,12 +39,8 @@ def get_unsloth_qualified_image(name: str, override: str | None = None) -> str: def get_tasks_image() -> str: - """CPU task steps (file_io, model_entity). - - Unsloth ships a single image, so the CPU task steps reuse the - ``nmp-unsloth-training`` image rather than a separate tasks image. - """ - return get_training_image() + """CPU task steps (file_io, model_entity) — shared ``nmp-customizer-tasks`` image.""" + return get_customizer_tasks_image(backend_override=config.tasks_image, image_registry=config.image_registry) def get_training_image() -> str: diff --git a/services/unsloth/src/nmp/unsloth/tasks/file_io/__init__.py b/services/unsloth/src/nmp/unsloth/tasks/file_io/__init__.py deleted file mode 100644 index 182d33c009..0000000000 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""File I/O task for unsloth customization jobs.""" - -from nmp.unsloth.tasks.file_io.run import run - -__all__ = ["run"] diff --git a/services/unsloth/src/nmp/unsloth/tasks/file_io/__main__.py b/services/unsloth/src/nmp/unsloth/tasks/file_io/__main__.py deleted file mode 100644 index cba1d22b48..0000000000 --- a/services/unsloth/src/nmp/unsloth/tasks/file_io/__main__.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import sys - -from nmp.unsloth.tasks.file_io.run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/unsloth/src/nmp/unsloth/tasks/model_entity/__init__.py b/services/unsloth/src/nmp/unsloth/tasks/model_entity/__init__.py deleted file mode 100644 index eae78497cc..0000000000 --- a/services/unsloth/src/nmp/unsloth/tasks/model_entity/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Model entity task for creating model entities after unsloth customization.""" - -from nmp.unsloth.tasks.model_entity.run import run - -__all__ = ["run"] diff --git a/services/unsloth/src/nmp/unsloth/tasks/model_entity/__main__.py b/services/unsloth/src/nmp/unsloth/tasks/model_entity/__main__.py deleted file mode 100644 index eb393313c9..0000000000 --- a/services/unsloth/src/nmp/unsloth/tasks/model_entity/__main__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Entry point for model_entity task. - -Usage: - python -m nmp.unsloth.tasks.model_entity -""" - -import sys - -from .run import run - -if __name__ == "__main__": - sys.exit(run()) diff --git a/services/unsloth/tests/test_compiler_validation_path.py b/services/unsloth/tests/test_compiler_validation_path.py index f48c002d27..51fe9fe0b8 100644 --- a/services/unsloth/tests/test_compiler_validation_path.py +++ b/services/unsloth/tests/test_compiler_validation_path.py @@ -97,3 +97,34 @@ async def test_training_step_gets_separate_validation_path_for_different_fileset download = next(s for s in job["steps"] if s["name"] == "model-and-dataset-download") assert len(download["config"]["download"]) == 3 + + +@pytest.mark.asyncio +async def test_upload_step_stamps_output_metadata() -> None: + from nmp.unsloth.app.jobs import compiler as compiler_mod + + original_fetch = compiler_mod.fetch_model_entity + compiler_mod.fetch_model_entity = AsyncMock( + return_value=types.SimpleNamespace( + workspace="default", + name="qwen3-1.7b", + fileset="default/qwen3-1.7b", + trust_remote_code=False, + ), + ) + try: + job = await platform_job_config_compiler( + workspace="default", + job_spec=_spec(validation_path=None), + sdk=MagicMock(), + ) + finally: + compiler_mod.fetch_model_entity = original_fetch + + upload = next(s for s in job["steps"] if s["name"] == "model-upload") + assert upload["config"]["upload"][0]["metadata"] == { + "model": "default/qwen3-1.7b", + "finetuning_type": "lora", + "save_method": "lora", + "output_type": "adapter", + } diff --git a/services/unsloth/tests/test_images.py b/services/unsloth/tests/test_images.py index 1b1b640566..0314937630 100644 --- a/services/unsloth/tests/test_images.py +++ b/services/unsloth/tests/test_images.py @@ -7,9 +7,9 @@ import nmp.customization_common.service.images as shared_images import nmp.unsloth.images as unsloth_images import pytest +from nmp.customization_common.service.images import CUSTOMIZER_TASKS_IMAGE_NAME from nmp.unsloth.config import UnslothConfig from nmp.unsloth.images import ( - TASKS_IMAGE_NAME, TRAINING_IMAGE_NAME, get_tasks_image, get_training_image, @@ -32,9 +32,10 @@ def test_default_unsloth_images_use_platform_registry(monkeypatch, platform_conf tasks = get_tasks_image() expected_training = f"{platform_config.image_registry}/{TRAINING_IMAGE_NAME}:{platform_config.image_tag}" + expected_tasks = f"{platform_config.image_registry}/{CUSTOMIZER_TASKS_IMAGE_NAME}:{platform_config.image_tag}" assert training == expected_training - assert tasks == expected_training # unsloth ships one image: tasks reuse the training image - assert TASKS_IMAGE_NAME.count("/") == 0 # single repo segment, no nested paths + assert tasks == expected_tasks + assert CUSTOMIZER_TASKS_IMAGE_NAME.count("/") == 0 # single repo segment, no nested paths def test_unsloth_image_registry_override(monkeypatch, platform_config): @@ -50,16 +51,15 @@ def test_unsloth_image_registry_override(monkeypatch, platform_config): ) -def test_unsloth_training_image_override_used_for_all_steps(monkeypatch, platform_config): - # A single training-image override drives both the training step and the - # CPU task steps, since unsloth reuses the training image for tasks. +def test_unsloth_tasks_image_override(monkeypatch, platform_config): monkeypatch.setattr( unsloth_images, "config", UnslothConfig( + tasks_image="my-registry/nemo-platform-dev/nmp-customizer-tasks:dev", training_image="my-registry/nemo-platform-dev/nmp-unsloth-training:dev", ), ) + assert get_tasks_image() == "my-registry/nemo-platform-dev/nmp-customizer-tasks:dev" assert get_training_image() == "my-registry/nemo-platform-dev/nmp-unsloth-training:dev" - assert get_tasks_image() == "my-registry/nemo-platform-dev/nmp-unsloth-training:dev" diff --git a/services/unsloth/tests/test_model_entity.py b/services/unsloth/tests/test_model_entity.py index 5a817cceec..d5231e19d3 100644 --- a/services/unsloth/tests/test_model_entity.py +++ b/services/unsloth/tests/test_model_entity.py @@ -38,7 +38,7 @@ def _make_job_ctx(workspace: str = "default"): def _make_runner(sdk): - from nmp.unsloth.tasks.model_entity.run import ModelEntityRunner + from nmp.customization_common.tasks.model_entity.run import ModelEntityRunner return ModelEntityRunner(sdk=sdk, job_ctx=_make_job_ctx()) @@ -58,7 +58,7 @@ def _raise_runner_conflict() -> None: """ import sys - run_mod = sys.modules["nmp.unsloth.tasks.model_entity.run"] + run_mod = sys.modules["nmp.customization_common.tasks.model_entity.run"] raise run_mod.ConflictError.__new__(run_mod.ConflictError, "already exists") @@ -78,19 +78,19 @@ def _model_entity(*, workspace: str = "default", name: str = "base", spec: objec class TestSanitizeName: def test_lowercases_and_replaces_invalid_chars(self) -> None: - from nmp.unsloth.tasks.model_entity.run import sanitize_name + from nmp.customization_common.tasks.model_entity.run import sanitize_name assert sanitize_name("sft-cfg", "Qwen/Qwen3-0.6B") == "sft-cfg-qwen-qwen3-0.6b" def test_collapses_consecutive_hyphens(self) -> None: - from nmp.unsloth.tasks.model_entity.run import sanitize_name + from nmp.customization_common.tasks.model_entity.run import sanitize_name # "/" is not in the allowed set, so each "/" becomes "-", then # the consecutive-hyphen collapse fires. assert sanitize_name("p", "a//b") == "p-a-b" def test_caps_length_below_60_and_strips_trailing_hyphen(self) -> None: - from nmp.unsloth.tasks.model_entity.run import sanitize_name + from nmp.customization_common.tasks.model_entity.run import sanitize_name # 59-char limit accounts for the "-v1" the backend appends. long_name = "a" * 80 @@ -105,7 +105,7 @@ def test_caps_length_below_60_and_strips_trailing_hyphen(self) -> None: class TestCreateFullEntity: - @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + @patch("nmp.customization_common.tasks.model_entity.run.client_from_platform") def test_creates_model_entity_for_full_sft(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig @@ -133,7 +133,7 @@ def test_creates_model_entity_for_full_sft(self, mock_cfp) -> None: assert deploy_target is new_me assert result is not None - @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + @patch("nmp.customization_common.tasks.model_entity.run.client_from_platform") def test_conflict_falls_back_to_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig @@ -160,7 +160,7 @@ def test_conflict_falls_back_to_update(self, mock_cfp) -> None: assert update_call.kwargs["name"] == "trained-model" assert update_call.kwargs["workspace"] == "default" - @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + @patch("nmp.customization_common.tasks.model_entity.run.client_from_platform") def test_missing_fileset_raises_creation_error(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityCreationError, ModelEntityTaskConfig @@ -187,7 +187,7 @@ def test_missing_fileset_raises_creation_error(self, mock_cfp) -> None: class TestCreateAdapter: - @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + @patch("nmp.customization_common.tasks.model_entity.run.client_from_platform") def test_creates_adapter_for_lora(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig @@ -213,7 +213,7 @@ def test_creates_adapter_for_lora(self, mock_cfp) -> None: sdk.models.adapters.create.assert_called_once() assert deploy_target is base_me - @patch("nmp.unsloth.tasks.model_entity.run.client_from_platform") + @patch("nmp.customization_common.tasks.model_entity.run.client_from_platform") def test_adapter_conflict_falls_back_to_update(self, mock_cfp) -> None: from nmp.customization_common.schemas.file_io import FileSetRef from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig, PEFTConfig diff --git a/tests/agentic-use/customizer-lora-job-cli/environment/setup-env.sh b/tests/agentic-use/customizer-lora-job-cli/environment/setup-env.sh index a77b73f944..7fa0b973af 100644 --- a/tests/agentic-use/customizer-lora-job-cli/environment/setup-env.sh +++ b/tests/agentic-use/customizer-lora-job-cli/environment/setup-env.sh @@ -6,7 +6,7 @@ source /app/image-env.sh echo '=== Pre-pulling nmp-automodel job images ===' if command -v docker &> /dev/null && [ -S /var/run/docker.sock ]; then TRAINING_IMAGE="${NMP_IMAGE_REGISTRY}/nmp-automodel-training:${NMP_IMAGE_TAG}" - TASKS_IMAGE="${NMP_IMAGE_REGISTRY}/nmp-automodel-tasks:${NMP_IMAGE_TAG}" + TASKS_IMAGE="${NMP_IMAGE_REGISTRY}/nmp-customizer-tasks:${NMP_IMAGE_TAG}" docker pull "$TRAINING_IMAGE" && echo "Pulled ${TRAINING_IMAGE}" || \ echo "WARNING: Failed to pull ${TRAINING_IMAGE} — bake and push with BASE_TAG_AUTOMODEL=${BASE_TAG_AUTOMODEL}" diff --git a/tests/agentic-use/customizer-lora-job-cli/instruction.md b/tests/agentic-use/customizer-lora-job-cli/instruction.md index 85942dc4da..d3f1966b60 100644 --- a/tests/agentic-use/customizer-lora-job-cli/instruction.md +++ b/tests/agentic-use/customizer-lora-job-cli/instruction.md @@ -87,7 +87,7 @@ The task is complete when: ## Notes -- LoRA/SFT training uses the **automodel** contributor (`nmp-automodel-training` / `nmp-automodel-tasks` images), not the legacy customizer automodel path +- LoRA/SFT training uses the **automodel** contributor (`nmp-automodel-training` / `nmp-customizer-tasks` images), not the legacy customizer automodel path - Model reference format: `workspace/model-entity-name` (e.g., `lora-training-workspace/smollm-135m`) - Dataset reference format in job JSON: `workspace/fileset-name` inside `dataset.training` - Jobs may take a few minutes depending on dataset size and GPU availability diff --git a/tests/smoke_gpu/conftest.py b/tests/smoke_gpu/conftest.py index 734a72dcf7..3009a990a2 100644 --- a/tests/smoke_gpu/conftest.py +++ b/tests/smoke_gpu/conftest.py @@ -5,7 +5,7 @@ def pytest_configure(config): config.addinivalue_line("markers", "smoke_gpu_tasks: Import smoke tests for the nmp-gpu-tasks image") config.addinivalue_line( - "markers", "smoke_nmp_automodel_tasks: Import smoke tests for the nmp/automodel-tasks image" + "markers", "smoke_nmp_customizer_tasks: Import smoke tests for the nmp-customizer-tasks image" ) config.addinivalue_line( "markers", "smoke_nmp_automodel_training: Import smoke tests for the nmp/automodel-training image" diff --git a/tests/smoke_gpu/test_customizer_automodel.py b/tests/smoke_gpu/test_customizer_automodel.py index e72783e8b0..339053e5d6 100644 --- a/tests/smoke_gpu/test_customizer_automodel.py +++ b/tests/smoke_gpu/test_customizer_automodel.py @@ -19,43 +19,31 @@ import pytest -@pytest.mark.smoke_nmp_automodel_tasks @pytest.mark.smoke_nmp_automodel_training def test_torch_importable(): import torch # noqa: F401 -@pytest.mark.smoke_nmp_automodel_tasks @pytest.mark.smoke_nmp_automodel_training def test_transformers_importable(): import transformers # noqa: F401 -@pytest.mark.smoke_nmp_automodel_tasks @pytest.mark.smoke_nmp_automodel_training def test_mamba_ssm_importable(): import mamba_ssm # noqa: F401 -@pytest.mark.smoke_nmp_automodel_tasks @pytest.mark.smoke_nmp_automodel_training def test_causal_conv1d_importable(): import causal_conv1d # noqa: F401 -@pytest.mark.smoke_nmp_automodel_tasks @pytest.mark.smoke_nmp_automodel_training def test_bitsandbytes_importable(): import bitsandbytes # noqa: F401 -@pytest.mark.smoke_nmp_automodel_tasks -def test_nmp_automodel_tasks_importable(): - from nmp.automodel.tasks import file_io # noqa: F401 - from nmp.automodel.tasks.model_entity import __main__ as model_entity_main # noqa: F401 - from nmp.core.models.tasks.model_spec import __main__ as model_spec_main # noqa: F401 - - @pytest.mark.smoke_nmp_automodel_training def test_nmp_automodel_training_importable(): import nemo_automodel # noqa: F401 diff --git a/tests/smoke_gpu/test_customizer_tasks.py b/tests/smoke_gpu/test_customizer_tasks.py new file mode 100644 index 0000000000..2b53ea84d4 --- /dev/null +++ b/tests/smoke_gpu/test_customizer_tasks.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""nmp-customizer-tasks image import smoke tests. + +Built as part of the docker-bake.hcl bake group (smoke-test stage) and run +on a CPU runner — no GPU hardware required. + +Two failure classes are caught at .so load time, before any GPU device is touched: + + ModuleNotFoundError — package missing from the image (e.g. excluded from + a tar layer without a compensating COPY command) + + ImportError — CUDA extension .so has an undefined symbol; the wheel + was compiled against a different PyTorch version than + the one installed (ABI mismatch) +""" + +import pytest + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_torch_importable(): + import torch # noqa: F401 + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_transformers_importable(): + import transformers # noqa: F401 + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_accelerate_importable(): + import accelerate # noqa: F401 + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_mamba_ssm_importable(): + import mamba_ssm # noqa: F401 + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_causal_conv1d_importable(): + import causal_conv1d # noqa: F401 + + +@pytest.mark.smoke_nmp_customizer_tasks +def test_nmp_customizer_tasks_importable(): + from nmp.core.models.sidecars.adapters.main import run as lora_sidecar_run # noqa: F401 + from nmp.core.models.tasks.model_spec import __main__ as model_spec_main # noqa: F401 + from nmp.customization_common.tasks import file_io # noqa: F401 + from nmp.customization_common.tasks.file_io import __main__ as file_io_main # noqa: F401 + from nmp.customization_common.tasks.model_entity import __main__ as model_entity_main # noqa: F401 diff --git a/uv.lock b/uv.lock index 907be070f0..ee0e10ca9b 100644 --- a/uv.lock +++ b/uv.lock @@ -4415,8 +4415,8 @@ source = { editable = "plugins/nemo-deployments" } dependencies = [ { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "httpx", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nemo-platform", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] @@ -4444,8 +4444,8 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.115" }, { name = "httpx", specifier = ">=0.27" }, { name = "kubernetes", marker = "extra == 'k8s'", specifier = ">=30.1.0" }, - { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nemo-platform-sdk", editable = "sdk/python/nemo-platform" }, { name = "pydantic", specifier = ">=2.10.6" }, ] provides-extras = ["docker", "k8s"]