xmanager-slurm 0.4.19__py3-none-any.whl

xm_slurm/experimental/parameter_controller.py

@@ -0,0 +1,206 @@
+import asyncio
+import base64
+import dataclasses
+import enum
+import functools
+import logging
+import typing as tp
+import zlib
+
+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 = tp.ParamSpec("P")
+T = tp.TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+async def _monitor_parameter_controller(
+    aux_unit: SlurmAuxiliaryUnit,
+    local_parameter_controller_coro: tp.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: tp.Mapping[str, str] | None = None,
+) -> tp.Callable[
+    [
+        tp.Callable[tp.Concatenate[SlurmExperiment, P], T]
+        | tp.Callable[tp.Concatenate[SlurmExperiment, P], tp.Awaitable[T]],
+    ],
+    tp.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: tp.Callable[tp.Concatenate[SlurmExperiment, P], T]
+        | tp.Callable[tp.Concatenate[SlurmExperiment, P], tp.Awaitable[T]],
+    ) -> tp.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 | tp.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 | tp.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/filesystems.py

@@ -0,0 +1,129 @@
+import abc
+import asyncio
+import os
+import typing as tp
+
+import aiofile
+import asyncssh
+import typing_extensions as tpe
+import wrapt
+
+
+class AsyncFileIO(tp.Protocol):
+    async def write(self, buffer: str | bytes, /) -> int: ...
+    async def read(self, size: int = -1, /) -> bytes: ...
+    async def seek(self, offset: int, /) -> int: ...
+    async def tell(self) -> int: ...
+    async def __aenter__(self) -> tpe.Self: ...
+    async def __aexit__(self, *args: tp.Any, **kwargs: tp.Any) -> None: ...
+
+
+class AsyncFileSystem(tp.Protocol):
+    async def open(
+        self,
+        path: os.PathLike[str] | str,
+        mode: tp.Literal["r", "w", "rb", "wb"],
+        *,
+        encoding: str = "utf-8",
+    ) -> AsyncFileIO: ...
+
+    async def read(
+        self, path: os.PathLike[str] | str, *, size: int = -1, offset: int = 0
+    ) -> bytes: ...
+    async def write(
+        self, path: os.PathLike[str] | str, data: str | bytes, *, offset: int = 0
+    ) -> None: ...
+
+    async def exists(self, path: os.PathLike[str] | str) -> bool: ...
+    async def size(self, path: os.PathLike[str] | str) -> int | None: ...
+    async def makedirs(
+        self, path: os.PathLike[str] | str, mode: int = 511, exist_ok: bool = False
+    ) -> None: ...
+
+
+class AbstractAsyncFileSystem(AsyncFileSystem, abc.ABC):
+    @abc.abstractmethod
+    async def open(
+        self,
+        path: os.PathLike[str] | str,
+        mode: tp.Literal["r", "w", "rb", "wb"],
+        *,
+        encoding: str = "utf-8",
+    ) -> AsyncFileIO: ...
+
+    async def read(self, path: os.PathLike[str] | str, *, size: int = -1, offset: int = 0) -> bytes:
+        async with await self.open(path, "rb") as f:
+            await f.seek(offset)
+            return await f.read(size)
+
+    async def write(
+        self, path: os.PathLike[str] | str, data: str | bytes, *, offset: int = 0
+    ) -> None:
+        async with await self.open(path, "wb") as f:
+            await f.seek(offset)
+            await f.write(data)
+
+
+class AsyncLocalFileIO(wrapt.ObjectProxy):
+    async def seek(self, offset: int, /) -> int:
+        await asyncio.to_thread(self.__wrapped__.seek, offset)
+        return await asyncio.to_thread(self.__wrapped__.tell)
+
+    async def tell(self) -> int:
+        return await asyncio.to_thread(self.__wrapped__.tell)
+
+    async def __aenter__(self) -> tpe.Self:
+        return AsyncLocalFileIO(await self.__wrapped__.__aenter__())
+
+    async def __aexit__(self, *args: tp.Any, **kwargs: tp.Any) -> None:
+        return await self.__wrapped__.__aexit__(*args, **kwargs)
+
+
+class AsyncLocalFileSystem(AbstractAsyncFileSystem):
+    def __init__(self): ...
+
+    async def open(
+        self,
+        path: os.PathLike[str] | str,
+        mode: tp.Literal["r", "w", "rb", "wb"],
+        *,
+        encoding: str = "utf-8",
+    ) -> AsyncFileIO:
+        return AsyncLocalFileIO(aiofile.async_open(os.fspath(path), mode=mode, encoding=encoding))  # type: ignore
+
+    async def exists(self, path: os.PathLike[str] | str) -> bool:
+        return await asyncio.to_thread(os.path.exists, os.fspath(path))
+
+    async def size(self, path: os.PathLike[str] | str) -> int | None:
+        return await asyncio.to_thread(os.path.getsize, os.fspath(path))
+
+    async def makedirs(
+        self, path: os.PathLike[str] | str, mode: int = 0o777, exist_ok: bool = False
+    ) -> None:
+        return await asyncio.to_thread(os.makedirs, os.fspath(path), mode=mode, exist_ok=exist_ok)
+
+
+class AsyncSSHFileSystem(AbstractAsyncFileSystem):
+    def __init__(self, client: asyncssh.SFTPClient):
+        self._client = client
+
+    async def open(
+        self,
+        path: os.PathLike[str] | str,
+        mode: tp.Literal["r", "w", "rb", "wb"],
+        *,
+        encoding: str = "utf-8",
+    ) -> AsyncFileIO:
+        return await self._client.open(os.fspath(path), mode, encoding=encoding)  # type: ignore
+
+    async def exists(self, path: os.PathLike[str] | str) -> bool:
+        return await self._client.exists(os.fspath(path))
+
+    async def size(self, path: os.PathLike[str] | str) -> int | None:
+        return (await self._client.stat(os.fspath(path))).size
+
+    async def makedirs(
+        self, path: os.PathLike[str] | str, mode: int = 0o777, exist_ok: bool = False
+    ) -> None:
+        attrs = asyncssh.SFTPAttrs(permissions=mode)
+        return await self._client.makedirs(os.fspath(path), attrs=attrs, exist_ok=exist_ok)

