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

xm_slurm/experimental/parameter_controller.py

@@ -18,6 +18,8 @@ from xm_slurm.experiment import SlurmAuxiliaryUnit, SlurmExperiment
 P = ParamSpec("P")
 T = TypeVar("T")
 
+logger = logging.getLogger(__name__)
+
 
 async def _monitor_parameter_controller(
     aux_unit: SlurmAuxiliaryUnit,
@@ -33,16 +35,16 @@ async def _monitor_parameter_controller(
         try:
             _ = future.result()
         except asyncio.CancelledError:
-            logging.info("Local parameter controller was cancelled, resuming on remote controller.")
+            logger.info("Local parameter controller was cancelled, resuming on remote controller.")
             pass
         except Exception:
-            logging.error("Local parameter controller failed, stopping remote controller.")
+            logger.error("Local parameter controller failed, stopping remote controller.")
             aux_unit.stop(
                 mark_as_failed=True, mark_as_completed=False, message="Local controller failed."
             )
             raise
         else:
-            logging.info(
+            logger.info(
                 "Local parameter controller finished before remote controller started, "
                 "stopping remote controller."
             )
@@ -56,33 +58,33 @@ async def _monitor_parameter_controller(
         interval=poll_interval,
     )
     async def wait_for_remote_controller() -> status.SlurmWorkUnitStatusEnum:
-        logging.info("Waiting for remote parameter controller to start.")
+        logger.info("Waiting for remote parameter controller to start.")
         if local_controller_finished.is_set():
             return status.SlurmWorkUnitStatusEnum.COMPLETED
         return (await aux_unit.get_status()).status
 
-    logging.info("Monitoring remote parameter controller.")
+    logger.info("Monitoring remote parameter controller.")
     # TODO(jfarebro): make get_status() more resiliant to errors when initially scheduling.
     # We run into issues if we call get_status() too quickly when Slurm hasn't ingested the job.
     await asyncio.sleep(15)
     match await wait_for_remote_controller():
         case status.SlurmWorkUnitStatusEnum.RUNNING:
-            logging.info("Remote parameter controller started.")
+            logger.info("Remote parameter controller started.")
             local_parameter_controller.cancel("Remote parameter controller started.")
         case status.SlurmWorkUnitStatusEnum.COMPLETED:
             if local_parameter_controller.done():
-                logging.info("Local parameter controller finished, stopping remote controller.")
+                logger.info("Local parameter controller finished, stopping remote controller.")
                 aux_unit.stop(
                     mark_as_completed=True, message="Local parameter controller finished."
                 )
             else:
-                logging.info("Remote parameter controller finished, stopping local controller.")
+                logger.info("Remote parameter controller finished, stopping local controller.")
                 local_parameter_controller.cancel()
         case status.SlurmWorkUnitStatusEnum.FAILED:
-            logging.error("Remote parameter controller failed, stopping local controller.")
+            logger.error("Remote parameter controller failed, stopping local controller.")
             local_parameter_controller.cancel()
         case status.SlurmWorkUnitStatusEnum.CANCELLED:
-            logging.info("Remote parameter controller was cancelled, stopping local controller.")
+            logger.info("Remote parameter controller was cancelled, stopping local controller.")
             local_parameter_controller.cancel()
         case status.SlurmWorkUnitStatusEnum.PENDING:
             raise RuntimeError("Remote parameter controller is still pending, invalid state.")

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,29 +12,31 @@ import shlex
 import shutil
 import subprocess
 import tempfile
-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__)
+
+
+def _hash_digest(obj: Hashable) -> str:
+    return hashlib.sha256(repr(obj).encode()).hexdigest()
 
-class LocalDockerClient(DockerClient):
-    """Build Docker images locally."""
 
+class DockerClient:
     class Builder(enum.Enum):
         BUILDKIT = enum.auto()
         BUILDAH = enum.auto()
@@ -47,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}")
 
@@ -115,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,
@@ -150,28 +202,27 @@ class LocalDockerClient(DockerClient):
             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
@@ -184,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"]
             )
@@ -205,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,7 +7,7 @@ import re
 import select
 import shutil
 import subprocess
-from typing import Callable, Concatenate, Hashable, ParamSpec, Sequence, TypeVar
+from typing import ParamSpec, Sequence, TypeVar
 
 from xmanager import xm
 
@@ -20,17 +17,7 @@ T = TypeVar("T")
 P = ParamSpec("P")
 ReturnT = TypeVar("ReturnT")
 
-
-def hash_digest(obj: Hashable) -> str:
-    # TODO(jfarebro): Need a better way to hash these objects
-    # 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,17 +29,6 @@ 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_ESCAPE_SEQUENCES_REGEX = re.compile(
     rb"\x1b\[\?25[hl]"  # Matches cursor show/hide commands (CSI ?25h and CSI ?25l)
