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

xm_slurm/experimental/parameter_controller.py

@@ -0,0 +1,202 @@
+import asyncio
+import base64
+import dataclasses
+import enum
+import functools
+import logging
+import zlib
+from typing import Awaitable, Callable, Concatenate, Coroutine, Mapping, ParamSpec, TypeVar
+
+import backoff
+import cloudpickle
+from xmanager import xm
+
+import xm_slurm
+from xm_slurm import job_blocks, status
+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,
+    local_parameter_controller_coro: Coroutine[None, None, T],
+    *,
+    poll_interval: float = 30.0,
+) -> None:
+    local_controller_finished = asyncio.Event()
+    local_parameter_controller = asyncio.create_task(local_parameter_controller_coro)
+
+    @local_parameter_controller.add_done_callback
+    def _(future: asyncio.Task[T]) -> None:
+        try:
+            _ = future.result()
+        except asyncio.CancelledError:
+            logger.info("Local parameter controller was cancelled, resuming on remote controller.")
+            pass
+        except Exception:
+            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:
+            logger.info(
+                "Local parameter controller finished before remote controller started, "
+                "stopping remote controller."
+            )
+            local_controller_finished.set()
+            aux_unit.stop(mark_as_completed=True, message="Local parameter controller finished.")
+
+    @backoff.on_predicate(
+        backoff.constant,
+        lambda aux_unit_status: aux_unit_status is status.SlurmWorkUnitStatusEnum.PENDING,
+        jitter=None,
+        interval=poll_interval,
+    )
+    async def wait_for_remote_controller() -> status.SlurmWorkUnitStatusEnum:
+        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
+
+    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:
+            logger.info("Remote parameter controller started.")
+            local_parameter_controller.cancel("Remote parameter controller started.")
+        case status.SlurmWorkUnitStatusEnum.COMPLETED:
+            if local_parameter_controller.done():
+                logger.info("Local parameter controller finished, stopping remote controller.")
+                aux_unit.stop(
+                    mark_as_completed=True, message="Local parameter controller finished."
+                )
+            else:
+                logger.info("Remote parameter controller finished, stopping local controller.")
+                local_parameter_controller.cancel()
+        case status.SlurmWorkUnitStatusEnum.FAILED:
+            logger.error("Remote parameter controller failed, stopping local controller.")
+            local_parameter_controller.cancel()
+        case status.SlurmWorkUnitStatusEnum.CANCELLED:
+            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.")
+
+
+class ParameterControllerMode(enum.Enum):
+    AUTO = enum.auto()
+    REMOTE_ONLY = enum.auto()
+    # TODO(jfarebro): is it possible to get LOCAL_ONLY?
+    # We'd need to have a dummy job type as we need to return a JobType?
+
+
+def parameter_controller(
+    *,
+    executable: xm.Executable,
+    executor: xm.Executor,
+    controller_mode: ParameterControllerMode = ParameterControllerMode.AUTO,
+    controller_name: str = "parameter_controller",
+    controller_args: xm.UserArgs | None = None,
+    controller_env_vars: Mapping[str, str] | None = None,
+) -> Callable[
+    [
+        Callable[Concatenate[SlurmExperiment, P], T]
+        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
+    ],
+    Callable[P, xm.AuxiliaryUnitJob],
+]:
+    """Converts a function to a controller which can be added to an experiment.
+
+    Calling the wrapped function would return an xm.JobGenerator which would run
+    it as auxiliary unit on the specified executor.
+
+    Args:
+        executable: An executable that has a Python entrypoint with all the necesarry dependencies.
+        executor: The executor to launch the controller job on.
+        controller_name: Name of the parameter controller job.
+        controller_args: Mapping of flag names and values to be used by the XM
+          client running inside the parameter controller job.
+        controller_env_vars: Mapping of env variable names and values to be passed
+          to the parameter controller job.
+
+    Returns:
+        A decorator to be applied to the function.
+    """
+
+    def decorator(
+        f: Callable[Concatenate[SlurmExperiment, P], T]
+        | Callable[Concatenate[SlurmExperiment, P], Awaitable[T]],
+    ) -> Callable[P, xm.AuxiliaryUnitJob]:
+        @functools.wraps(f)
+        def make_controller(*args: P.args, **kwargs: P.kwargs) -> xm.AuxiliaryUnitJob:
+            # Modify the function to read the experiment from the API so that it can be pickled.
+
+            async def job_generator(aux_unit: SlurmAuxiliaryUnit) -> None:
+                experiment_id = aux_unit.experiment.experiment_id
+
+                async def local_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                    if asyncio.iscoroutinefunction(f):
+                        return await f(aux_unit.experiment, *args, **kwargs)
+                    else:
+                        return f(aux_unit.experiment, *args, **kwargs)
+
+                async def remote_controller(*args: P.args, **kwargs: P.kwargs) -> T | Awaitable[T]:
+                    async with xm_slurm.get_experiment(experiment_id=experiment_id) as exp:
+                        if asyncio.iscoroutinefunction(f):
+                            return await f(exp, *args, **kwargs)
+                        else:
+                            return f(exp, *args, **kwargs)
+
+                remote_controller_serialized = base64.urlsafe_b64encode(
+                    zlib.compress(
+                        cloudpickle.dumps(
+                            functools.partial(remote_controller, *args, **kwargs),
+                        )
+                    )
+                )
+
+                parameter_controller_executable = dataclasses.replace(
+                    executable,
+                    args=xm.merge_args(
+                        job_blocks.get_args_for_python_entrypoint(
+                            xm.ModuleName("xm_slurm.scripts._cloudpickle")
+                        ),
+                        xm.SequentialArgs.from_collection({
+                            "cloudpickled_fn": remote_controller_serialized.decode("ascii"),
+                        }),
+                        xm.SequentialArgs.from_collection(controller_args),
+                    ),
+                    env_vars=controller_env_vars or {},
+                )
+
+                await aux_unit.add(
+                    xm.Job(
+                        executor=executor,
+                        executable=parameter_controller_executable,
+                        name=controller_name,
+                    )
+                )
+
+                # Launch local parameter controller and monitor for when it starts running
+                # so we can kill the local controller.
+                if controller_mode is ParameterControllerMode.AUTO:
+                    aux_unit._create_task(
+                        _monitor_parameter_controller(aux_unit, local_controller(*args, **kwargs))
+                    )
+
+            return xm.AuxiliaryUnitJob(
+                job_generator,
+                importance=xm.Importance.HIGH,
+                termination_delay_secs=0,  # TODO: add support for termination delay.?
+            )
+
+        return make_controller
+
+    return decorator

