xmanager-slurm 0.3.0__py3-none-any.whl

xm_slurm/__init__.py

@@ -0,0 +1,44 @@
+import logging
+
+from xm_slurm.executables import Dockerfile, DockerImage
+from xm_slurm.executors import Slurm, SlurmSpec
+from xm_slurm.experiment import (
+    Artifact,
+    create_experiment,
+    get_current_experiment,
+    get_current_work_unit,
+    get_experiment,
+)
+from xm_slurm.packageables import (
+    conda_container,
+    docker_container,
+    docker_image,
+    mamba_container,
+    pdm_container,
+    python_container,
+)
+from xm_slurm.resources import JobRequirements, ResourceQuantity, ResourceType
+
+logging.getLogger("asyncssh").setLevel(logging.WARN)
+logging.getLogger("httpx").setLevel(logging.WARN)
+
+__all__ = [
+    "Artifact",
+    "conda_container",
+    "create_experiment",
+    "docker_container",
+    "docker_image",
+    "Dockerfile",
+    "DockerImage",
+    "get_current_experiment",
+    "get_current_work_unit",
+    "get_experiment",
+    "JobRequirements",
+    "mamba_container",
+    "pdm_container",
+    "python_container",
+    "ResourceQuantity",
+    "ResourceType",
+    "Slurm",
+    "SlurmSpec",
+]

xm_slurm/api.py

