Add option to symlink from remote dir in packager

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

Author: hemildesai

src/nemo_run/core/execution/slurm.py

@@ -43,7 +43,7 @@
 from nemo_run.core.packaging.git import GitArchivePackager
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
 from nemo_run.core.tunnel.callback import Callback
-from nemo_run.core.tunnel.client import LocalTunnel, SSHConfigFile, SSHTunnel, Tunnel
+from nemo_run.core.tunnel.client import LocalTunnel, PackagingJob, SSHConfigFile, SSHTunnel, Tunnel
 from nemo_run.core.tunnel.server import TunnelMetadata, server_dir
 from nemo_run.devspace.base import DevSpace
 
@@ -388,7 +388,7 @@ def __post_init__(self):
             self.wait_time_for_group_job = 0
 
     def info(self) -> str:
-        return f"{self.__class__.__qualname__} on {self.tunnel._key}"
+        return f"{self.__class__.__qualname__} on {self.tunnel.key}"
 
     def alloc(self, job_name="interactive"):
         self.job_name = f"{self.job_name_prefix}{job_name}"
@@ -544,6 +544,21 @@ def package(self, packager: Packager, job_name: str):
             )
             return
 
+        if packager.symlink_from_remote_dir:
+            logger.info(
+                f"Packager {packager} is configured to symlink from remote dir. Skipping packaging."
+            )
+            if type(packager) is Packager:
+                self.tunnel.packaging_jobs[job_name] = PackagingJob(symlink=False)
+                return
+
+            self.tunnel.packaging_jobs[job_name] = PackagingJob(
+                symlink=True,
+                src_path=packager.symlink_from_remote_dir,
+                dst_path=os.path.join(self.tunnel.job_dir, Path(self.job_dir).name, "code"),
+            )
+            return
+
         assert self.experiment_id, "Executor not assigned to an experiment."
         if isinstance(packager, GitArchivePackager):
             output = subprocess.run(
@@ -573,7 +588,12 @@ def package(self, packager: Packager, job_name: str):
                 f"tar -xvzf {local_pkg} -C {local_code_extraction_path} --ignore-zeros", hide=True
             )
 
