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

xm_slurm/experiment.py

@@ -2,23 +2,31 @@ import asyncio
 import collections.abc
 import contextvars
 import dataclasses
+import datetime as dt
 import functools
 import inspect
 import json
+import logging
 import os
-import typing
+import traceback
+import typing as tp
 from concurrent import futures
-from typing import Any, Awaitable, Callable, Mapping, MutableSet, Sequence
 
+import more_itertools as mit
+from rich.console import ConsoleRenderable
 from xmanager import xm
-from xmanager.xm import async_packager, id_predictor
+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, execution, executors
+from xm_slurm import api, config, dependencies, execution, executors
 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__)
+
 _current_job_array_queue = contextvars.ContextVar[
     asyncio.Queue[tuple[xm.JobGroup, asyncio.Future]] | None
 ]("_current_job_array_queue", default=None)
@@ -26,7 +34,7 @@ _current_job_array_queue = contextvars.ContextVar[
 
 def _validate_job(
     job: xm.JobType,
-    args_view: Mapping[str, Any],
+    args_view: JobArgs | tp.Mapping[str, JobArgs],
 ) -> None:
     if not args_view:
         return
@@ -49,7 +57,7 @@ def _validate_job(
             )
 
         if isinstance(job, xm.JobGroup) and key in job.jobs:
-            _validate_job(job.jobs[key], expanded)
+            _validate_job(job.jobs[key], tp.cast(JobArgs, expanded))
         elif key not in allowed_keys:
             raise ValueError(f"Only `args` and `env_vars` are supported for args on job {job!r}.")
 
@@ -60,7 +68,7 @@ class Artifact:
     uri: str
 
     def __hash__(self) -> int:
-        return hash(self.name)
+        return hash((type(self), self.name))
 
 
 class ContextArtifacts(UserSet[Artifact]):
@@ -68,7 +76,7 @@ class ContextArtifacts(UserSet[Artifact]):
         self,
         owner: "SlurmExperiment | SlurmExperimentUnit",
         *,
-        artifacts: Sequence[Artifact],
+        artifacts: tp.Sequence[Artifact],
     ):
         super().__init__(
             artifacts,
@@ -117,50 +125,194 @@ class SlurmExperimentUnitMetadataContext:
 class SlurmExperimentUnit(xm.ExperimentUnit):
     """ExperimentUnit is a collection of semantically associated `Job`s."""
 
-    experiment: "SlurmExperiment"
+    experiment: "SlurmExperiment"  # type: ignore
 
     def __init__(
         self,
         experiment: xm.Experiment,
-        create_task: Callable[[Awaitable[Any]], futures.Future[Any]],
-        args: Mapping[str, Any] | None,
+        create_task: tp.Callable[[tp.Awaitable[tp.Any]], futures.Future[tp.Any]],
+        args: JobArgs | tp.Mapping[str, JobArgs] | None,
         role: xm.ExperimentUnitRole,
+        identity: str = "",
     ) -> None:
-        super().__init__(experiment, create_task, args, role)
+        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=[]),
         )
 
