xmanager-slurm 0.3.0__py3-none-any.whl

xm_slurm/packaging/docker/local.py

@@ -0,0 +1,206 @@
+import base64
+import dataclasses
+import enum
+import json
+import logging
+import os
+import pathlib
+import shlex
+import shutil
+import subprocess
+import tempfile
+import typing
+from typing import Sequence
+
+from xmanager import xm
+
+from xm_slurm.executables import (
+    Dockerfile,
+    ImageURI,
+    RemoteImage,
+    RemoteRepositoryCredentials,
+)
+from xm_slurm.executors import SlurmSpec
+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
+
+
+class LocalDockerClient(DockerClient):
+    """Build Docker images locally."""
+
+    class Builder(enum.Enum):
+        BUILDKIT = enum.auto()
+        BUILDAH = enum.auto()
+
+    def __init__(self):
+        if "XM_DOCKER_CLIENT" in os.environ:
+            client_call = shlex.split(os.environ["XM_DOCKER_CLIENT"])
+        elif shutil.which("docker"):
+            client_call = ["docker"]
+        elif shutil.which("podman"):
+            client_call = ["podman"]
+        else:
+            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
+        )
+        if backend_version.stdout.startswith("github.com/docker/buildx"):
+            self._builder = LocalDockerClient.Builder.BUILDKIT
+        else:
+            raise NotImplementedError(f"Unsupported Docker build backend: {backend_version}")
+
+        self._credentials_cache: dict[str, RemoteRepositoryCredentials] = {}
+
+    def credentials(self, hostname: str) -> RemoteRepositoryCredentials | None:
+        """Fetch credentials from the local Docker configuration."""
+        if hostname in self._credentials_cache:
+            return self._credentials_cache[hostname]
+
+        def _parse_docker_credentials(helper: str) -> RemoteRepositoryCredentials | None:
+            """Parse credentials from a Docker credential helper."""
+            if not shutil.which(f"docker-credential-{helper}"):
+                return None
+            returncode, output = subprocess.getstatusoutput(
+                f"echo {hostname} | docker-credential-{helper} get",
+            )
+
+            if returncode == 0:
+                credentials = json.loads(output)
+                return RemoteRepositoryCredentials(
+                    username=credentials["Username"], password=credentials["Secret"]
+                )
+            return None
+
+        def _parse_credentials_from_config(
+            config_path: pathlib.Path
+        ) -> RemoteRepositoryCredentials | None:
+            """Parse credentials from the Docker configuration file."""
+            if not config_path.exists():
+                return None
+            config = json.loads(config_path.read_text())
+
+            # Attempt to parse from the global credential store
+            if (creds_store := config.get("credsStore", None)) and (
+                credentials := _parse_docker_credentials(creds_store)
+            ):
+                self._credentials_cache[hostname] = credentials
+                return credentials
+            # Attempt to parse from the credential helper for this registry
+            if creds_helper := config.get("credHelpers", {}):
+                for registry, helper in creds_helper.items():
+                    registry = ImageURI(registry)
+                    if registry.domain == hostname and (
+                        credentials := _parse_docker_credentials(helper)
+                    ):
+                        self._credentials_cache[hostname] = credentials
+                        return credentials
+            # Last resort: attempt to parse raw auth
+            if auths := config.get("auths", None):
+                for registry, metadata in auths.items():
+                    registry = ImageURI(registry)
+                    if registry.domain == hostname:
+                        auth = base64.b64decode(metadata["auth"]).decode("utf-8")
+                        username, password = auth.split(":")
+                        credentials = RemoteRepositoryCredentials(username, password)
+                        self._credentials_cache[hostname] = credentials
+                        return credentials
+            return None
+
+        # Attempt to parse credentials from the Docker or Podman configuration
+        match self._builder:
+            case LocalDockerClient.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:
+                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
+
+    def bake(
+        self,
+        *,
+        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
+        )
+
+        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
+            hcl = self._bake_template.render(
+                executables=executors_by_executables,
+                hash=packaging_utils.hash_digest,
+            )
+            hcl_file.write_text(hcl)
+            logging.debug(hcl)
+
+            try:
+                command = DockerBakeCommand(
+                    targets=list(
+                        set(
+                            [
+                                packaging_utils.hash_digest(target.value.executable_spec)
+                                for target in targets
+                            ]
+                        )
+                    ),
+                    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
+            else:
+                metadata = json.loads(metadata_file.read_text())
+
+        images = []
+        for target in targets:
+            assert isinstance(target.value.executable_spec, Dockerfile)
+            assert isinstance(target.value.executor_spec, SlurmSpec)
+            assert target.value.executor_spec.tag
+
+            executable_metadata = metadata[
+                packaging_utils.hash_digest(target.value.executable_spec)
+            ]
+            uri = ImageURI(target.value.executor_spec.tag).with_digest(
+                executable_metadata["containerimage.digest"]
+            )
+
+            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=self.credentials(uri.domain),
+                    ),
+                )
+            )
+
+        return images

