Sync job code in local tunnel for Slurm Ray job

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

Author: hemildesai

nemo_run/run/ray/slurm.py

@@ -1106,13 +1106,14 @@ def start(
         # ------------------------------------------------------------------
         # Ship *workdir* over to the remote side (or package via packager)
         # ------------------------------------------------------------------
+        cluster_dir = os.path.join(self.executor.tunnel.job_dir, self.name)
         remote_workdir: Optional[str] = None
 
         if workdir:
-            if isinstance(self.executor.tunnel, SSHTunnel):
-                # Rsync workdir honouring .gitignore
-                remote_workdir = os.path.join(self.executor.tunnel.job_dir, self.name, "code")
-                if not dryrun:
+            remote_workdir = os.path.join(cluster_dir, "code")
+            if not dryrun:
+                if isinstance(self.executor.tunnel, SSHTunnel):
+                    # Rsync workdir honouring .gitignore
                     self.executor.tunnel.connect()
                     assert self.executor.tunnel.session is not None, (
                         "Tunnel session is not connected"
@@ -1123,11 +1124,22 @@ def start(
                         remote_workdir,
                         rsync_opts="--filter=':- .gitignore'",
                     )
-            else:
-                remote_workdir = workdir
+                else:
+                    os.makedirs(remote_workdir, exist_ok=True)
+                    subprocess.run(
+                        [
+                            "rsync",
+                            "-pthrvz",
+                            "--filter=:- .gitignore",
+                            f"{os.path.join(workdir, '')}",
+                            remote_workdir,
+                        ],
+                        check=True,
+                    )
         elif self.executor.packager is not None:
             # Use the packager to create an archive which we then extract on the
             # submission host and optionally rsync to the target.
+            remote_workdir = os.path.join(cluster_dir, "code")
             if not dryrun:
                 if isinstance(self.executor.tunnel, SSHTunnel):
                     package_dir = tempfile.mkdtemp(prefix="nemo_packager_")
@@ -1157,7 +1169,6 @@ def start(
                 )
 
                 if isinstance(self.executor.tunnel, SSHTunnel):
-                    remote_workdir = os.path.join(self.executor.tunnel.job_dir, self.name, "code")
                     self.executor.tunnel.connect()
                     assert self.executor.tunnel.session is not None, (
                         "Tunnel session is not connected"
@@ -1169,7 +1180,17 @@ def start(
                         rsync_opts="--filter=':- .gitignore'",
                     )
                 else:
-                    remote_workdir = local_code_extraction_path
+                    os.makedirs(remote_workdir, exist_ok=True)
+                    subprocess.run(
+                        [
+                            "rsync",
+                            "-pthrvz",
+                            "--filter=:- .gitignore",
+                            f"{os.path.join(local_code_extraction_path, '')}",
+                            remote_workdir,
+                        ],
+                        check=True,
+                    )
 
         assert remote_workdir is not None, "workdir could not be determined"