+    def add(  # type: ignore
+        self,
+        job: xm.JobType,
+        args: JobArgs | tp.Mapping[str, JobArgs] | None = None,
+        *,
+        dependency: dependencies.SlurmJobDependency | None = None,
+        identity: str = "",
+    ) -> tp.Awaitable[None]:
+        # Prioritize the identity given directly to the work unit at work unit
+        # creation time, as opposed to the identity passed when adding jobs to it as
+        # this is more consistent between job generator work units and regular work
+        # units.
+        identity = self.identity or identity
+
+        job = job_operators.shallow_copy_job_type(job)  # type: ignore
+        if args is not None:
+            core._apply_args(job, args)
+        job_operators.populate_job_names(job)  # type: ignore
+
+        def launch_job(job: xm.Job) -> tp.Awaitable[None]:
+            core._current_experiment.set(self.experiment)
+            core._current_experiment_unit.set(self)
+            return self._launch_job_group(
+                xm.JobGroup(**{job.name: job}),  # type: ignore
+                core._work_unit_arguments(job, self._args),
+                dependency=dependency,
+                identity=identity,
+            )
+
+        def launch_job_group(group: xm.JobGroup) -> tp.Awaitable[None]:
+            core._current_experiment.set(self.experiment)
+            core._current_experiment_unit.set(self)
+            return self._launch_job_group(
+                group,
+                core._work_unit_arguments(group, self._args),
+                dependency=dependency,
+                identity=identity,
+            )
+
+        def launch_job_generator(
+            job_generator: xm.JobGeneratorType,
+        ) -> tp.Awaitable[None]:
+            if not inspect.iscoroutinefunction(job_generator) and not inspect.iscoroutinefunction(
+                getattr(job_generator, "__call__")
+            ):
+                raise ValueError(
+                    "Job generator must be an async function. Signature needs to be "
+                    "`async def job_generator(work_unit: xm.WorkUnit) -> None:`"
+                )
+            core._current_experiment.set(self.experiment)
+            core._current_experiment_unit.set(self)
+            coroutine = job_generator(self, **(args or {}))
+            assert coroutine is not None
+            return coroutine
+
+        def launch_job_config(job_config: xm.JobConfig) -> tp.Awaitable[None]:
+            core._current_experiment.set(self.experiment)
+            core._current_experiment_unit.set(self)
+            return self._launch_job_config(
+                job_config, dependency, tp.cast(JobArgs, args) or {}, identity
+            )
+
+        job_awaitable: tp.Awaitable[tp.Any]
+        match job:
+            case xm.Job() as job:
+                job_awaitable = launch_job(job)
+            case xm.JobGroup() as job_group:
+                job_awaitable = launch_job_group(job_group)
+            case job_generator if xm_job_blocks.is_job_generator(job):
+                job_awaitable = launch_job_generator(job_generator)  # type: ignore
+            case xm.JobConfig() as job_config:
+                job_awaitable = launch_job_config(job_config)
+            case _:
+                raise TypeError(f"Unsupported job type: {job!r}")
+
+        launch_task = self._create_task(job_awaitable)
+        self._launch_tasks.append(launch_task)
+        return asyncio.wrap_future(launch_task)
+
+    async def _launch_job_group(  # type: ignore
+        self,
+        job_group: xm.JobGroup,
+        args_view: tp.Mapping[str, JobArgs],
+        *,
+        dependency: dependencies.SlurmJobDependency | None,
+        identity: str,
+    ) -> None:
+        del job_group, dependency, args_view, identity
+        raise NotImplementedError
+
+    async def _launch_job_config(  # type: ignore
+        self,
+        job_config: xm.JobConfig,
+        dependency: dependencies.SlurmJobDependency | None,
+        args_view: JobArgs,
+        identity: str,
+    ) -> None:
+        del job_config, dependency, args_view, identity
+        raise NotImplementedError
+
+    @tp.overload
+    async def _submit_jobs_for_execution(
+        self,
+        job: xm.Job,
+        dependency: dependencies.SlurmJobDependency | None,
+        args_view: JobArgs,
+        identity: str | None = ...,
+    ) -> execution.SlurmHandle: ...
+
+    @tp.overload
+    async def _submit_jobs_for_execution(
+        self,
+        job: xm.JobGroup,
+        dependency: dependencies.SlurmJobDependency | None,
+        args_view: tp.Mapping[str, JobArgs],
+        identity: str | None = ...,
+    ) -> execution.SlurmHandle: ...
+
+    @tp.overload
     async def _submit_jobs_for_execution(
         self,
-        job: xm.Job | xm.JobGroup,
-        args_view: Mapping[str, Any],
-        identity: str | None = None,
-    ) -> execution.SlurmHandle:
+        job: xm.Job,
+        dependency: dependencies.SlurmJobDependency | None,
+        args_view: tp.Sequence[JobArgs],
+        identity: str | None = ...,
+    ) -> list[execution.SlurmHandle]: ...
+
+    async def _submit_jobs_for_execution(self, job, dependency, args_view, identity=None):
         return await execution.launch(
             job=job,
+            dependency=dependency,
             args=args_view,
             experiment_id=self.experiment_id,
             identity=identity,
         )
 
     def _ingest_launched_jobs(self, job: xm.JobType, handle: execution.SlurmHandle) -> None:
+        self._execution_handles.append(handle)
+
+        def _ingest_job(job: xm.Job) -> None:
+            if not isinstance(self._role, xm.WorkUnitRole):
+                return
+            assert isinstance(self, SlurmWorkUnit)
+            assert job.name is not None
+            api.client().insert_job(
+                self.experiment_id,
+                self.work_unit_id,
+                api.SlurmJobModel(
+                    name=job.name,
+                    slurm_job_id=handle.slurm_job.job_id,
+                    slurm_ssh_config=handle.ssh.serialize(),
+                ),
+            )
+
         match job:
             case xm.JobGroup() as job_group:
                 for job in job_group.jobs.values():
+                    assert isinstance(job, xm.Job)
+                    _ingest_job(job)
                     self._launched_jobs.append(
                         xm.LaunchedJob(
                             name=job.name,  # type: ignore
-                            address=str(handle.job_id),
+                            address=str(handle.slurm_job.job_id),
                         )
                     )
             case xm.Job():
+                _ingest_job(job)
                 self._launched_jobs.append(
                     xm.LaunchedJob(
                         name=handle.job.name,  # type: ignore
-                        address=str(handle.job_id),
+                        address=str(handle.slurm_job.job_id),
                     )
                 )
 
@@ -187,85 +339,116 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
 
         self.experiment._create_task(_stop_awaitable())
 
-    async def get_status(self) -> SlurmWorkUnitStatus:
+    async def get_status(self) -> SlurmWorkUnitStatus:  # type: ignore
         states = await asyncio.gather(*[handle.get_state() for handle in self._execution_handles])
         return SlurmWorkUnitStatus.aggregate(states)
 
