xmanager-slurm 0.3.2__py3-none-any.whl → 0.4.1__py3-none-any.whl

xm_slurm/constants.py

@@ -0,0 +1,19 @@
+import re
+
+SLURM_JOB_ID_REGEX = re.compile(
+    r"^(?P<jobid>\d+)(?:(?:\+(?P<componentid>\d+))|(?:_(?P<arraytaskid>\d+)))?$"
+)
+
+IMAGE_URI_REGEX = re.compile(
+    r"^(?P<scheme>(?:[^:]+://)?)?(?P<domain>[^/]+)(?P<path>/[^:]*)?(?::(?P<tag>[^@]+))?@?(?P<digest>.+)?$"
+)
+
+DOMAIN_NAME_REGEX = re.compile(
+    r"^(?!-)(?!.*--)[A-Za-z0-9-]{1,63}(?<!-)(\.[A-Za-z0-9-]{1,63})*(\.[A-Za-z]{2,})$"
+)
+IPV4_REGEX = re.compile(
+    r"^((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$"
+)
+IPV6_REGEX = re.compile(
+    r"^(([0-9a-fA-F]{1,4}:){7}([0-9a-fA-F]{1,4}|:)|(([0-9a-fA-F]{1,4}:){1,7}|:):(([0-9a-fA-F]{1,4}:){1,6}|:)|(([0-9a-fA-F]{1,4}:){1,6}|:):(([0-9a-fA-F]{1,4}:){1,5}|:)|(([0-9a-fA-F]{1,4}:){1,5}|:):(([0-9a-fA-F]{1,4}:){1,4}|:)|(([0-9a-fA-F]{1,4}:){1,4}|:):(([0-9a-fA-F]{1,4}:){1,3}|:)|(([0-9a-fA-F]{1,4}:){1,3}|:):(([0-9a-fA-F]{1,4}:){1,2}|:)|(([0-9a-fA-F]{1,4}:){1,2}|:):([0-9a-fA-F]{1,4}|:)|([0-9a-fA-F]{1,4}|:):([0-9a-fA-F]{1,4}|:)(:([0-9a-fA-F]{1,4}|:)){1,6})$"
+)

xm_slurm/contrib/clusters/__init__.py

@@ -1,13 +1,20 @@
+import datetime as dt
+import logging
 import os
 
+from xmanager import xm
+
 from xm_slurm import config, resources
 from xm_slurm.contrib.clusters import drac
+from xm_slurm.executors import Slurm
 
 # ComputeCanada alias
 cc = drac
 
 __all__ = ["drac", "mila", "cc"]
 
+logger = logging.getLogger(__name__)
+
 
 def mila(
     *,
@@ -19,18 +26,39 @@ 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",
         }
 
+    def validate(job: xm.Job) -> None:
+        assert isinstance(job.executor, Slurm)
+
+        wants_requeue_with_grace_period = (
+            job.executor.requeue and job.executor.timeout_signal_grace_period > dt.timedelta(0)
+        )
+        partition = job.executor.partition or "main"
+
+        if wants_requeue_with_grace_period and (
+            partition is None or not partition.endswith("-grace")
+        ):
+            logger.warning(
+                f"Job {job.name} wants requeue with grace period, but partition `{partition}` does not end with '-grace'. "
+                "Mila Cluster requires you specify a grace partition. "
+                "This may result in the job not being requeued properly."
+            )
+
     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 +67,20 @@ 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",
         },
+        validate=validate,
     )

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/dependencies.py

