slurm_scheduler, dir_workspace: add isolated workspaces for Slurm

Author: d4l3k

docs/source/workspace.rst

@@ -9,6 +9,8 @@ torchx.workspace
 .. autoclass:: Workspace
   :members:
 
+.. autofunction:: walk_workspace
+
 torchx.workspace.docker_workspace
 #######################################
 
@@ -19,3 +21,13 @@ torchx.workspace.docker_workspace
 .. autoclass:: DockerWorkspace
   :members:
   :private-members: _update_app_images, _push_images
+
+torchx.workspace.dir_workspace
+#######################################
+
+
+.. automodule:: torchx.workspace.dir_workspace
+.. currentmodule:: torchx.workspace.dir_workspace
+
+.. autoclass:: DirWorkspace
+  :members:

scripts/slurmtest.sh

@@ -15,6 +15,8 @@ DIR="$BASE_DIR/project"
 mkdir "$DIR"
 cd "$DIR"
 
+JOB_DIR="$BASE_DIR/job"
+
 # shellcheck disable=SC1091
 source /opt/slurm/etc/slurm.sh
 sbatch --version
@@ -32,10 +34,12 @@ partition=compute
 time=10
 comment=hello
 nomem=true
+job_dir=$JOB_DIR
 EOT
 
 cat <<EOT > main.py
-print("hello world!")
+import sys
+print("hello world!", file=sys.stderr)
 EOT
 
 APP_ID="$(torchx run --wait --log --scheduler slurm dist.ddp -j 2x1 --script main.py)"

torchx/runner/api.py

@@ -283,7 +283,7 @@ def dryrun(
                 logger.info(
                     f"Building workspace: {workspace} for role[0]: {role.name}, image: {old_img}"
                 )
-                sched.build_workspace_and_update_role(role, workspace)
+                sched.build_workspace_and_update_role(role, workspace, cfg)
                 logger.info("Done building workspace")
                 if old_img != role.image:
                     logger.info(f"New image: {role.image} built from workspace")

torchx/schedulers/api.py

@@ -109,7 +109,7 @@ def submit(
             sched = self
             assert isinstance(sched, Workspace)
             role = app.roles[0]
-            sched.build_workspace_and_update_role(role, workspace)
+            sched.build_workspace_and_update_role(role, workspace, cfg)
         dryrun_info = self.submit_dryrun(app, cfg)
         return self.schedule(dryrun_info)
 

torchx/schedulers/slurm_scheduler.py

@@ -34,7 +34,9 @@
     macros,
     runopts,
 )
+from torchx.workspace.dir_workspace import DirWorkspace
 
+SLURM_JOB_DIRS = ".torchxslurmjobdirs"
 
 SLURM_STATES: Mapping[str, AppState] = {
     "BOOT_FAIL": AppState.FAILED,
@@ -166,6 +168,7 @@ class SlurmBatchRequest:
 
     cmd: List[str]
     replicas: Dict[str, SlurmReplicaRequest]
+    job_dir: Optional[str]
 
     def materialize(self) -> str:
         """
@@ -200,7 +203,7 @@ def materialize(self) -> str:
         return script
 
 
-class SlurmScheduler(Scheduler):
+class SlurmScheduler(Scheduler, DirWorkspace):
     """
     SlurmScheduler is a TorchX scheduling interface to slurm. TorchX expects
     that slurm CLI tools are locally installed and job accounting is enabled.
@@ -254,11 +257,8 @@ class SlurmScheduler(Scheduler):
                 Partial support. SlurmScheduler will return job and replica
                 status but does not provide the complete original AppSpec.
             workspaces: |
-                Partial support. Typical Slurm usage is from a shared NFS mount
-                so code will automatically be updated on the workers.
-                SlurmScheduler does not support programmatic patching via
-                WorkspaceScheduler.
-
+                If ``job_dir`` is specified the DirWorkspace will create a new
+                isolated directory with a snapshot of the workspace.
     """
 
     def __init__(self, session_name: str) -> None:
@@ -276,7 +276,9 @@ def run_opts(self) -> runopts:
             "time",
             type_=str,
             default=None,
-            help="The maximum time the job is allowed to run for.",
+            help='The maximum time the job is allowed to run for. Formats: \
+            "minutes", "minutes:seconds", "hours:minutes:seconds", "days-hours", \
+            "days-hours:minutes" or "days-hours:minutes:seconds"',
         )
         opts.add(
             "nomem",
@@ -304,25 +306,45 @@ def run_opts(self) -> runopts:
             type_=str,
             help="What events to mail users on.",
         )
+        opts.add(
+            "job_dir",
+            type_=str,
+            help="""The directory to place the job code and outputs. The
+            directory must not exist and will be created. To enable log
+            iteration, jobs will be tracked in ``.torchxslurmjobdirs``.
+            """,
+        )
         return opts
 
     def schedule(self, dryrun_info: AppDryRunInfo[SlurmBatchRequest]) -> str:
         req = dryrun_info.request
+        job_dir = req.job_dir
         with tempfile.TemporaryDirectory() as tmpdir:
             script = req.materialize()
-            path = os.path.join(tmpdir, "job.sh")
+            path = os.path.join(job_dir or tmpdir, "torchx-sbatch.sh")
 
             with open(path, "w") as f:
                 f.write(script)
 
-            cmd = req.cmd + [path]
+            cmd = req.cmd
+            if job_dir is not None:
+                cmd += [f"--chdir={job_dir}"]
+            cmd += [path]
 
             p = subprocess.run(cmd, stdout=subprocess.PIPE, check=True)
-            return p.stdout.decode("utf-8").strip()
+            job_id = p.stdout.decode("utf-8").strip()
+
+            if job_dir is not None:
+                _save_job_dir(job_id, job_dir)
+
+            return job_id
 
     def _submit_dryrun(
         self, app: AppDef, cfg: Mapping[str, CfgVal]
     ) -> AppDryRunInfo[SlurmBatchRequest]:
+        job_dir = cfg.get("job_dir")
+        assert job_dir is None or isinstance(job_dir, str), "job_dir must be str"
+
         replicas = {}
         for role in app.roles:
             for replica_id in range(role.num_replicas):
@@ -344,6 +366,7 @@ def _submit_dryrun(
         req = SlurmBatchRequest(
             cmd=cmd,
             replicas=replicas,
+            job_dir=job_dir,
         )
         return AppDryRunInfo(req, repr)
 
@@ -435,6 +458,9 @@ def log_iter(
             )
 
         log_file = f"slurm-{app_id}-{role_name}-{k}.{extension}"
+        job_dirs = _get_job_dirs()
+        if app_id in job_dirs:
+            log_file = os.path.join(job_dirs[app_id], log_file)
 
         return LogIterator(
             app_id, regex or ".*", log_file, self, should_tail=should_tail
@@ -445,3 +471,19 @@ def create_scheduler(session_name: str, **kwargs: Any) -> SlurmScheduler:
     return SlurmScheduler(
         session_name=session_name,
     )
+
+def _save_job_dir(job_id: str, job_dir: str) -> None:
+    with open(SLURM_JOB_DIRS, "at") as f:
+        f.write(f"{job_id} = {job_dir}\n")
+
+def _get_job_dirs() -> Mapping[str, str]:
+    with open(SLURM_JOB_DIRS, "rt") as f:
+        lines = f.readlines()
+    out = {}
+    for line in lines:
+        first, _, second = line.partition("=")
+        if not first or not second:
+            continue
+        out[first.strip] = second.strip
+    return out
+

torchx/schedulers/test/api_test.py

@@ -69,7 +69,9 @@ def run_opts(self) -> runopts:
         def resolve_resource(self, resource: Union[str, Resource]) -> Resource:
             return NULL_RESOURCE
 
-        def build_workspace_and_update_role(self, role: Role, workspace: str) -> None:
+        def build_workspace_and_update_role(
+            self, role: Role, workspace: str, cfg: Mapping[str, CfgVal]
+        ) -> None:
             role.image = workspace
 
     def test_invalid_run_cfg(self) -> None:

torchx/workspace/api.py

@@ -5,8 +5,16 @@
 # LICENSE file in the root directory of this source tree.
 
 import abc
+import posixpath
+import fnmatch
+from typing import TYPE_CHECKING, Iterable, Mapping, Tuple
 
-from torchx.specs import Role
+from torchx.specs import Role, CfgVal
+
+if TYPE_CHECKING:
+    from fsspec import AbstractFileSystem
+
+TORCHX_IGNORE = ".torchxignore"
 
 
 class Workspace(abc.ABC):
@@ -25,7 +33,9 @@ class Workspace(abc.ABC):
     """
 
     @abc.abstractmethod
-    def build_workspace_and_update_role(self, role: Role, workspace: str) -> None:
+    def build_workspace_and_update_role(
+        self, role: Role, workspace: str, cfg: Mapping[str, CfgVal]
+    ) -> None:
         """
         Builds the specified ``workspace`` with respect to ``img``
         and updates the ``role`` to reflect the built workspace artifacts.
@@ -36,3 +46,53 @@ def build_workspace_and_update_role(self, role: Role, workspace: str) -> None:
         Note: this method mutates the passed ``role``.
         """
         ...
+
+
+def _ignore(s: str, patterns: Iterable[str]) -> bool:
+    match = False
+    for pattern in patterns:
+        if pattern.startswith("!") and fnmatch.fnmatch(s, pattern[1:]):
+            match = False
+        elif fnmatch.fnmatch(s, pattern):
+            match = True
+    return match
+
+
+def walk_workspace(
+    fs: "AbstractFileSystem", path: str, ignore_name: str = TORCHX_IGNORE,
+) -> Iterable[Tuple[str, Iterable[str], Mapping[str, Mapping[str, object]]]]:
+    """
+    walk_workspace walks the filesystem path and applies the ignore rules
+    specified via ``ignore_name``.
+    This follows the rules for ``.dockerignore``.
+    https://docs.docker.com/engine/reference/builder/#dockerignore-file
+    """
+    ignore_patterns = []
+    ignore_path = posixpath.join(path, ignore_name)
+    if fs.exists(ignore_path):
+        with fs.open(ignore_path, "rt") as f:
+            lines = f.readlines()
+        for line in lines:
+            line, _, _ = line.partition("#")
+            line = line.strip()
+            if len(line) == 0 or line == ".":
+                continue
+            ignore_patterns.append(line)
+
+    for dir, dirs, files in fs.walk(path, detail=True):
+        assert isinstance(dir, str), "path must be str"
+        relpath = posixpath.relpath(dir, path)
+        if _ignore(relpath, ignore_patterns):
+            continue
+        dirs = [
+            d for d in dirs if not _ignore(posixpath.join(relpath, d), ignore_patterns)
+        ]
+        files = {
+            file: info
+            for file, info in files.items()
+            if not _ignore(
+                posixpath.join(relpath, file) if relpath != "." else file,
+                ignore_patterns,
+            )
+        }
+        yield dir, dirs, files

torchx/workspace/dir_workspace.py

@@ -0,0 +1,42 @@
+from typing import Mapping
+import os
+import posixpath
+import shutil
+
+from torchx.workspace.api import Workspace, walk_workspace
+from torchx.specs import Role, CfgVal
+
+import fsspec
+
+class DirWorkspace(Workspace):
+    def build_workspace_and_update_role(self, role: Role, workspace: str, cfg: Mapping[str, CfgVal]) -> None:
+        """
+        Creates a new directory specified by ``job_dir`` from the workspace. Role
+        image fields will be set to the ``job_dir``.
+
+        Any files listed in the ``.torchxignore`` folder will be skipped.
+        """
+        job_dir = cfg.get("job_dir")
+        if job_dir is None:
+            return
+        assert isinstance(job_dir, str), "job_dir must be str"
+
+        os.mkdir(job_dir)
+        _copy_to_dir(workspace, job_dir)
+        role.image = job_dir
+
+
+def _copy_to_dir(workspace: str, target: str) -> None:
+    fs, path = fsspec.core.url_to_fs(workspace)
+    assert isinstance(path, str), "path must be str"
+
+    for dir, dirs, files in walk_workspace(fs, path):
+        assert isinstance(dir, str), "path must be str"
+        relpath = posixpath.relpath(dir, path)
+        for file, info in files.items():
+            filepath = posixpath.join(
+                target,
+                posixpath.join(relpath, file) if relpath != "." else file,
+            )
+            with fs.open(info["name"], "rb") as src, fsspec.open(filepath, "wb") as dst:
+                shutil.copyfileobj(src, dst)