@@ -86,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()

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/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

xm_slurm/templates/slurm/fragments/monitor.bash.j2

@@ -24,6 +24,11 @@ __xm_slurm_wait_for_children() {
       echo "INFO: Received requeue exit code {{ requeue_exit_code }} from job ${job}. Requeing Slurm job ${JOB_ID} after ${SLURM_RESTART_COUNT-0} restarts." >&2
       scontrol requeue "${JOB_ID}"
       exit {{ requeue_exit_code }}
+    elif [ "${JOB_EXIT_CODE}" -ne 0 ]; then
+      echo "ERROR: Job ${job} exited with code ${JOB_EXIT_CODE}." >&2
+      exit "${JOB_EXIT_CODE}"
+    else
+      echo "INFO: Job ${job} exited successfully." >&2
     fi
   done
 }

xm_slurm/templates/slurm/job-array.bash.j2

@@ -2,8 +2,7 @@
 {% block directives %}
 {{ super() -}}
 #SBATCH --array=0-{{ args | length - 1 }}
-#SBATCH --output=xm-%j-%a.stdout
-#SBATCH --error=xm-%j-%a.stderr
+#SBATCH --output=slurm-%A_%a.out
 {% endblock directives %}
 
 {% block bootstrap %}

xm_slurm/templates/slurm/job.bash.j2

@@ -2,8 +2,7 @@
 {% block directives %}
 #SBATCH --open-mode=append
 #SBATCH --export=NONE
-#SBATCH --output=xm-%j.stdout
-#SBATCH --error=xm-%j.stderr
+#SBATCH --output=slurm-%j.out
 #SBATCH --comment="{'xid': {{ experiment_id }}}"
 {% if cluster.account and not job.executor.account %}
 #SBATCH --account={{ cluster.account }}
@@ -16,8 +15,10 @@
 {% endif %}
 {% if identity %}
 #SBATCH --job-name=xm[{{ experiment_id }}.{{ identity }}]
-#SBATCH --dependency=singleton
 {% else %}
+{% if dependency %}
+#SBATCH {{ dependency.to_directive() }}
+{% endif %}
 #SBATCH --job-name=xm[{{ experiment_id }}]
 {% endif %}
 {% for directive in job.executor.to_directives() %}

xm_slurm/types.py

@@ -0,0 +1,23 @@
+import typing as tp
+
+InstanceT_contra = tp.TypeVar("InstanceT_contra", contravariant=True)
+GetterT_co = tp.TypeVar("GetterT_co", covariant=True)
+SetterT_co = tp.TypeVar("SetterT_co", contravariant=True)
+
+
+class Descriptor(tp.Protocol[GetterT_co, SetterT_co]):
+    def __set_name__(self, owner: tp.Type[tp.Any], name: str) -> None: ...
+
+    @tp.overload
+    def __get__(
+        self, instance: InstanceT_contra, owner: tp.Type[InstanceT_contra] | None = None
+    ) -> GetterT_co: ...
+
+    @tp.overload
+    def __get__(self, instance: None, owner: tp.Type[InstanceT_contra]) -> GetterT_co: ...
+
+    def __get__(
+        self, instance: InstanceT_contra | None, owner: tp.Type[InstanceT_contra] | None = None
+    ) -> GetterT_co: ...
+
+    def __set__(self, instance: tp.Any, value: SetterT_co) -> None: ...

