xmanager-slurm 0.3.2__py3-none-any.whl → 0.4.1__py3-none-any.whl

xm_slurm/packaging/{docker/local.py → docker.py}

@@ -1,6 +1,9 @@
 import base64
+import collections.abc
 import dataclasses
 import enum
+import functools
+import hashlib
 import json
 import logging
 import os
@@ -9,35 +12,36 @@ import shlex
 import shutil
 import subprocess
 import tempfile
-import typing
-from typing import Sequence
+from typing import Hashable, Literal, Mapping, Sequence
 
+import jinja2 as j2
 from xmanager import xm
 
 from xm_slurm.executables import (
     Dockerfile,
+    DockerImage,
     ImageURI,
     RemoteImage,
     RemoteRepositoryCredentials,
 )
 from xm_slurm.executors import SlurmSpec
+from xm_slurm.packaging import registry
 from xm_slurm.packaging import utils as packaging_utils
-from xm_slurm.packaging.docker.abc import (
-    DockerBakeCommand,
-    DockerClient,
-    DockerVersionCommand,
-)
 from xm_slurm.packaging.registry import IndexedContainer
 
+logger = logging.getLogger(__name__)
+
 
-class LocalDockerClient(DockerClient):
-    """Build Docker images locally."""
+def _hash_digest(obj: Hashable) -> str:
+    return hashlib.sha256(repr(obj).encode()).hexdigest()
 
+
+class DockerClient:
     class Builder(enum.Enum):
         BUILDKIT = enum.auto()
         BUILDAH = enum.auto()
 
-    def __init__(self):
+    def __init__(self) -> None:
         if "XM_DOCKER_CLIENT" in os.environ:
             client_call = shlex.split(os.environ["XM_DOCKER_CLIENT"])
         elif shutil.which("docker"):
@@ -48,12 +52,11 @@ class LocalDockerClient(DockerClient):
             raise RuntimeError("No Docker client found.")
         self._client_call = client_call
 
-        version_command = DockerVersionCommand()
         backend_version = packaging_utils.run_command(
-            xm.merge_args(self._client_call, version_command.to_args()), return_stdout=True
+            xm.merge_args(self._client_call, ["buildx", "version"]), return_stdout=True
         )
         if backend_version.stdout.startswith("github.com/docker/buildx"):
-            self._builder = LocalDockerClient.Builder.BUILDKIT
+            self._builder = DockerClient.Builder.BUILDKIT
         else:
             raise NotImplementedError(f"Unsupported Docker build backend: {backend_version}")
 
@@ -80,7 +83,7 @@ class LocalDockerClient(DockerClient):
             return None
 
         def _parse_credentials_from_config(
-            config_path: pathlib.Path
+            config_path: pathlib.Path,
         ) -> RemoteRepositoryCredentials | None:
             """Parse credentials from the Docker configuration file."""
             if not config_path.exists():
@@ -116,21 +119,69 @@ class LocalDockerClient(DockerClient):
 
         # Attempt to parse credentials from the Docker or Podman configuration
         match self._builder:
-            case LocalDockerClient.Builder.BUILDKIT:
+            case DockerClient.Builder.BUILDKIT:
                 docker_config_path = (
                     pathlib.Path(os.environ.get("DOCKER_CONFIG", "~/.docker")).expanduser()
                     / "config.json"
                 )
                 return _parse_credentials_from_config(docker_config_path)
-            case LocalDockerClient.Builder.BUILDAH:
+            case DockerClient.Builder.BUILDAH:
                 podman_config_path = (
                     pathlib.Path(os.environ.get("XDG_CONFIG_HOME", "~/.config")).expanduser()
                     / "containers"
                     / "auth.json"
                 )
                 return _parse_credentials_from_config(podman_config_path)
