fix: split training dependency profiles (#1459)

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

docker/Dockerfile.nmp-customizer-tasks

@@ -79,6 +79,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # NGC system-site copies inherited from nvcr.io/nvidia/pytorch:26.07-py3.
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
+    --reinstall-package wandb \
     --overrides /app/docker/customizer/preserve_base_torch.txt \
     "aiohttp>=3.14.1,<4" \
     "grpcio>=1.81.1,<2" \
@@ -89,7 +90,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "tornado>=6.5.7,<7" \
     "urllib3>=2.7.0,<3" \
     "GitPython>=3.1.58,<4" \
-    "wandb==0.28.2"
+    "wandb>=0.28.2"
 
 # CVE cleanup for unused/stale packages from the NGC base image. The task entry
 # points import from /opt/venv; scanners still report the shadowed system copies.

docker/Dockerfile.nmp-unsloth-training

@@ -26,7 +26,7 @@
 #       also missing (common on newer CUDA stacks), logging "FA2 = False /
 #       Xformers = None".
 #   2. Editable installs of the platform glue (nemo-platform SDK, plugin,
-#      nmp-common, nmp-unsloth) from the in-repo workspace slice.
+#      nmp-common, nmp-unsloth[integrations]) from the in-repo workspace slice.
 #
 # Publish target: nmp-unsloth-training
 # Default tag:    `local` (override via BAKE_TAG at build time).
@@ -164,15 +164,17 @@ RUN chown -R ${USER_UID}:${USER_GID} /app
 
 # Step 2: install the platform glue editably. No [unsloth] extra here — that's
 # what we just installed in step 1, and re-triggering it would force the
-# resolver to re-evaluate the whole HF stack.
+# resolver to re-evaluate the whole HF stack. The [integrations] extra is
+# training-runtime only; nmp-api installs nmp-unsloth without extras.
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
+    --reinstall-package wandb \
     --overrides /opt/docker/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/unsloth"
+    -e "/app/services/unsloth[integrations]"
 
 # hf-transfer can pull huggingface-hub 1.x — install then re-pin hub + transformers.
 RUN --mount=type=cache,target=/root/.cache/uv \

docker/Dockerfile.safe-synthesizer-tasks

@@ -122,7 +122,7 @@ RUN printf '%s\n' \
         "${SAFE_SYNTHESIZER_RUNTIME_PACKAGE}" \
         > /tmp/safe-synthesizer-runtime.txt && \
     printf '%s\n' \
-        wandb==0.28.2 \
+        'wandb>=0.28.2' \
         'aiohttp>=3.14.3,<4' \
         'cryptography>=50.0.0,<51' \
         'pyarrow>=23.0.1,<24' \

docker/automodel/Dockerfile.nmp-automodel-base

@@ -106,6 +106,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     rm -rf /opt/venv/lib/python*/site-packages/vllm \
     /opt/venv/lib/python*/site-packages/vllm-*.dist-info && \
     uv pip install \
+    --reinstall-package wandb \
     "aiohttp>=3.13.3,<4" \
     "black>=26.3.1,<27" \
     "GitPython>=3.1.58,<4" \
@@ -125,7 +126,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "urllib3>=2.7.0,<3" \
     "mlflow-skinny>=3.13.0,<3.14.0" \
     "grpcio>=1.81.1,<2" \
-    "wandb==0.28.2"
+    "wandb>=0.28.2"
 
 ARG NMP_COLLECT_SOURCES=0
 ENV NMP_COLLECT_SOURCES=${NMP_COLLECT_SOURCES}

docker/automodel/no_override_requirements.txt

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

docker/rl/Dockerfile.nmp-rl-base

@@ -260,7 +260,7 @@ uv pip install --python /opt/nemo_rl_venv/bin/python "opensandbox>=0.1.9" "tenac
 
 # W&B ships a scanner-visible Go binary. Keep it above the fixed go-git/grpc
 # and Go stdlib floors even when NeMo-RL's lock has an older wheel.
-uv pip install --python /opt/nemo_rl_venv/bin/python "wandb==0.28.2"
+uv pip install --python /opt/nemo_rl_venv/bin/python "wandb>=0.28.2"
 
 # NeMo-Gym in the BASE venv, for the DRIVER - not a duplicate of the nemo_gym extra above.
 # Sandboxed (mode B) GRPO imports nemo_rl.environments.sandbox in the driver process:
@@ -440,7 +440,7 @@ fi
 for py in /opt/ray_venvs/*/bin/python /opt/gym_venvs/*/bin/python; do
     [ -x "${py}" ] || continue
     if "${py}" -c 'import importlib.util, sys; sys.exit(0 if importlib.util.find_spec("wandb") else 1)'; then
-        uv pip install --python "${py}" "wandb==0.28.2"
+        uv pip install --python "${py}" --upgrade-package wandb "wandb>=0.28.2"
     fi
 done
 for d in "${UV_CACHE_DIR}"/archive-v0/*/; do

docker/rl/Dockerfile.nmp-rl-training

@@ -42,9 +42,11 @@ RUN ln -sfn /opt/nemo_rl_venv /opt/venv
 RUN UV_CACHE_DIR=/root/.cache/uv \
     uv-glue pip freeze --python /opt/nemo_rl_venv/bin/python --exclude-editable > /tmp/base-python-freeze.txt
 
-# Editable platform glue WITH dependency resolution. The from-source base ships clean package
-# metadata, so uv can resolve the glue's real deps (typer, sniffio, ngcsdk, opentelemetry-*,
-# sqlalchemy, kubernetes, fastapi, ...) instead of us pinning them by hand.
+# Editable platform glue WITH dependency resolution. The [integrations] extra is
+# training-runtime only; nmp-api installs nmp-rl without extras. The from-source
+# base ships clean package metadata, so uv can resolve the glue's real deps
+# (typer, sniffio, ngcsdk, opentelemetry-*, sqlalchemy, kubernetes, fastapi, ...)
+# instead of us pinning them by hand.
 # To stop uv moving the pinned RL / vLLM / mcore / Gym stack while it does that, freeze the base venv
 # first and pass it as --overrides (base versions win over whatever the glue requests). Packages the
 # base doesn't have are still installed fresh. `| grep '==' ` drops git/URL installs (TE, flash-attn,
@@ -64,7 +66,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         -e /app/packages/nemo_platform_plugin \
         -e /app/packages/nmp_common \
         -e /app/packages/nmp_customization_common \
-        -e /app/services/rl
+        -e "/app/services/rl[integrations]"
 
 ARG NMP_COLLECT_SOURCES=0
 ENV NMP_COLLECT_SOURCES=${NMP_COLLECT_SOURCES}

docker/unsloth/no_override_requirements.txt

@@ -11,4 +11,4 @@ accelerate; sys_platform == 'never'
 safetensors; sys_platform == 'never'
 
 # Keep W&B's bundled wandb-core binary on a build with patched go-git/go-billy.
-wandb==0.28.2
+wandb>=0.28.2

e2e/test_nemo_agents_execute_job.py

@@ -239,7 +239,6 @@ def test_fabric_agent_invocation_job_runs_and_saves_results(sdk: NeMoPlatform, w
         artifacts_members = _tar_member_names(artifacts)
         assert _tar_contains(artifacts_members, "adapter-invocation.json")
         assert _tar_contains(artifacts_members, "stdout.txt")
-        assert _tar_contains(artifacts_members, "stderr.txt")
 
         # The run result captures platform-normalized Fabric RunResult details
         run_result = json.loads(_download_execute_job_result(sdk, workspace, job_name, "fabric_run_result"))