+    async def logs(
+        self,
+        *,
+        num_lines: int = 10,
+        block_size: int = 1024,
+        wait: bool = True,
+        follow: bool = False,
+    ) -> tp.AsyncGenerator[ConsoleRenderable, None]:
+        assert len(self._execution_handles) == 1, "Only one job handle is supported for logs."
+        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
+        ):
+            yield log
+
+    @property
     def launched_jobs(self) -> list[xm.LaunchedJob]:
         return self._launched_jobs
 
     @property
-    def context(self) -> SlurmExperimentUnitMetadataContext:
+    def context(self) -> SlurmExperimentUnitMetadataContext:  # type: ignore
         return self._context
 
+    def after_started(
+        self, *, time: dt.timedelta | None = None
+    ) -> dependencies.SlurmJobDependencyAfter:
+        return dependencies.SlurmJobDependencyAfter(self._execution_handles, time=time)
+
+    def after_finished(self) -> dependencies.SlurmJobDependencyAfterAny:
+        return dependencies.SlurmJobDependencyAfterAny(self._execution_handles)
+
+    def after_completed(self) -> dependencies.SlurmJobDependencyAfterOK:
+        return dependencies.SlurmJobDependencyAfterOK(self._execution_handles)
+
+    def after_failed(self) -> dependencies.SlurmJobDependencyAfterNotOK:
+        return dependencies.SlurmJobDependencyAfterNotOK(self._execution_handles)
+
 
 class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
     def __init__(
         self,
         experiment: "SlurmExperiment",
-        create_task: Callable[[Awaitable[Any]], futures.Future],
-        args: Mapping[str, Any],
+        create_task: tp.Callable[[tp.Awaitable[tp.Any]], futures.Future],
+        args: JobArgs | tp.Mapping[str, JobArgs] | None,
         role: xm.ExperimentUnitRole,
         work_unit_id_predictor: id_predictor.Predictor,
+        identity: str = "",
     ) -> None:
-        super().__init__(experiment, create_task, args, role)
+        super().__init__(experiment, create_task, args, role, identity=identity)
         self._work_unit_id_predictor = work_unit_id_predictor
         self._work_unit_id = self._work_unit_id_predictor.reserve_id()
-        api.client().insert_work_unit(
-            self.experiment_id,
-            api.WorkUnitPatchModel(
-                wid=self.work_unit_id,
-                identity=self.identity,
-                args=json.dumps(args),
-            ),
-        )
 
-    def _ingest_handle(self, handle: execution.SlurmHandle) -> None:
-        self._execution_handles.append(handle)
-        api.client().insert_job(
-            self.experiment_id,
-            self.work_unit_id,
-            api.SlurmJobModel(
-                name=self.experiment_unit_name,
-                slurm_job_id=handle.job_id,  # type: ignore
-                slurm_cluster=json.dumps({
-                    "host": handle.ssh_connection_options.host,
-                    "username": handle.ssh_connection_options.username,
-                    "port": handle.ssh_connection_options.port,
-                    "config": handle.ssh_connection_options.config.get_options(False),
-                }),
-            ),
-        )
+    def _get_existing_handle(self, job: xm.JobGroup) -> execution.SlurmHandle | None:
+        job_name = mit.one(job.jobs.keys())
+        for handle in self._execution_handles:
+            if handle.job_name == job_name:
+                return handle
+        return None
 
