fix(customizer): repair container execution regressions

packages/nmp_customization_common/src/nmp/customization_common/service/context.py

@@ -14,20 +14,20 @@
 
 from nmp.common.entities.constants import DEFAULT_WORKSPACE
 from nmp.common.jobs.constants import (
-    DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH,
     NEMO_JOB_ATTEMPT_ID_ENVVAR,
     NEMO_JOB_ID_ENVVAR,
-    NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR,
     NEMO_JOB_STEP_ENVVAR,
     NEMO_JOB_TASK_ENVVAR,
     NEMO_JOB_WORKSPACE_ENVVAR,
-    PERSISTENT_JOB_STORAGE_PATH_ENVVAR,
 )
 from nmp.customization_common.service.constants import (
-    DEFAULT_JOB_STORAGE_PATH,
     NMP_FILES_URL_ENVVAR,
     NMP_JOBS_URL_ENVVAR,
 )
+from nmp.customization_common.service.path_utils import (
+    get_job_step_config_path_from_env,
+    get_job_storage_path_from_env,
+)
 
 DEFAULT_JOB_ID = "unknown-job-id"
 DEFAULT_ATTEMPT_ID = "attempt-0"
@@ -97,8 +97,6 @@ def from_env(cls) -> Self:
             task=os.environ.get(NEMO_JOB_TASK_ENVVAR, DEFAULT_TASK),
             jobs_url=os.environ.get(NMP_JOBS_URL_ENVVAR),
             files_url=os.environ.get(NMP_FILES_URL_ENVVAR),
-            storage_path=Path(os.environ.get(PERSISTENT_JOB_STORAGE_PATH_ENVVAR, DEFAULT_JOB_STORAGE_PATH)),
-            config_path=Path(
-                os.environ.get(NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR, DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH)
-            ),
+            storage_path=get_job_storage_path_from_env(),
+            config_path=get_job_step_config_path_from_env(),
         )

packages/nmp_customization_common/src/nmp/customization_common/service/path_utils.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Path helpers for customization job storage mounts."""
+
+import os
+from pathlib import Path
+
+from nmp.common.jobs.constants import (
+    DEFAULT_JOB_STORAGE_PATH,
+    DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH,
+    NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR,
+    PERSISTENT_JOB_STORAGE_PATH_ENVVAR,
+)
+
+CURRENT_PERSISTENT_JOB_STORAGE_PATH_ENVVAR = "NEMO_JOB_PERSISTENT_JOB_STORAGE_PATH"
+LEGACY_PERSISTENT_JOB_STORAGE_PATH_ENVVARS = ("NMP_JOB_PERSISTENT_JOB_STORAGE_PATH",)
+
+CURRENT_NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR = "NEMO_JOB_STEP_CONFIG_FILE_PATH"
+LEGACY_NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVARS = ("NMP_JOB_STEP_CONFIG_FILE_PATH",)
+
+CURRENT_JOB_STORAGE_PATH = Path("/var/run/scratch/job")
+LEGACY_JOB_STORAGE_PATHS = (Path("/run/scratch/job"),)
+
+
+def _first_env_path(names: tuple[str, ...], default: str) -> Path:
+    for name in names:
+        value = os.environ.get(name)
+        if value:
+            return Path(value)
+    return Path(default)
+
+
+def get_job_storage_path_from_env() -> Path:
+    """Return the mounted persistent job storage path.
+
+    The explicit current env name is checked before the imported constant so a
+    task image still works if another installed package has an older constant.
+    """
+
+    return _first_env_path(
+        (
+            CURRENT_PERSISTENT_JOB_STORAGE_PATH_ENVVAR,
+            PERSISTENT_JOB_STORAGE_PATH_ENVVAR,
+            *LEGACY_PERSISTENT_JOB_STORAGE_PATH_ENVVARS,
+        ),
+        DEFAULT_JOB_STORAGE_PATH,
+    )
+
+
+def get_job_step_config_path_from_env() -> Path:
+    """Return the mounted step config file path."""
+
+    return _first_env_path(
+        (
+            CURRENT_NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR,
+            NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR,
+            *LEGACY_NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVARS,
+        ),
+        DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH,
+    )
+
+
+def _known_job_storage_roots(storage_path: Path) -> tuple[Path, ...]:
+    roots = (
+        storage_path,
+        Path(DEFAULT_JOB_STORAGE_PATH),
+        CURRENT_JOB_STORAGE_PATH,
+        *LEGACY_JOB_STORAGE_PATHS,
+    )
+    return tuple(dict.fromkeys(roots))
+
+
+def remap_job_storage_path(storage_path: Path, user_path: str | Path) -> Path:
+    """Remap known absolute job-storage roots to the mounted storage path.
+
+    Older customizer task configs used ``/run/scratch/job`` while the current
+    Jobs runner mounts persistent storage at ``/var/run/scratch/job``. This
+    keeps absolute paths produced by either side usable when API and task images
+    are briefly out of sync. Unknown absolute paths are returned unchanged so
+    callers can reject them with their normal traversal checks.
+    """
+
+    raw_path = Path(user_path)
+    if not raw_path.is_absolute():
+        return raw_path
+
+    for root in _known_job_storage_roots(storage_path):
+        try:
+            return storage_path / raw_path.relative_to(root)
+        except ValueError:
+            continue
+    return raw_path

packages/nmp_customization_common/src/nmp/customization_common/tasks/file_io_utils.py

@@ -28,6 +28,7 @@
     PathTraversalError,
     ProgressReportError,
 )
+from nmp.customization_common.service.path_utils import remap_job_storage_path
 
 logger = logging.getLogger(__name__)
 
@@ -166,7 +167,10 @@ def validate_safe_path(base_path: Path, user_path: str) -> Path:
         PathTraversalError: If the resolved path would escape base_path.
     """
     resolved_base = base_path.resolve()
