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

xm_slurm/__init__.py

@@ -10,6 +10,7 @@ from xm_slurm.experiment import (
     get_current_work_unit,
     get_experiment,
 )
+from xm_slurm.job_blocks import JobArgs
 from xm_slurm.packageables import (
     conda_container,
     docker_container,
@@ -34,13 +35,14 @@ __all__ = [
     "get_current_experiment",
     "get_current_work_unit",
     "get_experiment",
+    "JobArgs",
     "JobRequirements",
     "mamba_container",
-    "uv_container",
     "python_container",
     "ResourceQuantity",
     "ResourceType",
     "Slurm",
-    "SlurmSpec",
     "SlurmExperiment",
+    "SlurmSpec",
+    "uv_container",
 ]

xm_slurm/api.py

@@ -322,7 +322,7 @@ class XManagerSqliteAPI(XManagerAPI):
             db_path = Path(os.environ["XM_SLURM_STATE_DIR"]) / "db.sqlite3"
         else:
             db_path = Path.home() / ".local" / "state" / "xm-slurm" / "db.sqlite3"
-        logging.debug("Looking for db at: ", db_path)
+        logger.debug("Looking for db at: ", db_path)
         db_path.parent.mkdir(parents=True, exist_ok=True)
         engine = create_engine(f"sqlite:///{db_path}")
         Base.metadata.create_all(engine)

xm_slurm/config.py

@@ -5,9 +5,10 @@ import getpass
 import json
 import os
 import pathlib
-from typing import Literal, Mapping, NamedTuple
+from typing import Callable, Literal, Mapping, NamedTuple
 
 import asyncssh
+from xmanager import xm
 
 from xm_slurm import constants
 
@@ -137,7 +138,7 @@ class SlurmSSHConfig:
         )
 
     def __hash__(self):
-        return hash((self.host, self.host_public_key, self.user, self.port))
+        return hash((type(self), self.host, self.host_public_key, self.user, self.port))
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -176,6 +177,9 @@ class SlurmClusterConfig:
 
     features: Mapping["xm_slurm.FeatureType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
 
+    # Function to validate the Slurm executor config
+    validate: Callable[[xm.Job], None] | None = None
+
     def __post_init__(self) -> None:
         for src, dst in self.mounts.items():
             if not isinstance(src, (str, os.PathLike)):
@@ -194,6 +198,7 @@ class SlurmClusterConfig:
 
     def __hash__(self):
         return hash((
+            type(self),
             self.ssh,
             self.cwd,
             self.prolog,

xm_slurm/constants.py

@@ -1,5 +1,9 @@
 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>.+)?$"
 )

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(
     *,
@@ -24,6 +31,23 @@ def mila(
             "/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",
         ssh=config.SlurmSSHConfig(
@@ -58,4 +82,5 @@ def mila(
             resources.FeatureType.NVIDIA_MIG: "mig",
             resources.FeatureType.NVIDIA_NVLINK: "nvlink",
         },
+        validate=validate,
     )

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,11 @@
 import dataclasses
 import pathlib
-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)
@@ -32,22 +33,22 @@ class Dockerfile(xm.ExecutableSpec):
     target: str | None = None
 
     # SSH sockets/keys for the docker build step.
-    ssh: Sequence[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:
@@ -58,6 +59,7 @@ class Dockerfile(xm.ExecutableSpec):
 
     def __hash__(self) -> int:
         return hash((
+            type(self),
             self.dockerfile,
             self.context,
             self.target,
@@ -90,7 +92,7 @@ class DockerImage(xm.ExecutableSpec):
         return self.image
 
     def __hash__(self) -> int:
-        return hash((self.image, self.workdir))
+        return hash((type(self), self.image, self.workdir))
 
 
 @dataclasses.dataclass
@@ -148,6 +150,7 @@ class ImageURI:
 
     def __hash__(self) -> int:
         return hash((
+            type(self),
             self.scheme,
             self.domain,
             self.path,
@@ -161,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
@@ -192,18 +196,19 @@ 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())),