@@ -0,0 +1,261 @@
+import dataclasses
+import functools
+import importlib.util
+import logging
+import os
+import time
+import typing
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentPatchModel:
+    title: str | None = None
+    description: str | None = None
+    note: str | None = None
+    tags: list[str] | None = None
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class SlurmJobModel:
+    name: str
+    slurm_job_id: int
+    slurm_cluster: str
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ArtifactModel:
+    name: str
+    uri: str
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class WorkUnitPatchModel:
+    wid: int
+    identity: str | None
+    args: str | None = None
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class WorkUnitModel(WorkUnitPatchModel):
+    jobs: list[SlurmJobModel] = dataclasses.field(default_factory=list)
+    artifacts: list[ArtifactModel] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentModel:
+    title: str
+    description: str | None
+    note: str | None
+    tags: list[str] | None
+
+    work_units: list[WorkUnitModel]
+    artifacts: list[ArtifactModel]
+
+
+class XManagerAPI:
+    def get_experiment(self, xid: int) -> ExperimentModel:
+        del xid
+        raise NotImplementedError("`get_experiment` is not implemented without a storage backend.")
+
+    def delete_experiment(self, experiment_id: int) -> None:
+        del experiment_id
+        logger.debug("`delete_experiment` is not implemented without a storage backend.")
+
+    def insert_experiment(self, experiment: ExperimentPatchModel) -> int:
+        del experiment
+        logger.debug("`insert_experiment` is not implemented without a storage backend.")
+        return int(time.time() * 10**3)
+
+    def update_experiment(self, experiment_id: int, experiment_patch: ExperimentPatchModel) -> None:
+        del experiment_id, experiment_patch
+        logger.debug("`update_experiment` is not implemented without a storage backend.")
+
+    def insert_job(self, experiment_id: int, work_unit_id: int, job: SlurmJobModel) -> None:
+        del experiment_id, work_unit_id, job
+        logger.debug("`insert_job` is not implemented without a storage backend.")
+
+    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitPatchModel) -> None:
+        del experiment_id, work_unit
+        logger.debug("`insert_work_unit` is not implemented without a storage backend.")
+
+    def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
+        del experiment_id, work_unit_id, name
+        logger.debug("`delete_work_unit_artifact` is not implemented without a storage backend.")
+
+    def insert_work_unit_artifact(
+        self, experiment_id: int, work_unit_id: int, artifact: ArtifactModel
+    ) -> None:
+        del experiment_id, work_unit_id, artifact
+        logger.debug("`insert_work_unit_artifact` is not implemented without a storage backend.")
+
+    def delete_experiment_artifact(self, experiment_id: int, name: str) -> None:
+        del experiment_id, name
+        logger.debug("`delete_experiment_artifact` is not implemented without a storage backend.")
+
+    def insert_experiment_artifact(self, experiment_id: int, artifact: ArtifactModel) -> None:
+        del experiment_id, artifact
+        logger.debug("`insert_experiment_artifact` is not implemented without a storage backend.")
+
+
+class XManagerWebAPI(XManagerAPI):
+    def __init__(self, base_url: str, token: str):
+        if importlib.util.find_spec("xm_slurm_api_client") is None:
+            raise ImportError("xm_slurm_api_client not found.")
+
+        from xm_slurm_api_client import AuthenticatedClient  # type: ignore
+        from xm_slurm_api_client import models as _models  # type: ignore
+
+        self.models = _models
+        self.client = AuthenticatedClient(
+            base_url,
+            token=token,
+            raise_on_unexpected_status=True,
+            verify_ssl=False,
+        )
+
+    def get_experiment(self, xid: int) -> ExperimentModel:
+        from xm_slurm_api_client.api.experiment import (  # type: ignore
+            get_experiment as _get_experiment,
+        )
+
+        experiment: Any = _get_experiment.sync(xid, client=self.client)  # type: ignore
+        wus = []
+        for wu in experiment.work_units:
+            jobs = []
+            for job in wu.jobs:
+                jobs.append(SlurmJobModel(**job.dict()))
+            artifacts = []
+            for artifact in wu.artifacts:
+                artifacts.append(ArtifactModel(**artifact.dict()))
+            wus.append(
+                WorkUnitModel(
+                    wid=wu.wid,
+                    identity=wu.identity,
+                    args=wu.args,
+                    jobs=jobs,
+                    artifacts=artifacts,
+                )
+            )
+
+        artifacts = []
+        for artifact in experiment.artifacts:
+            artifacts.append(ArtifactModel(**artifact.dict()))
+
+        return ExperimentModel(
+            title=experiment.title,
+            description=experiment.description,
+            note=experiment.note,
+            tags=experiment.tags,
+            work_units=wus,
+            artifacts=artifacts,
+        )
+
+    def delete_experiment(self, experiment_id: int) -> None:
+        from xm_slurm_api_client.api.experiment import (  # type: ignore
+            delete_experiment as _delete_experiment,
+        )
+
+        _delete_experiment.sync(experiment_id, client=self.client)
+
+    def insert_experiment(self, experiment: ExperimentPatchModel) -> int:
+        from xm_slurm_api_client.api.experiment import (  # type: ignore
+            insert_experiment as _insert_experiment,
+        )
+
+        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."
+        experiment_response = _insert_experiment.sync(
+            client=self.client,
+            body=self.models.Experiment(title=experiment.title),
+        )
+        return typing.cast(int, experiment_response["xid"])  # type: ignore
+
+    def update_experiment(self, experiment_id: int, experiment_patch: ExperimentPatchModel) -> None:
+        from xm_slurm_api_client.api.experiment import (  # type: ignore
+            update_experiment as _update_experiment,
+        )
+
+        m = self.models.ExperimentPatch(**dataclasses.asdict(experiment_patch))
+
+        _update_experiment.sync(
+            experiment_id,
+            client=self.client,
+            body=self.models.ExperimentPatch(**dataclasses.asdict(experiment_patch)),
+        )
+
+    def insert_job(self, experiment_id: int, work_unit_id: int, job: SlurmJobModel) -> None:
+        from xm_slurm_api_client.api.job import insert_job as _insert_job  # type: ignore
+
+        _insert_job.sync(
+            experiment_id,
+            work_unit_id,
+            client=self.client,
+            body=self.models.SlurmJob(**dataclasses.asdict(job)),
+        )
+
+    def insert_work_unit(self, experiment_id: int, work_unit: WorkUnitPatchModel) -> None:
+        from xm_slurm_api_client.api.work_unit import (  # type: ignore
+            insert_work_unit as _insert_work_unit,
+        )
+
+        _insert_work_unit.sync(
+            experiment_id,
+            client=self.client,
+            body=self.models.WorkUnit(**dataclasses.asdict(work_unit)),
+        )
+
+    def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
+        from xm_slurm_api_client.api.artifact import (  # type: ignore
+            delete_work_unit_artifact as _delete_work_unit_artifact,
+        )
+
+        _delete_work_unit_artifact.sync(experiment_id, work_unit_id, name, client=self.client)
+
+    def insert_work_unit_artifact(
+        self, experiment_id: int, work_unit_id: int, artifact: ArtifactModel
+    ) -> None:
+        from xm_slurm_api_client.api.artifact import (  # type: ignore
+            insert_work_unit_artifact as _insert_work_unit_artifact,
+        )
+
+        _insert_work_unit_artifact.sync(
+            experiment_id,
+            work_unit_id,
+            client=self.client,
+            body=self.models.Artifact(**dataclasses.asdict(artifact)),
+        )
+
+    def delete_experiment_artifact(self, experiment_id: int, name: str) -> None: ...
+
+    def insert_experiment_artifact(self, experiment_id: int, artifact: ArtifactModel) -> None:
+        from xm_slurm_api_client.api.artifact import (  # type: ignore
+            insert_experiment_artifact as _insert_experiment_artifact,
+        )
+
+        _insert_experiment_artifact.sync(
+            experiment_id,
+            client=self.client,
+            body=self.models.Artifact(**dataclasses.asdict(artifact)),
+        )
+
+
+@functools.cache
+def client() -> XManagerAPI:
+    if importlib.util.find_spec("xm_slurm_api_client") is not None:
+        if (base_url := os.environ.get("XM_SLURM_API_BASE_URL")) is not None and (
+            token := os.environ.get("XM_SLURM_API_TOKEN")
+        ) is not None:
+            return XManagerWebAPI(base_url=base_url, token=token)
+        else:
+            logger.warn(
+                "XM_SLURM_API_BASE_URL and XM_SLURM_API_TOKEN not set. "
+                "Disabling XManager API client."
+            )
+
+    logger.debug("xm_slurm_api_client not found... skipping logging to the API.")
+    return XManagerAPI()