-            case _:
-                return None
+
+    @functools.cached_property
+    def _bake_template(self) -> j2.Template:
+        template_loader = j2.PackageLoader("xm_slurm", "templates/docker")
+        template_env = j2.Environment(loader=template_loader, trim_blocks=True, lstrip_blocks=False)
+
+        return template_env.get_template("docker-bake.hcl.j2")
+
+    def _bake_args(
+        self,
+        *,
+        targets: str | Sequence[str] | None = None,
+        builder: str | None = None,
+        files: str | os.PathLike[str] | Sequence[os.PathLike[str] | str] | None = None,
+        load: bool = False,
+        cache: bool = True,
+        print: bool = False,
+        pull: bool = False,
+        push: bool = False,
+        metadata_file: str | os.PathLike[str] | None = None,
+        progress: Literal["auto", "plain", "tty"] = "auto",
+        set: Mapping[str, str] | None = None,
+    ) -> xm.SequentialArgs:
+        files = files
+        if files is None:
+            files = []
+        if not isinstance(files, collections.abc.Sequence):
+            files = [files]
+
+        targets = targets
+        if targets is None:
+            targets = []
+        elif isinstance(targets, str):
+            targets = [targets]
+        assert isinstance(targets, collections.abc.Sequence)
+
+        return xm.merge_args(
+            ["buildx", "bake"],
+            [f"--progress={progress}"],
+            [f"--builder={builder}"] if builder else [],
+            [f"--metadata-file={metadata_file}"] if metadata_file else [],
+            ["--print"] if print else [],
+            ["--push"] if push else [],
+            ["--pull"] if pull else [],
+            ["--load"] if load else [],
+            ["--no-cache"] if not cache else [],
+            [f"--file={file}" for file in files],
+            [f"--set={key}={value}" for key, value in set.items()] if set else [],
+            targets,
+        )
 
     def bake(
         self,
@@ -138,39 +189,40 @@ class LocalDockerClient(DockerClient):
         targets: Sequence[IndexedContainer[xm.Packageable]],
     ) -> list[IndexedContainer[RemoteImage]]:
         executors_by_executables = packaging_utils.collect_executors_by_executable(targets)
-        executors_by_executables = typing.cast(
-            dict[Dockerfile, list[SlurmSpec]], executors_by_executables
-        )
+        for executable, executors in executors_by_executables.items():
+            assert isinstance(
+                executable, Dockerfile
+            ), "All executables must be Dockerfiles when building Docker images."
+            assert all(
+                isinstance(executor, SlurmSpec) and executor.tag for executor in executors
+            ), "All executors must be SlurmSpecs with tags when building Docker images."
 
         with tempfile.TemporaryDirectory() as tempdir:
             hcl_file = pathlib.Path(tempdir) / "docker-bake.hcl"
             metadata_file = pathlib.Path(tempdir) / "metadata.json"
 
             # Write HCL and bake it
+            # TODO(jfarebro): Need a better way to hash the executables
             hcl = self._bake_template.render(
                 executables=executors_by_executables,
-                hash=packaging_utils.hash_digest,
+                hash=_hash_digest,
             )
             hcl_file.write_text(hcl)
-            logging.debug(hcl)
+            logger.debug(hcl)
 
             try:
