xmanager-slurm 0.4.19__py3-none-any.whl

xm_slurm/contrib/clusters/drac.py

@@ -0,0 +1,242 @@
+import os
+import typing as tp
+
+from xm_slurm import config
+from xm_slurm.resources import FeatureType, ResourceType
+
+__all__ = [
+    "fir",
+    "narval",
+    "nibi",
+    "rorqual",
+    "killarney",
+    "tamia",
+    "vulcan",
+]
+
+
+def _drac_cluster(
+    *,
+    name: str,
+    host: str,
+    port: int = 22,
+    robot_host: str | None = None,
+    robot_port: int | None = None,
+    public_key: config.PublicKey,
+    user: str | None = None,
+    account: str | None = None,
+    modules: list[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:
+        mounts = {
+            "/scratch/$USER": "/scratch",
+            # TODO: move these somewhere common to all cluster configs.
+            "/home/$USER/.ssh": "/home/$USER/.ssh",
+            "/home/$USER/.local/state/xm-slurm": "/xm-slurm-state",
+        }
+    mounts = dict(mounts) | {"/dev/infiniband": "/dev/infiniband"}
+
+    endpoints = []
+    if robot_host is not None and robot_host != host:
+        endpoints.append(config.Endpoint(robot_host, robot_port))
+    endpoints.append(config.Endpoint(host, port))
+    endpoints = tuple(endpoints)
+
+    return config.SlurmClusterConfig(
+        name=name,
+        ssh=config.SSHConfig(user=user, endpoints=endpoints, public_key=public_key),
+        account=account,
+        proxy=proxy,
+        runtime=config.ContainerRuntime.APPTAINER,
+        prolog=f"module load apptainer {' '.join(modules) if modules else ''}".rstrip(),
+        host_environment={
+            "XDG_DATA_HOME": "$SLURM_TMPDIR/.local",
+            "APPTAINER_CACHEDIR": "$SCRATCH/.apptainer",
+            "APPTAINER_TMPDIR": "$SLURM_TMPDIR",
+            "APPTAINER_LOCALCACHEDIR": "$SLURM_TMPDIR",
+        },
+        container_environment={
+            "SCRATCH": "/scratch",
+            "XM_SLURM_STATE_DIR": "/xm-slurm-state",
+        },
+        mounts=mounts,
+        resources=resources or {},
+        features=features or {},
+    )
+
+
+def narval(
+    *,
+    user: str | None = None,
+    account: 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 = []
+    if proxy != "submission-host":
+        modules.append("httpproxy")
+
+    return _drac_cluster(
+        name="narval",
+        host="narval.alliancecan.ca",
+        robot_host="robot.narval.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAILFxB0spH5RApc43sBx0zOxo1ARVH0ezU+FbQH95FW+h",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        proxy=proxy,
+        modules=modules,
+        resources={ResourceType.A100: "a100"},
+        features={
+            FeatureType.NVIDIA_MIG: "a100mig",
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
+    )
+
+
+def rorqual(
+    *,
+    user: str | None = None,
+    account: 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 Rorqual Cluster (https://docs.alliancecan.ca/wiki/Rorqual/en)."""
+    modules = []
+    if proxy != "submission-host":
+        modules.append("httpproxy")
+
+    return _drac_cluster(
+        name="rorqual",
+        host="rorqual.alliancecan.ca",
+        robot_host="robot.rorqual.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAINME5e9bifKZbuKKOQSpe3xrvC4g1b0QLMYj+AXBQGJe",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        proxy=proxy,
+        modules=modules,
+        resources={ResourceType.H100: "h100"},
+        features={
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
+    )
+
+
+def fir(
+    *,
+    user: str | None = None,
+    account: 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(
+        name="fir",
+        host="fir.alliancecan.ca",
+        robot_host="robot.fir.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIJtenyJz+inwobvlJntWYFNu+ANcVWNcOHRKcEN6zmDo",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.H100: "h100"},
+    )
+
+
+def nibi(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC Nibi Cluster (https://docs.alliancecan.ca/wiki/Nibi/en)."""
+    return _drac_cluster(
+        name="nibi",
+        host="nibi.alliancecan.ca",
+        robot_host="robot.nibi.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIEcmFoQZr6+KUHm/zm/BJpnNIlME7GytMxbHgfAUfoQX",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.H100: "h100"},
+    )
+
+
+def killarney(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC (PAICE) Killarney Cluster (https://docs.alliancecan.ca/wiki/Killarney/en)."""
+    return _drac_cluster(
+        name="killarney",
+        host="killarney.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIGlzaBBtvhJsSr23rMoY41gy8Svv1IOct8TBRH9CGuJf",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.L40S: "l40s", ResourceType.H100: "h100"},
+    )
+
+
+def tamia(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC (PAICE) Tamia Cluster (https://docs.alliancecan.ca/wiki/Tamia/en)."""
+    return _drac_cluster(
+        name="tamia",
+        host="tamia.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIN2wL9wOa0VveA/2l2ky/OhPsQfYtKuX99dyNnUTSYeU",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.H100: "h100"},
+    )
+
+
+def vulcan(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: tp.Mapping[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC (PAICE) Vulcan Cluster (https://docs.alliancecan.ca/wiki/Vulcan/en)."""
+    return _drac_cluster(
+        name="vulcan",
+        host="vulcan.alliancecan.ca",
+        public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIMuIj6T45HqVeJgRotH9Qq46FzidekS2lXkD7FOTltnC",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.L40S: "l40s"},
+    )

xm_slurm/dependencies.py

@@ -0,0 +1,171 @@
+import abc
+import dataclasses
+import datetime as dt
+import typing as tp
+
+
+class SlurmDependencyException(Exception): ...
+
+
+NoChainingException = SlurmDependencyException(
+    "Slurm only supports chaining dependencies with the same logical operator. "
+    "For example, `dep1 & dep2 | dep3` is not supported but `dep1 & dep2 & dep3` is."
+)
+
+
+class SlurmJobDependency(abc.ABC):
+    @abc.abstractmethod
+    def to_dependency_str(self) -> str: ...
+
+    def to_directive(self) -> str:
+        return f"--dependency={self.to_dependency_str()}"
+
+    def __and__(self, other_dependency: "SlurmJobDependency") -> "SlurmJobDependencyAND":
+        if isinstance(self, SlurmJobDependencyOR):
+            raise NoChainingException
+        return SlurmJobDependencyAND(self, other_dependency)
+
+    def __or__(self, other_dependency: "SlurmJobDependency") -> "SlurmJobDependencyOR":
+        if isinstance(other_dependency, SlurmJobDependencyAND):
+            raise NoChainingException
+        return SlurmJobDependencyOR(self, other_dependency)
+
+    def flatten(self) -> tuple["SlurmJobDependency", ...]:
+        if isinstance(self, SlurmJobDependencyAND) or isinstance(self, SlurmJobDependencyOR):
+            return self.first_dependency.flatten() + self.second_dependency.flatten()
+        return (self,)
+
+    def traverse(
+        self, mapper: tp.Callable[["SlurmJobDependency"], "SlurmJobDependency"]
+    ) -> "SlurmJobDependency":
+        if isinstance(self, SlurmJobDependencyAND) or isinstance(self, SlurmJobDependencyOR):
+            return type(self)(
+                first_dependency=self.first_dependency.traverse(mapper),
+                second_dependency=self.second_dependency.traverse(mapper),
+            )
+        return mapper(self)
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAND(SlurmJobDependency):
+    first_dependency: SlurmJobDependency
+    second_dependency: SlurmJobDependency
+
+    def to_dependency_str(self) -> str:
+        return f"{self.first_dependency.to_dependency_str()},{self.second_dependency.to_dependency_str()}"
+
+    def __or__(self, other_dependency: SlurmJobDependency):
+        del other_dependency
+        raise NoChainingException
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.first_dependency, self.second_dependency))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyOR(SlurmJobDependency):
+    first_dependency: SlurmJobDependency
+    second_dependency: SlurmJobDependency
+
+    def to_dependency_str(self) -> str:
+        return f"{self.first_dependency.to_dependency_str()}?{self.second_dependency.to_dependency_str()}"
+
+    def __and__(self, other_dependency: SlurmJobDependency):
+        del other_dependency
+        raise NoChainingException
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.first_dependency, self.second_dependency))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfter(SlurmJobDependency):
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+    time: dt.timedelta | None = None
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+        if self.time is not None and self.time.total_seconds() % 60 != 0:
+            raise SlurmDependencyException("Time must be specified in exact minutes")
+
+    def to_dependency_str(self) -> str:
+        directive = "after"
+
+        for handle in self.handles:
+            directive += f":{handle.slurm_job.job_id}"
+            if self.time is not None:
+                directive += f"+{self.time.total_seconds() // 60:.0f}"
+        return directive
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterAny(SlurmJobDependency):
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afterany"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterNotOK(SlurmJobDependency):
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afternotok"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobDependencyAfterOK(SlurmJobDependency):
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        return ":".join(["afterok"] + [handle.slurm_job.job_id for handle in self.handles])
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))
+
+
+@dataclasses.dataclass(frozen=True)
+class SlurmJobArrayDependencyAfterOK(SlurmJobDependency):
+    handles: tp.Sequence["xm_slurm.execution.SlurmHandle[SlurmJob]"]  # type: ignore # noqa: F821
+
+    def __post_init__(self):
+        if len(self.handles) == 0:
+            raise SlurmDependencyException("Dependency doesn't have any handles.")
+
+    def to_dependency_str(self) -> str:
+        job_ids = []
+        for handle in self.handles:
+            job = handle.slurm_job
+            if job.is_array_job:
+                job_ids.append(job.array_job_id)
+            elif job.is_heterogeneous_job:
+                job_ids.append(job.het_job_id)
+            else:
+                job_ids.append(job.job_id)
+        return ":".join(["aftercorr"] + job_ids)
+
+    def __hash__(self) -> int:
+        return hash((type(self),) + tuple([handle.slurm_job for handle in self.handles]))

xm_slurm/executables.py

@@ -0,0 +1,215 @@
+import dataclasses
+import os
+import pathlib
+import typing as tp
+
+from xmanager import xm
+
+from xm_slurm import constants
+from xm_slurm.types import Descriptor
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Dockerfile(xm.ExecutableSpec):
+    """A specification describing a Dockerfile to build.
+
+    Args:
+        dockerfile: The path to the Dockerfile.
+        context: The path to the Docker context.
+        target: The Docker build target.
+        ssh: A list of docker SSH sockets/keys.
+        build_args: Build arguments to docker.
+        cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
+        labels: The container labels.
+        platforms: The target platform.
+    """
+
+    # Dockerfile
+    dockerfile: pathlib.Path
+    # Docker context
+    context: pathlib.Path
+
+    # Docker build target
+    target: str | None = None
+
+    # SSH sockets/keys for the docker build step.
+    ssh: tp.Sequence[str] = dataclasses.field(default_factory=list)
+
+    # Build arguments to docker
+    build_args: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+
+    # --cache-from field in BuildKit
+    cache_from: tp.Sequence[str] = dataclasses.field(default_factory=list)
+
+    # Container labels
+    labels: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+
+    # Target platform
+    platforms: tp.Sequence[str] = dataclasses.field(default_factory=lambda: ["linux/amd64"])
+
+    @property
+    def name(self) -> str:
+        name = self.dockerfile.stem
+        if self.target is not None:
+            name = f"{name}-{self.target}"
+        return name
+
+    def __hash__(self) -> int:
+        return hash((
+            type(self),
+            self.dockerfile,
+            self.context,
+            self.target,
+            tuple(sorted(self.ssh)),
+            tuple(sorted(self.build_args.items())),
+            tuple(sorted(self.cache_from)),
+            tuple(sorted(self.labels.items())),
+            tuple(sorted(self.platforms)),
+        ))
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DockerImage(xm.ExecutableSpec):
+    """A specification describing a pre-built Docker image.
+
+    Args:
+        image: The remote image URI.
+        workdir: The working directory in container.
+
+    """
+
+    image: str
+
+    # Working directory in container
+    workdir: pathlib.Path | None = None
+
+    @property
+    def name(self) -> str:
+        return self.image
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.image, self.workdir))
+
+
+@dataclasses.dataclass
+class ImageURI:
+    image: dataclasses.InitVar[str]
+
+    scheme: str | None = dataclasses.field(init=False, default=None)
+    domain: str = dataclasses.field(init=False)
+    path: str = dataclasses.field(init=False)
+    tag: str | None = dataclasses.field(init=False, default=None)
+    digest: str | None = dataclasses.field(init=False, default=None)
+
+    def __post_init__(self, image: str):
+        match = constants.IMAGE_URI_REGEX.match(image)
+        if not match:
+            raise ValueError(f"Invalid OCI image URI: {image}")
+        groups = {k: v for k, v in match.groupdict().items() if v is not None}
+        for k, v in groups.items():
+            setattr(self, k, v)
+
+        if self.tag is None and self.digest is None:
+            self.tag = "latest"
+
+    @property
+    def locator(self) -> str:
+        """Unique locator for this image.
+
+        Locator will return the digest if it exists otherwise the tag format.
+        If neither are present, it will raise an AssertionError.
+        """
+        if self.digest is not None:
+            return f"@{self.digest}"
+        assert self.tag is not None
+        return f":{self.tag}"
+
+    @property
+    def url(self) -> str:
+        """URL for this image without the locator."""
+        return f"{self.origin}{self.path}"
+
+    @property
+    def origin(self) -> str:
+        return f"{self.scheme}{self.domain}"
+
+    def with_tag(self, tag: str) -> "ImageURI":
+        self.tag = tag
+        return self
+
+    def with_digest(self, digest: str) -> "ImageURI":
+        self.digest = digest
+        return self
+
+    def __str__(self) -> str:
+        return self.format("{url}{locator}")
+
+    def __hash__(self) -> int:
+        return hash((
+            type(self),
+            self.scheme,
+            self.domain,
+            self.path,
+            self.tag,
+            self.digest,
+        ))
+
+    def format(self, format: str) -> str:
+        fields = {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
+        fields |= {"locator": self.locator, "url": self.url}
+        return format.format(**fields)
+
+
+class ImageDescriptor(Descriptor[ImageURI, str | ImageURI]):
+    def __set_name__(self, owner: type, name: str):
+        del owner
+        self.image = f"_{name}"
+
+    def __get__(self, instance: object | None, owner: tp.Type[object] | None = None) -> ImageURI:
+        del owner
+        return getattr(instance, self.image)
+
+    def __set__(self, instance: object, value: str | ImageURI):
+        _setattr = object.__setattr__ if not hasattr(instance, self.image) else setattr
+        if isinstance(value, str):
+            value = ImageURI(value)
+        _setattr(instance, self.image, value)
+
+
+class RemoteRepositoryCredentials(tp.NamedTuple):
+    username: str
+    password: str
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)  # type: ignore
+class RemoteImage(xm.Executable):
+    # Remote base image
+    image: Descriptor[ImageURI, str | ImageURI] = ImageDescriptor()
+
+    workdir: os.PathLike[str] | str
+    entrypoint: xm.SequentialArgs
+
+    # Container arguments
+    args: xm.SequentialArgs = dataclasses.field(default_factory=xm.SequentialArgs)
+    # Container environment variables
+    env_vars: tp.Mapping[str, str] = dataclasses.field(default_factory=dict)
+
+    # Remote repository credentials
+    credentials: RemoteRepositoryCredentials | None = None
+
+    @property
+    def name(self) -> str:  # type: ignore
+        return str(self.image)
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                type(self),
+                self.image,
+                self.workdir,
+                tuple(sorted(self.entrypoint.to_list())),
+                tuple(sorted(self.args.to_list())),
+                tuple(sorted(self.env_vars.items())),
+                self.credentials,
+            ),
+        )