xm_slurm/batching.py

@@ -0,0 +1,139 @@
+import asyncio
+import collections
+import dataclasses
+import inspect
+import time
+import types
+from typing import Any, Callable, Coroutine, Generic, ParamSpec, Sequence, TypeVar
+
+T = TypeVar("T", contravariant=True)
+R = TypeVar("R", covariant=True)
+P = 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: Sequence[inspect.BoundArguments]
+) -> inspect.BoundArguments:
+    """Stacks bound arguments into a single bound arguments object."""
+    stacked_args = collections.OrderedDict()
+    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(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: Callable[..., Coroutine[None, None, 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()
+
+    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: Any, objtype: type) -> Any:
+        del objtype
+        if isinstance(self.fn, staticmethod):
+            return self.__call__
+        else:
+            return types.MethodType(self, obj)
+
+    @property
+    def __func__(self) -> Callable[..., Coroutine[None, None, Sequence[R]]]:
+        return self.fn
+
+    @property
+    def __wrapped__(self) -> Callable[..., Coroutine[None, None, 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) -> asyncio.Future[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()
+        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,162 @@
+import dataclasses
+import enum
+import functools
+import getpass
+import os
+import pathlib
+from typing import Literal, Mapping, NamedTuple
+
+import asyncssh
+
+
+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: 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 PublicKey(NamedTuple):
+    algorithm: str
+    key: str
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class SlurmClusterConfig:
+    name: str
+
+    host: str
+    host_public_key: PublicKey | None = None
+    user: str | None = None
+    port: int | None = None
+
+    # 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: Literal["submission-host"] | str | None = None
+
+    runtime: ContainerRuntime
+
+    # Environment variables
+    environment: Mapping[str, str] = dataclasses.field(default_factory=dict)
+
+    # Mounts
+    mounts: Mapping[os.PathLike[str] | str, os.PathLike[str] | str] = dataclasses.field(
+        default_factory=dict
+    )
+
+    # Resource mapping
+    resources: Mapping[str, "xm_slurm.ResourceType"] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
+
+    def __post_init__(self) -> None:
+        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}")
+
+    @functools.cached_property
+    def ssh_known_hosts(self) -> asyncssh.SSHKnownHosts | None:
+        if self.host_public_key is None:
+            return None
+
+        return asyncssh.import_known_hosts(
+            f"[{self.host}]:{self.port} {self.host_public_key.algorithm} {self.host_public_key.key}"
+        )
+
+    @functools.cached_property
+    def ssh_config(self) -> asyncssh.config.SSHConfig:
+        ssh_config_paths = []
+        if (ssh_config := pathlib.Path.home() / ".ssh" / "config").exists():
+            ssh_config_paths.append(ssh_config)
+        if (xm_ssh_config := os.environ.get("XM_SLURM_SSH_CONFIG")) and (
+            xm_ssh_config := pathlib.Path(xm_ssh_config).expanduser()
+        ).exists():
+            ssh_config_paths.append(xm_ssh_config)
+
+        config = asyncssh.config.SSHClientConfig.load(
+            None,
+            ssh_config_paths,
+            True,
+            getpass.getuser(),
+            self.user or (),
+            self.host or (),
+            self.port or (),
+        )
+
+        if config.get("Hostname") is None:
+            raise RuntimeError(
+                f"Failed to parse hostname from host `{self.host}` using SSH configs: {', '.join(map(str, ssh_config_paths))}"
+            )
+        if config.get("User") is None:
+            raise RuntimeError(
+                f"Failed to parse user from SSH configs: {', '.join(map(str, ssh_config_paths))}"
+            )
+
+        return config
+
+    @functools.cached_property
+    def ssh_connection_options(self) -> asyncssh.SSHClientConnectionOptions:
+        options = asyncssh.SSHClientConnectionOptions(config=None)
+        options.prepare(last_config=self.ssh_config, known_hosts=self.ssh_known_hosts)
+        return options
+
+    def __hash__(self):
+        return hash((
+            self.host,
+            self.user,
+            self.port,
+            self.cwd,
+            self.prolog,
+            self.epilog,
+            self.account,
+            self.partition,
+            self.qos,
+            self.proxy,
+            self.runtime,
+            frozenset(self.environment.items()),
+        ))

xm_slurm/console.py

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

xm_slurm/contrib/clusters/__init__.py

@@ -0,0 +1,52 @@
+import os
+
+from xm_slurm import config, resources
+from xm_slurm.contrib.clusters import drac
+
+# ComputeCanada alias
+cc = drac
+
+__all__ = ["drac", "mila", "cc"]
+
+
+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",
+            "/network/archive/${USER:0:1}/$USER": "/archive",
+        }
+
+    return config.SlurmClusterConfig(
+        name="mila",
+        user=user,
+        host="login.server.mila.quebec",
+        host_public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIBTPCzWRkwYDr/cFb4d2uR6rFlUtqfH3MoLMXPpJHK0n",
+        ),
+        port=2222,
+        runtime=config.ContainerRuntime.SINGULARITY,
+        partition=partition,
+        prolog="module load singularity",
+        environment={
+            "SINGULARITY_CACHEDIR": "$SCRATCH/.apptainer",
+            "SINGULARITY_TMPDIR": "$SLURM_TMPDIR",
+            "SINGULARITY_LOCALCACHEDIR": "$SLURM_TMPDIR",
+            "SCRATCH": "/scratch",
+            "ARCHIVE": "/archive",
+        },
+        mounts=mounts,
+        resources={
+            "rtx8000": resources.ResourceType.RTX8000,
+            "v100": resources.ResourceType.V100,
+            "a100": resources.ResourceType.A100,
+            "a100l": resources.ResourceType.A100_80GIB,
+            "a6000": resources.ResourceType.A6000,
+        },
+    )