xmanager-slurm 0.4.14__py3-none-any.whl → 0.4.15__py3-none-any.whl

xm_slurm/__init__.py

@@ -18,7 +18,7 @@ from xm_slurm.packageables import (
     python_container,
     uv_container,
 )
-from xm_slurm.resources import JobRequirements, ResourceQuantity, ResourceType
+from xm_slurm.resources import JobRequirements, ResourceQuantity, ResourceType, Topology
 
 logging.getLogger("asyncssh").setLevel(logging.WARN)
 logging.getLogger("httpx").setLevel(logging.WARN)
@@ -42,5 +42,6 @@ __all__ = [
     "Slurm",
     "SlurmExperiment",
     "SlurmSpec",
+    "Topology",
     "uv_container",
 ]

xm_slurm/batching.py

@@ -4,11 +4,11 @@ import dataclasses
 import inspect
 import time
 import types
-from typing import Any, Callable, Coroutine, Generic, ParamSpec, Sequence, TypeVar
+import typing as tp
 
-T = TypeVar("T", contravariant=True)
-R = TypeVar("R", covariant=True)
-P = ParamSpec("P")
+T = tp.TypeVar("T", contravariant=True)
+R = tp.TypeVar("R", covariant=True)
+P = tp.ParamSpec("P")
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -18,10 +18,10 @@ class Request:
 
 
 def stack_bound_arguments(
-    signature: inspect.Signature, bound_arguments: Sequence[inspect.BoundArguments]
+    signature: inspect.Signature, bound_arguments: tp.Sequence[inspect.BoundArguments]
 ) -> inspect.BoundArguments:
     """Stacks bound arguments into a single bound arguments object."""
-    stacked_args = collections.OrderedDict[str, Any]()
+    stacked_args = collections.OrderedDict[str, tp.Any]()
     for bound_args in bound_arguments:
         for name, value in bound_args.arguments.items():
             stacked_args.setdefault(name, [])
@@ -29,7 +29,7 @@ def stack_bound_arguments(
     return inspect.BoundArguments(signature, stacked_args)
 
 
-class batch(Generic[R]):
+class batch(tp.Generic[R]):
     __slots__ = (
         "fn",
         "signature",
@@ -44,7 +44,7 @@ class batch(Generic[R]):
 
     def __init__(
         self,
-        fn: Callable[..., Coroutine[None, None, Sequence[R]]],
+        fn: tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]],
         /,
         *,
         max_batch_size: int,
@@ -102,7 +102,7 @@ class batch(Generic[R]):
 
         return batch
 
-    def __get__(self, obj: Any, objtype: type) -> Any:
+    def __get__(self, obj: tp.Any, objtype: tp.Type[tp.Any]) -> tp.Any:
         del objtype
         if isinstance(self.fn, staticmethod):
             return self.__call__
@@ -110,11 +110,11 @@ class batch(Generic[R]):
             return types.MethodType(self, obj)
 
     @property
-    def __func__(self) -> Callable[..., Coroutine[None, None, Sequence[R]]]:
+    def __func__(self) -> tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]]:
         return self.fn
 
     @property
-    def __wrapped__(self) -> Callable[..., Coroutine[None, None, Sequence[R]]]:
+    def __wrapped__(self) -> tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]]:
         return self.fn
 
     @property

xm_slurm/config.py

@@ -5,7 +5,7 @@ import getpass
 import json
 import os
 import pathlib
-from typing import Callable, Literal, Mapping, NamedTuple
+import typing as tp
 
 import asyncssh
 from xmanager import xm