-        self.tunnel.packaging_jobs.add(job_name)
+        self.tunnel.packaging_jobs[job_name] = PackagingJob(
+            symlink=False,
+            dst_path=None
+            if type(packager) is Packager
+            else os.path.join(self.tunnel.job_dir, Path(self.job_dir).name, "code"),
+        )
 
     def parse_deps(self) -> list[str]:
         """

src/nemo_run/core/packaging/base.py

@@ -16,7 +16,7 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-
+from typing import Optional
 
 from nemo_run.config import ConfigurableMixin
 
@@ -45,6 +45,10 @@ class Packager(ConfigurableMixin):
     #: Uses component or executor specific debug flags if set to True.
     debug: bool = False
 
+    #: Symlinks the package from the provided remote dir.
+    #: Only applicable when using SlurmExecutor at the moment.
+    symlink_from_remote_dir: Optional[str] = None
+
     def package(self, path: Path, job_dir: str, name: str) -> str: ...
 
     def setup(self):

src/nemo_run/core/tunnel/client.py

@@ -57,15 +57,25 @@ def authentication_handler(title, instructions, prompt_list):
     return [getpass.getpass(str(pr[0])) for pr in prompt_list]
 
 
+@dataclass(kw_only=True)
+class PackagingJob:
+    symlink: bool = False
+    src_path: Optional[str] = None
+    dst_path: Optional[str] = None
+
+    def symlink_cmd(self):
+        return f"ln -s {self.src_path} {self.dst_path}"
+
+
 @dataclass(kw_only=True)
 class Tunnel(ABC):
     job_dir: str
     host: str
     user: str
 
     def __post_init__(self):
-        self._key = f"{self.user}@{self.host}"
-        self._packaging_jobs = set()
+        self.key = f"{self.user}@{self.host}"
+        self._packaging_jobs: dict[str, PackagingJob] = {}
 
     @property
     def packaging_jobs(self):

src/nemo_run/run/experiment.py

@@ -645,9 +645,17 @@ def run(
         for tunnel in self.tunnels.values():
             if isinstance(tunnel, SSHTunnel):
                 tunnel.connect()
-                assert tunnel.session, f"SSH tunnel {tunnel._key} failed to connect."
+                assert tunnel.session, f"SSH tunnel {tunnel.key} failed to connect."
                 rsync(tunnel.session, self._exp_dir, os.path.dirname(tunnel.job_dir))
 
+            symlink_cmds = []
+            for packaging_job in tunnel.packaging_jobs.values():
+                if packaging_job.symlink:
+                    symlink_cmds.append(packaging_job.symlink_cmd())
+
+            if symlink_cmds:
+                tunnel.run(" && ".join(symlink_cmds))
+
         return self._run_dag(detach=detach, tail_logs=tail_logs, executors=executors)
 
     def _run_dag(self, detach: bool, tail_logs: bool, executors: set[Executor]):

src/nemo_run/run/torchx_backend/packaging.py

@@ -82,11 +82,6 @@ def package(
 
             args.append(fn_or_script_filename)
         else:
-            args += [
-                "-p",
-                _serialize(executor.packager.to_config()),
-            ]
-
             args.append(_serialize(fn_or_script))
 
         role_args = default_cmd + args

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

@@ -72,14 +72,14 @@ def _initialize_tunnel(self, tunnel: SSHTunnel | LocalTunnel):
             return
 
         experiment = run_experiment._current_experiment.get(None)
-        if experiment and tunnel._key in experiment.tunnels:
-            self.tunnel = experiment.tunnels[tunnel._key]
+        if experiment and tunnel.key in experiment.tunnels:
+            self.tunnel = experiment.tunnels[tunnel.key]
             return
 
         self.tunnel = tunnel
 
         if experiment:
-            experiment.tunnels[tunnel._key] = self.tunnel
+            experiment.tunnels[tunnel.key] = self.tunnel
 
     def _submit_dryrun(self, app: AppDef, cfg: Executor) -> AppDryRunInfo[Any]:  # type: ignore
         assert isinstance(cfg, SlurmExecutor), f"{cfg.__class__} not supported for slurm scheduler."

test/run/torchx_backend/test_packaging.py

@@ -71,8 +71,6 @@ def test_package_partial(mock_executor):
         "nemo_run.core.runners.fdl_runner",
         "-n",
         "test",
-        "-p",
-        "eJzdVE1Lw0AQ_StlL60goelRrKCePAgevJUSNtlJunazG_ajGEr_uzvbRNPYtBU9eQnZYea9N59bopWy5Ga0JbauwP8QDTm5HpE11PiqaLamBegkJjtvVekbZNbsA1wlwNs7wV8oVd3glIo5EUyp48JyadAqaRlsASMgcwsl4i4W5EkyeJ9w_M6nV-jOIHUFWS4RDyxl1FLvKp0QGMp4Zn-pAyEOZfS4_jTFY3l8JjL7B4lgOKOpHw-r6Qa08QPU-v2gUzlnTECUGJ1FmZI5L7p6HlqS15bjuZXSG6h7a_UEw-bjXCZKJ5kwYxysk-wSSpVoJz21hmi_CFwWUUoNdHW8NCtCdr4cB2RUF64EaRN89hkP96xdpmEMS4vTEM0aDCOsuLFK1-dBZh5kaGz-tk_YqM6JuXQwOq3psz3uLcMTEG5JrwYCaGDYUOHQkFNhIEizq4BAbvFO3kXNITpRnsNynlmE3REOj45mj5MpJ7_d2rh78OLu0YgvWby4X_AYydCTK4mKp9E08sI-AMy28Wg=",
         "eJzdlEtPwzAMgP_KlMuGhKp1RwRIcOOAxIHbNEVuk21haVKlzrRq2n8nDuto9wSxE5eqduLPdvxYM2ctsrvemmFdyvDDnJyy2x5byJok4Yui5iAET9kmqG32IXOsvix8qWXQt6y_MWW9BRVWeB1VmVcalalIa6CIusiIZIWyIO54zF6MkKuBou_D8IauA5vc9roHaTzI2GRCTiSCAIRgb7zWxBMqxz8GR4hubHu-rpr3sTx2iYz-QSJkLiALPYMOltJV0vHm3i8qNVVCaJnwyuVJbs1UzdrxPDdO3hsfr00oe132hOgGZPbQnxpuHc911aemOusdrcvnK55BvpBGJCgr5GUQYKZMJ5Dd5LBN7N2WO3AzX0iDnMR9n935a2bsNANhdh6xHYTThLmqQlb1ZcgoQE41znUrFfu-tXp-2htGFpY7b464ewOHCvSZLoC9F9ASIn4J2pMiDf8l4DxasnvanI9J2EwHL5tdAI2OgTICnXzdbjUuTNLmCD_QSR04ufXmYIOn7Y25E0Zb4eLkpgf1SskbXVXWUMjDZJiEyD4BcFMKTQ==",
     ]
 

test/test_api.py

@@ -16,8 +16,9 @@
 from dataclasses import dataclass
 from unittest.mock import Mock
 
-import nemo_run as run
 import pytest
+
+import nemo_run as run
 from nemo_run.api import dryrun_fn
 
 
@@ -117,7 +118,10 @@ def test_dryrun_fn_with_executor(self, capsys, configured_fn):
 
         captured = capsys.readouterr()
         assert "Dry run for task test.test_api:some_fn" in captured.out
-        assert "LocalExecutor(packager=Packager(debug=False)" in captured.out
+        assert (
+            "LocalExecutor(packager=Packager(debug=False, symlink_from_remote_dir=None)"
+            in captured.out
+        )
 
     def test_dryrun_fn_with_build(self, mocker, configured_fn):
         build_mock = Mock()