fix

Signed-off-by: Hemil Desai <hemild@nvidia.com>

Author: hemildesai

nemo_run/config.py

@@ -468,13 +468,33 @@ def get_name(self):
             return os.path.basename(self.path)
 
     def to_command(
-        self, with_entrypoint: bool = False, filename: Optional[str] = None, is_local: bool = False
+        self,
+        with_entrypoint: bool = False,
+        filename: Optional[str] = None,
+        is_local: bool = False,
+        substitute_rundir_path: Optional[str] = None,
     ) -> list[str]:
+        """Convert the script to a command.
+
+        Args:
+            with_entrypoint: If True, prepend the entrypoint to the command.
+            filename: If provided, write the inline script to this file.
+            is_local: If True, use the local filename in the command.
+            substitute_rundir_path: If provided, substitute /{RUNDIR_NAME} paths
+                with this path in the inline script content. Used for non-container
+                mode where container paths need to be replaced with actual cluster paths.
+        """
         if self.inline:
             if filename:
                 os.makedirs(os.path.dirname(filename), exist_ok=True)
+                inline_content = self.inline
+                # Substitute /{RUNDIR_NAME} paths if specified (non-container mode)
+                if substitute_rundir_path is not None:
+                    inline_content = inline_content.replace(
+                        f"/{RUNDIR_NAME}", substitute_rundir_path
+                    )
                 with open(filename, "w") as f:
-                    f.write("#!/usr/bin/bash\n" + self.inline)
+                    f.write("#!/usr/bin/bash\n" + inline_content)
 
                 if is_local:
                     cmd = [filename]

nemo_run/core/execution/slurm.py