xm_slurm/packaging/registry.py

@@ -0,0 +1,45 @@
+import dataclasses
+from typing import Callable, Generic, ParamSpec, Sequence, Type, TypeVar
+
+from xmanager import xm
+
+T_co = TypeVar("T_co", covariant=True)
+P = ParamSpec("P")
+ExecutableSpecT = TypeVar("ExecutableSpecT", bound=xm.ExecutableSpec)
+
+
+@dataclasses.dataclass(frozen=True)
+class IndexedContainer(Generic[T_co]):
+    index: int
+    value: T_co
+
+
+RegistrationCallable = Callable[
+    [Sequence[IndexedContainer[xm.Packageable]]],
+    Sequence[IndexedContainer[xm.Executable]],
+]
+
+
+_REGISTRY: dict[Type[xm.ExecutableSpec], RegistrationCallable] = {}
+
+
+def register(
+    *typs: Type[ExecutableSpecT],
+) -> Callable[[RegistrationCallable], RegistrationCallable]:
+    def decorator(
+        registration_callable: RegistrationCallable,
+    ) -> RegistrationCallable:
+        global _REGISTRY
+        for typ in typs:
+            _REGISTRY[typ] = registration_callable
+        return registration_callable
+
+    return decorator
+
+
+def route(
+    typ: Type[ExecutableSpecT],
+    packageables: Sequence[IndexedContainer[xm.Packageable]],
+) -> Sequence[IndexedContainer[xm.Executable]]:
+    global _REGISTRY
+    return _REGISTRY[typ](packageables)

xm_slurm/packaging/router.py

@@ -0,0 +1,52 @@
+import collections
+import logging
+from typing import Sequence, Type
+
+from xmanager import xm
+
+from xm_slurm.console import console
+from xm_slurm.executors import SlurmSpec
+from xm_slurm.packaging import registry
+
+IndexedContainer = registry.IndexedContainer
+
+
+def package(
+    packageables: Sequence[xm.Packageable],
+) -> list[xm.Executable]:
+    """
+    Takes as input a list of packageables and returns a mapping of
+    `DockerTarget`'s to the latest digest of that image.
+    """
+    # Docker targets to be collected.
+    # These are a mapping from `DockerTarget` to the latest digest of the image.
+    targets_by_type = collections.defaultdict[
+        Type[xm.ExecutableSpec], list[IndexedContainer[xm.Packageable]]
+    ](list)
+
+    # Collect dockerfiles that need to be built locally
+    for index, packageable in enumerate(packageables):
+        if not isinstance(packageable.executor_spec, SlurmSpec):
+            raise ValueError(
+                f"Unsupported executor spec for packageable: {packageable}."
+                "xm_slurm only supports `xm_slurm.SlurmSpec`."
+            )
+        targets_by_type[type(packageable.executable_spec)].append(
+            IndexedContainer[xm.Packageable](index, packageable)
+        )
+
+    targets: list[IndexedContainer[xm.Executable]] = []
+    # 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.")
+            targets.extend(registry.route(executable_spec_type, targets_for_type))
+
+    console.print(
+        f"[magenta]:package: Finished packaging [bold]{len(targets)} executable"
+        f"{'s' if len(targets) > 1 else ''}[/bold]."
+    )
+
+    assert len(targets) == len(packageables), "Number of targets must match packageables"
+    targets = sorted(targets, key=lambda t: t.index)
+    return [target.value for target in targets]

