xmanager-slurm 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

xm_slurm/experiment.py

@@ -18,12 +18,11 @@ from xmanager import xm
 from xmanager.xm import async_packager, core, id_predictor, job_operators
 from xmanager.xm import job_blocks as xm_job_blocks
 
-from xm_slurm import api, config, dependencies, execution, executors
+from xm_slurm import api, config, dependencies, execution, executors, metadata_context
 from xm_slurm.console import console
 from xm_slurm.job_blocks import JobArgs
 from xm_slurm.packaging import router
 from xm_slurm.status import SlurmWorkUnitStatus
-from xm_slurm.utils import UserSet
 
 logger = logging.getLogger(__name__)
 
@@ -62,66 +61,6 @@ def _validate_job(
             raise ValueError(f"Only `args` and `env_vars` are supported for args on job {job!r}.")
 
 
-@dataclasses.dataclass(kw_only=True, frozen=True)
-class Artifact:
-    name: str
-    uri: str
-
-    def __hash__(self) -> int:
-        return hash((type(self), self.name))
-
-
-class ContextArtifacts(UserSet[Artifact]):
-    def __init__(
-        self,
-        owner: "SlurmExperiment | SlurmExperimentUnit",
-        *,
-        artifacts: tp.Sequence[Artifact],
-    ):
-        super().__init__(
-            artifacts,
-            on_add=self._on_add_artifact,
-            on_remove=self._on_remove_artifact,
-            on_discard=self._on_remove_artifact,
-        )
-        self._owner = owner
-        self._create_task = self._owner._create_task
-
-    def _on_add_artifact(self, artifact: Artifact) -> None:
-        match self._owner:
-            case SlurmExperiment():
-                api.client().insert_experiment_artifact(
-                    self._owner.experiment_id,
-                    api.ArtifactModel(
-                        name=artifact.name,
-                        uri=artifact.uri,
-                    ),
-                )
-            case SlurmWorkUnit():
-                api.client().insert_work_unit_artifact(
-                    self._owner.experiment_id,
-                    self._owner.work_unit_id,
-                    api.ArtifactModel(
-                        name=artifact.name,
-                        uri=artifact.uri,
-                    ),
-                )
-
-    def _on_remove_artifact(self, artifact: Artifact) -> None:
-        match self._owner:
-            case SlurmExperiment():
-                api.client().delete_experiment_artifact(self._owner.experiment_id, artifact.name)
-            case SlurmWorkUnit():
-                api.client().delete_work_unit_artifact(
-                    self._owner.experiment_id, self._owner.work_unit_id, artifact.name
-                )
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class SlurmExperimentUnitMetadataContext:
-    artifacts: ContextArtifacts
-
-
 class SlurmExperimentUnit(xm.ExperimentUnit):
     """ExperimentUnit is a collection of semantically associated `Job`s."""
 
@@ -138,8 +77,9 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
         super().__init__(experiment, create_task, args, role, identity=identity)
         self._launched_jobs: list[xm.LaunchedJob] = []
         self._execution_handles: list[execution.SlurmHandle] = []
-        self._context = SlurmExperimentUnitMetadataContext(
-            artifacts=ContextArtifacts(owner=self, artifacts=[]),
+        self._context = metadata_context.SlurmExperimentUnitMetadataContext(
+            self,
+            artifacts=metadata_context.SlurmContextArtifacts(owner=self, artifacts=[]),
         )
 
     def add(  # type: ignore
@@ -289,7 +229,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
             api.client().insert_job(
                 self.experiment_id,
                 self.work_unit_id,
-                api.SlurmJobModel(
+                api.models.SlurmJob(
                     name=job.name,
                     slurm_job_id=handle.slurm_job.job_id,
                     slurm_ssh_config=handle.ssh.serialize(),
@@ -351,7 +291,12 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
         wait: bool = True,
         follow: bool = False,
     ) -> tp.AsyncGenerator[ConsoleRenderable, None]:
-        assert len(self._execution_handles) == 1, "Only one job handle is supported for logs."
+        if not self._execution_handles:
+            raise ValueError(f"No execution handles found for experiment unit {self!r}")
+        elif len(self._execution_handles) > 1:
+            raise ValueError(f"Multiple execution handles found for experiment unit {self!r}")
+        assert len(self._execution_handles) == 1
+
         handle = self._execution_handles[0]  # TODO(jfarebro): interleave?
         async for log in handle.logs(
             num_lines=num_lines, block_size=block_size, wait=wait, follow=follow
@@ -363,7 +308,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
         return self._launched_jobs
 
     @property
-    def context(self) -> SlurmExperimentUnitMetadataContext:  # type: ignore
+    def context(self) -> metadata_context.SlurmExperimentUnitMetadataContext:  # type: ignore
         return self._context
 
     def after_started(
@@ -427,7 +372,7 @@ class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
             api.client().update_work_unit(
                 self.experiment_id,
                 self.work_unit_id,
-                api.ExperimentUnitPatchModel(args=json.dumps(args_view), identity=None),
+                api.models.ExperimentUnitPatch(args=json.dumps(args_view), identity=None),
             )
 
         async with self._work_unit_id_predictor.submit_id(self.work_unit_id):  # type: ignore
@@ -488,112 +433,10 @@ class SlurmAuxiliaryUnit(SlurmExperimentUnit):
         return f"<SlurmAuxiliaryUnit {self.experiment_unit_name}>"
 
 
-class SlurmExperimentContextAnnotations:
-    def __init__(
-        self,
-        experiment: "SlurmExperiment",
-        *,
-        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 = UserSet(
-            tags or set(),
-            on_add=self._on_tag_added,
-            on_remove=self._on_tag_removed,
-            on_discard=self._on_tag_removed,
-        )
-        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.ExperimentPatchModel(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.ExperimentPatchModel(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.ExperimentPatchModel(note=value),
-        )
-
-    @property
-    def tags(self) -> tp.MutableSet[str]:
-        return self._tags
-
-    @tags.setter
-    def tags(self, tags: set[str]) -> None:
-        # TODO(jfarebro): Create custom tag collection
-        # and set it here, we need this so we can hook add and remove
-        # to mutate the database transparently
-        self._tags = UserSet(tags, on_add=self._on_tag_added, on_remove=self._on_tag_removed)
-        api.client().update_experiment(
-            self._experiment.experiment_id,
-            api.ExperimentPatchModel(tags=list(self._tags)),
-        )
-
-    def _on_tag_added(self, tag: str) -> None:
-        del tag
-        api.client().update_experiment(
-            self._experiment.experiment_id,
-            api.ExperimentPatchModel(tags=list(self._tags)),
-        )
-
-    def _on_tag_removed(self, tag: str) -> None:
-        del tag
-        api.client().update_experiment(
-            self._experiment.experiment_id,
-            api.ExperimentPatchModel(tags=list(self._tags)),
-        )
-
-
-class SlurmExperimentContextArtifacts(ContextArtifacts):
-    def add_graphviz_config(self, config: str) -> None:
-        self.add(Artifact(name="GRAPHVIZ", uri=f"graphviz://{config}"))
-
-    def add_python_config(self, config: str) -> None:
-        self.add(Artifact(name="PYTHON", uri=config))
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class SlurmExperimentMetadataContext:
-    annotations: SlurmExperimentContextAnnotations
-    artifacts: ContextArtifacts
-
-
 class SlurmExperiment(xm.Experiment):
     _id: int
     _experiment_units: list[SlurmExperimentUnit]
-    _experiment_context: SlurmExperimentMetadataContext
+    _experiment_context: metadata_context.SlurmExperimentMetadataContext
     _work_unit_count: int
     _async_packager = async_packager.AsyncPackager(router.package)
 
@@ -605,12 +448,13 @@ class SlurmExperiment(xm.Experiment):
         super().__init__()
         self._id = experiment_id
         self._experiment_units = []
-        self._experiment_context = SlurmExperimentMetadataContext(
-            annotations=SlurmExperimentContextAnnotations(
+        self._experiment_context = metadata_context.SlurmExperimentMetadataContext(
+            self,
+            annotations=metadata_context.SlurmExperimentContextAnnotations(
                 experiment=self,
                 title=experiment_title,
             ),
-            artifacts=ContextArtifacts(self, artifacts=[]),
+            artifacts=metadata_context.SlurmContextArtifacts(self, artifacts=[]),
         )
         self._work_unit_count = 0
 
@@ -1005,7 +849,7 @@ class SlurmExperiment(xm.Experiment):
 
             api.client().insert_work_unit(
                 self.experiment_id,
-                api.WorkUnitModel(
+                api.models.WorkUnitPatch(
                     wid=work_unit.work_unit_id,
                     identity=work_unit.identity,
                     args=json.dumps(args),
@@ -1080,7 +924,7 @@ class SlurmExperiment(xm.Experiment):
         return self.context.annotations.title
 
     @property
-    def context(self) -> SlurmExperimentMetadataContext:  # type: ignore
+    def context(self) -> metadata_context.SlurmExperimentMetadataContext:  # type: ignore
         return self._experiment_context
 
     @property
@@ -1099,7 +943,9 @@ class SlurmExperiment(xm.Experiment):
 
 def create_experiment(experiment_title: str) -> SlurmExperiment:
     """Create Experiment."""
-    experiment_id = api.client().insert_experiment(api.ExperimentPatchModel(title=experiment_title))
+    experiment_id = api.client().insert_experiment(
+        api.models.ExperimentPatch(title=experiment_title)
+    )
     return SlurmExperiment(experiment_title=experiment_title, experiment_id=experiment_id)
 
 
@@ -1114,11 +960,11 @@ def get_experiment(experiment_id: int) -> SlurmExperiment:
     # Populate annotations
     experiment.context.annotations.description = experiment_model.description or ""
     experiment.context.annotations.note = experiment_model.note or ""
-    experiment.context.annotations.tags = set(experiment_model.tags or [])
+    experiment.context.annotations.tags = experiment_model.tags or []
 
     # Populate artifacts
     for artifact in experiment_model.artifacts:
-        experiment.context.artifacts.add(Artifact(name=artifact.name, uri=artifact.uri))
+        experiment.context.artifacts[artifact.name] = artifact.uri
 
     # Populate work units
     for wu_model in experiment_model.work_units:
@@ -1145,7 +991,7 @@ def get_experiment(experiment_id: int) -> SlurmExperiment:
 
         # Populate artifacts for each work unit
         for artifact in wu_model.artifacts:
-            work_unit.context.artifacts.add(Artifact(name=artifact.name, uri=artifact.uri))
+            work_unit.context.artifacts[artifact.name] = artifact.uri
 
         experiment._experiment_units.append(work_unit)
         experiment._work_unit_count += 1

xm_slurm/filesystem.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/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")