{xmanager_slurm-0.4.0.dist-info → xmanager_slurm-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: xmanager-slurm
-Version: 0.4.0
+Version: 0.4.1
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>

xmanager_slurm-0.4.1.dist-info/RECORD

@@ -0,0 +1,44 @@
+xm_slurm/__init__.py,sha256=Ld2w7ofLlTieWOHP_Jb3f48-qtVQBjFXynxUm9WF8mc,1116
+xm_slurm/api.py,sha256=LeGgHz82t8Oay0Z1Ourv9-r-DBur3lhCUTnmmGhGFY4,18502
+xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
+xm_slurm/config.py,sha256=GLLEkRLJxQW0urmHCLmwq_4ECmimEBQFl8Nz62SIo78,6787
+xm_slurm/console.py,sha256=UpMqeJ0C8i0pkue1AHnnyyX0bFJ9zZeJ7HBR6yhuA8A,54
+xm_slurm/constants.py,sha256=zefVtlFdflgSolie5g_rVxWV-Zpydxapchm3y0a2FDc,999
+xm_slurm/dependencies.py,sha256=-5gN_tpfs3dOA7H5_MIHO2ratb7F5Pm_yjkR5rZcgI8,6421
+xm_slurm/executables.py,sha256=S3z8jSDL6AdyGYpzy_cCs03Mj0vgA4ZTqIe8APYor3E,6469
+xm_slurm/execution.py,sha256=i2oYH5RS-mHsHPwFDFZvo5qCudbgqBML-Hzq6DPNItw,25721
+xm_slurm/executors.py,sha256=fMtxGUCi4vEKmb_p4JEpqPUTh7L_f1LcR_TamMLAWNg,4667
+xm_slurm/experiment.py,sha256=trHapcYxPNKofzSqu7KZawML59tZ8FVjoEZYe2Wal7w,44521
+xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
+xm_slurm/packageables.py,sha256=YZFTL6UWx9A_zyztTy1THUlj3pW1rA0cBPHJxD1LOJk,12884
+xm_slurm/resources.py,sha256=EaYDATVudrEDPKKdSZoWgfqPiidc6DMjIctmzLQmiH0,5683
+xm_slurm/status.py,sha256=WTWiDHi-ZHtwHRnDP0cGa-27zTSm6LkA-GCKsN-zBgg,6916
+xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
+xm_slurm/utils.py,sha256=ESjOkGT7bRSzIeZrUtZplSHP4oaH6VZ92y2woYdcyKM,2239
+xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xm_slurm/contrib/clusters/__init__.py,sha256=vugR50D9fPJQN5bTd7cSArDGrA6pC-YJHMXrEyvr_Uw,2980
+xm_slurm/contrib/clusters/drac.py,sha256=tJeQFWFIpeZ1gD3j6AAJssNoLSiDkB-3lz1_ObnkRhc,5905
+xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
+xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
+xm_slurm/packaging/docker.py,sha256=TA8-QG09EdhF4K1ixrEboVFarF9LcURNHhzUXL-7Iqg,11518
+xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CUE,1148
+xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
+xm_slurm/packaging/utils.py,sha256=KI5s32rNTCfgwzY_7Ghck27jHKvKg5sl5_NEEqJbJqI,3999
+xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
+xm_slurm/scripts/cli.py,sha256=ZXqYOs8X23TYDdKxvV-wIa-0mTfpxSl4_Pli6TiKI7s,1435
+xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=ClsFpj91Mr1VfA8L6eqBG3HQz0Z8VenF6mEfmAhQgUo,1498
+xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgCdBXjzqslFJdZI,753
+xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
+xm_slurm/templates/docker/uv.Dockerfile,sha256=kYD32oUS1jUaARsNV1o6EFnIfLCNh5GMmck27b-5NRU,969
+xm_slurm/templates/slurm/job-array.bash.j2,sha256=iYtGMRDXgwwc2_8E3v4a30f3fKuq4zWgZHkxCXJ9iXc,567
+xm_slurm/templates/slurm/job-group.bash.j2,sha256=UkjfBE7jg9mepcUWaHZEAjkiXsIM1j_sLxLzxkteD-Y,1120
+xm_slurm/templates/slurm/job.bash.j2,sha256=v0xGYzagDdWW6Tg44qobGJLNSUP1Cf4CcekrPibYdrE,1864
+xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=HYqYhXsTv8TCed5UaGCZVGIYsqxSKHcnPyNNTHWNvxc,1279
+xm_slurm/templates/slurm/fragments/proxy.bash.j2,sha256=VJLglZo-Nvx9R-qe3rHTxr07CylTQ6Z9NwBzvIpAZrA,814
+xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=dMntzelhs8DqKyIpO9S6wzMfH2PDevmgvyjCW8Xc2dY,3222
+xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=xKXYFvQvazMx0PgvmlRXR6eecoiBUl8y52dIzQtWkBE,1469
+xmanager_slurm-0.4.1.dist-info/METADATA,sha256=3mT4XIm8evv-5qw7oney4nYn3IasIA_l1rWz86XNOY8,954
+xmanager_slurm-0.4.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+xmanager_slurm-0.4.1.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.1.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.1.dist-info/RECORD,,

xmanager_slurm-0.4.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+xm = xm_slurm.scripts.cli:main

xm_slurm/packaging/docker/__init__.py

@@ -1,69 +0,0 @@
-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
-DOCKER_CLIENT_PROVIDER = flags.DEFINE_enum(
-    "xm_docker_client", "docker", ["docker"], "Docker image build client."
-)
-
-IndexedContainer = registry.IndexedContainer
-
-
-@functools.cache
-def docker_client() -> DockerClient:
-    match DOCKER_CLIENT_PROVIDER.value:
-        case "docker":
-            from xm_slurm.packaging.docker.local import LocalDockerClient
-
-            return LocalDockerClient()
-        case _:
-            raise ValueError(f"Unknown build client: {DOCKER_CLIENT_PROVIDER.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