Fix dependencies in Slurm Executor and local inline scripts (#93)

* Fix dependencies in Slurm Executor

Signed-off-by: Hemil Desai <[email protected]>

* Optimize

Signed-off-by: Hemil Desai <[email protected]>

* Fix local inline scripts

Signed-off-by: Hemil Desai <[email protected]>

---------

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

src/nemo_run/config.py

@@ -425,15 +425,19 @@ def get_name(self):
             return os.path.basename(self.path)
 
     def to_command(
-        self, with_entrypoint: bool = False, filename: Optional[str] = None
+        self, with_entrypoint: bool = False, filename: Optional[str] = None, is_local: bool = False
     ) -> list[str]:
         if self.inline:
             if filename:
                 os.makedirs(os.path.dirname(filename), exist_ok=True)
                 with open(filename, "w") as f:
                     f.write("#!/usr/bin/bash\n" + self.inline)
 
-                cmd = [os.path.join(f"/{RUNDIR_NAME}", SCRIPTS_DIR, Path(filename).name)]
+                if is_local:
+                    cmd = [filename]
+                else:
+                    cmd = [os.path.join(f"/{RUNDIR_NAME}", SCRIPTS_DIR, Path(filename).name)]
+
                 if with_entrypoint:
                     cmd = [self.entrypoint] + cmd
 

src/nemo_run/run/torchx_backend/packaging.py

@@ -23,6 +23,7 @@
 
 from nemo_run.config import SCRIPTS_DIR, Partial, Script
 from nemo_run.core.execution.base import Executor, FaultTolerance, Torchrun
+from nemo_run.core.execution.local import LocalExecutor
 from nemo_run.core.serialization.yaml import YamlSerializer
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
 from nemo_run.run.torchx_backend.components import ft_launcher, torchrun
@@ -119,7 +120,8 @@ def package(
 
         args = fn_or_script.args
         role_args = fn_or_script.to_command(
-            filename=os.path.join(executor.job_dir, SCRIPTS_DIR, f"{name}.sh")
+            filename=os.path.join(executor.job_dir, SCRIPTS_DIR, f"{name}.sh"),
+            is_local=True if isinstance(executor, LocalExecutor) else False,
         )
         m = fn_or_script.path if fn_or_script.m else None
         no_python = fn_or_script.entrypoint != "python"

src/nemo_run/run/torchx_backend/schedulers/slurm.py

@@ -144,24 +144,31 @@ def _submit_dryrun(self, app: AppDef, cfg: Executor) -> AppDryRunInfo[Any]:  # t
     def schedule(self, dryrun_info: AppDryRunInfo[SlurmBatchRequest]) -> str:  # type: ignore
         # Setup
         req = dryrun_info.request
-        slurm_cfg = dryrun_info.request.slurm_config
-        assert slurm_cfg.experiment_id, "Executor not assigned to experiment."
+        slurm_executor = dryrun_info.request.slurm_config
+        assert slurm_executor.experiment_id, "Executor not assigned to experiment."
 
-        job_dir = slurm_cfg.job_dir
-        tunnel = slurm_cfg.tunnel
+        job_dir = slurm_executor.job_dir
+        tunnel = slurm_executor.tunnel
         assert tunnel, f"Tunnel required for {self.__class__}"
         assert job_dir, f"Need to provide job_dir for {self.__class__}"
 
         self._initialize_tunnel(tunnel)
         assert self.tunnel, f"Cannot initialize tunnel {tunnel}"
 
-        dst_path = os.path.join(self.tunnel.job_dir, f"{slurm_cfg.job_name}_sbatch.sh")
+        dst_path = os.path.join(self.tunnel.job_dir, f"{slurm_executor.job_name}_sbatch.sh")
+
+        if slurm_executor.dependencies:
+            cmd = ["sbatch", "--requeue", "--parsable"]
+            slurm_deps = slurm_executor.parse_deps()
+            cmd.append(f"--dependency={slurm_executor.dependency_type}:{':'.join(slurm_deps)}")
+            req.cmd = cmd
+
         # Run sbatch script
         req.cmd += [dst_path]
         job_id = self.tunnel.run(" ".join(req.cmd)).stdout.strip()
 
         # Save metadata
-        _save_job_dir(job_id, job_dir, tunnel, slurm_cfg.job_details.ls_term)
+        _save_job_dir(job_id, job_dir, tunnel, slurm_executor.job_details.ls_term)
         return job_id
 
     def _cancel_existing(self, app_id: str) -> None: