xmanager-slurm 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

xm_slurm/contrib/clusters/__init__.py

@@ -19,18 +19,22 @@ def mila(
     if mounts is None:
         mounts = {
             "/network/scratch/${USER:0:1}/$USER": "/scratch",
-            "/network/archive/${USER:0:1}/$USER": "/archive",
+            # TODO: move these somewhere common to all cluster configs.
+            "/home/mila/${USER:0:1}/$USER/.local/state/xm-slurm": "/xm-slurm-state",
+            "/home/mila/${USER:0:1}/$USER/.ssh": "/home/mila/${USER:0:1}/$USER/.ssh",
         }
 
     return config.SlurmClusterConfig(
         name="mila",
-        user=user,
-        host="login.server.mila.quebec",
-        host_public_key=config.PublicKey(
-            "ssh-ed25519",
-            "AAAAC3NzaC1lZDI1NTE5AAAAIBTPCzWRkwYDr/cFb4d2uR6rFlUtqfH3MoLMXPpJHK0n",
+        ssh=config.SlurmSSHConfig(
+            user=user,
+            host="login.server.mila.quebec",
+            host_public_key=config.PublicKey(
+                "ssh-ed25519",
+                "AAAAC3NzaC1lZDI1NTE5AAAAIBTPCzWRkwYDr/cFb4d2uR6rFlUtqfH3MoLMXPpJHK0n",
+            ),
+            port=2222,
         ),
-        port=2222,
         runtime=config.ContainerRuntime.SINGULARITY,
         partition=partition,
         prolog="module load singularity",
@@ -39,14 +43,19 @@ def mila(
             "SINGULARITY_TMPDIR": "$SLURM_TMPDIR",
             "SINGULARITY_LOCALCACHEDIR": "$SLURM_TMPDIR",
             "SCRATCH": "/scratch",
-            "ARCHIVE": "/archive",
+            # TODO: move this somewhere common to all cluster configs.
+            "XM_SLURM_STATE_DIR": "/xm-slurm-state",
         },
         mounts=mounts,
         resources={
-            "rtx8000": resources.ResourceType.RTX8000,
-            "v100": resources.ResourceType.V100,
-            "a100": resources.ResourceType.A100,
-            "a100l": resources.ResourceType.A100_80GIB,
-            "a6000": resources.ResourceType.A6000,
+            resources.ResourceType.RTX8000: "rtx8000",
+            resources.ResourceType.V100: "v100",
+            resources.ResourceType.A100: "a100",
+            resources.ResourceType.A100_80GIB: "a100l",
+            resources.ResourceType.A6000: "a6000",
+        },
+        features={
+            resources.FeatureType.NVIDIA_MIG: "mig",
+            resources.FeatureType.NVIDIA_NVLINK: "nvlink",
         },
     )

xm_slurm/contrib/clusters/drac.py

@@ -2,7 +2,7 @@ import os
 from typing import Literal
 
 from xm_slurm import config
-from xm_slurm.resources import ResourceType
+from xm_slurm.resources import FeatureType, ResourceType
 
 __all__ = ["narval", "beluga", "cedar", "graham"]
 
@@ -18,18 +18,26 @@ def _drac_cluster(
     modules: list[str] | None = None,
     proxy: Literal["submission-host"] | str | None = None,
     mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
-    resources: dict[str, ResourceType] | None = None,
+    resources: dict[ResourceType, str] | None = None,
+    features: dict[FeatureType, str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Cluster."""
     if mounts is None:
-        mounts = {"/scratch/$USER": "/scratch"}
+        mounts = {
+            "/scratch/$USER": "/scratch",
+            # TODO: move these somewhere common to all cluster configs.
+            "/home/$USER/.ssh": "/home/$USER/.ssh",
+            "/home/$USER/.local/state/xm-slurm": "/xm-slurm-state",
+        }
 
     return config.SlurmClusterConfig(
         name=name,
-        user=user,
-        host=host,
-        host_public_key=host_public_key,
-        port=port,
+        ssh=config.SlurmSSHConfig(
+            user=user,
+            host=host,
+            host_public_key=host_public_key,
+            port=port,
+        ),
         account=account,
         proxy=proxy,
         runtime=config.ContainerRuntime.APPTAINER,
@@ -40,9 +48,12 @@ def _drac_cluster(
             "APPTAINER_LOCALCACHEDIR": "$SLURM_TMPDIR",
             "_XDG_DATA_HOME": "$SLURM_TMPDIR/.local",
             "SCRATCH": "/scratch",
+            # TODO: move this somewhere common to all cluster configs.
+            "XM_SLURM_STATE_DIR": "/xm-slurm-state",
         },
         mounts=mounts,
         resources=resources or {},
+        features=features or {},
     )
 
 
@@ -70,7 +81,11 @@ def narval(
         mounts=mounts,
         proxy=proxy,
         modules=modules,
-        resources={"a100": ResourceType.A100},
+        resources={ResourceType.A100: "a100"},
+        features={
+            FeatureType.NVIDIA_MIG: "a100mig",
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
     )
 
 
@@ -98,7 +113,10 @@ def beluga(
         mounts=mounts,
         proxy=proxy,
         modules=modules,
-        resources={"tesla_v100-sxm2-16gb": ResourceType.V100},
+        resources={ResourceType.V100: "tesla_v100-sxm2-16gb"},
+        features={
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
     )
 
 
@@ -120,9 +138,9 @@ def cedar(
         account=account,
         mounts=mounts,
         resources={
-            "v100l": ResourceType.V100_32GIB,
-            "p100": ResourceType.P100,
-            "p100l": ResourceType.P100_16GIB,
+            ResourceType.V100_32GIB: "v100l",
+            ResourceType.P100: "p100",
+            ResourceType.P100_16GIB: "p100l",
         },
     )
 
@@ -147,10 +165,10 @@ def graham(
         mounts=mounts,
         proxy=proxy,
         resources={
-            "v100": ResourceType.V100,
-            "p100": ResourceType.P100,
-            "a100": ResourceType.A100,
-            "a5000": ResourceType.A5000,
+            ResourceType.V100: "v100",
+            ResourceType.P100: "p100",
+            ResourceType.A100: "a100",
+            ResourceType.A5000: "a5000",
         },
     )
 

xm_slurm/executables.py

@@ -1,10 +1,11 @@
 import dataclasses
 import pathlib
-import re
 from typing import Mapping, NamedTuple, Sequence
 
 from xmanager import xm
 
+from xm_slurm import constants
+
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Dockerfile(xm.ExecutableSpec):
@@ -14,6 +15,7 @@ class Dockerfile(xm.ExecutableSpec):
         dockerfile: The path to the Dockerfile.
         context: The path to the Docker context.
         target: The Docker build target.
+        ssh: A list of docker SSH sockets/keys.
         build_args: Build arguments to docker.
         cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
         workdir: The working directory in container.
@@ -29,6 +31,9 @@ class Dockerfile(xm.ExecutableSpec):
     # Docker build target
     target: str | None = None
 
+    # SSH sockets/keys for the docker build step.
+    ssh: Sequence[str] = dataclasses.field(default_factory=list)
+
     # Build arguments to docker
     build_args: Mapping[str, str] = dataclasses.field(default_factory=dict)
 
@@ -56,6 +61,7 @@ class Dockerfile(xm.ExecutableSpec):
             self.dockerfile,
             self.context,
             self.target,
+            tuple(sorted(self.ssh)),
             tuple(sorted(self.build_args.items())),
             tuple(sorted(self.cache_from)),
             self.workdir,
@@ -87,11 +93,6 @@ class DockerImage(xm.ExecutableSpec):
         return hash((self.image, self.workdir))
 
 
-_IMAGE_URI_REGEX = re.compile(
-    r"^(?P<scheme>(?:[^:]+://)?)?(?P<domain>[^/]+)(?P<path>/[^:]*)?(?::(?P<tag>[^@]+))?@?(?P<digest>.+)?$"
-)
-
-
 @dataclasses.dataclass
 class ImageURI:
     image: dataclasses.InitVar[str]
@@ -103,7 +104,7 @@ class ImageURI:
     digest: str | None = dataclasses.field(init=False, default=None)
 
     def __post_init__(self, image: str):
-        match = _IMAGE_URI_REGEX.match(image)
+        match = constants.IMAGE_URI_REGEX.match(image)
         if not match:
             raise ValueError(f"Invalid OCI image URI: {image}")
         groups = {k: v for k, v in match.groupdict().items() if v is not None}
@@ -199,3 +200,14 @@ class RemoteImage(xm.Executable):
     @property
     def name(self) -> str:
         return str(self.image)
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.image,
+                self.workdir,
+                tuple(sorted(self.args.to_list())),
+                tuple(sorted(self.env_vars.items())),
+                self.credentials,
+            ),
+        )

xm_slurm/execution.py

@@ -13,12 +13,14 @@ from typing import Any, Mapping, Sequence
 import asyncssh
 import backoff
 import jinja2 as j2
+import more_itertools as mit
 from asyncssh.auth import KbdIntPrompts, KbdIntResponse
 from asyncssh.misc import MaybeAwait
 from xmanager import xm
 
 from xm_slurm import batching, config, executors, status
 from xm_slurm.console import console
+from xm_slurm.job_blocks import JobArgs
 
 SlurmClusterConfig = config.SlurmClusterConfig
 ContainerRuntime = config.ContainerRuntime
@@ -163,8 +165,9 @@ class _BatchedSlurmHandle:
 class SlurmHandle(_BatchedSlurmHandle):
     """A handle for referring to the launched container."""
 
-    ssh_connection_options: asyncssh.SSHClientConnectionOptions
+    ssh: config.SlurmSSHConfig
     job_id: str
+    job_name: str  # XManager job name associated with this handle
 
     def __post_init__(self):
         if re.match(r"^\d+(_\d+|\+\d+)?$", self.job_id) is None:
@@ -180,10 +183,10 @@ class SlurmHandle(_BatchedSlurmHandle):
         return await self.get_state()
 
     async def stop(self) -> None:
-        await self._batched_cancel(self.ssh_connection_options, self.job_id)
+        await self._batched_cancel(self.ssh.connection_options, self.job_id)
 
     async def get_state(self) -> status.SlurmJobState:
-        return await self._batched_get_state(self.ssh_connection_options, self.job_id)
+        return await self._batched_get_state(self.ssh.connection_options, self.job_id)
 
 
 @functools.cache
@@ -208,7 +211,7 @@ def get_template_env(container_runtime: ContainerRuntime) -> j2.Environment:
         case ContainerRuntime.PODMAN:
             runtime_template = template_env.get_template("runtimes/podman.bash.j2")
         case _:
-            raise NotImplementedError
+            raise NotImplementedError(f"Container runtime {container_runtime} is not implemented.")
     # Update our global env with the runtime template's exported globals
     template_env.globals.update(runtime_template.module.__dict__)
 
@@ -216,16 +219,17 @@ def get_template_env(container_runtime: ContainerRuntime) -> j2.Environment:
 
 
 class Client:
-    def __init__(self):
-        self._connections: dict[
+    def __init__(self) -> None:
+        self._connections = dict[
             asyncssh.SSHClientConnectionOptions, asyncssh.SSHClientConnection
-        ] = {}
+        ]()
         self._connection_lock = asyncio.Lock()
 
     @backoff.on_exception(backoff.expo, asyncssh.Error, max_tries=5, max_time=60.0)
     async def _setup_remote_connection(self, conn: asyncssh.SSHClientConnection) -> None:
         # Make sure the xm-slurm state directory exists
-        await conn.run("mkdir -p ~/.local/state/xm-slurm", check=True)
+        async with conn.start_sftp_client() as sftp_client:
+            await sftp_client.makedirs(".local/state/xm-slurm", exist_ok=True)
 
     async def connection(
         self,
@@ -238,7 +242,24 @@ class Client:
                     await self._setup_remote_connection(conn)
                     self._connections[options] = conn
                 except asyncssh.misc.PermissionDenied as ex:
-                    raise RuntimeError(f"Permission denied connecting to {options.host}") from ex
+                    raise SlurmExecutionError(
+                        f"Permission denied connecting to {options.host}"
+                    ) from ex
+                except asyncssh.misc.ConnectionLost as ex:
+                    raise SlurmExecutionError(f"Connection lost to host {options.host}") from ex
+                except asyncssh.misc.HostKeyNotVerifiable as ex:
+                    raise SlurmExecutionError(
+                        f"Cannot verify the public key for host {options.host}"
+                    ) from ex
+                except asyncssh.misc.KeyExchangeFailed as ex:
+                    raise SlurmExecutionError(
+                        f"Failed to exchange keys with host {options.host}"
+                    ) from ex
+                except asyncssh.Error as ex:
+                    raise SlurmExecutionError(
+                        f"SSH connection error when connecting to {options.host}"
+                    ) from ex
+
         return self._connections[options]
 
     @backoff.on_exception(backoff.expo, asyncssh.Error, max_tries=5, max_time=60.0)
@@ -340,8 +361,8 @@ class Client:
         self,
         *,
         cluster: SlurmClusterConfig,
-        job: xm.Job | xm.JobGroup,
-        args: Mapping[str, Any] | None,
+        job: xm.JobGroup,
+        args: Mapping[str, JobArgs] | None,
         experiment_id: int,
         identity: str | None = ...,
     ) -> SlurmHandle: ...
@@ -351,21 +372,24 @@ class Client:
         self,
         *,
         cluster: SlurmClusterConfig,
-        job: xm.Job | xm.JobGroup,
-        args: Sequence[Mapping[str, Any]],
+        job: xm.Job,
+        args: Sequence[JobArgs],
         experiment_id: int,
         identity: str | None = ...,
-    ) -> Sequence[SlurmHandle]: ...
+    ) -> list[SlurmHandle]: ...
 
+    @typing.overload
     async def launch(
         self,
         *,
         cluster: SlurmClusterConfig,
-        job: xm.Job | xm.JobGroup,
-        args: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None,
+        job: xm.Job,
+        args: JobArgs,
         experiment_id: int,
-        identity: str | None = None,
-    ) -> SlurmHandle | Sequence[SlurmHandle]:
+        identity: str | None = ...,
+    ) -> SlurmHandle: ...
+
+    async def launch(self, *, cluster, job, args, experiment_id, identity=None):
         # Construct template
         template = await self.template(
             job=job,
@@ -379,7 +403,7 @@ class Client:
         # Hash submission script
         template_hash = hashlib.blake2s(template.encode()).hexdigest()[:8]
 
-        conn = await self.connection(cluster.ssh_connection_options)
+        conn = await self.connection(cluster.ssh.connection_options)
         async with conn.start_sftp_client() as sftp:
             # Write the submission script to the cluster
             # TODO(jfarebro): SHOULD FIND A WAY TO GET THE HOME DIRECTORY
@@ -392,9 +416,9 @@ class Client:
 
         # Construct and run command on the cluster
         command = f"sbatch --chdir .local/state/xm-slurm/{experiment_id} --parsable submission-script-{template_hash}.sh"
-        result = await self.run(cluster.ssh_connection_options, command)
+        result = await self.run(cluster.ssh.connection_options, command)
         if result.returncode != 0:
-            raise RuntimeError(f"Failed to schedule job on {cluster.host}: {result.stderr}")
+            raise RuntimeError(f"Failed to schedule job on {cluster.ssh.host}: {result.stderr}")
 
         assert isinstance(result.stdout, str)
         slurm_job_id, *_ = result.stdout.split(",")
@@ -405,26 +429,40 @@ class Client:
             f"[cyan]{cluster.name}[/cyan] "
         )
 
+        # If we scheduled an array job make sure to return a list of handles
+        # The indexing is always sequential in 0, 1, ..., n - 1
         if isinstance(job, xm.Job) and isinstance(args, collections.abc.Sequence):
+            assert job.name is not None
             return [
                 SlurmHandle(
-                    ssh_connection_options=cluster.ssh_connection_options,
+                    ssh=cluster.ssh,
                     job_id=f"{slurm_job_id}_{array_index}",
+                    job_name=job.name,
                 )
                 for array_index in range(len(args))
             ]
-
-        return SlurmHandle(
-            ssh_connection_options=cluster.ssh_connection_options,
-            job_id=slurm_job_id,
-        )
+        elif isinstance(job, xm.Job):
+            assert job.name is not None
+            return SlurmHandle(ssh=cluster.ssh, job_id=slurm_job_id, job_name=job.name)
+        elif isinstance(job, xm.JobGroup):
+            # TODO: make this work for actual job groups.
+            job = mit.one(job.jobs.values())
+            assert isinstance(job, xm.Job)
+            assert job.name is not None
+            return SlurmHandle(ssh=cluster.ssh, job_id=slurm_job_id, job_name=job.name)
+        else:
+            raise ValueError(f"Unsupported job type: {type(job)}")
+
+    def __del__(self):
+        for conn in self._connections.values():
+            conn.close()
 
 
 @typing.overload
 async def launch(
     *,
-    job: xm.Job | xm.JobGroup,
-    args: Mapping[str, Any],
+    job: xm.JobGroup,
+    args: Mapping[str, JobArgs],
     experiment_id: int,
     identity: str | None = ...,
 ) -> SlurmHandle: ...
@@ -433,22 +471,32 @@ async def launch(
 @typing.overload
 async def launch(
     *,
-    job: xm.Job | xm.JobGroup,
-    args: Sequence[Mapping[str, Any]],
+    job: xm.Job,
+    args: Sequence[JobArgs],
     experiment_id: int,
     identity: str | None = ...,
-) -> Sequence[SlurmHandle]: ...
+) -> list[SlurmHandle]: ...
+
+
+@typing.overload
+async def launch(
+    *,
+    job: xm.Job,
+    args: JobArgs,
+    experiment_id: int,
+    identity: str | None = ...,
+) -> SlurmHandle: ...
 
 
 async def launch(
     *,
     job: xm.Job | xm.JobGroup,
-    args: Mapping[str, Any] | Sequence[Mapping[str, Any]],
+    args: Mapping[str, JobArgs] | Sequence[JobArgs] | JobArgs,
     experiment_id: int,
     identity: str | None = None,
-) -> SlurmHandle | Sequence[SlurmHandle]:
+) -> SlurmHandle | list[SlurmHandle]:
     match job:
-        case xm.Job():
+        case xm.Job() as job:
             if not isinstance(job.executor, executors.Slurm):
                 raise ValueError("Job must have a Slurm executor")
             job_requirements = job.executor.requirements
@@ -459,7 +507,7 @@ async def launch(
             return await get_client().launch(
                 cluster=cluster,
                 job=job,
-                args=args,
+                args=typing.cast(JobArgs | Sequence[JobArgs], args),
                 experiment_id=experiment_id,
                 identity=identity,
             )
@@ -482,8 +530,8 @@ async def launch(
 
             return await get_client().launch(
                 cluster=job_group_clusters.pop(),
-                job=job,
-                args=args,
+                job=job_group,
+                args=typing.cast(Mapping[str, JobArgs], args),
                 experiment_id=experiment_id,
                 identity=identity,
             )