xm_slurm/job_blocks.py

@@ -0,0 +1,21 @@
+import typing as tp
+
+from xmanager import xm
+
+
+class JobArgs(tp.TypedDict, total=False):
+    args: xm.UserArgs
+    env_vars: tp.Mapping[str, str]
+
+
+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/metadata_context.py

@@ -0,0 +1,253 @@
+import collections.abc
+import typing as tp
+
+from xmanager import xm
+
+from xm_slurm import api
+
+
+class SlurmContextArtifacts(collections.abc.MutableMapping[str, str]):
+    def __init__(
+        self,
+        owner: xm.Experiment | xm.ExperimentUnit,
+        *,
+        artifacts: tp.Sequence[api.models.Artifact],
+    ):
+        self._data = {artifact.name: artifact.uri for artifact in artifacts}
+        self._owner = owner
+        self._create_task = self._owner._create_task
+
+    def add(self, name: str, uri: str) -> None:
+        artifact = api.models.Artifact(name=name, uri=uri)
+        match self._owner:
+            case xm.Experiment():
+                api.client().insert_experiment_artifact(self._owner.experiment_id, artifact)
+            case xm.WorkUnit():
+                api.client().insert_work_unit_artifact(
+                    self._owner.experiment_id, self._owner.work_unit_id, artifact
+                )
+
+        self._data[name] = uri
+
+    def remove(self, name: str) -> None:
+        match self._owner:
+            case xm.Experiment():
+                api.client().delete_experiment_artifact(self._owner.experiment_id, name)
+            case xm.WorkUnit():
+                api.client().delete_work_unit_artifact(
+                    self._owner.experiment_id, self._owner.work_unit_id, name
+                )
+
+    def __setitem__(self, name: str, uri: str) -> None:
+        self.add(name, uri)
+
+    def __delitem__(self, name: str) -> None:
+        self.remove(name)
+
+    def __getitem__(self, name: str) -> str:
+        return self._data[name]
+
+    def __iter__(self) -> tp.Iterator[str]:
+        return iter(self._data)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+
+class SlurmExperimentAnnotationTags(collections.abc.MutableSet[str]):
+    def __init__(self, experiment: xm.Experiment, *, tags: tp.Iterable[str]):
+        self._experiment = experiment
+        # Use a dict to ensure order is preserved
+        self._tags = dict.fromkeys(tags)
+
+    def add(self, tag: str) -> None:
+        self._tags[tag] = None
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(tags=list(self._tags)),
+        )
+
+    def remove(self, tag: str) -> None:
+        self.discard(tag)
+
+    def discard(self, tag: str) -> None:
+        self._tags.pop(tag)
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(tags=list(self._tags)),
+        )
+
+    def __contains__(self, tag: str) -> bool:
+        return tag in self._tags
+
+    def __iter__(self) -> tp.Iterator[str]:
+        return iter(self._tags)
+
+    def __len__(self) -> int:
+        return len(self._tags)
+
+
+class SlurmExperimentUnitMetadataContext:
+    def __init__(
+        self,
+        experiment_unit: xm.ExperimentUnit,
+        *,
+        artifacts: SlurmContextArtifacts,
+    ):
+        self._experiment_unit = experiment_unit
+        self._artifacts = artifacts
+
+    @property
+    def artifacts(self) -> SlurmContextArtifacts:
+        return self._artifacts
+
+    @artifacts.setter
+    def artifacts(self, artifacts: SlurmContextArtifacts) -> None:
+        del artifacts
+        raise ValueError("The artifacts object is immutable.")
+
+
+class SlurmExperimentContextAnnotations:
+    def __init__(
+        self,
+        experiment: xm.Experiment,
+        *,
+        title: str,
+        tags: set[str] | None = None,
+        description: str | None = None,
+        note: str | None = None,
+    ):
+        self._experiment = experiment
+        self._create_task = self._experiment._create_task
+        self._title = title
+        self._tags = SlurmExperimentAnnotationTags(experiment, tags=tags or [])
+        self._description = description or ""
+        self._note = note or ""
+
+    @property
+    def title(self) -> str:
+        return self._title
+
+    @title.setter
+    def title(self, value: str) -> None:
+        self._title = value
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(title=value),
+        )
+
+    @property
+    def description(self) -> str:
+        return self._description
+
+    @description.setter
+    def description(self, value: str) -> None:
+        self._description = value
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(description=value),
+        )
+
+    @property
+    def note(self) -> str:
+        return self._note
+
+    @note.setter
+    def note(self, value: str) -> None:
+        self._note = value
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(note=value),
+        )
+
+    @property
+    def tags(self) -> SlurmExperimentAnnotationTags:
+        return self._tags
+
+    @tags.setter
+    def tags(self, tags: tp.Iterable[str]) -> None:
+        self._tags = SlurmExperimentAnnotationTags(self._experiment, tags=tags)
+        api.client().update_experiment(
+            self._experiment.experiment_id,
+            api.models.ExperimentPatch(tags=list(self._tags)),
+        )
+
+
+class SlurmExperimentMetadataContext:
+    def __init__(
+        self,
+        experiment: xm.Experiment,
+        *,
+        annotations: SlurmExperimentContextAnnotations,
+        artifacts: SlurmContextArtifacts,
+    ):
+        self._experiment = experiment
+        self._annotations = annotations
+        self._artifacts = artifacts
+
+        self._graphviz_config = None
+        self._python_config = None
+
+    @property
+    def annotations(self) -> SlurmExperimentContextAnnotations:
+        return self._annotations
+
+    @annotations.setter
+    def annotations(self, annotations: SlurmExperimentContextAnnotations) -> None:
+        del annotations
+        raise ValueError("The annotations object is immutable.")
+
+    @property
+    def artifacts(self) -> SlurmContextArtifacts:
+        return self._artifacts
+
+    @artifacts.setter
+    def artifacts(self, artifacts: SlurmContextArtifacts) -> None:
+        del artifacts
+        raise ValueError("The artifacts object is immutable.")
+
+    @property
+    def graphviz_config(self) -> str | None:
+        return self._graphviz_config
+
+    @graphviz_config.setter
+    def graphviz_config(self, config: str | None) -> None:
+        self._graphviz_config = config
+        match config:
+            case None:
+                api.client().delete_experiment_config_artifact(
+                    self._experiment.experiment_id, "GRAPHVIZ"
+                )
+            case str():
+                api.client().insert_experiment_config_artifact(
+                    self._experiment.experiment_id,
+                    api.models.ConfigArtifact(name="GRAPHVIZ", uri=f"graphviz://{config}"),
+                )
+
+    @graphviz_config.deleter
+    def graphviz_config(self) -> None:
+        self._graphviz_config = None
+        api.client().delete_experiment_config_artifact(self._experiment.experiment_id, "GRAPHVIZ")
+
+    @property
+    def python_config(self) -> str | None:
+        return self._python_config
+
+    @python_config.setter
+    def python_config(self, config: str | None) -> None:
+        self._python_config = config
+        match config:
+            case None:
+                api.client().delete_experiment_config_artifact(
+                    self._experiment.experiment_id, "PYTHON"
+                )
+            case str():
+                api.client().insert_experiment_config_artifact(
+                    self._experiment.experiment_id,
+                    api.models.ConfigArtifact(name="PYTHON", uri=config),
+                )
+
+    @python_config.deleter
+    def python_config(self) -> None:
+        self._python_config = None
+        api.client().delete_experiment_config_artifact(self._experiment.experiment_id, "PYTHON")