xmanager-slurm 0.4.19__py3-none-any.whl

xm_slurm/packageables.py

@@ -0,0 +1,309 @@
+import importlib.resources as resources
+import pathlib
+import sys
+import typing as tp
+
+from xmanager import xm
+
+from xm_slurm import job_blocks, utils
+from xm_slurm.executables import Dockerfile, DockerImage
+from xm_slurm.executors import SlurmSpec
+
+
+def docker_image(
+    *,
+    image: str,
+    args: xm.UserArgs | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
+) -> xm.Packageable:
+    """Creates a packageable for a pre-built Docker image.
+
+    Args:
+        image: The remote image URI.
+        args: The user arguments to pass to the executable.
+        env_vars: The environment variables to pass to the executable.
+
+    Returns: A packageable for a pre-built Docker image.
+    """
+    return xm.Packageable(
+        executor_spec=SlurmSpec(),
+        executable_spec=DockerImage(image=image),
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=env_vars or {},
+    )
+
+
+def docker_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    dockerfile: pathlib.Path | None = None,
+    context: pathlib.Path | None = None,
+    target: str | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
+    build_args: tp.Mapping[str, str] | None = None,
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    args: xm.UserArgs | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
+) -> xm.Packageable:
+    """Creates a Docker container packageable from a dockerfile.
+
+    Args:
+        executor_spec: The executor specification, where will the container be stored at.
+        dockerfile: The path to the dockerfile.
+        context: The path to the docker context.
+        target: The docker build target.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
+        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.
+        args: The user arguments to pass to the executable.
+        env_vars: The environment variables to pass to the executable.
+
+    Returns: A Docker container packageable.
+    """
+    if context is None:
+        context = utils.find_project_root()
+    context = context.resolve()
+    if dockerfile is None:
+        dockerfile = context / "Dockerfile"
+    dockerfile = dockerfile.resolve()
+
+    if ssh is None:
+        ssh = []
+    elif ssh is True:
+        ssh = ["default"]
+
+    if cache_from is None and isinstance(executor_spec, SlurmSpec):
+        cache_from = executor_spec.tag
+    if cache_from is None:
+        cache_from = []
+    elif isinstance(cache_from, str):
+        cache_from = [cache_from]
+
+    return xm.Packageable(
+        executor_spec=executor_spec,
+        executable_spec=Dockerfile(
+            dockerfile=dockerfile,
+            context=context,
+            target=target,
+            ssh=ssh,
+            build_args=build_args or {},
+            cache_from=cache_from,
+            labels=labels or {},
+        ),
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=env_vars or {},
+    )
+
+
+def python_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    entrypoint: xm.ModuleName | xm.CommandList,
+    context: pathlib.Path | None = None,
+    requirements: pathlib.Path | None = None,
+    base_image: str = "docker.io/python:{major}.{minor}-slim",
+    extra_system_packages: tp.Sequence[str] = (),
+    extra_python_packages: tp.Sequence[str] = (),
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
+    args: xm.UserArgs | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
+) -> xm.Packageable:
+    """Creates a Python container from a base image using pip from a `requirements.txt` file.
+
+    NOTE: The base image will use the Python version of the current interpreter.
+    NOTE: uv is used to install packages from `requirements`.
+
+    Args:
+        executor_spec: The executor specification, where will the container be stored at.
+        entrypoint: The entrypoint to run in the container.
+        context: The path to the docker context.
+        requirements: The path to the pip requirements file.
+        base_image: The base image to use. NOTE: The base image must contain the Python runtime.
+        extra_system_packages: Additional system packages to install. NOTE: These are installed via `apt-get`.
+        extra_python_packages: Additional Python packages to install. NOTE: These are installed via `uv pip`.
+        cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
+        labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
+        args: The user arguments to pass to the executable.
+        env_vars: The environment variables to pass to the executable.
+
+    Returns: A Python container packageable.
+    """
+    entrypoint_args = job_blocks.get_args_for_python_entrypoint(entrypoint)
+    args = xm.merge_args(entrypoint_args, args or {})
+
+    if context is None:
+        context = utils.find_project_root()
+    context = context.resolve()
+    if requirements is None:
+        requirements = context / "requirements.txt"
+    requirements = requirements.resolve()
+    if not requirements.exists():
+        raise ValueError(f"Pip requirements `{requirements}` doesn't exist.")
+    if not requirements.is_relative_to(context):
+        raise ValueError(
+            f"Pip requirements `{requirements}` must be relative to context: `{context}`"
+        )
+
+    with resources.as_file(
+        resources.files("xm_slurm.templates").joinpath("docker/python.Dockerfile")
+    ) as dockerfile:
+        return docker_container(
+            executor_spec=executor_spec,
+            dockerfile=dockerfile,
+            context=context,
+            ssh=ssh,
+            build_args={
+                "PIP_REQUIREMENTS": requirements.relative_to(context).as_posix(),
+                "EXTRA_SYSTEM_PACKAGES": " ".join(extra_system_packages),
+                "EXTRA_PYTHON_PACKAGES": " ".join(extra_python_packages),
+                "BASE_IMAGE": base_image.format_map({
+                    "major": sys.version_info.major,
+                    "minor": sys.version_info.minor,
+                    "micro": sys.version_info.micro,
+                }),
+            },
+            cache_from=cache_from,
+            labels=labels,
+            args=args,
+            env_vars=env_vars,
+        )
+
+
+def mamba_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    entrypoint: xm.ModuleName | xm.CommandList,
+    context: pathlib.Path | None = None,
+    environment: pathlib.Path | None = None,
+    base_image: str = "gcr.io/distroless/base-debian10",
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
+    args: xm.UserArgs | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
+) -> xm.Packageable:
+    """Creates a Conda container from a base image using mamba from a `environment.yml` file.
+
+    Note: The base image *doesn't* need to contain the Python runtime.
+
+    Args:
+        executor_spec: The executor specification, where will the container be stored at.
+        entrypoint: The entrypoint to run in the container.
+        context: The path to the docker context.
+        environment: The path to the conda environment file.
+        base_image: The base image to use.
+        cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
+        labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
+        args: The user arguments to pass to the executable.
+        env_vars: The environment variables to pass to the executable.
+
+    Returns: A Conda container packageable.
+    """
+    entrypoint_args = job_blocks.get_args_for_python_entrypoint(entrypoint)
+    args = xm.merge_args(entrypoint_args, args or {})
+
+    if context is None:
+        context = utils.find_project_root()
+    context = context.resolve()
+    if environment is None:
+        environment = context / "environment.yml"
+    environment = environment.resolve()
+    if not environment.exists():
+        raise ValueError(f"Conda environment manifest `{environment}` doesn't exist.")
+    if not environment.is_relative_to(context):
+        raise ValueError(
+            f"Conda environment manifest `{environment}` must be relative to context: `{context}`"
+        )
+
+    with resources.as_file(
+        resources.files("xm_slurm.templates").joinpath("docker/mamba.Dockerfile")
+    ) as dockerfile:
+        return docker_container(
+            executor_spec=executor_spec,
+            dockerfile=dockerfile,
+            context=context,
+            ssh=ssh,
+            build_args={
+                "CONDA_ENVIRONMENT": environment.relative_to(context).as_posix(),
+                "BASE_IMAGE": base_image,
+            },
+            cache_from=cache_from,
+            labels=labels,
+            args=args,
+            env_vars=env_vars,
+        )
+
+
+conda_container = mamba_container
+
+
+def uv_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    entrypoint: xm.ModuleName | xm.CommandList,
+    context: pathlib.Path | None = None,
+    base_image: str = "docker.io/python:{major}.{minor}-slim-bookworm",
+    extra_system_packages: tp.Sequence[str] = (),
+    extra_python_packages: tp.Sequence[str] = (),
+    cache_from: str | tp.Sequence[str] | None = None,
+    labels: tp.Mapping[str, str] | None = None,
+    ssh: tp.Sequence[str] | tp.Literal[True] | None = None,
+    args: xm.UserArgs | None = None,
+    env_vars: tp.Mapping[str, str] | None = None,
+) -> xm.Packageable:
+    """Creates a Python container from a base image using uv from a `uv.lock` file.
+
+    Args:
+        executor_spec: The executor specification, where will the container be stored at.
+        entrypoint: The entrypoint to run in the container.
+        context: The path to the docker context.
+        base_image: The base image to use. NOTE: The base image must contain the Python runtime.
+        extra_system_packages: Additional system packages to install. NOTE: These are installed via `apt-get`.
+        extra_python_packages: Additional Python packages to install. NOTE: These are installed via `uv pip`.
+        cache_from: Where to pull the BuildKit cache from. See `--cache-from` in `docker build`.
+        labels: The container labels.
+        ssh: A list of SSH sockets/keys for the docker build step or `True` to use the default SSH agent.
+        args: The user arguments to pass to the executable.
+        env_vars: The environment variables to pass to the executable.
+
+    Returns: A Python container packageable.
+    """
+    entrypoint_args = job_blocks.get_args_for_python_entrypoint(entrypoint)
+    args = xm.merge_args(entrypoint_args, args or {})
+
+    if context is None:
+        context = utils.find_project_root()
+    context = context.resolve()
+    if not (context / "pyproject.toml").exists():
+        raise ValueError(f"Python project file `{context / 'pyproject.toml'}` doesn't exist.")
+    if not (context / "uv.lock").exists():
+        raise ValueError(f"UV lock file `{context / 'uv.lock'}` doesn't exist.")
+
+    with resources.as_file(
+        resources.files("xm_slurm.templates").joinpath("docker/uv.Dockerfile")
+    ) as dockerfile:
+        return docker_container(
+            executor_spec=executor_spec,
+            dockerfile=dockerfile,
+            context=context,
+            ssh=ssh,
+            build_args={
+                "EXTRA_SYSTEM_PACKAGES": " ".join(extra_system_packages),
+                "EXTRA_PYTHON_PACKAGES": " ".join(extra_python_packages),
+                "BASE_IMAGE": base_image.format_map({
+                    "major": sys.version_info.major,
+                    "minor": sys.version_info.minor,
+                    "micro": sys.version_info.micro,
+                }),
+            },
+            cache_from=cache_from,
+            labels=labels,
+            args=args,
+            env_vars=env_vars,
+        )