-    resolved_path = (base_path / user_path).resolve()
+    candidate_path = remap_job_storage_path(base_path, user_path)
+    if not candidate_path.is_absolute():
+        candidate_path = base_path / candidate_path
+    resolved_path = candidate_path.resolve()
 
     if not resolved_path.is_relative_to(resolved_base):
         raise PathTraversalError(

packages/nmp_customization_common/tests/tasks/test_file_io_utils.py

@@ -3,7 +3,8 @@
 
 from pathlib import Path
 
-from nmp.customization_common.tasks.file_io_utils import list_local_files
+import pytest
+from nmp.customization_common.tasks.file_io_utils import list_local_files, validate_safe_path
 
 
 def test_list_local_files_skips_unreadable_entries(tmp_path: Path) -> None:
@@ -34,3 +35,23 @@ def test_list_local_files_single_file(tmp_path: Path) -> None:
     assert len(files) == 1
     assert files[0].path == "weights.bin"
     assert files[0].size == 3
+
+
+def test_validate_safe_path_remaps_legacy_job_storage_root(tmp_path: Path) -> None:
+    storage = tmp_path / "job"
+    target = storage / "output_model"
+    target.mkdir(parents=True)
+
+    resolved = validate_safe_path(storage, "/run/scratch/job/output_model")
+
+    assert resolved == target.resolve()
+
+
+def test_validate_safe_path_rejects_unknown_absolute_path(tmp_path: Path) -> None:
+    from nmp.customization_common.schemas.file_io import PathTraversalError
+
+    storage = tmp_path / "job"
+    storage.mkdir()
+
+    with pytest.raises(PathTraversalError):
+        validate_safe_path(storage, "/etc/passwd")

services/automodel/src/nmp/automodel/app/jobs/compiler.py

@@ -381,9 +381,17 @@ async def platform_job_config_compiler(
     workspace: str,
     job_spec: CustomizationJobOutput,
     sdk: AsyncNeMoPlatform,
+    *,
+    job_name: str | None = None,
+    profile: str | None = None,
 ) -> PlatformJobSpec:
     """Compile canonical job spec into a four-step PlatformJobSpec."""
+    del job_name  # reserved for future scheduling decisions
     transformed_spec = job_spec
+    if profile is not None and transformed_spec.training.execution_profile is None:
+        transformed_spec = transformed_spec.model_copy(
+            update={"training": transformed_spec.training.model_copy(update={"execution_profile": profile})},
+        )
     logger.info("Compiling Automodel job to PlatformJobSpec: %s", transformed_spec.model_dump_json(indent=2))
 
     try:
@@ -394,6 +402,7 @@ async def platform_job_config_compiler(
     # output is a required field in CustomizationJobOutput
     cpu_resources = _get_cpu_resources()
     base_env = _get_base_environment()
+    task_profile = transformed_spec.training.execution_profile or config.default_training_execution_profile
 
     # Fetch the primary model entity
     me = await fetch_model_entity(transformed_spec.model, workspace, sdk)
@@ -447,17 +456,11 @@ async def platform_job_config_compiler(
     trust_remote_code = me.trust_remote_code or False
     model_entity_config = _build_model_entity_config(workspace, transformed_spec, trust_remote_code)
 
-    cpu_profile = (
-        transformed_spec.training.execution_profile
-        if transformed_spec.training.execution_profile is not None
-        else config.default_training_execution_profile
-    )
-
     steps = [
         # Step 1: Download model and dataset files from Files service
         PlatformJobStep(
             name="model-and-dataset-download",
-            executor=_cpu_tasks_executor(FILE_IO_TASK_COMMAND, cpu_resources, cpu_profile),
+            executor=_cpu_tasks_executor(FILE_IO_TASK_COMMAND, cpu_resources, task_profile),
             environment=base_env,
             config=file_io_download_config.model_dump(mode="json"),
         ),
@@ -471,14 +474,14 @@ async def platform_job_config_compiler(
         # Step 3: Upload customized model
         PlatformJobStep(
             name="model-upload",
-            executor=_cpu_tasks_executor(FILE_IO_TASK_COMMAND, cpu_resources, cpu_profile),
+            executor=_cpu_tasks_executor(FILE_IO_TASK_COMMAND, cpu_resources, task_profile),
             environment=base_env,
             config=file_io_upload_config.model_dump(mode="json"),
         ),
         # Step 4: Create model entity
         PlatformJobStep(
             name="model-entity-creation",
-            executor=_cpu_tasks_executor(MODEL_ENTITY_TASK_COMMAND, cpu_resources, cpu_profile),
+            executor=_cpu_tasks_executor(MODEL_ENTITY_TASK_COMMAND, cpu_resources, task_profile),
             environment=base_env,
             config=model_entity_config.model_dump(mode="json"),
         ),

services/automodel/src/nmp/automodel/compile.py

@@ -25,14 +25,12 @@ async def platform_job_config_compiler(
     """Compile Automodel job spec (plugin or legacy shape) to PlatformJobSpec."""
     if not isinstance(job_spec, CustomizationJobOutput):
         job_spec = automodel_spec_to_compiler_output(job_spec)
-    if profile and job_spec.training.execution_profile is None:
-        job_spec = job_spec.model_copy(
-            update={"training": job_spec.training.model_copy(update={"execution_profile": profile})},
-        )
     return await _compile_canonical(
         workspace,
         job_spec,
         sdk,
+        job_name=job_name,
+        profile=profile,
     )
 
 

services/automodel/src/nmp/automodel/tasks/training/runner.py

@@ -19,6 +19,7 @@
 import yaml
 from nmp.automodel.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME
 from nmp.customization_common.service.context import NMPJobContext
+from nmp.customization_common.service.path_utils import remap_job_storage_path
 
 from .backends.backend import AUTOMODEL_CONFIG_FILENAME, AutomodelBackend
 from .distributed import DistributedContext
@@ -116,7 +117,36 @@ def _get_barrier_dir(self) -> Path:
 
     def _load_config(self, config_path: Path) -> TrainingStepConfig:
         with open(config_path) as f:
-            return TrainingStepConfig.model_validate(json.load(f))
+            config = TrainingStepConfig.model_validate(json.load(f))
+        return self._normalize_storage_paths(config)
+
+    def _normalize_storage_paths(self, config: TrainingStepConfig) -> TrainingStepConfig:
+        storage_path = self._job_ctx.storage_path
+
+        model = config.model.model_copy(
+            update={"path": str(remap_job_storage_path(storage_path, config.model.path))},
+        )
+        dataset = config.dataset.model_copy(
+            update={"path": str(remap_job_storage_path(storage_path, config.dataset.path))},
+        )
+        training = config.training
+        if training.kd is not None:
+            teacher_model = training.kd.teacher_model.model_copy(
+                update={"path": str(remap_job_storage_path(storage_path, training.kd.teacher_model.path))},
+            )
+            training = training.model_copy(
+                update={"kd": training.kd.model_copy(update={"teacher_model": teacher_model})},
+            )
+
+        return config.model_copy(
+            update={
+                "model": model,
+                "dataset": dataset,
+                "training": training,
+                "workspace_path": str(remap_job_storage_path(storage_path, config.workspace_path)),
+                "output_path": str(remap_job_storage_path(storage_path, config.output_path)),
+            },
+        )
 
     def _get_library_config_path(self) -> Path:
         return self._workspace_path / AUTOMODEL_CONFIG_FILENAME

services/automodel/tests/tasks/training/test_runner.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for the Automodel training runner."""
+
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock
+
+sys.modules["nemo_automodel"] = MagicMock()
+sys.modules["nemo_automodel._transformers"] = MagicMock()
+sys.modules["nemo_automodel._transformers.registry"] = MagicMock()
+
+from nmp.automodel.entities.values import TrainingType  # noqa: E402
+from nmp.automodel.tasks.training.runner import TrainingRunner  # noqa: E402
+from nmp.automodel.tasks.training.schemas import DistillationConfig, ModelConfig, TrainingStepConfig  # noqa: E402
+from nmp.customization_common.service.context import NMPJobContext  # noqa: E402
+
+
+def _job_context(storage_path: Path) -> NMPJobContext:
+    return NMPJobContext(
+        workspace="default",
+        job_id="job-1",
+        attempt_id="attempt-0",
+        step="training",
+        task="task-1",
+        jobs_url=None,
+        files_url=None,
+        storage_path=storage_path,
+        config_path=storage_path / "config.json",
+    )
+
+
+def _config() -> TrainingStepConfig:
+    return TrainingStepConfig(
+        model=ModelConfig(path="/run/scratch/job/model"),
+        dataset=TrainingStepConfig.DatasetConfig(path="/run/scratch/job/dataset"),
+        training=TrainingStepConfig.TrainingConfig(
+            training_type=TrainingType.DISTILLATION,
+            kd=DistillationConfig(teacher_model=ModelConfig(path="/run/scratch/job/teacher_model")),
+        ),
+        schedule=TrainingStepConfig.ScheduleConfig(),
+        batch=TrainingStepConfig.BatchConfig(),
+        optimizer=TrainingStepConfig.OptimizerConfig(),
+        parallelism=TrainingStepConfig.ParallelismConfig(),
+        output_model="trained-model",
+        workspace_path="/run/scratch/job/training",
+        output_path="/run/scratch/job/output_model",
+    )
+
+
+def test_normalizes_legacy_job_storage_paths_to_runtime_mount(tmp_path: Path) -> None:
+    storage = tmp_path / "job"
+    runner = TrainingRunner.__new__(TrainingRunner)
+    runner._job_ctx = _job_context(storage)
+
+    normalized = runner._normalize_storage_paths(_config())
+
+    assert normalized.model.path == str(storage / "model")
+    assert normalized.dataset.path == str(storage / "dataset")
+    assert normalized.training.kd is not None
+    assert normalized.training.kd.teacher_model.path == str(storage / "teacher_model")
+    assert normalized.workspace_path == str(storage / "training")
+    assert normalized.output_path == str(storage / "output_model")

services/automodel/tests/test_compiler.py

@@ -445,3 +445,15 @@ def _step_image(step: Any) -> str:
     assert _step_image(steps[1]) == get_training_image()
     assert _step_image(steps[2]) == get_tasks_image()
     assert _step_image(steps[3]) == get_tasks_image()
+
+
+@pytest.mark.asyncio
+async def test_platform_job_config_compiler_applies_profile_to_task_steps(mock_sdk, monkeypatch):
+    monkeypatch.setattr(
+        "nmp.automodel.app.jobs.compiler.fetch_model_entity",
+        AsyncMock(return_value=_make_mock_model_entity()),
+    )
+
+    spec = await platform_job_config_compiler(_make_job_output(), "default", mock_sdk, profile="custom-gpu")
+
+    assert [step.executor.profile for step in spec.steps] == ["custom-gpu"] * 4