xmanager-slurm 0.4.19__py3-none-any.whl

xm_slurm/api/web/client.py

@@ -0,0 +1,173 @@
+import dataclasses
+
+import backoff
+import httpx
+
+from xm_slurm.api import models
+from xm_slurm.api.abc import XManagerAPI
+
+# Define which exceptions should trigger a retry
+RETRY_EXCEPTIONS = (
+    httpx.ConnectError,
+    httpx.ConnectTimeout,
+    httpx.ReadTimeout,
+    httpx.WriteTimeout,
+    httpx.NetworkError,
+)
+
+
+# Common backoff decorator for all API calls
+def with_backoff(f):
+    return backoff.on_exception(
+        backoff.expo,
+        RETRY_EXCEPTIONS,
+        max_tries=3,  # Maximum number of attempts
+        max_time=30,  # Maximum total time to try in seconds
+        jitter=backoff.full_jitter,  # Add jitter to prevent thundering herd
+    )(f)
+
+
+class XManagerWebAPI(XManagerAPI):
+    def __init__(self, base_url: str, token: str):
+        self.base_url = base_url.rstrip("/")
+        self.client = httpx.Client(headers={"Authorization": f"Bearer {token}"}, verify=False)
+
+    def _make_url(self, path: str) -> str:
+        return f"{self.base_url}{path}"
+
+    @with_backoff
+    def get_experiment(self, xid: int) -> models.Experiment:
+        response = self.client.get(self._make_url(f"/experiment/{xid}"))
+        response.raise_for_status()
+        data = response.json()
+        # Construct work units with nested jobs and artifacts
+        work_units = []
+        for wu_data in data.pop("work_units", []):
+            # Build jobs for this work unit
+            jobs = [
+                models.SlurmJob(
+                    name=job["name"],
+                    slurm_job_id=job["slurm_job_id"],
+                    slurm_ssh_config=job["slurm_ssh_config"],
+                )
+                for job in wu_data.pop("jobs", [])
+            ]
+
+            # Build artifacts for this work unit
+            artifacts = [
+                models.Artifact(name=artifact["name"], uri=artifact["uri"])
+                for artifact in wu_data.pop("artifacts", [])
+            ]
+
+            # Create work unit with its jobs and artifacts
+            wu_data["jobs"] = jobs
+            wu_data["artifacts"] = artifacts
+            work_units.append(models.WorkUnit(**wu_data))
+
+        # Build experiment artifacts
+        artifacts = [
+            models.Artifact(name=artifact["name"], uri=artifact["uri"])
+            for artifact in data.pop("artifacts", [])
+        ]
+
+        return models.Experiment(**data, work_units=work_units, artifacts=artifacts)
+
+    @with_backoff
+    def delete_experiment(self, experiment_id: int) -> None:
+        response = self.client.delete(self._make_url(f"/experiment/{experiment_id}"))
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_experiment(self, experiment: models.ExperimentPatch) -> int:
+        assert experiment.title is not None, "Title must be set in the experiment model."
+        assert (
+            experiment.description is None and experiment.note is None and experiment.tags is None
+        ), "Only title should be set in the experiment model."
+
+        response = self.client.put(
+            self._make_url("/experiment"), json=dataclasses.asdict(experiment)
+        )
+        response.raise_for_status()
+        return int(response.json()["xid"])
+
+    @with_backoff
+    def update_experiment(
+        self, experiment_id: int, experiment_patch: models.ExperimentPatch
+    ) -> None:
+        response = self.client.patch(
+            self._make_url(f"/experiment/{experiment_id}"),
+            json=dataclasses.asdict(experiment_patch),
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_work_unit(self, experiment_id: int, work_unit: models.WorkUnitPatch) -> None:
+        response = self.client.put(
+            self._make_url(f"/experiment/{experiment_id}/wu"),
+            json=dataclasses.asdict(work_unit),
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_job(self, experiment_id: int, work_unit_id: int, job: models.SlurmJob) -> None:
+        response = self.client.put(
+            self._make_url(f"/experiment/{experiment_id}/wu/{work_unit_id}/job"),
+            json=dataclasses.asdict(job),
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_work_unit_artifact(
+        self, experiment_id: int, work_unit_id: int, artifact: models.Artifact
+    ) -> None:
+        response = self.client.put(
+            self._make_url(f"/experiment/{experiment_id}/wu/{work_unit_id}/artifact"),
+            json=dataclasses.asdict(artifact),
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
+        response = self.client.delete(
+            self._make_url(f"/experiment/{experiment_id}/wu/{work_unit_id}/artifact/{name}")
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def delete_experiment_artifact(self, experiment_id: int, name: str) -> None:
+        response = self.client.delete(
+            self._make_url(f"/experiment/{experiment_id}/artifact/{name}")
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_experiment_artifact(self, experiment_id: int, artifact: models.Artifact) -> None:
+        response = self.client.put(
+            self._make_url(f"/experiment/{experiment_id}/artifact"),
+            json=dataclasses.asdict(artifact),
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def insert_experiment_config_artifact(
+        self, experiment_id: int, artifact: models.ConfigArtifact
+    ) -> None:
+        response = self.client.put(
+            self._make_url(f"/experiment/{experiment_id}/config"), json=dataclasses.asdict(artifact)
+        )
+        response.raise_for_status()
+
+    @with_backoff
+    def delete_experiment_config_artifact(self, experiment_id: int, name: str) -> None:
+        response = self.client.delete(self._make_url(f"/experiment/{experiment_id}/config/{name}"))
+        response.raise_for_status()
+
+    @with_backoff
+    def update_work_unit(
+        self, experiment_id: int, work_unit_id: int, patch: models.ExperimentUnitPatch
+    ) -> None:
+        response = self.client.patch(
+            self._make_url(f"/experiment/{experiment_id}/wu/{work_unit_id}"),
+            json=dataclasses.asdict(patch),
+        )
+        response.raise_for_status()

xm_slurm/batching.py

@@ -0,0 +1,139 @@
+import asyncio
+import collections
+import dataclasses
+import inspect
+import time
+import types
+import typing as tp
+
+T = tp.TypeVar("T", contravariant=True)
+R = tp.TypeVar("R", covariant=True)
+P = tp.ParamSpec("P")
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Request:
+    args: inspect.BoundArguments
+    future: asyncio.Future
+
+
+def stack_bound_arguments(
+    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, tp.Any]()
+    for bound_args in bound_arguments:
+        for name, value in bound_args.arguments.items():
+            stacked_args.setdefault(name, [])
+            stacked_args[name].append(value)
+    return inspect.BoundArguments(signature, stacked_args)
+
+
+class batch(tp.Generic[R]):
+    __slots__ = (
+        "fn",
+        "signature",
+        "max_batch_size",
+        "batch_timeout",
+        "loop",
+        "process_batch_task",
+        "queue",
+    )
+    __name__: str = "batch"
+    __qualname__: str = "batch"
+
+    def __init__(
+        self,
+        fn: tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]],
+        /,
+        *,
+        max_batch_size: int,
+        batch_timeout: float,
+    ) -> None:
+        self.fn = fn
+        self.signature = inspect.signature(fn)
+
+        self.max_batch_size = max_batch_size
+        self.batch_timeout = batch_timeout
+
+        self.loop: asyncio.AbstractEventLoop | None = None
+        self.process_batch_task: asyncio.Task | None = None
+
+        self.queue = asyncio.Queue[Request]()
+
+    async def _process_batch(self):
+        assert self.loop is not None
+        while not self.loop.is_closed():
+            batch = await self._wait_for_batch()
+            assert len(batch) > 0
+
+            bound_args = stack_bound_arguments(self.signature, [request.args for request in batch])
+            futures = [request.future for request in batch]
+
+            results_future = self.fn(*bound_args.args, *bound_args.kwargs)
+
+            try:
+                results = await results_future
+
+                for result, future in zip(results, futures):
+                    future.set_result(result)
+            except Exception as e:
+                for future in futures:
+                    future.set_exception(e)
+
+    async def _wait_for_batch(self) -> list[Request]:
+        batch = [await self.queue.get()]
+
+        batch_start_time = time.time()
+        while True:
+            remaining_batch_time_s = max(self.batch_timeout - (time.time() - batch_start_time), 0)
+
+            try:
+                request = await asyncio.wait_for(self.queue.get(), timeout=remaining_batch_time_s)
+                batch.append(request)
+            except asyncio.TimeoutError:
+                break
+
+            if (
+                time.time() - batch_start_time >= self.batch_timeout
+                or len(batch) >= self.max_batch_size
+            ):
+                break
+
+        return batch
+
+    def __get__(self, obj: tp.Any, objtype: tp.Type[tp.Any]) -> tp.Any:
+        del objtype
+        if isinstance(self.fn, staticmethod):
+            return self.__call__
+        else:
+            return types.MethodType(self, obj)
+
+    @property
+    def __func__(self) -> tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]]:
+        return self.fn
+
+    @property
+    def __wrapped__(self) -> tp.Callable[..., tp.Coroutine[None, None, tp.Sequence[R]]]:
+        return self.fn
+
+    @property
+    def __isabstractmethod__(self) -> bool:
+        """Return whether the wrapped function is abstract."""
+        return getattr(self.fn, "__isabstractmethod__", False)
+
+    @property
+    def _is_coroutine(self) -> bool:
+        # TODO(jfarebro): py312 adds inspect.markcoroutinefunction
+        # until then this is just a hack
+        return asyncio.coroutines._is_coroutine  # type: ignore
+
+    async def __call__(self, *args, **kwargs) -> R:
+        if self.loop is None and self.process_batch_task is None:
+            self.loop = asyncio.get_event_loop()
+            self.process_batch_task = self.loop.create_task(self._process_batch())
+
+        future = asyncio.Future[R]()
+        bound_args = self.signature.bind(*args, **kwargs)
+        self.queue.put_nowait(Request(args=bound_args, future=future))
+        return await future

xm_slurm/config.py

@@ -0,0 +1,189 @@
+import dataclasses
+import enum
+import functools
+import json
+import os
+import pathlib
+import typing as tp
+
+import asyncssh
+from xmanager import xm
+
+
+class ContainerRuntime(enum.Enum):
+    """The container engine to use."""
+
+    SINGULARITY = enum.auto()
+    APPTAINER = enum.auto()
+    DOCKER = enum.auto()
+    PODMAN = enum.auto()
+
+    @classmethod
+    def from_string(
+        cls, runtime: tp.Literal["singularity", "apptainer", "docker", "podman"]
+    ) -> "ContainerRuntime":
+        return {
+            "singularity": cls.SINGULARITY,
+            "apptainer": cls.APPTAINER,
+            "docker": cls.DOCKER,
+            "podman": cls.PODMAN,
+        }[runtime]
+
+    def __str__(self):
+        if self is self.SINGULARITY:
+            return "singularity"
+        elif self is self.APPTAINER:
+            return "apptainer"
+        elif self is self.DOCKER:
+            return "docker"
+        elif self is self.PODMAN:
+            return "podman"
+        else:
+            raise NotImplementedError
+
+
+class Endpoint(tp.NamedTuple):
+    hostname: str
+    port: int | None
+
+    def __str__(self) -> str:
+        if self.port is None or self.port == asyncssh.DEFAULT_PORT:
+            return self.hostname
+        return f"[{self.hostname}]:{self.port}"
+
+
+class PublicKey(tp.NamedTuple):
+    algorithm: str
+    key: str
+
+
+@dataclasses.dataclass
+class SSHConfig:
+    endpoints: tuple[Endpoint, ...]
+    public_key: PublicKey | None = None
+    user: str | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.endpoints, tuple):
+            raise TypeError(f"endpoints must be a tuple, not {type(self.endpoints)}")
+        if len(self.endpoints) == 0:
+            raise ValueError("endpoints must be a non-empty tuple")
+        if not all(isinstance(endpoint, Endpoint) for endpoint in self.endpoints):
+            raise TypeError(f"endpoints must be a tuple of strings, not {type(self.endpoints)}")
+
+        if not isinstance(self.user, str | None):
+            raise TypeError(f"user must be a string or None, not {type(self.user)}")
+        if not isinstance(self.public_key, PublicKey | None):
+            raise TypeError(
+                f"public_key must be a SSHPublicKey or None, not {type(self.public_key)}"
+            )
+
+    @functools.cached_property
+    def known_hosts(self) -> asyncssh.SSHKnownHosts | None:
+        if self.public_key is None:
+            return None
+
+        known_hosts = []
+        for endpoint in self.endpoints:
+            known_hosts.append(f"{endpoint!s} {self.public_key.algorithm} {self.public_key.key}")
+
+        return asyncssh.SSHKnownHosts("\n".join(known_hosts))
+
+    def serialize(self):
+        return json.dumps({
+            "endpoints": tuple(tuple(endpoint) for endpoint in self.endpoints),
+            "public_key": tuple(self.public_key),
+            "user": self.user,
+        })
+
+    @classmethod
+    def deserialize(cls, data):
+        data = json.loads(data)
+        return cls(
+            endpoints=tuple(Endpoint(*endpoint) for endpoint in data["endpoints"]),
+            public_key=PublicKey(*data["public_key"]) if data["public_key"] else None,
+            user=data["user"],
+        )
+
+    def __hash__(self):
+        return hash((
+            type(self),
+            *(tuple(endpoint) for endpoint in self.endpoints),
+            self.public_key,
+            self.user,
+        ))
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class SlurmClusterConfig:
+    name: str
+
+    ssh: SSHConfig
+
+    # Job submission directory
+    cwd: str | None = None
+
+    # Additional scripting
+    prolog: str | None = None
+    epilog: str | None = None
+
+    # Job scheduling
+    account: str | None = None
+    partition: str | None = None
+    qos: str | None = None
+
+    # If true, a reverse proxy is initiated via the submission host.
+    proxy: tp.Literal["submission-host"] | str | None = None
+
+    runtime: ContainerRuntime
+
+    # Environment variables
+    host_environment: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+    container_environment: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+
+    # Mounts
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] = dataclasses.field(
+        default_factory=dict
+    )
+
+    # Resource mapping
+    resources: tp.Mapping["xm_slurm.ResourceType", 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: tp.Callable[[xm.Job], None] | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.ssh, SSHConfig):
+            raise TypeError(f"ssh must be a SlurmSSHConfig, not {type(self.ssh)}")
+        for src, dst in self.mounts.items():
+            if not isinstance(src, (str, os.PathLike)):
+                raise TypeError(
+                    f"Mount source must be a string or path-like object, not {type(src)}"
+                )
+            if not isinstance(dst, (str, os.PathLike)):
+                raise TypeError(
+                    f"Mount destination must be a string or path-like object, not {type(dst)}"
+                )
+
+            if not pathlib.Path(src).is_absolute():
+                raise ValueError(f"Mount source must be an absolute path: {src}")
+            if not pathlib.Path(dst).is_absolute():
+                raise ValueError(f"Mount destination must be an absolute path: {dst}")
+
+    def __hash__(self):
+        return hash((
+            type(self),
+            self.ssh,
+            self.cwd,
+            self.prolog,
+            self.epilog,
+            self.account,
+            self.partition,
+            self.qos,
+            self.proxy,
+            self.runtime,
+            frozenset(self.host_environment.items()),
+            frozenset(self.container_environment.items()),
+        ))

xm_slurm/console.py

@@ -0,0 +1,3 @@
+from rich.console import Console
+
+console = Console()

xm_slurm/constants.py

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

xm_slurm/contrib/clusters/__init__.py

@@ -0,0 +1,67 @@
+import logging
+import os
+
+from xm_slurm import config, resources
+from xm_slurm.contrib.clusters import drac
+
+# ComputeCanada alias
+cc = drac
+
+__all__ = ["drac", "mila", "cc"]
+
+logger = logging.getLogger(__name__)
+
+
+def mila(
+    *,
+    user: str | None = None,
+    partition: str | None = None,
+    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """Mila Cluster (https://docs.mila.quebec/)."""
+    if mounts is None:
+        mounts = {
+            "/network/scratch/${USER:0:1}/$USER": "/scratch",
+            # TODO: move these somewhere common to all cluster configs.
+            "/home/mila/${USER:0:1}/$USER/.local/state/xm-slurm": "/xm-slurm-state",
+            "/home/mila/${USER:0:1}/$USER/.ssh": "/home/mila/${USER:0:1}/$USER/.ssh",
+        }
+    mounts = dict(mounts) | {"/dev/infiniband": "/dev/infiniband"}
+
+    return config.SlurmClusterConfig(
+        name="mila",
+        ssh=config.SSHConfig(
+            user=user,
+            endpoints=(config.Endpoint("login.server.mila.quebec", 2222),),
+            public_key=config.PublicKey(
+                "ssh-ed25519",
+                "AAAAC3NzaC1lZDI1NTE5AAAAIBTPCzWRkwYDr/cFb4d2uR6rFlUtqfH3MoLMXPpJHK0n",
+            ),
+        ),
+        runtime=config.ContainerRuntime.SINGULARITY,
+        partition=partition,
+        prolog="module load singularity",
+        host_environment={
+            "SINGULARITY_CACHEDIR": "$SCRATCH/.apptainer",
+            "SINGULARITY_TMPDIR": "$SLURM_TMPDIR",
+            "SINGULARITY_LOCALCACHEDIR": "$SLURM_TMPDIR",
+        },
+        container_environment={
+            "SCRATCH": "/scratch",
+            "XM_SLURM_STATE_DIR": "/xm-slurm-state",
+        },
+        mounts=mounts,
+        resources={
+            resources.ResourceType.RTX8000: "rtx8000",
+            resources.ResourceType.V100: "v100",
+            resources.ResourceType.A100: "a100",
+            resources.ResourceType.A100_80GIB: "a100l",
+            resources.ResourceType.A6000: "a6000",
+            resources.ResourceType.L40S: "l40s",
+            resources.ResourceType.H100: "h100",
+        },
+        features={
+            resources.FeatureType.NVIDIA_MIG: "mig",
+            resources.FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
+    )