-    async def _launch_job_group(
+    async def _launch_job_group(  # type: ignore
         self,
         job: xm.JobGroup,
-        args_view: Mapping[str, Any],
+        args_view: tp.Mapping[str, JobArgs],
+        *,
+        dependency: dependencies.SlurmJobDependency | None,
         identity: str,
     ) -> None:
         global _current_job_array_queue
         _validate_job(job, args_view)
+        future = asyncio.Future[execution.SlurmHandle]()
+
+        # If we already have a handle for this job, we don't need to submit it again.
+        # We'll just resolve the future with the existing handle.
+        # Otherwise we'll add callbacks to ingest the handle and the launched jobs.
+        if existing_handle := self._get_existing_handle(job):
+            future.set_result(existing_handle)
+        else:
+            future.add_done_callback(
+                lambda handle: self._ingest_launched_jobs(job, handle.result())
+            )
+
+            api.client().update_work_unit(
+                self.experiment_id,
+                self.work_unit_id,
+                api.ExperimentUnitPatchModel(args=json.dumps(args_view), identity=None),
+            )
 
-        future = asyncio.Future()
         async with self._work_unit_id_predictor.submit_id(self.work_unit_id):  # type: ignore
             # If we're scheduling as part of a job queue (i.e., the queue is set on the context)
-            # then we'll insert the job and current future that'll get resolved to the
-            # proper handle.
+            # then we'll insert the job and future that'll get resolved to the proper handle
+            # when the Slurm job array is scheduled.
             if job_array_queue := _current_job_array_queue.get():
                 job_array_queue.put_nowait((job, future))
-            # Otherwise we'll resolve the future with the scheduled job immediately
-            else:
-                # Set the result inside of the context manager so we don't get out-of-order
-                # id scheduling...
-                future.set_result(
-                    await self._submit_jobs_for_execution(job, args_view, identity=identity)
+            # Otherwise, we're scheduling a single job and we'll submit it for execution.
+            # If the future is already done, i.e., the handle is already resolved, we don't need
+            # to submit the job again.
+            elif not future.done():
+                handle = await self._submit_jobs_for_execution(
+                    job, dependency, args_view, identity=identity
                 )
+                future.set_result(handle)
 
         # Wait for the job handle, this is either coming from scheduling the job array
-        # or from the single job above.
-        handle = await future
-        self._ingest_handle(handle)
-        self._ingest_launched_jobs(job, handle)
+        # or from a single job submission. If an existing handle was found, this will be
+        # a no-op.
+        await future
 
     @property
     def experiment_unit_name(self) -> str:
@@ -282,20 +465,19 @@ class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
 class SlurmAuxiliaryUnit(SlurmExperimentUnit):
     """An auxiliary unit operated by the Slurm backend."""
 
-    def _ingest_handle(self, handle: execution.SlurmHandle) -> None:
-        del handle
-        console.print("[red]Auxiliary units do not currently support ingestion.[/red]")
-
-    async def _launch_job_group(
+    async def _launch_job_group(  # type: ignore
         self,
-        job: xm.Job | xm.JobGroup,
-        args_view: Mapping[str, Any],
+        job: xm.JobGroup,
+        args_view: tp.Mapping[str, JobArgs],
+        *,
+        dependency: dependencies.SlurmJobDependency | None,
         identity: str,
     ) -> None:
         _validate_job(job, args_view)
 
-        slurm_handle = await self._submit_jobs_for_execution(job, args_view, identity=identity)
-        self._ingest_handle(slurm_handle)
+        slurm_handle = await self._submit_jobs_for_execution(
+            job, dependency, args_view, identity=identity
+        )
         self._ingest_launched_jobs(job, slurm_handle)
 
     @property
@@ -365,7 +547,7 @@ class SlurmExperimentContextAnnotations:
         )
 
     @property
-    def tags(self) -> MutableSet[str]:
+    def tags(self) -> tp.MutableSet[str]:
         return self._tags
 
     @tags.setter
@@ -432,75 +614,85 @@ class SlurmExperiment(xm.Experiment):
         )
         self._work_unit_count = 0
 
-    @typing.overload
-    def add(
+    @tp.overload
+    def add(  # type: ignore
         self,
         job: xm.AuxiliaryUnitJob,
-        args: Mapping[str, Any] | None = ...,
+        args: JobArgs | tp.Mapping[str, JobArgs] | None = ...,
         *,
-        identity: str = "",
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
+        dependency: dependencies.SlurmJobDependency | None = ...,
+        identity: str = ...,
+    ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None = ...,
+        job: xm.JobGroup,
+        args: tp.Mapping[str, JobArgs] | None = ...,
         *,
-        role: xm.WorkUnitRole = ...,
-        identity: str = "",
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency | None = ...,
+        identity: str = ...,
     ) -> asyncio.Future[SlurmWorkUnit]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None,
+        job: xm.Job | xm.JobGeneratorType,
+        args: tp.Sequence[JobArgs],
         *,
-        role: xm.ExperimentUnitRole,
-        identity: str = "",
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
-
-    @typing.overload
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency
+        | tp.Sequence[dependencies.SlurmJobDependency]
+        | None = ...,
+        identity: str = ...,
+    ) -> asyncio.Future[tp.Sequence[SlurmWorkUnit]]: ...
+
+    @tp.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None = ...,
+        job: xm.Job | xm.JobGeneratorType | xm.JobConfig,
+        args: JobArgs | None = ...,
         *,
-        role: xm.ExperimentUnitRole,
-        identity: str = "",
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency | None = ...,
+        identity: str = ...,
+    ) -> asyncio.Future[SlurmWorkUnit]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
-        job: xm.Job | xm.JobGeneratorType,
-        args: Sequence[Mapping[str, Any]],
+        job: xm.JobType,
         *,
-        role: xm.WorkUnitRole = ...,
-        identity: str = "",
-    ) -> asyncio.Future[Sequence[SlurmWorkUnit]]: ...
+        role: xm.AuxiliaryUnitRole,
+        dependency: dependencies.SlurmJobDependency | None = ...,
+        identity: str = ...,
+    ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
 