@@ -369,12 +369,12 @@ def merge(
         main_executor.run_as_group = True
 
         if main_executor.het_group_indices:
-            assert (
-                main_executor.heterogeneous
-            ), "heterogeneous must be True if het_group_indices is provided"
-            assert (
-                len(main_executor.het_group_indices) == num_tasks
-            ), "het_group_indices must be the same length as the number of tasks"
+            assert main_executor.heterogeneous, (
+                "heterogeneous must be True if het_group_indices is provided"
+            )
+            assert len(main_executor.het_group_indices) == num_tasks, (
+                "het_group_indices must be the same length as the number of tasks"
+            )
             assert all(
                 x <= y
                 for x, y in zip(
@@ -858,9 +858,9 @@ def materialize(self) -> str:
 
         sbatch_flags = []
         if self.executor.heterogeneous:
-            assert (
-                len(self.jobs) == len(self.executor.resource_group)
-            ), f"Number of jobs {len(self.jobs)} must match number of resource group requests {len(self.executor.resource_group)}.\nIf you are just submitting a single job, make sure that heterogeneous=False in the executor."
+            assert len(self.jobs) == len(self.executor.resource_group), (
+                f"Number of jobs {len(self.jobs)} must match number of resource group requests {len(self.executor.resource_group)}.\nIf you are just submitting a single job, make sure that heterogeneous=False in the executor."
+            )
             final_group_index = len(self.executor.resource_group) - 1
             if self.executor.het_group_indices:
                 final_group_index = self.executor.het_group_indices.index(
@@ -870,9 +870,9 @@ def materialize(self) -> str:
             for i in range(len(self.executor.resource_group)):
                 resource_req = self.executor.resource_group[i]
                 if resource_req.het_group_index is not None:
-                    assert (
-                        self.executor.resource_group[i - 1].het_group_index is not None
-                    ), "het_group_index must be set for all requests in resource_group"
+                    assert self.executor.resource_group[i - 1].het_group_index is not None, (
+                        "het_group_index must be set for all requests in resource_group"
+                    )
                     if (
                         i > 0
                         and resource_req.het_group_index

nemo_run/run/torchx_backend/packaging.py

@@ -15,6 +15,7 @@
 
 import logging
 import os
+from pathlib import Path
 from typing import Iterator, Optional, Type, Union
 
 import fiddle as fdl
@@ -120,9 +121,23 @@ def _get_details_from_script(fn_or_script: Script, serialize_configs: bool):
                 log.warning(f"Failed saving yaml configs due to: {e}")
 
         args = fn_or_script.args
+
+        # For SlurmExecutor without container, substitute /{RUNDIR_NAME} paths
+        # with actual cluster paths in inline scripts
+        substitute_rundir_path = None
+        if (
+            isinstance(executor, SlurmExecutor)
+            and executor.container_image is None
+            and executor.tunnel is not None
+        ):
+            substitute_rundir_path = os.path.join(
+                executor.tunnel.job_dir, Path(executor.job_dir).name
+            )
+
         role_args = fn_or_script.to_command(
             filename=os.path.join(executor.job_dir, SCRIPTS_DIR, f"{name}.sh"),
             is_local=True if isinstance(executor, LocalExecutor) else False,
+            substitute_rundir_path=substitute_rundir_path,
         )
         m = fn_or_script.path if fn_or_script.m else None
         no_python = fn_or_script.entrypoint != "python"

test/run/torchx_backend/test_packaging.py

@@ -14,16 +14,22 @@
 # limitations under the License.
 
 from dataclasses import dataclass
+from unittest.mock import MagicMock
 
 import pytest
 from torchx import specs
 
-from nemo_run.config import Partial, Script
+from nemo_run.config import RUNDIR_NAME, Partial, Script
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.launcher import FaultTolerance, Torchrun
 from nemo_run.core.execution.local import LocalExecutor
+from nemo_run.core.execution.slurm import SlurmExecutor
 from nemo_run.core.packaging.base import Packager
-from nemo_run.run.torchx_backend.packaging import merge_executables, package
+from nemo_run.core.tunnel.client import LocalTunnel
+from nemo_run.run.torchx_backend.packaging import (
+    merge_executables,
+    package,
+)
 
 
 @dataclass(kw_only=True)
@@ -301,3 +307,132 @@ def test_merge_executables():
     assert len(merged_app_def.roles) == 2
     assert merged_app_def.roles[0].name == "role1"
     assert merged_app_def.roles[1].name == "role2"
+
+
+class TestPackagingNonContainerMode:
+    """Tests for non-container mode path substitution in packaging."""
+
+    def test_package_script_inline_with_slurm_non_container_mode(self, tmp_path):
+        """Test that inline scripts have /nemo_run paths substituted in non-container mode."""
+        # Create a SlurmExecutor without container
+        executor = SlurmExecutor(
+            account="test_account",
+            partition="gpu",
+            nodes=1,
+            ntasks_per_node=8,
+            container_image=None,  # Non-container mode
+        )
+        executor.job_dir = str(tmp_path / "test-job")
+        executor.experiment_id = "exp-123"
+
+        # Mock tunnel
+        tunnel = MagicMock(spec=LocalTunnel)
+        tunnel.job_dir = "/remote/experiments/exp-123"
+        executor.tunnel = tunnel
+
+        # Create an inline script with /nemo_run paths
+        fn_or_script = Script(
+            inline=f"cd /{RUNDIR_NAME}/code && python /{RUNDIR_NAME}/scripts/run.py"
+        )
+
+        app_def = package(
+            name="test",
+            fn_or_script=fn_or_script,
+            executor=executor,
+        )
+
+        assert app_def.name == "test"
+        assert len(app_def.roles) == 1
+
+        # Read the generated script file and verify paths were substituted
+        script_file = tmp_path / "test-job" / "scripts" / "test.sh"
+        assert script_file.exists()
+
+        content = script_file.read_text()
+        actual_job_dir = "/remote/experiments/exp-123/test-job"
+
+        # Should NOT contain /nemo_run paths
+        assert f"/{RUNDIR_NAME}" not in content
+        # Should contain the actual job directory path
+        assert f"{actual_job_dir}/code" in content
+        assert f"{actual_job_dir}/scripts/run.py" in content
+
+    def test_package_script_inline_with_slurm_container_mode(self, tmp_path):
+        """Test that inline scripts preserve /nemo_run paths in container mode."""
+        # Create a SlurmExecutor with container
+        executor = SlurmExecutor(
+            account="test_account",
+            partition="gpu",
+            nodes=1,
+            ntasks_per_node=8,
+            container_image="nvcr.io/nvidia/pytorch:24.01-py3",
+        )
+        executor.job_dir = str(tmp_path / "test-job")
+        executor.experiment_id = "exp-123"
+
+        # Mock tunnel
+        tunnel = MagicMock(spec=LocalTunnel)
+        tunnel.job_dir = "/remote/experiments/exp-123"
+        executor.tunnel = tunnel
+
+        # Create an inline script with /nemo_run paths
+        fn_or_script = Script(
+            inline=f"cd /{RUNDIR_NAME}/code && python /{RUNDIR_NAME}/scripts/run.py"
+        )
+
+        app_def = package(
+            name="test",
+            fn_or_script=fn_or_script,
+            executor=executor,
+        )
+
+        assert app_def.name == "test"
+        assert len(app_def.roles) == 1
+
+        # Read the generated script file and verify paths were NOT substituted
+        script_file = tmp_path / "test-job" / "scripts" / "test.sh"
+        assert script_file.exists()
+
+        content = script_file.read_text()
+
+        # Should contain /nemo_run paths (not substituted)
+        assert f"/{RUNDIR_NAME}/code" in content
+        assert f"/{RUNDIR_NAME}/scripts/run.py" in content
+
+    def test_package_script_path_not_affected_by_non_container_mode(self, tmp_path):
+        """Test that path-based scripts are not affected by non-container mode substitution."""
+        # Create a SlurmExecutor without container
+        executor = SlurmExecutor(
+            account="test_account",
+            partition="gpu",
+            nodes=1,
+            ntasks_per_node=8,
+            container_image=None,  # Non-container mode
+        )
+        executor.job_dir = str(tmp_path / "test-job")
+        executor.experiment_id = "exp-123"
+
+        # Mock tunnel
+        tunnel = MagicMock(spec=LocalTunnel)
+        tunnel.job_dir = "/remote/experiments/exp-123"
+        executor.tunnel = tunnel
+
+        # Create a path-based script (not inline)
+        fn_or_script = Script(
+            path="test.py",
+            args=["--config", f"/{RUNDIR_NAME}/configs/config.yaml"],
+        )
+
+        app_def = package(
+            name="test",
+            fn_or_script=fn_or_script,
+            executor=executor,
+        )
+
+        assert app_def.name == "test"
+        assert len(app_def.roles) == 1
+        role = app_def.roles[0]
+
+        # Path-based scripts don't write files, so args should remain unchanged
+        # (the substitution only affects inline script file content)
+        assert role.args == ["test.py", "--config", f"/{RUNDIR_NAME}/configs/config.yaml"]

test/test_config.py

@@ -391,6 +391,45 @@ def test_inline_script(self):
             "\"echo 'test'\"",
         ]
 
+    def test_to_command_substitute_rundir_path(self, tmp_path):
+        """Test that substitute_rundir_path substitutes /nemo_run paths in inline scripts."""
+        from nemo_run.config import RUNDIR_NAME
+
+        script = Script(inline=f"cd /{RUNDIR_NAME}/code && python /{RUNDIR_NAME}/scripts/run.py")
+        filename = str(tmp_path / "test_script.sh")
+
+        script.to_command(
+            filename=filename,
+            substitute_rundir_path="/remote/experiments/exp-123/test-job",
+        )
+
+        # Read the file and verify paths were substituted
+        with open(filename) as f:
+            content = f.read()
+
+        assert f"/{RUNDIR_NAME}" not in content
+        assert "/remote/experiments/exp-123/test-job/code" in content
+        assert "/remote/experiments/exp-123/test-job/scripts/run.py" in content
+
+    def test_to_command_without_substitute_rundir_path(self, tmp_path):
+        """Test that paths are NOT substituted when substitute_rundir_path is None."""
+        from nemo_run.config import RUNDIR_NAME
+
+        script = Script(inline=f"cd /{RUNDIR_NAME}/code && python /{RUNDIR_NAME}/scripts/run.py")
+        filename = str(tmp_path / "test_script.sh")
+
+        script.to_command(
+            filename=filename,
+            substitute_rundir_path=None,  # Default, no substitution
+        )
+
+        # Read the file and verify paths were NOT substituted
+        with open(filename) as f:
+            content = f.read()
+
+        assert f"/{RUNDIR_NAME}/code" in content
+        assert f"/{RUNDIR_NAME}/scripts/run.py" in content
+
 
 class TestGetUnderlyingTypes:
     def test_simple_type(self):