xm_slurm/job_blocks.py

@@ -1,6 +1,13 @@
+from typing import Mapping, TypedDict
+
 from xmanager import xm
 
 
+class JobArgs(TypedDict, total=False):
+    args: xm.UserArgs
+    env_vars: Mapping[str, str]
+
+
 def get_args_for_python_entrypoint(
     entrypoint: xm.ModuleName | xm.CommandList,
 ) -> xm.SequentialArgs:

xm_slurm/packageables.py

@@ -1,7 +1,7 @@
 import importlib.resources as resources
 import pathlib
 import sys
-from typing import Mapping, Sequence
+from typing import Literal, Mapping, Sequence
 
 import immutabledict
 from xmanager import xm
@@ -29,8 +29,8 @@ def docker_image(
     return xm.Packageable(
         executor_spec=SlurmSpec(),
         executable_spec=DockerImage(image=image),
-        args=args,
-        env_vars=env_vars,
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=dict(env_vars),
     )
 
 
@@ -40,7 +40,7 @@ def docker_container(
     dockerfile: pathlib.Path | None = None,
     context: pathlib.Path | None = None,
     target: str | None = None,
-    ssh: list[str] | None = None,
+    ssh: Sequence[str] | Literal[True] | None = None,
     build_args: Mapping[str, str] = immutabledict.immutabledict(),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
@@ -55,7 +55,7 @@ def docker_container(
         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.
+        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.
@@ -71,8 +71,12 @@ def docker_container(
     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:
@@ -92,8 +96,8 @@ def docker_container(
             workdir=workdir,
             labels=labels,
         ),
-        args=args,
-        env_vars=env_vars,
+        args=xm.SequentialArgs.from_collection(args),
+        env_vars=dict(env_vars),
     )
 
 
@@ -104,14 +108,18 @@ def python_container(
     context: pathlib.Path | None = None,
     requirements: pathlib.Path | None = None,
     base_image: str = "docker.io/python:{major}.{minor}-slim",
+    extra_system_packages: Sequence[str] = (),
+    extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     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.
+    NOTE: uv is used to install packages from `requirements`.
 
     Args:
         executor_spec: The executor specification, where will the container be stored at.
@@ -119,8 +127,11 @@ def python_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.
 
@@ -149,11 +160,11 @@ def python_container(
             executor_spec=executor_spec,
             dockerfile=dockerfile,
             context=context,
+            ssh=ssh,
             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),
+                "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,
@@ -178,6 +189,7 @@ def mamba_container(
     base_image: str = "gcr.io/distroless/base-debian10",
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
     env_vars: Mapping[str, str] = immutabledict.immutabledict(),
 ) -> xm.Packageable:
@@ -193,6 +205,7 @@ def mamba_container(
         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.
 
@@ -221,6 +234,7 @@ def mamba_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,
@@ -237,26 +251,32 @@ def mamba_container(
 conda_container = mamba_container
 
 
-def pdm_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",
+    base_image: str = "docker.io/python:{major}.{minor}-slim-bookworm",
+    extra_system_packages: Sequence[str] = (),
+    extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
     labels: Mapping[str, str] = immutabledict.immutabledict(),
+    ssh: Sequence[str] | Literal[True] | None = None,
     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.
+    """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.
 
@@ -268,20 +288,22 @@ def pdm_container(
     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.")
+    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/pdm.Dockerfile")
+        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={
-                "PYTHON_MAJOR": str(sys.version_info.major),
-                "PYTHON_MINOR": str(sys.version_info.minor),
-                "PYTHON_MICRO": str(sys.version_info.micro),
+                "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,
@@ -291,7 +313,7 @@ def pdm_container(
             cache_from=cache_from,
             labels=labels,
             # We must specify the workdir manually for apptainer support
-            workdir=pathlib.Path("/workspace/src"),
+            workdir=pathlib.Path("/workspace"),
             args=args,
             env_vars=env_vars,
         )