xm_slurm/packaging/__init__.py

@@ -0,0 +1,8 @@
+# First register our built-in packaging methods
+import xm_slurm.packaging.docker  # noqa: F401
+from xm_slurm.packaging import registry, router
+
+package = router.package
+register = registry.register
+
+__all__ = ["package", "register"]

xm_slurm/packaging/docker.py

@@ -0,0 +1,348 @@
+import base64
+import collections.abc
+import dataclasses
+import enum
+import functools
+import hashlib
+import json
+import logging
+import os
+import pathlib
+import shlex
+import shutil
+import tempfile
+import typing as tp
+
+import jinja2 as j2
+from xmanager import xm
+
+from xm_slurm import utils
+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.registry import IndexedContainer
+
+logger = logging.getLogger(__name__)
+
+
+def _hash_digest(obj: tp.Hashable) -> str:
+    return hashlib.sha256(repr(obj).encode()).hexdigest()
+
+
+class DockerClient:
+    class Builder(enum.Enum):
+        BUILDKIT = enum.auto()
+        BUILDAH = enum.auto()
+
+    def __init__(self) -> None:
+        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
+
+        backend_version = utils.run_command(
+            xm.merge_args(self._client_call, ["buildx", "version"]), return_stdout=True
+        )
+        if backend_version.stdout.startswith("github.com/docker/buildx"):
+            self._builder = DockerClient.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
+            result = utils.run_command(
+                [f"docker-credential-{helper}", "get"],
+                stdin=hostname,
+                return_stdout=True,
+            )
+
+            if result.returncode == 0:
+                credentials = json.loads(result.stdout)
+                return RemoteRepositoryCredentials(
+                    username=str.strip(credentials["Username"]),
+                    password=str.strip(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(
+                            str.strip(username),
+                            str.strip(password),
+                        )
+                        self._credentials_cache[hostname] = credentials
+                        return credentials
+            return None
+
+        # Attempt to parse credentials from the Docker or Podman configuration
+        match self._builder:
+            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 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)
+
+    def inspect(self, image: ImageURI, element: str) -> dict[str, tp.Any]:
+        output = utils.run_command(
+            xm.merge_args(
+                self._client_call,
+                ["buildx", "imagetools", "inspect"],
+                ["--format", f"{{{{json .{element}}}}}"],
+                [str(image)],
+            ),
+            check=True,
+            return_stdout=True,
+        )
+        return json.loads(output.stdout.strip().strip("'"))
+
+    @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 | tp.Sequence[str] | None = None,
+        builder: str | None = None,
+        files: str | os.PathLike[str] | tp.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: tp.Literal["auto", "plain", "tty"] = "auto",
+        set: tp.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, *, targets: tp.Sequence[IndexedContainer[xm.Packageable]]
+    ) -> list[IndexedContainer[RemoteImage]]:
+        executors_by_executables = packaging_utils.collect_executors_by_executable(targets)
+        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=_hash_digest,
+            )
+            hcl_file.write_text(hcl)
+            logger.debug(hcl)
+
+            try:
+                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=False,
+                        push=True,
+                    ),
+                )
+                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[_hash_digest(target.value.executable_spec)]
+            uri = ImageURI(target.value.executor_spec.tag).with_digest(
+                executable_metadata["containerimage.digest"]
+            )
+            config = self.inspect(uri, "Image.Config")
+            if "WorkingDir" not in config:
+                raise ValueError(
+                    "Docker image does not have a working directory. "
+                    "To support all runtimes, we need to set a working directory. "
+                    "Please set `WORKDIR` in the `Dockerfile`."
+                )
+            if "Entrypoint" not in config:
+                raise ValueError(
+                    "Docker image does not have an entrypoint. "
+                    "To support all runtimes, we need to set an entrypoint. "
+                    "Please set `ENTRYPOINT` in the `Dockerfile`."
+                )
+
+            images.append(
+                dataclasses.replace(
+                    target,
+                    value=RemoteImage(  # type: ignore
+                        image=str(uri),
+                        workdir=config["WorkingDir"],
+                        entrypoint=xm.SequentialArgs.from_collection(config["Entrypoint"]),
+                        args=target.value.args,
+                        env_vars=target.value.env_vars,
+                        credentials=self.credentials(uri.domain),
+                    ),
+                )
+            )
+
+        return images
+
+
+@functools.cache
+def docker_client() -> DockerClient:
+    return DockerClient()
+
+
+@registry.register(Dockerfile)
+def _(
+    targets: tp.Sequence[IndexedContainer[xm.Packageable]],
+) -> list[IndexedContainer[RemoteImage]]:
+    return docker_client().bake(targets=targets)
+
+
+@registry.register(DockerImage)
+def _(
+    targets: tp.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)
+
+        config = client.inspect(uri, "Image.Config")
+        if "WorkingDir" not in config:
+            raise ValueError(
+                "Docker image does not have a working directory. "
+                "To support all runtimes, we need to set a working directory. "
+                "Please set `WORKDIR` in the `Dockerfile`."
+            )
+        if "Entrypoint" not in config:
+            raise ValueError(
+                "Docker image does not have an entrypoint. "
+                "To support all runtimes, we need to set an entrypoint. "
+                "Please set `ENTRYPOINT` in the `Dockerfile`."
+            )
+
+        images.append(
+            dataclasses.replace(
+                target,
+                value=RemoteImage(  # type: ignore
+                    image=str(uri),
+                    workdir=config["WorkingDir"],
+                    entrypoint=xm.SequentialArgs.from_collection(config["Entrypoint"]),
+                    args=target.value.args,
+                    env_vars=target.value.env_vars,
+                    credentials=client.credentials(hostname=uri.domain),
+                ),
+            )
+        )
+
+    return images