xm_slurm/packaging/utils.py

@@ -0,0 +1,202 @@
+import collections
+import concurrent.futures
+import functools
+import hashlib
+import logging
+import os
+import pathlib
+import pty
+import re
+import select
+import shutil
+import subprocess
+import typing
+from typing import Callable, Concatenate, Hashable, Literal, ParamSpec, Sequence, TypeVar
+
+from xmanager import xm
+
+from xm_slurm.packaging.registry import IndexedContainer
+
+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()
+
+
+def collect_executors_by_executable(
+    targets: Sequence[IndexedContainer[xm.Packageable]],
+) -> dict[xm.ExecutableSpec, set[xm.ExecutorSpec]]:
+    executors_by_executable = collections.defaultdict(set)
+    for target in targets:
+        executors_by_executable[target.value.executable_spec].add(target.value.executor_spec)
+    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(
+    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)
+    rb"|\x1b\[u"  # Matches cursor restore position (CSI u)
+    rb"|\x1b\[2J"  # Matches clear screen (CSI 2J)
+    rb"|\x1b\[K"  # Matches clear line (CSI K)
+)
+
+
+@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,
+    tty: bool = False,
+    cwd: str | os.PathLike[str] | None = None,
+    check: bool = False,
+    return_stdout: bool = False,
+    return_stderr: bool = False,
+) -> subprocess.CompletedProcess[str] | subprocess.CompletedProcess[None]:
+    if isinstance(args, xm.SequentialArgs):
+        args = args.to_list()
+    args = list(args)
+
+    executable = shutil.which(args[0])
+    if not executable:
+        raise RuntimeError(f"Couldn't find executable {args[0]}")
+    executable = pathlib.Path(executable)
+
+    subprocess_env = os.environ.copy() | (env if env else {})
+    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)}")
+
+    stdout_master, stdout_slave = pty.openpty()
+    stderr_master, stderr_slave = pty.openpty()
+
+    stdout_data, stderr_data = b"", b""
+    with subprocess.Popen(
+        executable=executable,
+        args=args,
+        shell=False,
+        text=True,
+        bufsize=0,
+        stdout=stdout_slave,
+        stderr=stderr_slave,
+        start_new_session=True,
+        close_fds=True,
+        cwd=cwd,
+        env=subprocess_env,
+    ) as process:
+        os.close(stdout_slave)
+        os.close(stderr_slave)
+
+        fds = [stdout_master, stderr_master]
+        while fds:
+            rlist, _, _ = select.select(fds, [], [])
+            for fd in rlist:
+                try:
+                    data = os.read(fd, 1024)
+                except OSError:
+                    data = None
+
+                if not data:
+                    os.close(fd)
+                    fds.remove(fd)
+                    continue
+
+                data = re.sub(cursor_commands_regex, b"", data)
+
+                if fd == stdout_master:
+                    if return_stdout:
+                        stdout_data += data
+                    if tty:
+                        os.write(pty.STDOUT_FILENO, data)
+                elif fd == stderr_master:
+                    if return_stderr:
+                        stderr_data += data
+                    if tty:
+                        os.write(pty.STDERR_FILENO, data)
+                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
+
+    retcode = process.poll()
+    assert retcode is not None
+
+    if check and retcode:
+        raise subprocess.CalledProcessError(retcode, process.args)
+    return subprocess.CompletedProcess(
+        process.args,
+        retcode,
+        stdout=stdout,
+        stderr=stderr,
+    )

xm_slurm/resources.py