-    def add(
+    def add(  # type: ignore
         self,
         job: xm.JobType,
-        args: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
+        args: JobArgs
+        | tp.Mapping[str, JobArgs]
+        | tp.Sequence[tp.Mapping[str, tp.Any]]
+        | None = None,
         *,
-        role: xm.ExperimentUnitRole = xm.WorkUnitRole(),
+        role: xm.ExperimentUnitRole | None = None,
+        dependency: dependencies.SlurmJobDependency
+        | tp.Sequence[dependencies.SlurmJobDependency]
+        | None = None,
         identity: str = "",
     ) -> (
-        asyncio.Future[SlurmExperimentUnit]
+        asyncio.Future[SlurmAuxiliaryUnit]
         | asyncio.Future[SlurmWorkUnit]
-        | asyncio.Future[Sequence[SlurmWorkUnit]]
+        | asyncio.Future[tp.Sequence[SlurmWorkUnit]]
     ):
+        if role is None:
+            role = xm.WorkUnitRole()
+
         if isinstance(args, collections.abc.Sequence):
             if not isinstance(role, xm.WorkUnitRole):
                 raise ValueError("Only `xm.WorkUnit`s are supported for job arrays.")
-            if identity:
-                raise ValueError(
-                    "Cannot set an identity on the root add call. "
-                    "Please use a job generator and set the identity within."
-                )
             if isinstance(job, xm.JobGroup):
                 raise ValueError(
                     "Job arrays over `xm.JobGroup`s aren't supported. "
@@ -509,19 +701,79 @@ class SlurmExperiment(xm.Experiment):
                 )
             assert isinstance(job, xm.Job) or inspect.iscoroutinefunction(job), "Invalid job type"
 
+            # Validate job & args
+            for trial in args:
+                _validate_job(job, trial)
+            args = tp.cast(tp.Sequence[JobArgs], args)
+
             return asyncio.wrap_future(
-                self._create_task(self._launch_job_array(job, args, role, identity))
+                self._create_task(self._launch_job_array(job, dependency, args, role, identity)),
+                loop=self._event_loop,
+            )
+        if not (isinstance(dependency, dependencies.SlurmJobDependency) or dependency is None):
+            raise ValueError("Invalid dependency type, expected a SlurmJobDependency or None")
+
+        if isinstance(job, xm.AuxiliaryUnitJob):
+            role = job.role
+        self._added_roles[type(role)] += 1
+
+        if self._should_reload_experiment_unit(role):
+            experiment_unit_future = self._get_experiment_unit(
+                self.experiment_id, identity, role, args
             )
         else:
-            return super().add(job, args, role=role, identity=identity)  # type: ignore
+            experiment_unit_future = self._create_experiment_unit(args, role, identity)
+
+        async def launch():
+            experiment_unit = await experiment_unit_future
+            try:
+                await experiment_unit.add(job, args, dependency=dependency, identity=identity)
+            except Exception as experiment_exception:
+                logger.error(
+                    "Stopping experiment unit (identity %r) after it failed with: %s",
+                    identity,
+                    experiment_exception,
+                )
+                try:
+                    if isinstance(job, xm.AuxiliaryUnitJob):
+                        experiment_unit.stop()
+                    else:
+                        experiment_unit.stop(
+                            mark_as_failed=True,
+                            message=f"Work unit creation failed. {traceback.format_exc()}",
+                        )
+                except Exception as stop_exception:  # pylint: disable=broad-except
+                    logger.error("Couldn't stop experiment unit: %s", stop_exception)
+                    raise
+            return experiment_unit
+
+        async def reload():
+            experiment_unit = await experiment_unit_future
+            try:
+                await experiment_unit.add(job, args, dependency=dependency, identity=identity)
+            except Exception as update_exception:
+                logging.error(
+                    "Could not reload the experiment unit: %s",
+                    update_exception,
+                )
+                raise
+            return experiment_unit
+
+        return asyncio.wrap_future(
+            self._create_task(reload() if self._should_reload_experiment_unit(role) else launch()),
+            loop=self._event_loop,
+        )
 
     async def _launch_job_array(
         self,
         job: xm.Job | xm.JobGeneratorType,
-        args: Sequence[Mapping[str, Any]],
+        dependency: dependencies.SlurmJobDependency
+        | tp.Sequence[dependencies.SlurmJobDependency]
+        | None,
+        args: tp.Sequence[JobArgs],
         role: xm.WorkUnitRole,
         identity: str = "",
-    ) -> Sequence[SlurmWorkUnit]:
+    ) -> tp.Sequence[SlurmWorkUnit]:
         global _current_job_array_queue
 
         # Create our job array queue and assign it to the current context
@@ -531,10 +783,13 @@ class SlurmExperiment(xm.Experiment):
         # For each trial we'll schedule the job
         # and collect the futures
         wu_futures = []
-        for trial in args:
-            wu_futures.append(super().add(job, args=trial, role=role, identity=identity))
+        for idx, trial in enumerate(args):
+            wu_futures.append(
+                self.add(
+                    job, args=trial, role=role, identity=f"{identity}_{idx}" if identity else ""
+                )
+            )
 
-        # TODO(jfarebro): Set a timeout here
         # We'll wait until XManager has filled the queue.
         # There are two cases here, either we were given an xm.Job
         # in which case this will be trivial and filled immediately.
@@ -543,7 +798,8 @@ class SlurmExperiment(xm.Experiment):
         while not job_array_queue.full():
             await asyncio.sleep(0.1)
 