-                command = DockerBakeCommand(
-                    targets=list(
-                        set(
-                            [
-                                packaging_utils.hash_digest(target.value.executable_spec)
-                                for target in targets
-                            ]
-                        )
+                bake_command = xm.merge_args(
+                    self._client_call,
+                    self._bake_args(
+                        targets=list(
+                            set([_hash_digest(target.value.executable_spec) for target in targets])
+                        ),
+                        files=[hcl_file],
+                        metadata_file=metadata_file,
+                        pull=True,
+                        push=True,
                     ),
-                    files=[hcl_file],
-                    metadata_file=metadata_file,
-                    pull=True,
-                    push=True,
                 )
-
-                bake_command = xm.merge_args(self._client_call, command.to_args())
                 packaging_utils.run_command(bake_command.to_list(), tty=True, check=True)
             except Exception as ex:
                 raise RuntimeError(f"Failed to build Dockerfiles: {ex}") from ex
@@ -183,9 +235,7 @@ class LocalDockerClient(DockerClient):
             assert isinstance(target.value.executor_spec, SlurmSpec)
             assert target.value.executor_spec.tag
 
-            executable_metadata = metadata[
-                packaging_utils.hash_digest(target.value.executable_spec)
-            ]
+            executable_metadata = metadata[_hash_digest(target.value.executable_spec)]
             uri = ImageURI(target.value.executor_spec.tag).with_digest(
                 executable_metadata["containerimage.digest"]
             )
@@ -204,3 +254,48 @@ class LocalDockerClient(DockerClient):
             )
 
         return images
+
+
+@functools.cache
+def docker_client() -> DockerClient:
+    return DockerClient()
+
+
+@registry.register(Dockerfile)
+def _(
+    targets: Sequence[IndexedContainer[xm.Packageable]],
+) -> list[IndexedContainer[RemoteImage]]:
+    return docker_client().bake(targets=targets)
+
+
+@registry.register(DockerImage)
+def _(
+    targets: Sequence[IndexedContainer[xm.Packageable]],
+) -> list[IndexedContainer[RemoteImage]]:
+    """Build Docker images, this is essentially a passthrough."""
+    images = []
+    client = docker_client()
+    for target in targets:
+        assert isinstance(target.value.executable_spec, DockerImage)
+        assert isinstance(target.value.executor_spec, SlurmSpec)
+        if target.value.executor_spec.tag is not None:
+            raise ValueError(
+                "Executable `DockerImage` should not be tagged via `SlurmSpec`. "
+                "The image URI is provided by the `DockerImage` itself."
+            )
+
+        uri = ImageURI(target.value.executable_spec.image)
+        images.append(
+            dataclasses.replace(
+                target,
+                value=RemoteImage(  # type: ignore
+                    image=str(uri),
+                    workdir=target.value.executable_spec.workdir,
+                    args=target.value.args,
+                    env_vars=target.value.env_vars,
+                    credentials=client.credentials(hostname=uri.domain),
+                ),
+            )
+        )
+
+    return images

xm_slurm/packaging/router.py

@@ -10,6 +10,8 @@ from xm_slurm.packaging import registry
 
 IndexedContainer = registry.IndexedContainer
 
+logger = logging.getLogger(__name__)
+
 
 def package(
     packageables: Sequence[xm.Packageable],
@@ -39,7 +41,7 @@ def package(
     # TODO(jfarebro): Could make this async as well...?
     with console.status("[magenta] :package: Packaging executables..."):
         for executable_spec_type, targets_for_type in targets_by_type.items():
-            logging.info(f"Packaging {len(targets_for_type)} {executable_spec_type!r} targets.")
+            logger.info(f"Packaging {len(targets_for_type)} {executable_spec_type!r} targets.")
             targets.extend(registry.route(executable_spec_type, targets_for_type))
 
     console.print(

xm_slurm/packaging/utils.py

@@ -1,7 +1,4 @@
 import collections
-import concurrent.futures
-import functools
-import hashlib
 import logging
 import os
 import pathlib
@@ -10,8 +7,7 @@ import re
 import select
 import shutil
 import subprocess
-import typing
-from typing import Callable, Concatenate, Hashable, Literal, ParamSpec, Sequence, TypeVar
+from typing import ParamSpec, Sequence, TypeVar
 
 from xmanager import xm
 
@@ -21,16 +17,7 @@ T = TypeVar("T")
 P = ParamSpec("P")
 ReturnT = TypeVar("ReturnT")
 
-
-def hash_digest(obj: Hashable) -> str:
-    # obj_hash = hash(obj)
-    # unsigned_obj_hash = obj_hash.from_bytes(
-    #     obj_hash.to_bytes((obj_hash.bit_length() + 7) // 8, "big", signed=True),
-    #     "big",
-    #     signed=False,
-    # )
-    # return hex(unsigned_obj_hash).removeprefix("0x")
-    return hashlib.sha256(repr(obj).encode()).hexdigest()
+logger = logging.getLogger(__name__)
 
 
 def collect_executors_by_executable(
@@ -42,19 +29,8 @@ def collect_executors_by_executable(
     return executors_by_executable
 
 
-def parallel_map(
-    f: Callable[Concatenate[T, P], ReturnT],
-) -> Callable[Concatenate[Sequence[T], P], list[ReturnT]]:
-    @functools.wraps(f)
-    def decorator(sequence: Sequence[T], *args: P.args, **kwargs: P.kwargs) -> list[ReturnT]:
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            return list(executor.map(lambda x: f(x, *args, **kwargs), sequence))
-
-    return decorator
-
-
 # Cursor commands to filter out from the command data stream
-cursor_commands_regex = re.compile(
+_CURSOR_ESCAPE_SEQUENCES_REGEX = re.compile(
     rb"\x1b\[\?25[hl]"  # Matches cursor show/hide commands (CSI ?25h and CSI ?25l)
     rb"|\x1b\[[0-9;]*[Hf]"  # Matches cursor position commands (CSI n;mH and CSI n;mf)
     rb"|\x1b\[s"  # Matches cursor save position (CSI s)
@@ -64,54 +40,6 @@ cursor_commands_regex = re.compile(
 )
 
 
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[False] = False,
-    return_stderr: Literal[False] = False,
-) -> subprocess.CompletedProcess[None]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[True] = True,
-    return_stderr: Literal[False] = False,
-) -> subprocess.CompletedProcess[str]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[False] = False,
-    return_stderr: Literal[True] = True,
-) -> subprocess.CompletedProcess[str]: ...
-
-
-@typing.overload
-def run_command(
-    args: Sequence[str] | xm.SequentialArgs,
-    env: dict[str, str] | None = ...,
-    tty: bool = ...,
-    cwd: str | os.PathLike[str] | None = ...,
-    check: bool = ...,
-    return_stdout: Literal[True] = True,
-    return_stderr: Literal[True] = True,
-) -> subprocess.CompletedProcess[str]: ...
-
-
 def run_command(
     args: Sequence[str] | xm.SequentialArgs,
     env: dict[str, str] | None = None,
@@ -120,7 +48,7 @@ def run_command(
     check: bool = False,
     return_stdout: bool = False,
     return_stderr: bool = False,
-) -> subprocess.CompletedProcess[str] | subprocess.CompletedProcess[None]:
+) -> subprocess.CompletedProcess[str]:
     if isinstance(args, xm.SequentialArgs):
         args = args.to_list()
     args = list(args)
@@ -134,8 +62,8 @@ def run_command(
     if executable.name == "docker" and args[1] == "buildx":
         subprocess_env |= {"DOCKER_CLI_EXPERIMENTAL": "enabled"}
 
-    logging.debug(f"env: {subprocess_env}")
-    logging.debug(f"command: {' '.join(args)}")
+    logger.debug(f"env: {subprocess_env}")
+    logger.debug(f"command: {' '.join(args)}")
 
     stdout_master, stdout_slave = pty.openpty()
     stderr_master, stderr_slave = pty.openpty()
@@ -171,7 +99,7 @@ def run_command(
                     fds.remove(fd)
                     continue
 
-                data = re.sub(cursor_commands_regex, b"", data)
+                data = _CURSOR_ESCAPE_SEQUENCES_REGEX.sub(b"", data)
 
                 if fd == stdout_master:
                     if return_stdout:
@@ -186,8 +114,8 @@ def run_command(
                 else:
                     raise RuntimeError("Unexpected file descriptor")
 
-    stdout = stdout_data.decode(errors="replace") if stdout_data else None
-    stderr = stderr_data.decode(errors="replace") if stderr_data else None
+    stdout = stdout_data.decode(errors="replace") if stdout_data else ""
+    stderr = stderr_data.decode(errors="replace") if stderr_data else ""
 
     retcode = process.poll()
     assert retcode is not None

xm_slurm/resources.py

@@ -36,20 +36,35 @@ class ResourceType(enum.IntEnum):
 
 
 AcceleratorType = set([
+    ResourceType.RTX8000,
     ResourceType.P4,
     ResourceType.P100,
+    ResourceType.P100_16GIB,
     ResourceType.V100,
+    ResourceType.V100_32GIB,
     ResourceType.A100,
     ResourceType.A100_80GIB,
+    ResourceType.A5000,
     ResourceType.A6000,
     ResourceType.H100,
-    ResourceType.GPU,
 ])
 
+assert AcceleratorType | {
+    ResourceType.CPU,
+    ResourceType.MEMORY,
+    ResourceType.DISK,
+    ResourceType.GPU,
+} == set(ResourceType.__members__.values()), "Resource types are not exhaustive."
+
 
 ResourceQuantity = int | float
 
 
+class FeatureType(enum.IntEnum):
+    NVIDIA_MIG = 1
+    NVIDIA_NVLINK = 2
+
+
 class JobRequirements:
     replicas: int
     location: str | None
@@ -80,8 +95,10 @@ class JobRequirements:
                 case str():
                     resource = resource_name
 
-            if resource in AcceleratorType or (
-                isinstance(resource, str) and resource.startswith("gpu")
+            if (
+                resource in AcceleratorType
+                or resource == ResourceType.GPU
+                or (isinstance(resource, str) and resource.startswith("gpu"))
             ):
                 if self.accelerator is not None:
                     raise ValueError("Accelerator already set.")
@@ -92,6 +109,8 @@ class JobRequirements:
             self.task_requirements[resource] = value
 
     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():
@@ -112,7 +131,12 @@ class JobRequirements:
                     directives.append(f"--gpus-per-task={value}")
                 case ResourceType() if resource in AcceleratorType:
                     assert isinstance(value, int), "Accelerator must be an integer"
-                    directives.append(f"--gpus-per-task={resource.name.lower()}:{value}")
+                    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}")
                 case str():
                     directives.append(f"--gres={resource}:{value}")
 

xm_slurm/scripts/_cloudpickle.py

@@ -0,0 +1,28 @@
+import base64
+import logging
+import zlib
+
+import cloudpickle
+from absl import app, flags
+from xmanager import xm
+
+CLOUDPICKLED_FN = flags.DEFINE_string(
+    "cloudpickled_fn", None, "Base64 encoded cloudpickled function", required=True
+)
+
+logger = logging.getLogger(__name__)
+
+
+@xm.run_in_asyncio_loop
+async def main(argv):
+    del argv
+
+    logger.info("Loading cloudpickled function...")
+    cloudpickled_fn = zlib.decompress(base64.urlsafe_b64decode(CLOUDPICKLED_FN.value))
+    function = cloudpickle.loads(cloudpickled_fn)
+    logger.info("Running cloudpickled function...")
+    await function()
+
+
+if __name__ == "__main__":
+    app.run(main)

xm_slurm/scripts/cli.py

@@ -0,0 +1,52 @@
+import argparse
+
+from xmanager import xm
+
+import xm_slurm
+from xm_slurm.console import console
+
+
+async def logs(
+    experiment_id: int,
+    wid: int,
+    *,
+    follow: bool = True,
+    num_lines: int = 10,
+    block_size: int = 1024,
+):
+    wu = xm_slurm.get_experiment(experiment_id).work_units()[wid]
+    async for log in wu.logs(num_lines=num_lines, block_size=block_size, wait=True, follow=follow):
+        console.print(log, end="\n")
+
+
+@xm.run_in_asyncio_loop
+async def main():
+    parser = argparse.ArgumentParser(description="XManager.")
+    subparsers = parser.add_subparsers(dest="subcommand", required=True)
+
+    logs_parser = subparsers.add_parser("logs", help="Display logs for a specific experiment.")
+    logs_parser.add_argument("xid", type=int, help="Experiment ID.")
+    logs_parser.add_argument("wid", type=int, help="Work Unit ID.")
+    logs_parser.add_argument(
+        "-n",
+        "--n-lines",
+        type=int,
+        default=50,
+        help="Number of lines to display from the end of the log file.",
+    )
+    logs_parser.add_argument(
+        "-f",
+        "--follow",
+        default=True,
+        action="store_true",
+        help="Follow the log file as it is updated.",
+    )
+
+    args = parser.parse_args()
+    match args.subcommand:
+        case "logs":
+            await logs(args.xid, args.wid, follow=args.follow, num_lines=args.n_lines)
+
+
+if __name__ == "__main__":
+    main()  # type: ignore

xm_slurm/status.py

@@ -106,6 +106,15 @@ SlurmFailedJobStates = set([
 ])
 SlurmCancelledJobStates = set([SlurmJobState.CANCELLED])
 
+assert (
+    SlurmPendingJobStates
+    | SlurmRunningJobStates
+    | SlurmActiveJobStates
+    | SlurmCompletedJobStates
+    | SlurmFailedJobStates
+    | SlurmCancelledJobStates
+) == set(SlurmJobState.__members__.values()), "Slurm job states are not exhaustive."
+
 
 class SlurmWorkUnitStatusEnum(enum.IntEnum):
     """Status of a local experiment job."""

xm_slurm/templates/docker/mamba.Dockerfile

@@ -1,7 +1,7 @@
 # syntax=docker/dockerfile:1.4
 ARG BASE_IMAGE=gcr.io/distroless/base-debian10
 
-FROM docker.io/mambaorg/micromamba:jammy as mamba
+FROM docker.io/mambaorg/micromamba:bookworm-slim as mamba
 ARG CONDA_ENVIRONMENT=environment.yml
 
 USER root
@@ -9,7 +9,9 @@ USER root
 COPY $CONDA_ENVIRONMENT /tmp/
 
 # Setup mamba environment
-RUN --mount=type=cache,target=/opt/conda/pkgs --mount=type=cache,target=/root/.cache/pip \
+RUN --mount=type=cache,target=/opt/conda/pkgs \
+    --mount=type=cache,target=/root/.cache/pip \
+    --mount=type=ssh \
     micromamba create --yes --always-copy --no-pyc --prefix /opt/env --file /tmp/environment.yml
 
 RUN find /opt/env/ -follow -type f -name '*.a' -delete && \

xm_slurm/templates/docker/python.Dockerfile

@@ -1,24 +1,32 @@
 # syntax=docker/dockerfile:1.4
-ARG BASE_IMAGE=docker.io/python:3.10-slim
+ARG BASE_IMAGE=docker.io/python:3.10-slim-bookworm
 FROM $BASE_IMAGE as builder
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+ARG EXTRA_SYSTEM_PACKAGES=""
+ARG EXTRA_PYTHON_PACKAGES=""
+
+ENV UV_PYTHON_DOWNLOADS=0
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
+
+WORKDIR /workspace
 
 RUN apt-get update \
     && apt-get install -y --no-install-recommends \
-    git \
+    git $EXTRA_SYSTEM_PACKAGES \
     && rm -rf /var/lib/apt/lists/*
 
 # Install and update necesarry global Python packages
-RUN pip install -U pip setuptools wheel pysocks
+RUN uv pip install --system pysocks $EXTRA_PYTHON_PACKAGES
 
 ARG PIP_REQUIREMENTS=requirements.txt
 
-RUN python -m venv --copies --upgrade --upgrade-deps --system-site-packages /venv
-COPY $PIP_REQUIREMENTS /tmp/requirements.txt
-RUN --mount=type=cache,target=/root/.cache/pip \
-    PIP_CACHE_DIR=/root/.cache/pip /venv/bin/pip install -r /tmp/requirements.txt \
-    && rm -rf /tmp/requirements.txt
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=$PIP_REQUIREMENTS,target=requirements.txt \
+    --mount=type=ssh \
+    uv pip install --system --requirement requirements.txt
 
 COPY --link . /workspace
-WORKDIR /workspace
 
-ENTRYPOINT [ "/venv/bin/python" ]
+ENTRYPOINT [ "python" ]