xmanager-slurm 0.3.0__py3-none-any.whl

xm_slurm/job_blocks.py

@@ -0,0 +1,14 @@
+from xmanager import xm
+
+
+def get_args_for_python_entrypoint(
+    entrypoint: xm.ModuleName | xm.CommandList,
+) -> xm.SequentialArgs:
+    match entrypoint:
+        case xm.ModuleName():
+            entrypoint_args = ["-m", entrypoint.module_name]
+        case xm.CommandList():
+            entrypoint_args = entrypoint.commands
+        case _:
+            raise TypeError(f"Invalid entrypoint type: {type(entrypoint)}")
+    return xm.SequentialArgs.from_collection(entrypoint_args)

xm_slurm/packageables.py

@@ -0,0 +1,292 @@
+import importlib.resources as resources
+import pathlib
+import sys
+from typing import Mapping, Sequence
+
+import immutabledict
+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: Mapping[str, str] = immutabledict.immutabledict(),
+) -> 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=args,
+        env_vars=env_vars,
+    )
+
+
+def docker_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    dockerfile: pathlib.Path | None = None,
+    context: pathlib.Path | None = None,
+    target: str | None = None,
+    build_args: Mapping[str, str] = immutabledict.immutabledict(),
+    cache_from: str | Sequence[str] | None = None,
+    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    workdir: pathlib.Path | None = None,
+    args: xm.UserArgs | None = None,
+    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+) -> 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.
+        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.
+        workdir: The working directory in container.
+        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 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,
+            build_args=build_args,
+            cache_from=cache_from,
+            workdir=workdir,
+            labels=labels,
+        ),
+        args=args,
+        env_vars=env_vars,
+    )
+
+
+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",
+    cache_from: str | Sequence[str] | None = None,
+    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    args: xm.UserArgs | None = None,
+    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+) -> 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.
+
+    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.
+        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 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,
+            build_args={
+                "PIP_REQUIREMENTS": requirements.relative_to(context).as_posix(),
+                "PYTHON_MAJOR": str(sys.version_info.major),
+                "PYTHON_MINOR": str(sys.version_info.minor),
+                "PYTHON_MICRO": str(sys.version_info.micro),
+                "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,
+            # We must specify the workdir manually for apptainer support
+            workdir=pathlib.Path("/workspace"),
+            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 | Sequence[str] | None = None,
+    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    args: xm.UserArgs | None = None,
+    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+) -> 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.
+        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,
+            build_args={
+                "CONDA_ENVIRONMENT": environment.relative_to(context).as_posix(),
+                "BASE_IMAGE": base_image,
+            },
+            cache_from=cache_from,
+            labels=labels,
+            # We must specify the workdir manually for apptainer support
+            workdir=pathlib.Path("/workspace"),
+            args=args,
+            env_vars=env_vars,
+        )
+
+
+conda_container = mamba_container
+
+
+def pdm_container(
+    *,
+    executor_spec: xm.ExecutorSpec,
+    entrypoint: xm.ModuleName | xm.CommandList,
+    context: pathlib.Path | None = None,
+    base_image: str = "docker.io/python:{major}.{minor}-slim",
+    cache_from: str | Sequence[str] | None = None,
+    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    args: xm.UserArgs | None = None,
+    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+) -> xm.Packageable:
+    """Creates a Python container from a base image using pdm from a `pdm.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.
+        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 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 / "pdm.lock").exists():
+        raise ValueError(f"PDM lockfile `{context / 'pdm.lock'}` doesn't exist.")
+
+    with resources.as_file(
+        resources.files("xm_slurm.templates").joinpath("docker/pdm.Dockerfile")
+    ) as dockerfile:
+        return docker_container(
+            executor_spec=executor_spec,
+            dockerfile=dockerfile,
+            context=context,
+            build_args={
+                "PYTHON_MAJOR": str(sys.version_info.major),
+                "PYTHON_MINOR": str(sys.version_info.minor),
+                "PYTHON_MICRO": str(sys.version_info.micro),
+                "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,
+            # We must specify the workdir manually for apptainer support
+            workdir=pathlib.Path("/workspace/src"),
+            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/__init__.py

@@ -0,0 +1,75 @@
+import dataclasses
+import functools
+from typing import Sequence
+
+from absl import flags
+from xmanager import xm
+
+from xm_slurm.executables import Dockerfile, DockerImage, ImageURI, RemoteImage
+from xm_slurm.executors import SlurmSpec
+from xm_slurm.packaging import registry
+from xm_slurm.packaging.docker.abc import DockerClient
+
+FLAGS = flags.FLAGS
+REMOTE_BUILD = flags.DEFINE_enum(
+    "xm_builder", "local", ["local", "gcp", "azure"], "Remote build provider."
+)
+
+IndexedContainer = registry.IndexedContainer
+
+
+@functools.cache
+def docker_client() -> DockerClient:
+    match REMOTE_BUILD.value:
+        case "local":
+            from xm_slurm.packaging.docker.local import LocalDockerClient
+
+            return LocalDockerClient()
+        case "gcp":
+            from xm_slurm.packaging.docker.cloud import GoogleCloudRemoteDockerClient
+
+            return GoogleCloudRemoteDockerClient()
+        case "azure":
+            raise NotImplementedError("Azure remote build is not yet supported.")
+        case _:
+            raise ValueError(f"Unknown remote build provider: {REMOTE_BUILD.value}")
+
+
+@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/docker/abc.py

@@ -0,0 +1,112 @@
+import abc
+import collections.abc
+import dataclasses
+import functools
+import os
+from typing import Literal, Mapping, Protocol, Sequence
+
+import jinja2 as j2
+from xmanager import xm
+
+from xm_slurm.executables import RemoteImage, RemoteRepositoryCredentials
+from xm_slurm.packaging.registry import IndexedContainer
+
+
+class DockerCommandProtocol(Protocol):
+    def to_args(self) -> xm.SequentialArgs: ...
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DockerBakeCommand(DockerCommandProtocol):
+    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
+
+    def to_args(self) -> xm.SequentialArgs:
+        files = self.files
+        if files is None:
+            files = []
+        if not isinstance(files, collections.abc.Sequence):
+            files = [files]
+
+        targets = self.targets
+        if targets is None:
+            targets = []
+        elif not isinstance(targets, collections.abc.Sequence):
+            targets = [targets]
+
+        return xm.merge_args(
+            ["buildx", "bake"],
+            [f"--progress={self.progress}"],
+            [f"--builder={self.builder}"] if self.builder else [],
+            [f"--metadata-file={self.metadata_file}"] if self.metadata_file else [],
+            ["--print"] if self.print else [],
+            ["--push"] if self.push else [],
+            ["--pull"] if self.pull else [],
+            ["--load"] if self.load else [],
+            ["--no-cache"] if not self.cache else [],
+            [f"--file={file}" for file in files],
+            [f"--set={key}={value}" for key, value in self.set.items()] if self.set else [],
+            targets,
+        )
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DockerPullCommand(DockerCommandProtocol):
+    image: str
+
+    def to_args(self) -> xm.SequentialArgs:
+        return xm.merge_args(["pull", self.image])
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DockerLoginCommand(DockerCommandProtocol):
+    server: str
+    username: str
+    password: str | None = None
+    password_stdin: bool = False
+
+    def __post_init__(self):
+        if self.password is None and not self.password_stdin:
+            raise ValueError("Either password or password_stdin must be set")
+        if self.password is not None and self.password_stdin:
+            raise ValueError("Only one of password or password_stdin must be set")
+
+    def to_args(self) -> xm.SequentialArgs:
+        return xm.merge_args(
+            ["login", "--username", self.username],
+            ["--password", self.password] if self.password else [],
+            ["--password-stdin"] if self.password_stdin else [],
+            [self.server],
+        )
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DockerVersionCommand(DockerCommandProtocol):
+    def to_args(self) -> xm.SequentialArgs:
+        return xm.merge_args(["buildx", "version"])
+
+
+class DockerClient(abc.ABC):
+    @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")
+
+    @abc.abstractmethod
+    def credentials(self, *, hostname: str) -> RemoteRepositoryCredentials | None: ...
+
+    @abc.abstractmethod
+    def bake(
+        self, *, targets: Sequence[IndexedContainer[xm.Packageable]]
+    ) -> list[IndexedContainer[RemoteImage]]: ...