@@ -0,0 +1,171 @@
+import abc
+import dataclasses
+import datetime as dt
+from typing import Callable, Sequence
+
+
+class SlurmDependencyException(Exception): ...
+
+
+NoChainingException = SlurmDependencyException(
+    "Slurm only supports chaining dependencies with the same logical operator. "
+    "For example, `dep1 & dep2 | dep3` is not supported but `dep1 & dep2 & dep3` is."
+)
+
+
+class SlurmJobDependency(abc.ABC):
+    @abc.abstractmethod
+    def to_dependency_str(self) -> str: ...
+
+    def to_directive(self) -> str:
+        return f"--dependency={self.to_dependency_str()}"
+
+    def __and__(self, other_dependency: "SlurmJobDependency") -> "SlurmJobDependencyAND":
+        if isinstance(self, SlurmJobDependencyOR):
+            raise NoChainingException
+        return SlurmJobDependencyAND(self, other_dependency)
+
+    def __or__(self, other_dependency: "SlurmJobDependency") -> "SlurmJobDependencyOR":
+        if isinstance(other_dependency, SlurmJobDependencyAND):
+            raise NoChainingException
+        return SlurmJobDependencyOR(self, other_dependency)
+
+    def flatten(self) -> tuple["SlurmJobDependency", ...]:
+        if isinstance(self, SlurmJobDependencyAND) or isinstance(self, SlurmJobDependencyOR):
+            return self.first_dependency.flatten() + self.second_dependency.flatten()
+        return (self,)
+
+    def traverse(
+        self, mapper: Callable[["SlurmJobDependency"], "SlurmJobDependency"]
+    ) -> "SlurmJobDependency":
+        if isinstance(self, SlurmJobDependencyAND) or isinstance(self, SlurmJobDependencyOR):
+            return type(self)(
+                first_dependency=self.first_dependency.traverse(mapper),
+                second_dependency=self.second_dependency.traverse(mapper),
+            )
+        return mapper(self)
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAND(SlurmJobDependency):
+    first_dependency: SlurmJobDependency
+    second_dependency: SlurmJobDependency
+
+    def to_dependency_str(self) -> str:
+        return f"{self.first_dependency.to_dependency_str()},{self.second_dependency.to_dependency_str()}"
+
+    def __or__(self, other_dependency: SlurmJobDependency):
+        del other_dependency
+        raise NoChainingException
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.first_dependency, self.second_dependency))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyOR(SlurmJobDependency):
+    first_dependency: SlurmJobDependency
+    second_dependency: SlurmJobDependency
+
+    def to_dependency_str(self) -> str:
+        return f"{self.first_dependency.to_dependency_str()}?{self.second_dependency.to_dependency_str()}"
+
+    def __and__(self, other_dependency: SlurmJobDependency):
+        del other_dependency
+        raise NoChainingException
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.first_dependency, self.second_dependency))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfter(SlurmJobDependency):
+    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    time: dt.timedelta | None = None
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+        if self.time is not None and self.time.total_seconds() % 60 != 0:
+            raise SlurmDependencyException("Time must be specified in exact minutes")
+
+    def to_dependency_str(self) -> str:
+        directive = "after"
+
+        for handle in self.handles:
+            directive += f":{handle.slurm_job.job_id}"
+            if self.time is not None:
+                directive += f"+{self.time.total_seconds() // 60:.0f}"
+        return directive
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterAny(SlurmJobDependency):
+    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afterany"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterNotOK(SlurmJobDependency):
+    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afternotok"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterOK(SlurmJobDependency):
+    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afterok"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobArrayDependencyAfterOK(SlurmJobDependency):
+    handles: Sequence["xm_slurm.execution.SlurmHandle[SlurmJob]"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        job_ids = []
+        for handle in self.handles:
+            job = handle.slurm_job
+            if job.is_array_job:
+                job_ids.append(job.array_job_id)
+            elif job.is_heterogeneous_job:
+                job_ids.append(job.het_job_id)
+            else:
+                job_ids.append(job.job_id)
+        return ":".join(["aftercorr"] + job_ids)
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))

xm_slurm/executables.py

@@ -1,10 +1,12 @@
 import dataclasses
 import pathlib
-import re
-from typing import Mapping, NamedTuple, Sequence
+import typing as tp
 
 from xmanager import xm
 
+from xm_slurm import constants
+from xm_slurm.types import Descriptor
+
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class Dockerfile(xm.ExecutableSpec):
@@ -31,22 +33,22 @@ class Dockerfile(xm.ExecutableSpec):
     target: str | None = None
 
     # SSH sockets/keys for the docker build step.
