Add experiment integration

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

Author: hemildesai

nemo_run/config.py

@@ -449,6 +449,8 @@ class Script(ConfigurableMixin):
     #: Whether to use ``python -m`` when executing via python.
     m: bool = False
 
+    metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
+
     def __post_init__(self):
         assert self.path or self.inline
         assert self.entrypoint, "Need to provide an entrypoint for script."

nemo_run/core/execution/slurm.py

@@ -748,10 +748,10 @@ def _as_sbatch_flag(key: str, value: Any) -> str:
 
 @dataclass(kw_only=True)
 class SlurmBatchRequest:
-    cmd: list[str]
+    launch_cmd: list[str]
     jobs: list[str]
     command_groups: list[list[str]]
-    slurm_config: SlurmExecutor
+    executor: SlurmExecutor
     max_retries: int
     setup: Optional[list[str]] = None
     extra_env: dict[str, str]
@@ -786,7 +786,7 @@ def materialize(self) -> str:
             In case an erroneous keyword argument is added, a list of all eligible parameters
             is printed, with their default values
         """
-        args = asdict(self.slurm_config)  # noqa: F821
+        args = asdict(self.executor)  # noqa: F821
         parameters = {
             k: v for k, v in args.items() if v is not None and k in SlurmExecutor.SBATCH_FLAGS
         }
@@ -800,18 +800,16 @@ def materialize(self) -> str:
         # add necessary parameters
         original_job_name: str = self.jobs[0]  # type: ignore
         job_name_prefix = (
-            self.slurm_config.job_name_prefix
-            if self.slurm_config.job_name_prefix
-            else f"{self.slurm_config.account}-{self.slurm_config.account.split('_')[-1]}."
+            self.executor.job_name_prefix
+            if self.executor.job_name_prefix
+            else f"{self.executor.account}-{self.executor.account.split('_')[-1]}."
         )
         job_name = f"{job_name_prefix}{original_job_name}"
         slurm_job_dir = (
-            self.slurm_config.tunnel.job_dir
-            if self.slurm_config.tunnel
-            else self.slurm_config.job_dir
+            self.executor.tunnel.job_dir if self.executor.tunnel else self.executor.job_dir
         )
-        job_directory_name = Path(self.slurm_config.job_dir).name
-        job_details = self.slurm_config.job_details
+        job_directory_name = Path(self.executor.job_dir).name
+        job_details = self.executor.job_details
 
         if not job_details.job_name:
             job_details.job_name = job_name
@@ -824,41 +822,41 @@ def materialize(self) -> str:
         stdout = str(job_details.stdout)
         stderr = str(job_details.stderr)
 
-        if self.slurm_config.array is not None:
+        if self.executor.array is not None:
             stdout = stdout.replace("%j", "%A_%a")
             stderr = stderr.replace("%j", "%A_%a")
         parameters["output"] = stdout.replace("%t", "0")
 
-        if not self.slurm_config.stderr_to_stdout:
+        if not self.executor.stderr_to_stdout:
             parameters["error"] = stderr.replace("%t", "0")
 
-        if self.slurm_config.additional_parameters is not None:
-            parameters.update(self.slurm_config.additional_parameters)
+        if self.executor.additional_parameters is not None:
+            parameters.update(self.executor.additional_parameters)
 
         # now create
-        sbatch_cmd = " ".join([shlex.quote(arg) for arg in self.cmd])
+        sbatch_cmd = " ".join([shlex.quote(arg) for arg in self.launch_cmd])
 
         sbatch_flags = []
-        if self.slurm_config.heterogeneous:
-            assert len(self.jobs) == len(self.slurm_config.resource_group), (
-                f"Number of jobs {len(self.jobs)} must match number of resource group requests {len(self.slurm_config.resource_group)}.\nIf you are just submitting a single job, make sure that heterogeneous=False in the executor."
+        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."
             )
-            final_group_index = len(self.slurm_config.resource_group) - 1
-            if self.slurm_config.het_group_indices:
-                final_group_index = self.slurm_config.het_group_indices.index(
-                    max(self.slurm_config.het_group_indices)
+            final_group_index = len(self.executor.resource_group) - 1
+            if self.executor.het_group_indices:
+                final_group_index = self.executor.het_group_indices.index(
+                    max(self.executor.het_group_indices)
                 )
 
-            for i in range(len(self.slurm_config.resource_group)):
-                resource_req = self.slurm_config.resource_group[i]
+            for i in range(len(self.executor.resource_group)):
+                resource_req = self.executor.resource_group[i]
                 if resource_req.het_group_index:
-                    assert self.slurm_config.resource_group[i - 1].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"
                     )
                     if (
                         i > 0
                         and resource_req.het_group_index
-                        == self.slurm_config.resource_group[i - 1].het_group_index
+                        == self.executor.resource_group[i - 1].het_group_index
                     ):
                         continue
 
@@ -887,33 +885,31 @@ def materialize(self) -> str:
             for k in sorted(parameters):
                 sbatch_flags.append(_as_sbatch_flag(k, parameters[k]))
 
-        if self.slurm_config.dependencies:
-            slurm_deps = self.slurm_config.parse_deps()
+        if self.executor.dependencies:
+            slurm_deps = self.executor.parse_deps()
             sbatch_flags.append(
                 _as_sbatch_flag(
-                    "dependency", f"{self.slurm_config.dependency_type}:{':'.join(slurm_deps)}"
+                    "dependency", f"{self.executor.dependency_type}:{':'.join(slurm_deps)}"
                 )
             )
 
         env_vars = []
-        full_env_vars = self.slurm_config.env_vars | self.extra_env
+        full_env_vars = self.executor.env_vars | self.extra_env
         for key, value in full_env_vars.items():
             env_vars.append(f"export {key.upper()}={value}")
 
         # commandline (this will run the function and args specified in the file provided as argument)
         # We pass --output and --error here, because the SBATCH command doesn't work as expected with a filename pattern
-        stderr_flags = [] if self.slurm_config.stderr_to_stdout else ["--error", stderr]
+        stderr_flags = [] if self.executor.stderr_to_stdout else ["--error", stderr]
 
         srun_commands = []
         group_env_vars = []
         srun_stdout = noquote(job_details.srun_stdout)
         stderr_flags = (
-            []
-            if self.slurm_config.stderr_to_stdout
-            else ["--error", noquote(job_details.srun_stderr)]
+            [] if self.executor.stderr_to_stdout else ["--error", noquote(job_details.srun_stderr)]
         )
         memory_measure_out = None
-        if self.slurm_config.memory_measure:
+        if self.executor.memory_measure:
             memory_measure_out = srun_stdout
 
         def get_container_flags(
@@ -937,10 +933,10 @@ def get_container_flags(
             return _container_flags
 
         for group_ind, command_group in enumerate(self.command_groups):
-            if self.slurm_config.run_as_group and len(self.slurm_config.resource_group) == len(
+            if self.executor.run_as_group and len(self.executor.resource_group) == len(
                 self.command_groups
             ):
-                resource_req = self.slurm_config.resource_group[group_ind]
+                resource_req = self.executor.resource_group[group_ind]
                 if not resource_req.job_details.job_name:
                     resource_req.job_details.job_name = f"{job_name_prefix}{self.jobs[group_ind]}"
 
@@ -952,7 +948,7 @@ def get_container_flags(
                 cmd_stdout = noquote(resource_req.job_details.srun_stdout)
                 cmd_stderr = (
                     []
-                    if self.slurm_config.stderr_to_stdout
+                    if self.executor.stderr_to_stdout
                     else [
                         "--error",
                         noquote(resource_req.job_details.srun_stderr),
@@ -980,20 +976,20 @@ def get_container_flags(
                 if cmd_stderr:
                     cmd_stderr[-1] = cmd_stderr[-1].replace(original_job_name, self.jobs[group_ind])
                 _container_flags = get_container_flags(
-                    base_mounts=self.slurm_config.container_mounts,
+                    base_mounts=self.executor.container_mounts,
                     src_job_dir=os.path.join(
                         slurm_job_dir,
                         job_directory_name,
                     ),
-                    container_image=self.slurm_config.container_image,
+                    container_image=self.executor.container_image,
                 )
                 _srun_args = ["--wait=60", "--kill-on-bad-exit=1"]
-                _srun_args.extend(self.slurm_config.srun_args or [])
+                _srun_args.extend(self.executor.srun_args or [])
 
-            if self.slurm_config.run_as_group and self.slurm_config.heterogeneous:
+            if self.executor.run_as_group and self.executor.heterogeneous:
                 het_group_index = (
-                    self.slurm_config.resource_group[group_ind].het_group_index
-                    if self.slurm_config.resource_group[group_ind].het_group_index is not None
+                    self.executor.resource_group[group_ind].het_group_index
+                    if self.executor.resource_group[group_ind].het_group_index is not None
                     else group_ind
                 )
                 het_group_flag = [f"--het-group={het_group_index}"]
@@ -1018,10 +1014,10 @@ def get_container_flags(
             )
             command = " ".join(command_group)
 
-            if self.slurm_config.run_as_group:
+            if self.executor.run_as_group:
                 srun_command = f"{srun_cmd} {command} & pids[{group_ind}]=$!"
                 if group_ind != len(self.command_groups) - 1:
-                    srun_command += f"\n\nsleep {self.slurm_config.wait_time_for_group_job}\n"
+                    srun_command += f"\n\nsleep {self.executor.wait_time_for_group_job}\n"
             else:
                 srun_command = f"{srun_cmd} {command}"
 
@@ -1033,15 +1029,14 @@ def get_container_flags(
             "max_retries": self.max_retries,
             "env_vars": env_vars,
             "head_node_ip_var": SlurmExecutor.HEAD_NODE_IP_VAR,
-            "setup_lines": self.slurm_config.setup_lines,
+            "setup_lines": self.executor.setup_lines,
             "memory_measure": memory_measure_out,
             "srun_commands": srun_commands,
             "group_env_vars": group_env_vars,
-            "heterogeneous": self.slurm_config.heterogeneous,
-            "run_as_group": self.slurm_config.run_as_group,
-            "monitor_group_job": self.slurm_config.run_as_group
-            and self.slurm_config.monitor_group_job,
-            "monitor_group_job_wait_time": self.slurm_config.monitor_group_job_wait_time,
+            "heterogeneous": self.executor.heterogeneous,
+            "run_as_group": self.executor.run_as_group,
+            "monitor_group_job": self.executor.run_as_group and self.executor.monitor_group_job,
+            "monitor_group_job_wait_time": self.executor.monitor_group_job_wait_time,
             "het_group_host_var": SlurmExecutor.HET_GROUP_HOST_VAR,
             "ft_enabled": self.launcher and isinstance(self.launcher, FaultTolerance),
         }
@@ -1060,7 +1055,7 @@ def get_container_flags(
         return sbatch_script
 
     def __repr__(self) -> str:
-        return f"""{" ".join(self.cmd + ["$SBATCH_SCRIPT"])}
+        return f"""{" ".join(self.launch_cmd + ["$SBATCH_SCRIPT"])}
 
 #----------------
 # SBATCH_SCRIPT

nemo_run/core/execution/utils.py

@@ -25,7 +25,7 @@ def fill_template(template_name: str, variables: dict, template_dir: Optional[st
     template_dir = template_dir or os.path.join(os.path.dirname(__file__), "templates")
     template_path = os.path.join(template_dir, template_name)
     if not os.path.exists(template_path):
-        raise FileNotFoundError(f'Template "{template_name}" does not exist.')
+        raise FileNotFoundError(f'Template "{template_path}" does not exist.')
     with open(template_path, "r", encoding="utf-8") as fin:
         template = fin.read()
 

nemo_run/run/ray/cluster.py

@@ -22,6 +22,8 @@
 from nemo_run.run.ray.kuberay import KubeRayCluster
 from nemo_run.run.ray.slurm import SlurmRayCluster
 
+USE_WITH_RAY_CLUSTER_KEY = "use_with_ray_cluster"
+
 
 @dataclass(kw_only=True)
 class RayCluster:

nemo_run/run/ray/slurm.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import json
 import logging
 import os
@@ -28,6 +29,7 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, TypeAlias, Union
 
+from nemo_run.config import RUNDIR_NAME, RUNDIR_SPECIAL_NAME
 from nemo_run.core.execution.slurm import SlurmExecutor, _as_sbatch_flag
 from nemo_run.core.execution.utils import fill_template
 from nemo_run.core.packaging.git import GitArchivePackager
@@ -43,11 +45,14 @@
 class SlurmRayRequest:
     name: str
     cluster_dir: str
-    template_path: str
+    template_name: str
+    template_dir: Optional[str] = None
     executor: SlurmExecutor
     pre_ray_start_commands: Optional[list[str]] = None
     command: Optional[str] = None
     workdir: Optional[str] = None
+    nemo_run_dir: Optional[str] = None
+    launch_cmd: list[str]
 
     @staticmethod
     def get_job_name(executor: SlurmExecutor, name: str) -> str:
@@ -125,8 +130,19 @@ def get_srun_flags(mounts: list[str], container_image: Optional[str]) -> str:
             else:
                 _srun_flags.append("--gres=gpu:8")
 
-            _srun_flags.append(f"--container-workdir={self.cluster_dir}")
-            _srun_flags += ["--container-mounts", ",".join(mounts)]
+            if self.nemo_run_dir:
+                new_mounts = copy.deepcopy(mounts)
+                for i, mount in enumerate(new_mounts):
+                    if mount.startswith(RUNDIR_SPECIAL_NAME):
+                        new_mounts[i] = mount.replace(RUNDIR_SPECIAL_NAME, self.nemo_run_dir, 1)
+
+                new_mounts.append(f"{self.nemo_run_dir}:/{RUNDIR_NAME}")
+            else:
+                new_mounts = mounts
+
+            _srun_flags += ["--container-mounts", ",".join(new_mounts)]
+            container_workdir = self.workdir or self.cluster_dir
+            _srun_flags.append(f"--container-workdir={container_workdir}")
 
             return " ".join(_srun_flags)
 
@@ -148,7 +164,12 @@ def get_srun_flags(mounts: list[str], container_image: Optional[str]) -> str:
         if self.pre_ray_start_commands:
             vars_to_fill["pre_ray_start_commands"] = "\n".join(self.pre_ray_start_commands)
 
-        sbatch_script = fill_template(self.template_path, vars_to_fill)
+        sbatch_script = fill_template(
+            self.template_name,
+            vars_to_fill,
+            template_dir=self.template_dir
+            or os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates"),
+        )
         return sbatch_script
 
     def __repr__(self) -> str:
@@ -254,13 +275,12 @@ def create_ray_cluster(
         ray_sbatch = SlurmRayRequest(
             name=name,
             cluster_dir=cluster_dir,
-            template_path=os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "templates", "ray.sub.j2"
-            ),
+            template_name="ray.sub.j2",
             executor=executor,
             pre_ray_start_commands=pre_ray_start_commands,
             command=command,
             workdir=workdir,
+            launch_cmd=["sbatch", "--requeue", "--parsable"],
         ).materialize()
 
         if dryrun:
@@ -321,6 +341,7 @@ def schedule_ray_job(
         command: str,
         workdir: Optional[str] = None,
         pre_ray_start_commands: Optional[list[str]] = None,
+        runtime_env_yaml: Optional[str] = None,
         dryrun: bool = False,
     ):
         remote_workdir = None

nemo_run/run/ray/templates/ray.sub.j2

@@ -289,10 +289,11 @@ echo "[INFO] Ray cluster information saved to $CLUSTER_DIR/ray_cluster_info.json
 # This driver process is responsible for launching a job on the Ray cluster
 CONTAINER_CWD=$(scontrol show job $SLURM_JOB_ID --json | jq -r '.jobs[].current_working_directory')
 # Define command to be empty by default
-COMMAND=""
+COMMAND="${COMMAND:-{{ command }}}"
+COMMAND_WORKDIR={{ command_workdir | default('$CONTAINER_CWD') }}
 
 if [[ -n "$COMMAND" ]]; then
-  srun --no-container-mount-home --gpus=0 --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-driver.log bash -c "$COMMAND"
+  srun --no-container-mount-home --gpus=0 --overlap --container-name=ray-head --container-workdir=$COMMAND_WORKDIR --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-job.log bash -c "$COMMAND"
 else
   echo "[INFO]: Ray Cluster is idled, run this on the slurm head node to get a shell to the head node:"
   cat <<EOF >$CLUSTER_DIR/scripts/${SLURM_JOB_ID}-attach.sh