@@ -23,7 +23,7 @@ class ContainerRuntime(enum.Enum):
 
     @classmethod
     def from_string(
-        cls, runtime: Literal["singularity", "apptainer", "docker", "podman"]
+        cls, runtime: tp.Literal["singularity", "apptainer", "docker", "podman"]
     ) -> "ContainerRuntime":
         return {
             "singularity": cls.SINGULARITY,
@@ -45,7 +45,7 @@ class ContainerRuntime(enum.Enum):
             raise NotImplementedError
 
 
-class PublicKey(NamedTuple):
+class PublicKey(tp.NamedTuple):
     algorithm: str
     key: str
 
@@ -172,26 +172,26 @@ class SlurmClusterConfig:
     qos: str | None = None
 
     # If true, a reverse proxy is initiated via the submission host.
-    proxy: Literal["submission-host"] | str | None = None
+    proxy: tp.Literal["submission-host"] | str | None = None
 
     runtime: ContainerRuntime
 
     # Environment variables
-    host_environment: Mapping[str, str] = dataclasses.field(default_factory=dict)
-    container_environment: Mapping[str, str] = dataclasses.field(default_factory=dict)
+    host_environment: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+    container_environment: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
 
     # Mounts
-    mounts: Mapping[os.PathLike[str] | str, os.PathLike[str] | str] = dataclasses.field(
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] = dataclasses.field(
         default_factory=dict
     )
 
     # Resource mapping
-    resources: Mapping["xm_slurm.ResourceType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
+    resources: tp.Mapping["xm_slurm.ResourceType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
 
-    features: Mapping["xm_slurm.FeatureType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
+    features: tp.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
+    validate: tp.Callable[[xm.Job], None] | None = None
 
     def __post_init__(self) -> None:
         for src, dst in self.mounts.items():

xm_slurm/contrib/clusters/drac.py

@@ -1,5 +1,5 @@
 import os
-from typing import Literal
+import typing as tp
 
 from xm_slurm import config
 from xm_slurm.resources import FeatureType, ResourceType
@@ -16,10 +16,10 @@ def _drac_cluster(
     user: str | None = None,
     account: str | None = None,
     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[ResourceType, str] | None = None,
-    features: dict[FeatureType, str] | None = None,
+    proxy: tp.Literal["submission-host"] | str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    resources: tp.Mapping[ResourceType, str] | None = None,
+    features: tp.Mapping[FeatureType, str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Cluster."""
     if mounts is None:
@@ -62,8 +62,8 @@ def narval(
     *,
     user: str | None = None,
     account: str | None = None,
-    proxy: Literal["submission-host"] | str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    proxy: tp.Literal["submission-host"] | str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Narval Cluster (https://docs.alliancecan.ca/wiki/Narval/en)."""
     modules = []
@@ -94,8 +94,8 @@ def beluga(
     *,
     user: str | None = None,
     account: str | None = None,
-    proxy: Literal["submission-host"] | str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    proxy: tp.Literal["submission-host"] | str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Beluga Cluster (https://docs.alliancecan.ca/wiki/B%C3%A9luga/en)."""
     modules = []
@@ -125,8 +125,8 @@ def rorqual(
     *,
     user: str | None = None,
     account: str | None = None,
-    proxy: Literal["submission-host"] | str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    proxy: tp.Literal["submission-host"] | str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Beluga Cluster (https://docs.alliancecan.ca/wiki/Rorqual/en)."""
     modules = []
@@ -155,7 +155,7 @@ def cedar(
     *,
     user: str | None = None,
     account: str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Cedar Cluster (https://docs.alliancecan.ca/wiki/Cedar/en)."""
     return _drac_cluster(
@@ -180,7 +180,7 @@ def fir(
     *,
     user: str | None = None,
     account: str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Fir Cluster (https://docs.alliancecan.ca/wiki/Fir/en)."""
     return _drac_cluster(
@@ -201,8 +201,8 @@ def graham(
     *,
     user: str | None = None,
     account: str | None = None,
-    proxy: Literal["submission-host"] | str | None = "submission-host",
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+    proxy: tp.Literal["submission-host"] | str | None = "submission-host",
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
 ) -> config.SlurmClusterConfig:
     """DRAC Cedar Cluster (https://docs.alliancecan.ca/wiki/Graham/en)."""
     return _drac_cluster(
@@ -223,17 +223,3 @@ def graham(
             ResourceType.A5000: "a5000",
         },
     )
-
-
-def all(
-    user: str | None = None,
-    account: str | None = None,
-    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
-) -> list[config.SlurmClusterConfig]:
-    """All DRAC clusters."""
-    return [
-        narval(user=user, account=account, mounts=mounts),
-        beluga(user=user, account=account, mounts=mounts),
-        cedar(user=user, account=account, mounts=mounts),
-        graham(user=user, account=account, mounts=mounts),
-    ]

xm_slurm/dependencies.py

@@ -1,7 +1,7 @@
 import abc
 import dataclasses
 import datetime as dt
-from typing import Callable, Sequence
+import typing as tp
 
 
 class SlurmDependencyException(Exception): ...
@@ -36,7 +36,7 @@ class SlurmJobDependency(abc.ABC):
         return (self,)
 
     def traverse(
-        self, mapper: Callable[["SlurmJobDependency"], "SlurmJobDependency"]
+        self, mapper: tp.Callable[["SlurmJobDependency"], "SlurmJobDependency"]
     ) -> "SlurmJobDependency":
         if isinstance(self, SlurmJobDependencyAND) or isinstance(self, SlurmJobDependencyOR):
             return type(self)(
@@ -80,7 +80,7 @@ class SlurmJobDependencyOR(SlurmJobDependency):
 
 @dataclasses.dataclass(frozen=True)
 class SlurmJobDependencyAfter(SlurmJobDependency):
-    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
     time: dt.timedelta | None = None
 
     def __post_init__(self):
@@ -104,7 +104,7 @@ class SlurmJobDependencyAfter(SlurmJobDependency):
 
 @dataclasses.dataclass(frozen=True)
 class SlurmJobDependencyAfterAny(SlurmJobDependency):
-    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
 
     def __post_init__(self):
         if len(self.handles) == 0:
@@ -119,7 +119,7 @@ class SlurmJobDependencyAfterAny(SlurmJobDependency):
 
 @dataclasses.dataclass(frozen=True)
 class SlurmJobDependencyAfterNotOK(SlurmJobDependency):
-    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
 
     def __post_init__(self):
         if len(self.handles) == 0:
@@ -134,7 +134,7 @@ class SlurmJobDependencyAfterNotOK(SlurmJobDependency):
 
 @dataclasses.dataclass(frozen=True)
 class SlurmJobDependencyAfterOK(SlurmJobDependency):
-    handles: Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
 
     def __post_init__(self):
         if len(self.handles) == 0:
@@ -149,7 +149,7 @@ class SlurmJobDependencyAfterOK(SlurmJobDependency):
 
 @dataclasses.dataclass(frozen=True)
 class SlurmJobArrayDependencyAfterOK(SlurmJobDependency):
-    handles: Sequence["xm_slurm.execution.SlurmHandle[SlurmJob]"]  # type: ignore # noqa: F821
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle[SlurmJob]"]  # type: ignore # noqa: F821
 
     def __post_init__(self):
         if len(self.handles) == 0:

xm_slurm/executors.py

@@ -1,10 +1,18 @@
+import collections.abc
 import dataclasses
 import datetime as dt
 import signal
+import typing as tp
 
 from xmanager import xm
 
-from xm_slurm import resources
+from xm_slurm import resources, utils
+
+ResourceBindType = tp.Literal[
+    resources.ResourceType.GPU,
+    resources.ResourceType.MEMORY,
+    resources.ResourceType.RAM,
+]
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -26,10 +34,19 @@ class Slurm(xm.Executor):
     Args:
         requirements: The requirements for the job.
         time: The maximum time to run the job.
+        switches: Maximum count of leaf switches desired for the job allocation.
+        switches_grace_period: Maximum time to wait for that number of switches.
+        bind: How to bind tasks to resource (memory, GPU, or generic resource).
         account: The account to charge the job to.
         partition: The partition to run the job in.
         qos: The quality of service to run the job with.
         priority: The priority of the job.
+        reservation: Allocate resources for the job from the named reservation.
+        exclusive: Allow sharing nodes with other running jobs.
+        oversubscribe: Allow over-subscribing resources with other running jobs.
+        overcommit: Allow sharing of allocated resources as if only one task per was requested.
+        nice: Run the job with an adjusted scheduling priority.
+        kill_on_invalid_dependencies: Whether to kill the job if it has invalid dependencies.
         timeout_signal: The signal to send to the job when it runs out of time.
         timeout_signal_grace_period: The time to wait before sending `timeout_signal`.
         requeue: Whether or not the job is eligible for requeueing.
@@ -41,12 +58,18 @@ class Slurm(xm.Executor):
     # Job requirements
     requirements: resources.JobRequirements
     time: dt.timedelta
+    bind: tp.Mapping[ResourceBindType | str, str | None] | None = None
 
     # Placement
     account: str | None = None
     partition: str | None = None
     qos: str | None = None
     priority: int | None = None
+    reservation: str | tp.Iterable[str] | None = None
+    exclusive: bool = False
+    oversubscribe: bool = False
+    overcommit: bool = False
+    nice: int | None = None
 
     # Job dependency handling
     kill_on_invalid_dependencies: bool = True
@@ -65,13 +88,28 @@ class Slurm(xm.Executor):
         return self.time - self.timeout_signal_grace_period
 
     def __post_init__(self) -> None:
-        if not isinstance(self.time, dt.timedelta):
-            raise TypeError(f"time must be a `datetime.timedelta`, got {type(self.time)}")
         if not isinstance(self.requirements, resources.JobRequirements):
             raise TypeError(
                 f"requirements must be a `xm_slurm.JobRequirements`, got {type(self.requirements)}. "
                 "If you're still using `xm.JobRequirements`, please update to `xm_slurm.JobRequirements`."
             )
+        if not isinstance(self.time, dt.timedelta):
+            raise TypeError(f"time must be a `datetime.timedelta`, got {type(self.time)}")
+        if self.bind is not None:
+            if not isinstance(self.bind, collections.abc.Mapping):
+                raise TypeError(f"bind must be a mapping, got {type(self.bind)}")
+            for resource, value in self.bind.items():
+                if resource not in (
+                    resources.ResourceType.GPU,
+                    resources.ResourceType.MEMORY,
+                    resources.ResourceType.RAM,
+                ) and not isinstance(resource, str):
+                    raise TypeError(
+                        f"bind resource must be a {resources.ResourceType.GPU.name}, {resources.ResourceType.MEMORY.name}, or {resources.ResourceType.RAM.name}, got {type(resource)}"
+                    )
+                if value is not None and not isinstance(value, str):
+                    raise TypeError(f"bind value must be None or a string, got {type(value)}")
+
         if not isinstance(self.timeout_signal, signal.Signals):
             raise TypeError(
                 f"termination_signal must be a `signal.Signals`, got {type(self.timeout_signal)}"
@@ -86,6 +124,10 @@ class Slurm(xm.Executor):
             )
         if self.requeue_on_exit_code == 0:
             raise ValueError("requeue_on_exit_code should not be 0 to avoid unexpected behavior.")
+        if self.exclusive and self.oversubscribe:
+            raise ValueError("exclusive and oversubscribe are mutually exclusive.")
+        if self.nice is not None and not (-2147483645 <= self.nice <= 2147483645):
+            raise ValueError(f"nice must be between -2147483645 and 2147483645, got {self.nice}")
 
     @classmethod
     def Spec(cls, tag: str | None = None) -> SlurmSpec:
@@ -96,10 +138,22 @@ class Slurm(xm.Executor):
         directives = self.requirements.to_directives()
 
         # Time
-        days = self.time.days
-        hours, remainder = divmod(self.time.seconds, 3600)
-        minutes, seconds = divmod(remainder, 60)
-        directives.append(f"--time={days}-{hours:02}:{minutes:02}:{seconds:02}")
+        directives.append(f"--time={utils.timestr_from_timedelta(self.time)}")
+
+        # Resource binding
+        if self.bind is not None:
+            for resource, value in self.bind.items():
+                if value is None:
+                    value = "none"
+                match resource:
+                    case resources.ResourceType.MEMORY | resources.ResourceType.RAM:
+                        directives.append(f"--mem-bind={value}")
+                    case resources.ResourceType.GPU:
+                        directives.append(f"--gpu-bind={value}")
+                    case str():
+                        directives.append(f"--tres-bind=gres/{resource}:{value}")
+                    case _:
+                        raise ValueError(f"Unsupported resource type {resource!r} for binding.")
 
         # Job dependency handling
         directives.append(
@@ -107,20 +161,36 @@ class Slurm(xm.Executor):
         )
 
         # Placement
-        if self.account:
+        if self.account is not None:
             directives.append(f"--account={self.account}")
-        if self.partition:
+        if self.partition is not None:
             directives.append(f"--partition={self.partition}")
-        if self.qos:
+        if self.qos is not None:
             directives.append(f"--qos={self.qos}")
-        if self.priority:
+        if self.priority is not None:
             directives.append(f"--priority={self.priority}")
+        if self.reservation is not None:
+            match self.reservation:
+                case str():
+                    directives.append(f"--reservation={self.reservation}")
+                case collections.abc.Iterable():
+                    directives.append(f"--reservation={','.join(self.reservation)}")
+                case _:
+                    raise ValueError(f"Invalid reservation type: {type(self.reservation)}")
+        if self.exclusive:
+            directives.append("--exclusive")
+        if self.oversubscribe:
+            directives.append("--oversubscribe")
+        if self.overcommit:
+            directives.append("--overcommit")
+        if self.nice is not None:
+            directives.append(f"--nice={self.nice}")
 
         # Job rescheduling
         directives.append(
             f"--signal={self.timeout_signal.name.removeprefix('SIG')}@{self.timeout_signal_grace_period.seconds}"
         )
-        if self.requeue and self.requeue_max_attempts > 0:
+        if self.requeue is not None and self.requeue_max_attempts > 0:
             directives.append("--requeue")
         else:
             directives.append("--no-requeue")

xm_slurm/experimental/parameter_controller.py

@@ -4,8 +4,8 @@ import dataclasses
 import enum
 import functools
 import logging
+import typing as tp
 import zlib
-from typing import Awaitable, Callable, Concatenate, Coroutine, Mapping, ParamSpec, TypeVar
 
 import backoff
 import cloudpickle
@@ -15,15 +15,15 @@ import xm_slurm
 from xm_slurm import job_blocks, status
 from xm_slurm.experiment import SlurmAuxiliaryUnit, SlurmExperiment
 
-P = ParamSpec("P")
-T = TypeVar("T")
+P = tp.ParamSpec("P")
+T = tp.TypeVar("T")
 
 logger = logging.getLogger(__name__)
 
 
 async def _monitor_parameter_controller(
     aux_unit: SlurmAuxiliaryUnit,
-    local_parameter_controller_coro: Coroutine[None, None, T],
+    local_parameter_controller_coro: tp.Coroutine[None, None, T],
     *,
     poll_interval: float = 30.0,
 ) -> None:
@@ -104,13 +104,13 @@ def parameter_controller(
     controller_mode: ParameterControllerMode = ParameterControllerMode.AUTO,
     controller_name: str = "parameter_controller",
     controller_args: xm.UserArgs | None = None,
-    controller_env_vars: Mapping[str, str] | None = None,
-) -> Callable[
+    controller_env_vars: tp.Mapping[str, str] | None = None,
+) -> tp.Callable[
     [
-        Callable[Concatenate[SlurmExperiment, P], T]
-        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
+        tp.Callable[tp.Concatenate[SlurmExperiment, P], T]
+        | tp.Callable[tp.Concatenate[SlurmExperiment, P], tp.Awaitable[T]],
     ],
-    Callable[P, xm.AuxiliaryUnitJob],
+    tp.Callable[P, xm.AuxiliaryUnitJob],
 ]:
     """Converts a function to a controller which can be added to an experiment.
 
@@ -131,9 +131,9 @@ def parameter_controller(
     """
 
     def decorator(
-        f: Callable[Concatenate[SlurmExperiment, P], T]
-        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
-    ) -> Callable[P, xm.AuxiliaryUnitJob]:
+        f: tp.Callable[tp.Concatenate[SlurmExperiment, P], T]
+        | tp.Callable[tp.Concatenate[SlurmExperiment, P], tp.Awaitable[T]],
+    ) -> tp.Callable[P, xm.AuxiliaryUnitJob]:
         @functools.wraps(f)
         def make_controller(*args: P.args, **kwargs: P.kwargs) -> xm.AuxiliaryUnitJob:
             # Modify the function to read the experiment from the API so that it can be pickled.
@@ -141,13 +141,17 @@ def parameter_controller(
             async def job_generator(aux_unit: SlurmAuxiliaryUnit) -> None:
                 experiment_id = aux_unit.experiment.experiment_id
 
-                async def local_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                async def local_controller(
+                    *args: P.args, **kwargs: P.kwargs
+                ) -> T | tp.Awaitable[T]:
                     if asyncio.iscoroutinefunction(f):
                         return await f(aux_unit.experiment, *args, **kwargs)
                     else:
                         return f(aux_unit.experiment, *args, **kwargs)
 
-                async def remote_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                async def remote_controller(
+                    *args: P.args, **kwargs: P.kwargs
+                ) -> T | tp.Awaitable[T]:
                     async with xm_slurm.get_experiment(experiment_id=experiment_id) as exp:
                         if asyncio.iscoroutinefunction(f):
                             return await f(exp, *args, **kwargs)

xm_slurm/job_blocks.py

@@ -1,11 +1,11 @@
-from typing import Mapping, TypedDict
+import typing as tp
 
 from xmanager import xm
 
 
-class JobArgs(TypedDict, total=False):
+class JobArgs(tp.TypedDict, total=False):
     args: xm.UserArgs
-    env_vars: Mapping[str, str]
+    env_vars: tp.Mapping[str, str]
 
 
 def get_args_for_python_entrypoint(

xm_slurm/packageables.py

@@ -1,7 +1,7 @@
 import importlib.resources as resources
 import pathlib
 import sys
-from typing import Literal, Mapping, Sequence
+import typing as tp
 
 from xmanager import xm
 
@@ -14,7 +14,7 @@ def docker_image(
     *,
     image: str,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a packageable for a pre-built Docker image.
 
@@ -39,12 +39,12 @@ def docker_container(
     dockerfile: pathlib.Path | None = None,
     context: pathlib.Path | None = None,
     target: str | None = None,
-    ssh: Sequence[str] | Literal[True] | None = None,
-    build_args: Mapping[str, str] | None = None,
-    cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
+    build_args: tp.Mapping[str, str] | None = None,
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Docker container packageable from a dockerfile.
 
@@ -104,13 +104,13 @@ def python_container(
     context: pathlib.Path | None = None,
     requirements: pathlib.Path | None = None,
     base_image: str = "docker.io/python:{major}.{minor}-slim",
-    extra_system_packages: Sequence[str] = (),
-    extra_python_packages: Sequence[str] = (),
-    cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] | None = None,
-    ssh: Sequence[str] | Literal[True] | None = None,
+    extra_system_packages: tp.Sequence[str] = (),
+    extra_python_packages: tp.Sequence[str] = (),
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Python container from a base image using pip from a `requirements.txt` file.
 
@@ -181,11 +181,11 @@ def mamba_container(
     context: pathlib.Path | None = None,
     environment: pathlib.Path | None = None,
     base_image: str = "gcr.io/distroless/base-debian10",
-    cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] | None = None,
-    ssh: Sequence[str] | Literal[True] | None = None,
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Conda container from a base image using mamba from a `environment.yml` file.
 
@@ -249,13 +249,13 @@ def uv_container(
     entrypoint: xm.ModuleName | xm.CommandList,
     context: pathlib.Path | None = None,
     base_image: str = "docker.io/python:{major}.{minor}-slim-bookworm",
-    extra_system_packages: Sequence[str] = (),
-    extra_python_packages: Sequence[str] = (),
-    cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] | None = None,
-    ssh: Sequence[str] | Literal[True] | None = None,
+    extra_system_packages: tp.Sequence[str] = (),
+    extra_python_packages: tp.Sequence[str] = (),
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Python container from a base image using uv from a `uv.lock` file.
 

xm_slurm/packaging/registry.py

@@ -1,31 +1,31 @@
 import dataclasses
-from typing import Callable, Generic, ParamSpec, Sequence, Type, TypeVar
+import typing as tp
 
 from xmanager import xm
 
-T_co = TypeVar("T_co", covariant=True)
-P = ParamSpec("P")
-ExecutableSpecT = TypeVar("ExecutableSpecT", bound=xm.ExecutableSpec)
+T_co = tp.TypeVar("T_co", covariant=True)
+P = tp.ParamSpec("P")
+ExecutableSpecT = tp.TypeVar("ExecutableSpecT", bound=xm.ExecutableSpec)
 
 
 @dataclasses.dataclass(frozen=True)
-class IndexedContainer(Generic[T_co]):
+class IndexedContainer(tp.Generic[T_co]):
     index: int
     value: T_co
 
 
-RegistrationCallable = Callable[
-    [Sequence[IndexedContainer[xm.Packageable]]],
-    Sequence[IndexedContainer[xm.Executable]],
+RegistrationCallable = tp.Callable[
+    [tp.Sequence[IndexedContainer[xm.Packageable]]],
+    tp.Sequence[IndexedContainer[xm.Executable]],
 ]
 
 
-_REGISTRY: dict[Type[xm.ExecutableSpec], RegistrationCallable] = {}
+_REGISTRY: dict[tp.Type[xm.ExecutableSpec], RegistrationCallable] = {}
 
 
 def register(
-    *typs: Type[ExecutableSpecT],
-) -> Callable[[RegistrationCallable], RegistrationCallable]:
+    *typs: tp.Type[ExecutableSpecT],
+) -> tp.Callable[[RegistrationCallable], RegistrationCallable]:
     def decorator(
         registration_callable: RegistrationCallable,
     ) -> RegistrationCallable:
@@ -38,8 +38,8 @@ def register(
 
 
 def route(
-    typ: Type[ExecutableSpecT],
-    packageables: Sequence[IndexedContainer[xm.Packageable]],
-) -> Sequence[IndexedContainer[xm.Executable]]:
+    typ: tp.Type[ExecutableSpecT],
+    packageables: tp.Sequence[IndexedContainer[xm.Packageable]],
+) -> tp.Sequence[IndexedContainer[xm.Executable]]:
     global _REGISTRY
     return _REGISTRY[typ](packageables)

xm_slurm/packaging/router.py

@@ -1,6 +1,6 @@
 import collections
 import logging
-from typing import Sequence, Type
+import typing as tp
 
 from xmanager import xm
 
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 
 def package(
-    packageables: Sequence[xm.Packageable],
+    packageables: tp.Sequence[xm.Packageable],
 ) -> list[xm.Executable]:
     """
     Takes as input a list of packageables and returns a mapping of
@@ -23,7 +23,7 @@ def package(
     # Docker targets to be collected.
     # These are a mapping from `DockerTarget` to the latest digest of the image.
     targets_by_type = collections.defaultdict[
-        Type[xm.ExecutableSpec], list[IndexedContainer[xm.Packageable]]
+        tp.Type[xm.ExecutableSpec], list[IndexedContainer[xm.Packageable]]
     ](list)
 
     # Collect dockerfiles that need to be built locally

xm_slurm/packaging/utils.py

@@ -1,20 +1,20 @@
 import collections
 import logging
-from typing import ParamSpec, Sequence, TypeVar
+import typing as tp
 
 from xmanager import xm
 
 from xm_slurm.packaging.registry import IndexedContainer
 
-T = TypeVar("T")
-P = ParamSpec("P")
-ReturnT = TypeVar("ReturnT")
+T = tp.TypeVar("T")
+P = tp.ParamSpec("P")
+ReturnT = tp.TypeVar("ReturnT")
 
 logger = logging.getLogger(__name__)
 
 
 def collect_executors_by_executable(
-    targets: Sequence[IndexedContainer[xm.Packageable]],
+    targets: tp.Sequence[IndexedContainer[xm.Packageable]],
 ) -> dict[xm.ExecutableSpec, set[xm.ExecutorSpec]]:
     executors_by_executable = collections.defaultdict(set)
     for target in targets:

xm_slurm/resources.py

@@ -1,9 +1,13 @@
+import builtins
+import collections.abc
+import datetime as dt
 import enum
 import itertools
 import math
-from typing import Mapping
+import re
+import typing as tp
 
-from xm_slurm import config
+from xm_slurm import config, utils
 
 
 class ResourceType(enum.IntEnum):
@@ -57,32 +61,123 @@ assert AcceleratorType | {
 } == set(ResourceType.__members__.values()), "Resource types are not exhaustive."
 
 
-ResourceQuantity = int | float
-
-
 class FeatureType(enum.IntEnum):
     NVIDIA_MIG = 1
     NVIDIA_NVLINK = 2
 
 
+class InvalidTopologyError(Exception):
+    """An unrecognized topology has been provided."""
+
+
+TOPOLOGY_REGEX = re.compile(r"^(?P<dims>[\d]+(?:x[\d]+)*)$")
+
+
+class Topology:
+    mesh: str
+    dimensions: list[int]
+    switches: int | None
+    switches_grace_period: dt.timedelta | None
+
+    def __init__(
+        self,
+        mesh: str,
+        /,
+        *,
+        switches: int | None = None,
+        switches_grace_period: dt.timedelta | None = None,
+    ):
+        mesh_match = TOPOLOGY_REGEX.fullmatch(mesh)
+        if not mesh_match:
+            raise InvalidTopologyError(f"Invalid topology mesh: {mesh!r}.")
+
+        self.mesh = mesh
+        self.dimensions = list(map(int, mesh_match.group("dims").split("x")))
+        if switches is not None:
+            assert (
+                isinstance(switches, int) and switches > 0
+            ), "Switches must be a positive integer."
+        self.switches = switches
+        if switches_grace_period is not None:
+            assert isinstance(
+                switches_grace_period, dt.timedelta
+            ), "Switches grace period must be a `datetime.timedelta`."
+        self.switches_grace_period = switches_grace_period
+
+    @property
+    def chip_count(self) -> int:
+        return math.prod(self.dimensions)
+
+    @property
+    def ndim(self) -> int:
+        return len(self.dimensions)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Topology):
+            return False
+        return (
+            self.mesh == other.mesh
+            and self.switches == other.switches
+            and self.switches_grace_period == other.switches_grace_period
+        )
+
+    def __hash__(self) -> int:
+        return hash((self.mesh, self.switches, self.switches_grace_period))
+
+
+ResourceQuantity = int | float | Topology
+
+
+def _parse_resource_quantity(
+    resource_name: ResourceType | str, value: ResourceQuantity
+) -> tuple[float, Topology | None]:
+    if isinstance(resource_name, ResourceType):
+        resource_name = resource_name.name
+    match value:
+        case Topology() as topology:
+            return topology.chip_count, topology
+        case builtins.str() as topology_str if (
+            "x" in topology_str and TOPOLOGY_REGEX.fullmatch(topology_str) is not None
+        ):
+            topology = Topology(topology_str)
+            return topology.chip_count, topology
+        case builtins.str() as num_str:
+            try:
+                value = float(num_str)
+                return int(value) if value.is_integer() else value, None
+            except ValueError as e:
+                raise ValueError(
+                    f"Couldn't parse resource quantity for {resource_name}. "
+                    f"{num_str!r} was given."
+                ) from e
+        case int() | float():
+            return value, None
+        case _:
+            raise ValueError(f"Invalid resource quantity: {value!r} for {resource_name!r}.")
+
+
 class JobRequirements:
     replicas: int
     location: str | None
     accelerator: ResourceType | None
-    cluster: config.SlurmClusterConfig | None = None
+    topology: Topology | None
+    cluster: config.SlurmClusterConfig
 
     def __init__(
         self,
         *,
-        resources: Mapping[ResourceType | str, ResourceQuantity] | None = None,
-        replicas: int = 1,
-        location: str | None = None,
-        cluster: config.SlurmClusterConfig | None = None,
+        resources: tp.Mapping[ResourceType | str, ResourceQuantity] | None = None,
+        replicas: int | None = None,
+        location: str | tp.Iterable[str] | None = None,
+        cluster: config.SlurmClusterConfig,
         **kw_resources: ResourceQuantity,
     ):
-        self.replicas = replicas or 1
+        if isinstance(location, collections.abc.Iterable) and not isinstance(location, str):
+            location = ",".join(location)
         self.location = location
+
         self.accelerator = None
+        self.topology = None
         self.cluster = cluster
 
         if resources is None:
@@ -90,6 +185,7 @@ class JobRequirements:
 
         self.task_requirements: dict[ResourceType | str, ResourceQuantity] = {}
         for resource_name, value in itertools.chain(resources.items(), kw_resources.items()):
+            quantity, topology = _parse_resource_quantity(resource_name, value)
             match resource_name:
                 case str() if resource_name.upper() in ResourceType.__members__:
                     resource = ResourceType[resource_name.upper()]
@@ -106,58 +202,132 @@ class JobRequirements:
                 if self.accelerator is not None:
                     raise ValueError("Accelerator already set.")
                 self.accelerator = resource  # type: ignore
+                self.topology = topology or Topology(f"{quantity:g}")
+            elif topology is not None:
+                raise ValueError(
+                    f"A topology was specified for a non-accelerator resource: {resource_name!r}."
+                )
 
             if resource in self.task_requirements:
                 raise ValueError(f"{resource} has been specified twice.")
-            self.task_requirements[resource] = value
+            self.task_requirements[resource] = quantity
+
+        if self.topology is not None and self.topology.ndim > 2:
+            raise ValueError("Topologies with more than 2 dimensions are not supported.")
+
+        if (
+            self.accelerator is not None
+            and self.topology is not None
+            and len(self.topology.dimensions) == 2
+        ):
+            if replicas is not None and replicas != self.topology.dimensions[1]:
+                raise ValueError(
+                    f"For multihost GPUs with topology {self.topology}, replicas should"
+                    f"be either None or {self.topology.dimensions[1]}. Found: "
+                    f"{replicas}"
+                )
+            replicas = self.topology.dimensions[1]
+
+        if replicas is not None and replicas <= 0:
+            raise ValueError(f"Replicas must be a positive integer, got {replicas!r}")
+        self.replicas = replicas or 1
 
     def to_directives(self) -> list[str]:
-        if self.cluster is None:
-            raise ValueError("Cannnot derive Slurm directives for requirements without a cluster.")
         directives = []
 
         for resource, value in self.task_requirements.items():
             match resource:
                 case ResourceType.EPHEMERAL_STORAGE | ResourceType.DISK:
-                    assert isinstance(value, int), "Disk space must be an integer"
+                    assert isinstance(
+                        value, int
+                    ), f"Disk space must be an integer, got {type(value)!r}"
                     directives.append(f"--tmp={math.ceil(value / 2**20)}M")
                 case ResourceType.MEMORY | ResourceType.RAM:
                     num_cpus = self.task_requirements.get(ResourceType.CPU, 1)
-                    assert isinstance(value, (int, float)), "Memory must be an integer or float"
-                    assert isinstance(num_cpus, int), "CPU must be an integer"
+                    assert isinstance(
+                        value, (int, float)
+                    ), f"Memory must be an integer or float, got {type(value)!r}"
+                    assert isinstance(
+                        num_cpus, int
+                    ), f"CPU must be an integer, got {type(num_cpus)!r}"
                     directives.append(f"--mem-per-cpu={math.ceil(value / num_cpus / 2**20)}M")
                 case ResourceType.CPU:
-                    assert isinstance(value, int), "CPU must be an integer"
+                    assert isinstance(value, int), f"CPU must be an integer, got {type(value)!r}"
                     directives.append(f"--cpus-per-task={value}")
                 case ResourceType.GPU:
-                    assert isinstance(value, int), "GPU must be an integer"
-                    directives.append(f"--gpus-per-task={value}")
+                    assert isinstance(value, int), f"GPU must be an integer, got {type(value)!r}"
+                    directives.append(f"--gpus={value}")
                 case ResourceType() if resource in AcceleratorType:
-                    assert isinstance(value, int), "Accelerator must be an integer"
+                    assert isinstance(
+                        value, int
+                    ), f"Accelerator must be an integer, got {type(value)!r}"
                     resource_type = self.cluster.resources.get(resource, None)
                     if resource_type is None:
                         raise ValueError(
                             f"Cluster {self.cluster.name} does not map resource type {resource!r}."
                         )
-                    directives.append(f"--gpus-per-task={resource_type}:{value}")
+                    directives.append(f"--gpus={resource_type}:{value}")
                 case str():
                     directives.append(f"--gres={resource}:{value}")
 
-        directives.append(f"--ntasks={self.replicas}")
         if self.location:
+            assert isinstance(
+                self.location, str
+            ), f"Location must be a string, got {type(self.location)!r}"
             directives.append(f"--nodelist={self.location}")
 
+        assert (
+            isinstance(self.replicas, int) and self.replicas > 0
+        ), f"Replicas must be a positive integer, got {self.replicas!r}"
+        directives.append(f"--ntasks={self.replicas}")
+
+        if self.topology is not None:
+            assert self.accelerator is not None, "Accelerator must be set."
+            match self.accelerator:
+                case ResourceType.GPU:
+                    directives.append(f"--gpus-per-task={self.topology.dimensions[0]}")
+                case ResourceType() if self.accelerator in AcceleratorType:
+                    resource_type = self.cluster.resources[self.accelerator]
+                    directives.append(
+                        f"--gpus-per-task={resource_type}:{self.topology.dimensions[0]}"
+                    )
+
+            if self.topology.switches is not None:
+                switches_timeout = (
+                    f"@{utils.timestr_from_timedelta(self.topology.switches_grace_period)}"
+                    if self.topology.switches_grace_period is not None
+                    else ""
+                )
+                directives.append(f"--switches={self.topology.switches}{switches_timeout}")
+
         return directives
 
     def replace(
         self,
-        cluster: config.SlurmClusterConfig | None,
+        replicas: int | None = None,
+        location: str | None = None,
+        cluster: config.SlurmClusterConfig | None = None,
         **kw_resources: ResourceQuantity,
     ) -> "JobRequirements":
+        # Merge kw_resources into existing task_requirements, removing conflicting enum keys
+        merged_resources = dict(self.task_requirements)
+
+        # Remove ResourceType keys that will be overridden by string keys in kw_resources
+        for key in list(merged_resources.keys()):
+            if isinstance(key, ResourceType) and any(
+                ResourceType[name.upper()] == key
+                for name in kw_resources
+                if name.upper() in ResourceType.__members__
+            ):
+                del merged_resources[key]
+
+        merged_resources.update(kw_resources)  # type: ignore
+
         return JobRequirements(
-            resources=self.task_requirements | kw_resources,  # type: ignore
-            replicas=self.replicas,
-            cluster=cluster or self.cluster,
+            resources=merged_resources,
+            replicas=replicas if replicas is not None else self.replicas,
+            location=location if location is not None else self.location,
+            cluster=cluster if cluster is not None else self.cluster,
         )
 
     def __repr__(self) -> str:
@@ -169,7 +339,7 @@ class JobRequirements:
             args.append(f"{resource.lower()}={value!r}")
 
         if self.replicas != 1:
-            args.append(f"replicas={self.replicas}")
+            args.append(f"replicas={self.replicas!r}")
 
         if self.cluster is not None:
             args.append(f"cluster={self.cluster!r}")

xm_slurm/status.py

@@ -2,7 +2,7 @@
 
 import enum
 import re
-from typing import Sequence
+import typing as tp
 
 from xmanager import xm
 
@@ -151,7 +151,7 @@ class SlurmWorkUnitStatus(xm.ExperimentUnitStatus):
     """Status of a Slurm experiment job."""
 
     @classmethod
-    def aggregate(cls, states: Sequence[SlurmJobState]) -> "SlurmWorkUnitStatus":
+    def aggregate(cls, states: tp.Sequence[SlurmJobState]) -> "SlurmWorkUnitStatus":
         """Aggregate a sequence of statuses into a single status."""
         assert len(states) > 0, "Cannot aggregate empty sequence of statuses."
         max_error_state: SlurmJobState | None = None

xm_slurm/templates/slurm/job-array.bash.j2

@@ -7,9 +7,9 @@
 
 {% block bootstrap %}
 srun \
+  --label \
   --unbuffered \
   --kill-on-bad-exit=0 \
-  --overlap \
   --export="ALL" \
   bash <<'SRUN_EOF' &
 set -Eeuxo pipefail

xm_slurm/templates/slurm/job-group.bash.j2

@@ -29,6 +29,7 @@
 {% block bootstrap %}
 {% for job in job_group.jobs.values() +%}
 srun \
+  --label \
   --unbuffered \
   --kill-on-bad-exit=0 \
   --export="ALL" \

xm_slurm/templates/slurm/job.bash.j2

@@ -59,9 +59,9 @@ export {{ key }}="{{ value }}"
 
 {% block bootstrap %}
 srun \
+  --label \
   --unbuffered \
   --kill-on-bad-exit=0 \
-  --overlap \
   --export="ALL" \
   bash <<'SRUN_EOF' &
 set -Eeuxo pipefail

xm_slurm/utils.py

@@ -1,4 +1,5 @@
 import asyncio
+import datetime as dt
 import functools
 import logging
 import os
@@ -186,3 +187,10 @@ def run_command(
         stdout=stdout,
         stderr=stderr,
     )
+
+
+def timestr_from_timedelta(time: dt.timedelta) -> str:
+    days = time.days
+    hours, remainder = divmod(time.seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    return f"{days}-{hours:02}:{minutes:02}:{seconds:02}"

{xmanager_slurm-0.4.14.dist-info → xmanager_slurm-0.4.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.14
+Version: 0.4.15
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>

{xmanager_slurm-0.4.14.dist-info → xmanager_slurm-0.4.15.dist-info}/RECORD

@@ -1,21 +1,21 @@
-xm_slurm/__init__.py,sha256=WgRn9HDYa5H3sfIH-HZu33liBOh98jM4GqcR349RaSY,1086
-xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
-xm_slurm/config.py,sha256=i9WuxjfGBgVoHHDNk3BvO7LCwpBTJeRnOiFvTU-FHrk,7330
+xm_slurm/__init__.py,sha256=VNbvBLbv5ccbPxQUpbiwgoo72qI3FrATTloevufstzY,1112
+xm_slurm/batching.py,sha256=ynbMRItxNtBn0SbkhHbrv5ugYuHeMw-7BP7a_-I6Oqg,4384
+xm_slurm/config.py,sha256=6RAdMTgteUJ_tmr3H_RE6MHtzda202_unU2S24iPHUE,7319
 xm_slurm/console.py,sha256=UpMqeJ0C8i0pkue1AHnnyyX0bFJ9zZeJ7HBR6yhuA8A,54
 xm_slurm/constants.py,sha256=zefVtlFdflgSolie5g_rVxWV-Zpydxapchm3y0a2FDc,999
-xm_slurm/dependencies.py,sha256=-5gN_tpfs3dOA7H5_MIHO2ratb7F5Pm_yjkR5rZcgI8,6421
+xm_slurm/dependencies.py,sha256=G-8vfmvSptZH6c_Ow51SwT84Dr6LI1clRj8F8wOUkiw,6421
 xm_slurm/executables.py,sha256=fGmrFBl-258bMn6ip5adYeM7xxUHAeIbDN9zD2FDGtY,6373
 xm_slurm/execution.py,sha256=mTy5u2oP2StIbGzjaSiGCUAwXuBFOiaJ5ephWoc25hI,31799
-xm_slurm/executors.py,sha256=bUgKcgtvf-nPGjcuHRzUAqD1r3_vwea_h-Y9MAB-Kqo,4887
+xm_slurm/executors.py,sha256=karM5u2UEG2IWi0z548_vasyBACrXGV675rCllJmwZw,8616
 xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
-xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
+xm_slurm/job_blocks.py,sha256=BFOOYgeodoGIQsB5PdC7SsOUou5aZx-1qbQ7lcqqylI,604
 xm_slurm/metadata_context.py,sha256=mksVRbVUuistL1uE7TC-fkW-Y69On52jN_svP1e1kiQ,7841
-xm_slurm/packageables.py,sha256=K6vNhLvASdnqsc8vXlT3h9cObJpC9Rbw93pUBvBwapQ,12209
-xm_slurm/resources.py,sha256=T7uje3E6oWbZSrsxykgW-40DE-Bvw_NWDM2qXbw2rgI,5740
-xm_slurm/status.py,sha256=WTWiDHi-ZHtwHRnDP0cGa-27zTSm6LkA-GCKsN-zBgg,6916
+xm_slurm/packageables.py,sha256=aEZUQpddfq4FK6h4f6kgGEI4XcOufhm68MjoDFOYR4U,12261
+xm_slurm/resources.py,sha256=aC8MzO_7fB9IAdTACvhwVOaNDjLOlWnCh428-8_IDYA,12322
+xm_slurm/status.py,sha256=JIBCJPOYsmeJOQbzdACXA2vTWK7g8YWWhzpGP79e7JE,6911
 xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
-xm_slurm/utils.py,sha256=xtFvktaxr0z65sTdu6HhOVfyo0OAB9t-EYXWcYrQQEU,5958
+xm_slurm/utils.py,sha256=9w98HlXF0U9cKKtoB8QtGm0CnB0MnnzBARKlbbVNNpU,6211
 xm_slurm/api/__init__.py,sha256=cyao3LZ3uLftu1wIv1aN7Qvsl6gYzYpkxeehTHZ0fA8,1089
 xm_slurm/api/abc.py,sha256=-lS2OndnOuEiwNdr8ccQKkwMd1iDmKMmkBOSTvo5H5w,1816
 xm_slurm/api/models.py,sha256=_INVh0j-4-rRs0WASyg4fNB6NF1L1nUeGgQ6-XnbwsM,1610
@@ -23,13 +23,13 @@ xm_slurm/api/sqlite/client.py,sha256=jAesCKDuYwnNcAxwJk_1b1TB8cT_QGbSjo1UE3mZjEQ
 xm_slurm/api/web/client.py,sha256=uO67Y7fnQ-w__Vm_A5BEuy7Qi8wQcWk3vIsBGEBkyfk,6261
 xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xm_slurm/contrib/clusters/__init__.py,sha256=XFCVnkThiU3_8uA_tUgDByOBanXNHrxDvfmuptmQ2KE,2214
-xm_slurm/contrib/clusters/drac.py,sha256=vY3dxrNUk12H9Gq-tuCcqo2YcdTGq-4LJnQF6DzD4_k,7431
-xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
+xm_slurm/contrib/clusters/drac.py,sha256=_iubsmzz5tK2KTaKqSuykS3IDtsdXqJ0MXep1THSJUM,7020
+xm_slurm/experimental/parameter_controller.py,sha256=IrFzq104LkZrhzuirit5GUZDXDvv2bBSYNMh3orsiPY,8518
 xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
 xm_slurm/packaging/docker.py,sha256=-DWcB9qqbeHmIEqyfF0-v6xOT25ae90u2x-QZ7kluOw,13579
-xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CUE,1148
-xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
-xm_slurm/packaging/utils.py,sha256=6EAb17zKQQeuyNY2EV9AoW1RvnDGrQwmIT9wtQEsC4c,632
+xm_slurm/packaging/registry.py,sha256=Hq56KhqsQRxgr_y1EQhcZORlnrs13xY5vDGge5WEgYU,1134
+xm_slurm/packaging/router.py,sha256=MLWvy-shJzRAx4YCi9z9Dj_PoWrXZO8T71DDlLOcjaM,2062
+xm_slurm/packaging/utils.py,sha256=KlU_GGkFH1Xu5VZkAMqRilmq6SV1iLai80beEZ3UQmw,616
 xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
 xm_slurm/scripts/cli.py,sha256=zzsQpvkx9VThAeQPM34iDK9wAWfCVCIIvLNI12UaMhw,2577
 xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=7qSJl2VN5poz-Hh8Gjo7--qR-k3lmfGtBu2mNbfG2uA,1499
@@ -37,16 +37,16 @@ xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgC
 xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
 xm_slurm/templates/docker/uv.Dockerfile,sha256=L2UJMX2c8waMdrRhiqPytQe3pTBu6u5PpMhJYsKkbEg,1040
 xm_slurm/templates/slurm/entrypoint.bash.j2,sha256=MRdSVwgGrgQdpEhqfkP35IidgsblrtVXB1YWzvE9hkk,666
-xm_slurm/templates/slurm/job-array.bash.j2,sha256=smxmSSzBEUHm6MJF-nYPVVjK6CLKrb1fRxF_tfrzAX8,552
-xm_slurm/templates/slurm/job-group.bash.j2,sha256=Cp8YhNOxYqaOkl4MFjQlcaLMGZwdDh97m8OGT5RWbAo,1101
-xm_slurm/templates/slurm/job.bash.j2,sha256=DrDipliaEfiHbq9vDfOdfD8zBVFLy1jjlvCV-9-6k9s,2086
+xm_slurm/templates/slurm/job-array.bash.j2,sha256=j7jkJjSbe39XvSTJ9rmK2oVnHdntElIhdS5PFpZzpFs,550
+xm_slurm/templates/slurm/job-group.bash.j2,sha256=vH5HwneVsVSHx6dPZwbLa4KT9NedRbrZ7cWNE5pXi-M,1113
+xm_slurm/templates/slurm/job.bash.j2,sha256=JnK0D8_3tVNpnvPwM5yL_rjLcjqhuHiCtolDjUGAwpk,2084
 xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=ri5FgoKs6_bQVf5DO8SL4rJf4UsLxV34aOV-OD8VWDU,2526
 xm_slurm/templates/slurm/fragments/proxy.bash.j2,sha256=VJLglZo-Nvx9R-qe3rHTxr07CylTQ6Z9NwBzvIpAZrA,814
 xm_slurm/templates/slurm/library/retry.bash,sha256=bLe59qvfWEk17rE1wZ4EHiHba3RvR2WWZPq-kSe8RAA,2164
 xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=v0LwHM-kBW8sJqVcVA2jYr1n44imDSZrJqmqlr5uTGc,1980
 xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=zWLsFEuVzOMSETOmv4A5ZCV4oQHwCipiR6wi79XVzNI,1188
-xmanager_slurm-0.4.14.dist-info/METADATA,sha256=T7xNy0jmrKhQemaDhCg9E-J64gWkabjJGpxDYgdsBx8,1007
-xmanager_slurm-0.4.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.14.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.14.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.14.dist-info/RECORD,,
+xmanager_slurm-0.4.15.dist-info/METADATA,sha256=xgJUFDConlb4R5W0cOK3xb2fuIR7x_tjNpOybMagT_A,1007
+xmanager_slurm-0.4.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.15.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.15.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.15.dist-info/RECORD,,