xmanager-slurm 0.4.13__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/execution.py

@@ -3,6 +3,8 @@ import collections.abc
 import dataclasses
 import functools
 import hashlib
+import importlib
+import importlib.resources
 import logging
 import operator
 import os
@@ -311,6 +313,14 @@ def get_template_env(runtime: ContainerRuntime) -> j2.Environment:
     template_env.globals["raise"] = _raise_template_exception
     template_env.globals["operator"] = operator
 
+    # Iterate over stdlib files and insert them into the template environment
+    stdlib = []
+    for file in importlib.resources.files("xm_slurm.templates.slurm.library").iterdir():
+        if not file.is_file() or not file.name.endswith(".bash"):
+            continue
+        stdlib.append(file.read_text())
+    template_env.globals["stdlib"] = stdlib
+
     entrypoint_template = template_env.get_template("entrypoint.bash.j2")
     template_env.globals.update(entrypoint_template.module.__dict__)
 

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(