-    ssh: list[str] = dataclasses.field(default_factory=list)
+    ssh: tp.Sequence[str] = dataclasses.field(default_factory=list)
 
     # Build arguments to docker
-    build_args: Mapping[str, str] = dataclasses.field(default_factory=dict)
+    build_args: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
 
     # --cache-from field in BuildKit
-    cache_from: Sequence[str] = dataclasses.field(default_factory=list)
+    cache_from: tp.Sequence[str] = dataclasses.field(default_factory=list)
 
     # Working directory in container
     workdir: pathlib.Path | None = None
 
     # Container labels
-    labels: Mapping[str, str] = dataclasses.field(default_factory=dict)
+    labels: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
 
     # Target platform
-    platforms: Sequence[str] = dataclasses.field(default_factory=lambda: ["linux/amd64"])
+    platforms: tp.Sequence[str] = dataclasses.field(default_factory=lambda: ["linux/amd64"])
 
     @property
     def name(self) -> str:
@@ -57,6 +59,7 @@ class Dockerfile(xm.ExecutableSpec):
 
     def __hash__(self) -> int:
         return hash((
+            type(self),
             self.dockerfile,
             self.context,
             self.target,
@@ -89,12 +92,7 @@ class DockerImage(xm.ExecutableSpec):
         return self.image
 
     def __hash__(self) -> int:
-        return hash((self.image, self.workdir))
-
-
-_IMAGE_URI_REGEX = re.compile(
-    r"^(?P<scheme>(?:[^:]+://)?)?(?P<domain>[^/]+)(?P<path>/[^:]*)?(?::(?P<tag>[^@]+))?@?(?P<digest>.+)?$"
-)
+        return hash((type(self), self.image, self.workdir))
 
 
 @dataclasses.dataclass
@@ -108,7 +106,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}
@@ -152,6 +150,7 @@ class ImageURI:
 
     def __hash__(self) -> int:
         return hash((
+            type(self),
             self.scheme,
             self.domain,
             self.path,
@@ -165,30 +164,31 @@ class ImageURI:
         return format.format(**fields)
 
 
-class ImageDescriptor:
+class ImageDescriptor(Descriptor[ImageURI, str | ImageURI]):
     def __set_name__(self, owner: type, name: str):
         del owner
         self.image = f"_{name}"
 
-    def __get__(self, instance: object, owner: type) -> ImageURI:
+    def __get__(self, instance: object | None, owner: tp.Type[object] | None = None) -> ImageURI:
         del owner
         return getattr(instance, self.image)
 
     def __set__(self, instance: object, value: str | ImageURI):
+        _setattr = object.__setattr__ if not hasattr(instance, self.image) else setattr
         if isinstance(value, str):
             value = ImageURI(value)
-        setattr(instance, self.image, value)
+        _setattr(instance, self.image, value)
 
 
-class RemoteRepositoryCredentials(NamedTuple):
+class RemoteRepositoryCredentials(tp.NamedTuple):
     username: str
     password: str
 
 
-@dataclasses.dataclass(kw_only=True)  # type: ignore
+@dataclasses.dataclass(frozen=True, kw_only=True)  # type: ignore
 class RemoteImage(xm.Executable):
     # Remote base image
-    image: ImageDescriptor = ImageDescriptor()
+    image: Descriptor[ImageURI, str | ImageURI] = ImageDescriptor()
 
     # Working directory in container
     workdir: pathlib.Path | None = None
@@ -196,11 +196,23 @@ class RemoteImage(xm.Executable):
     # Container arguments
     args: xm.SequentialArgs = dataclasses.field(default_factory=xm.SequentialArgs)
     # Container environment variables
-    env_vars: Mapping[str, str] = dataclasses.field(default_factory=dict)
+    env_vars: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
 
     # Remote repository credentials
     credentials: RemoteRepositoryCredentials | None = None
 
     @property
-    def name(self) -> str:
+    def name(self) -> str:  # type: ignore
         return str(self.image)
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                type(self),
+                self.image,
+                self.workdir,
+                tuple(sorted(self.args.to_list())),
+                tuple(sorted(self.env_vars.items())),
+                self.credentials,
+            ),
+        )