-        # All jobs have been resolved
+        # All jobs have been resolved so now we'll perform sanity checks
+        # to make sure we can infer the sweep
         executable, executor, name = None, None, None
         resolved_args, resolved_env_vars, resolved_futures = [], [], []
         while not job_array_queue.empty():
@@ -551,9 +807,13 @@ class SlurmExperiment(xm.Experiment):
             # that there's only a single job in this job group
             job_group_view, future = job_array_queue.get_nowait()
             assert isinstance(job_group_view, xm.JobGroup), "Expected a job group from xm"
-            _, job_view = job_group_view.jobs.popitem()
+            job_view = mit.one(
+                job_group_view.jobs.values(),
+                too_short=ValueError("Expected a single `xm.Job` in job group."),
+                too_long=ValueError("Only one `xm.Job` is supported for job arrays."),
+            )
 
-            if len(job_group_view.jobs) != 0 or not isinstance(job_view, xm.Job):
+            if not isinstance(job_view, xm.Job):
                 raise ValueError("Only `xm.Job` is supported for job arrays. ")
 
             if executable is None:
@@ -571,92 +831,223 @@ class SlurmExperiment(xm.Experiment):
             if job_view.name != name:
                 raise RuntimeError("Found multiple names in job array")
 
-            resolved_args.append(
-                set(xm.SequentialArgs.from_collection(job_view.args).to_dict().items())
-            )
+            resolved_args.append(xm.SequentialArgs.from_collection(job_view.args).to_list())
             resolved_env_vars.append(set(job_view.env_vars.items()))
             resolved_futures.append(future)
         assert executable is not None, "No executable found?"
         assert executor is not None, "No executor found?"
         assert isinstance(executor, executors.Slurm), "Only Slurm executors are supported."
-        assert executor.requirements.cluster is not None, "Cluster must be set on executor."
-
-        common_args: set = functools.reduce(lambda a, b: a & b, resolved_args, set())
+        assert (
+            executor.requirements.cluster is not None
+        ), "Cluster must be specified on requirements."
+
+        # XManager merges job arguments with keyword arguments with job arguments
+        # coming first. These are the arguments that may be common across all jobs
+        # so we can find the largest common prefix and remove them from each job.
+        common_args: list[str] = list(mit.longest_common_prefix(resolved_args))
         common_env_vars: set = functools.reduce(lambda a, b: a & b, resolved_env_vars, set())
 
         sweep_args = [
-            {
-                "args": dict(a.difference(common_args)),
-                "env_vars": dict(e.difference(common_env_vars)),
-            }
+            JobArgs(
+                args=functools.reduce(
+                    # Remove the common arguments from each job
+                    lambda args, to_remove: args.remove_args(to_remove),
+                    common_args,
+                    xm.SequentialArgs.from_collection(a),
+                ),
+                env_vars=dict(e.difference(common_env_vars)),
+            )
             for a, e in zip(resolved_args, resolved_env_vars)
         ]
 
+        # Dependency resolution
+        resolved_dependency = None
+        resolved_dependency_task_id_order = None
+        # one-to-one
+        if isinstance(dependency, collections.abc.Sequence):
+            if len(dependency) != len(wu_futures):
+                raise ValueError("Dependency list must be the same length as the number of trials.")
+            assert len(dependency) > 0, "Dependency list must not be empty."
+
+            # Convert any SlurmJobDependencyAfterOK to SlurmJobArrayDependencyAfterOK
+            # for any array jobs.
+            def _maybe_convert_afterok(
+                dep: dependencies.SlurmJobDependency,
+            ) -> dependencies.SlurmJobDependency:
+                if isinstance(dep, dependencies.SlurmJobDependencyAfterOK) and all([
+                    handle.slurm_job.is_array_job for handle in dep.handles
+                ]):
+                    return dependencies.SlurmJobArrayDependencyAfterOK([
+                        dataclasses.replace(
+                            handle,
+                            slurm_job=handle.slurm_job.array_job_id,
+                        )
+                        for handle in dep.handles
+                    ])
+                return dep
+
+            dependencies_converted = [dep.traverse(_maybe_convert_afterok) for dep in dependency]
+            dependency_sets = [set(dep.flatten()) for dep in dependencies_converted]
+            dependency_differences = functools.reduce(set.difference, dependency_sets, set())
+            # There should be NO differences between the dependencies of each trial after conversion.
+            if len(dependency_differences) > 0:
+                raise ValueError(
+                    f"Found variable dependencies across trials: {dependency_differences}. "
+                    "Slurm job arrays require the same dependencies across all trials. "
+                )
+            resolved_dependency = dependencies_converted[0]
+
+            # This is slightly annoying but we need to re-sort the sweep arguments in case the dependencies were passed
+            # in a different order than 1, 2, ..., N as the Job array can only have correspondance with the same task id.
+            original_array_dependencies = [
+                mit.one(
+                    filter(
+                        lambda dep: isinstance(dep, dependencies.SlurmJobDependencyAfterOK)
+                        and all([handle.slurm_job.is_array_job for handle in dep.handles]),
+                        deps.flatten(),
+                    )
+                )
+                for deps in dependency
+            ]
+            resolved_dependency_task_id_order = [
+                int(
+                    mit.one(
+                        functools.reduce(
+                            set.difference,
+                            [handle.slurm_job.array_task_id for handle in dep.handles],  # type: ignore
+                        )
+                    )
+                )
+                for dep in original_array_dependencies
+            ]
+            assert len(resolved_dependency_task_id_order) == len(sweep_args)
+            assert set(resolved_dependency_task_id_order) == set(range(len(sweep_args))), (
+                "Dependent job array tasks should have task ids 0, 1, ..., N. "
+                f"Found: {resolved_dependency_task_id_order}"
+            )
+        # one-to-many
+        elif isinstance(dependency, dependencies.SlurmJobDependency):
+            resolved_dependency = dependency
+        assert resolved_dependency is None or isinstance(
+            resolved_dependency, dependencies.SlurmJobDependency
+        ), "Invalid dependency type"
+
         # No support for sweep_env_vars right now.
         # We schedule the job array and then we'll resolve all the work units with
         # the handles Slurm gives back to us.
