Add support for docker

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

Author: hemildesai

src/nemo_run/config.py

@@ -46,6 +46,7 @@
 RECURSIVE_TYPES = (typing.Union, typing.Optional)
 NEMORUN_HOME = os.environ.get("NEMORUN_HOME", os.path.expanduser("~/.nemo_run"))
 RUNDIR_NAME = "nemo_run"
+RUNDIR_SPECIAL_NAME = "/$nemo_run"
 SCRIPTS_DIR = "scripts"
 
 

src/nemo_run/core/execution/slurm.py

@@ -29,7 +29,7 @@
 from rich.console import Console
 from rich.text import Text
 
-from nemo_run.config import RUNDIR_NAME
+from nemo_run.config import RUNDIR_NAME, RUNDIR_SPECIAL_NAME
 from nemo_run.core.execution.base import (
     Executor,
     ExecutorMacros,
@@ -866,8 +866,8 @@ def get_container_flags(
 
             new_mounts = copy.deepcopy(base_mounts)
             for i, mount in enumerate(new_mounts):
-                if mount.startswith("/$nemo_run"):
-                    new_mounts[i] = mount.replace("$nemo_run", src_job_dir)
+                if mount.startswith(RUNDIR_SPECIAL_NAME):
+                    new_mounts[i] = mount.replace(RUNDIR_SPECIAL_NAME, src_job_dir, 1)
 
             new_mounts.append(f"{src_job_dir}:/{RUNDIR_NAME}")
             _mount_arg = ",".join(new_mounts)

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

@@ -42,7 +42,7 @@
     is_terminal,
 )
 
-from nemo_run.config import RUNDIR_NAME
+from nemo_run.config import RUNDIR_NAME, RUNDIR_SPECIAL_NAME
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.docker import (
     DockerContainer,
@@ -120,6 +120,11 @@ def schedule(self, dryrun_info: AppDryRunInfo[DockerJobRequest]) -> str:  # type
                 log.warning(f"failed to pull image {image}, falling back to local: {e}")
 
         for container in req.containers:
+            for i, mount in enumerate(container.executor.volumes):
+                if mount.startswith(RUNDIR_SPECIAL_NAME):
+                    container.executor.volumes[i] = mount.replace(
+                        RUNDIR_SPECIAL_NAME, req.executor.job_dir, 1
+                    )
             container.executor.volumes.append(f"{req.executor.job_dir}:/{RUNDIR_NAME}")
 
         req.run(client=client)