@@ -0,0 +1,150 @@
+import enum
+import itertools
+import math
+from typing import Mapping
+
+import immutabledict
+
+from xm_slurm import config
+
+
+class ResourceType(enum.IntEnum):
+    CPU = 1
+
+    MEMORY = 2
+    RAM = 2
+
+    EPHEMERAL_STORAGE = 3
+    DISK = 3
+
+    GPU = 1000
+    RTX8000 = 1001
+    P4 = 1010
+
+    P100 = 1011
+    P100_16GIB = 1012
+
+    V100 = 1020
+    V100_32GIB = 1021
+
+    A100 = 1030
+    A100_80GIB = 1031
+    A5000 = 1032
+    A6000 = 1033
+
+    H100 = 1040
+
+
+AcceleratorType = set([
+    ResourceType.P4,
+    ResourceType.P100,
+    ResourceType.V100,
+    ResourceType.A100,
+    ResourceType.A100_80GIB,
+    ResourceType.A6000,
+    ResourceType.H100,
+    ResourceType.GPU,
+])
+
+
+ResourceQuantity = int | float
+
+
+class JobRequirements:
+    replicas: int
+    location: str | None
+    accelerator: ResourceType | None
+    cluster: config.SlurmClusterConfig | None = None
+
+    def __init__(
+        self,
+        *,
+        resources: Mapping[ResourceType | str, ResourceQuantity] = immutabledict.immutabledict(),
+        replicas: int = 1,
+        location: str | None = None,
+        cluster: config.SlurmClusterConfig | None = None,
+        **kw_resources: ResourceQuantity,
+    ):
+        self.replicas = replicas or 1
+        self.location = location
+        self.accelerator = None
+        self.cluster = cluster
+
+        self.task_requirements: dict[ResourceType | str, ResourceQuantity] = {}
+        for resource_name, value in itertools.chain(resources.items(), kw_resources.items()):
+            match resource_name:
+                case str() if resource_name.upper() in ResourceType.__members__:
+                    resource = ResourceType[resource_name.upper()]
+                case ResourceType():
+                    resource = resource_name
+                case str():
+                    resource = resource_name
+
+            if resource in AcceleratorType or (
+                isinstance(resource, str) and resource.startswith("gpu")
+            ):
+                if self.accelerator is not None:
+                    raise ValueError("Accelerator already set.")
+                self.accelerator = resource  # type: ignore
+
+            if resource in self.task_requirements:
+                raise ValueError(f"{resource} has been specified twice.")
+            self.task_requirements[resource] = value
+
+    def to_directives(self) -> list[str]:
+        directives = []
+
+        for resource, value in self.task_requirements.items():
+            match resource:
+                case ResourceType.EPHEMERAL_STORAGE | ResourceType.DISK:
+                    assert isinstance(value, int), "Disk space must be an integer"
+                    directives.append(f"--tmp={math.ceil(value / 2**20)}M")
+                case ResourceType.MEMORY | ResourceType.RAM:
+                    num_cpus = self.task_requirements.get(ResourceType.CPU, 1)
+                    assert isinstance(value, (int, float)), "Memory must be an integer or float"
+                    assert isinstance(num_cpus, int), "CPU must be an integer"
+                    directives.append(f"--mem-per-cpu={math.ceil(value / num_cpus / 2**20)}M")
+                case ResourceType.CPU:
+                    assert isinstance(value, int), "CPU must be an integer"
+                    directives.append(f"--cpus-per-task={value}")
+                case ResourceType.GPU:
+                    assert isinstance(value, int), "GPU must be an integer"
+                    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}")
+                case str():
+                    directives.append(f"--gres={resource}:{value}")
+
+        directives.append(f"--ntasks={self.replicas}")
+        if self.location:
+            directives.append(f"--nodelist={self.location}")
+
+        return directives
+
+    def replace(
+        self,
+        cluster: config.SlurmClusterConfig | None,
+        **kw_resources: ResourceQuantity,
+    ) -> "JobRequirements":
+        return JobRequirements(
+            resources=self.task_requirements | kw_resources,  # type: ignore
+            replicas=self.replicas,
+            cluster=cluster or self.cluster,
+        )
+
+    def __repr__(self) -> str:
+        args = []
+
+        for resource, value in self.task_requirements.items():
+            if isinstance(resource, ResourceType):
+                resource = resource.name
+            args.append(f"{resource.lower()}={value!r}")
+
+        if self.replicas != 1:
+            args.append(f"replicas={self.replicas}")
+
+        if self.cluster is not None:
+            args.append(f"cluster={self.cluster!r}")
+
+        return f'xm_slurm.JobRequirements({", ".join(args)})'