-        try:
-            handles = await execution.get_client().launch(
-                cluster=executor.requirements.cluster,
-                job=xm.Job(
-                    executable=executable,
-                    executor=executor,
-                    name=name,
-                    args=dict(common_args),
-                    env_vars=dict(common_env_vars),
-                ),
-                args=sweep_args,
-                experiment_id=self.experiment_id,
-                identity=identity,
+        # If we already have handles for all the work units, we don't need to submit the
+        # job array to SLURM.
+        num_resolved_handles = sum(future.done() for future in resolved_futures)
+        if num_resolved_handles == 0:
+            try:
+                handles = await execution.launch(
+                    job=xm.Job(
+                        executable=executable,
+                        executor=executor,
+                        name=name,
+                        args=xm.SequentialArgs.from_collection(common_args),
+                        env_vars=dict(common_env_vars),
+                    ),
+                    dependency=resolved_dependency,
+                    args=[
+                        sweep_args[resolved_dependency_task_id_order.index(i)]
+                        for i in range(len(sweep_args))
+                    ]
+                    if resolved_dependency_task_id_order
+                    else sweep_args,
+                    experiment_id=self.experiment_id,
+                    identity=identity,
+                )
+                if resolved_dependency_task_id_order:
+                    handles = [handles[i] for i in resolved_dependency_task_id_order]
+            except Exception as e:
+                for future in resolved_futures:
+                    future.set_exception(e)
+                raise
+            else:
+                for handle, future in zip(handles, resolved_futures):
+                    future.set_result(handle)
+        elif 0 < num_resolved_handles < len(resolved_futures):
+            raise RuntimeError(
+                "Some array job elements have handles, but some don't. This shouldn't happen."
             )
-        except Exception as e:
-            for future in resolved_futures:
-                future.set_exception(e)
-            raise
-        else:
-            for handle, future in zip(handles, resolved_futures):
-                future.set_result(handle)
 
         wus = await asyncio.gather(*wu_futures)
+
         _current_job_array_queue.set(None)
         return wus
 
-    def _create_experiment_unit(
+    def _get_work_unit_by_identity(self, identity: str) -> SlurmWorkUnit | None:
+        if identity == "":
+            return None
+        for unit in self._experiment_units:
+            if isinstance(unit, SlurmWorkUnit) and unit.identity == identity:
+                return unit
+        return None
+
+    def _create_experiment_unit(  # type: ignore
         self,
-        args: Mapping[str, Any],
+        args: JobArgs | tp.Mapping[str, JobArgs] | None,
         role: xm.ExperimentUnitRole,
         identity: str,
-    ) -> Awaitable[SlurmWorkUnit]:
-        del identity
-
-        def _create_work_unit(role: xm.WorkUnitRole) -> Awaitable[SlurmWorkUnit]:
+    ) -> tp.Awaitable[SlurmWorkUnit | SlurmAuxiliaryUnit]:
+        def _create_work_unit(role: xm.WorkUnitRole) -> tp.Awaitable[SlurmWorkUnit]:
             work_unit = SlurmWorkUnit(
                 self,
                 self._create_task,
                 args,
                 role,
                 self._work_unit_id_predictor,
+                identity=identity,
             )
             self._experiment_units.append(work_unit)
             self._work_unit_count += 1
 
-            future = asyncio.Future()
+            api.client().insert_work_unit(
+                self.experiment_id,
+                api.WorkUnitModel(
+                    wid=work_unit.work_unit_id,
+                    identity=work_unit.identity,
+                    args=json.dumps(args),
+                ),
+            )
+
+            future = asyncio.Future[SlurmWorkUnit]()
             future.set_result(work_unit)
             return future
 
+        def _create_auxiliary_unit(role: xm.AuxiliaryUnitRole) -> tp.Awaitable[SlurmAuxiliaryUnit]:
+            auxiliary_unit = SlurmAuxiliaryUnit(
+                self,
+                self._create_task,
+                args,
+                role,
+                identity=identity,
+            )
+            self._experiment_units.append(auxiliary_unit)
+            future = asyncio.Future[SlurmAuxiliaryUnit]()
+            future.set_result(auxiliary_unit)
+            return future
+
         match role:
             case xm.WorkUnitRole():
+                if (existing_unit := self._get_work_unit_by_identity(identity)) is not None:
+                    future = asyncio.Future[SlurmWorkUnit]()
+                    future.set_result(existing_unit)
+                    return future
                 return _create_work_unit(role)
+            case xm.AuxiliaryUnitRole():
+                return _create_auxiliary_unit(role)
             case _:
                 raise ValueError(f"Unsupported role {role}")
 
-    def _get_experiment_unit(
+    def _get_experiment_unit(  # type: ignore
         self,
         experiment_id: int,
         identity: str,
         role: xm.ExperimentUnitRole,
-        args: Mapping[str, Any] | None = None,
-    ) -> Awaitable[xm.ExperimentUnit]:
+        args: JobArgs | tp.Mapping[str, JobArgs] | None = None,
+    ) -> tp.Awaitable[SlurmExperimentUnit]:
         del experiment_id, identity, role, args
         raise NotImplementedError
 
@@ -672,7 +1063,7 @@ class SlurmExperiment(xm.Experiment):
         # If no work units were added, delete this experiment
         # This is to prevent empty experiments from being persisted
         # and cluttering the database.
-        if self.work_unit_count == 0:
+        if len(self._experiment_units) == 0:
             console.print(
                 f"[red]No work units were added to experiment `{self.experiment_title}`... deleting.[/red]"
             )
@@ -689,15 +1080,14 @@ class SlurmExperiment(xm.Experiment):
         return self.context.annotations.title
 
     @property
-    def context(self) -> SlurmExperimentMetadataContext:
+    def context(self) -> SlurmExperimentMetadataContext:  # type: ignore
         return self._experiment_context
 
     @property
     def work_unit_count(self) -> int:
         return self._work_unit_count
 
-    @property
-    def work_units(self) -> Mapping[int, SlurmWorkUnit]:
+    def work_units(self) -> dict[int, SlurmWorkUnit]:
         """Gets work units created via self.add()."""
         return {
             wu.work_unit_id: wu for wu in self._experiment_units if isinstance(wu, SlurmWorkUnit)
@@ -716,8 +1106,51 @@ def create_experiment(experiment_title: str) -> SlurmExperiment:
 def get_experiment(experiment_id: int) -> SlurmExperiment:
     """Get Experiment."""
     experiment_model = api.client().get_experiment(experiment_id)
-    # TODO(jfarebro): Fill in jobs and work units and annotations
-    return SlurmExperiment(experiment_title=experiment_model.title, experiment_id=experiment_id)
+    experiment = SlurmExperiment(
+        experiment_title=experiment_model.title, experiment_id=experiment_id
+    )
+    experiment._work_unit_id_predictor = id_predictor.Predictor(1)
+
+    # 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 [])
+
+    # Populate artifacts
+    for artifact in experiment_model.artifacts:
+        experiment.context.artifacts.add(Artifact(name=artifact.name, uri=artifact.uri))
+
+    # Populate work units
+    for wu_model in experiment_model.work_units:
+        work_unit = SlurmWorkUnit(
+            experiment=experiment,
+            create_task=experiment._create_task,
+            args=json.loads(wu_model.args) if wu_model.args else {},
+            role=xm.WorkUnitRole(),
+            identity=wu_model.identity or "",
+            work_unit_id_predictor=experiment._work_unit_id_predictor,
+        )
+        work_unit._work_unit_id = wu_model.wid
+
+        # Populate jobs for each work unit
+        for job_model in wu_model.jobs:
+            slurm_ssh_config = config.SlurmSSHConfig.deserialize(job_model.slurm_ssh_config)
+            handle = execution.SlurmHandle(
+                experiment_id=experiment_id,
+                ssh=slurm_ssh_config,
+                slurm_job=str(job_model.slurm_job_id),
+                job_name=job_model.name,
+            )
+            work_unit._execution_handles.append(handle)
+
+        # Populate artifacts for each work unit
+        for artifact in wu_model.artifacts:
+            work_unit.context.artifacts.add(Artifact(name=artifact.name, uri=artifact.uri))
+
+        experiment._experiment_units.append(work_unit)
+        experiment._work_unit_count += 1
+
+    return experiment
 
 
 @functools.cache
@@ -733,5 +1166,5 @@ def get_current_work_unit() -> SlurmWorkUnit | None:
         wid := os.environ.get("XM_SLURM_WORK_UNIT_ID")
     ):
         experiment = get_experiment(int(xid))
-        return experiment.work_units[int(wid)]
+        return experiment.work_units()[int(wid)]
     return None