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

xm_slurm/experiment.py

@@ -2,25 +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, config, 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)
@@ -28,7 +34,7 @@ _current_job_array_queue = contextvars.ContextVar[
 
 def _validate_job(
     job: xm.JobType,
-    args_view: JobArgs | Mapping[str, JobArgs],
+    args_view: JobArgs | tp.Mapping[str, JobArgs],
 ) -> None:
     if not args_view:
         return
@@ -51,7 +57,7 @@ def _validate_job(
             )
 
         if isinstance(job, xm.JobGroup) and key in job.jobs:
-            _validate_job(job.jobs[key], typing.cast(JobArgs, 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}.")
 
@@ -62,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]):
@@ -70,7 +76,7 @@ class ContextArtifacts(UserSet[Artifact]):
         self,
         owner: "SlurmExperiment | SlurmExperimentUnit",
         *,
-        artifacts: Sequence[Artifact],
+        artifacts: tp.Sequence[Artifact],
     ):
         super().__init__(
             artifacts,
@@ -124,8 +130,8 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
     def __init__(
         self,
         experiment: xm.Experiment,
-        create_task: Callable[[Awaitable[Any]], futures.Future[Any]],
-        args: JobArgs | 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:
@@ -136,25 +142,137 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
             artifacts=ContextArtifacts(owner=self, artifacts=[]),
         )
 
-    @typing.overload
+    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: ...
 
-    @typing.overload
+    @tp.overload
     async def _submit_jobs_for_execution(
         self,
         job: xm.JobGroup,
-        args_view: Mapping[str, JobArgs],
+        dependency: dependencies.SlurmJobDependency | None,
+        args_view: tp.Mapping[str, JobArgs],
         identity: str | None = ...,
     ) -> execution.SlurmHandle: ...
 
-    async def _submit_jobs_for_execution(self, job, args_view, identity=None):
+    @tp.overload
+    async def _submit_jobs_for_execution(
+        self,
+        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,
@@ -173,7 +291,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
                 self.work_unit_id,
                 api.SlurmJobModel(
                     name=job.name,
-                    slurm_job_id=handle.job_id,  # type: ignore
+                    slurm_job_id=handle.slurm_job.job_id,
                     slurm_ssh_config=handle.ssh.serialize(),
                 ),
             )
@@ -186,7 +304,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
                     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():
@@ -194,7 +312,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
                 self._launched_jobs.append(
                     xm.LaunchedJob(
                         name=handle.job.name,  # type: ignore
-                        address=str(handle.job_id),
+                        address=str(handle.slurm_job.job_id),
                     )
                 )
 
@@ -221,10 +339,25 @@ 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
@@ -233,13 +366,27 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
     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: JobArgs,
+        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 = "",
@@ -258,7 +405,9 @@ class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
     async def _launch_job_group(  # type: ignore
         self,
         job: xm.JobGroup,
-        args_view: Mapping[str, JobArgs],
+        args_view: tp.Mapping[str, JobArgs],
+        *,
+        dependency: dependencies.SlurmJobDependency | None,
         identity: str,
     ) -> None:
         global _current_job_array_queue
@@ -291,7 +440,9 @@ class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
             # 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, args_view, identity=identity)
+                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
@@ -317,12 +468,16 @@ class SlurmAuxiliaryUnit(SlurmExperimentUnit):
     async def _launch_job_group(  # type: ignore
         self,
         job: xm.JobGroup,
-        args_view: Mapping[str, JobArgs],
+        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)
+        slurm_handle = await self._submit_jobs_for_execution(
+            job, dependency, args_view, identity=identity
+        )
         self._ingest_launched_jobs(job, slurm_handle)
 
     @property
@@ -392,7 +547,7 @@ class SlurmExperimentContextAnnotations:
         )
 
     @property
-    def tags(self) -> MutableSet[str]:
+    def tags(self) -> tp.MutableSet[str]:
         return self._tags
 
     @tags.setter
@@ -459,87 +614,82 @@ 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 = ...,
         *,
+        dependency: dependencies.SlurmJobDependency | None = ...,
         identity: str = ...,
     ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
         job: xm.JobGroup,
-        args: Mapping[str, Mapping[str, Any]] | None,
+        args: tp.Mapping[str, JobArgs] | None = ...,
         *,
-        role: xm.WorkUnitRole = xm.WorkUnitRole(),
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency | None = ...,
         identity: str = ...,
     ) -> asyncio.Future[SlurmWorkUnit]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
-        job: xm.JobGroup,
-        args: Mapping[str, Mapping[str, Any]] | None,
+        job: xm.Job | xm.JobGeneratorType,
+        args: tp.Sequence[JobArgs],
         *,
-        role: xm.ExperimentUnitRole,
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency
+        | tp.Sequence[dependencies.SlurmJobDependency]
+        | None = ...,
         identity: str = ...,
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
+    ) -> asyncio.Future[tp.Sequence[SlurmWorkUnit]]: ...
 
-    @typing.overload
+    @tp.overload
     def add(
         self,
         job: xm.Job | xm.JobGeneratorType | xm.JobConfig,
-        args: Mapping[str, Any] | None,
+        args: JobArgs | None = ...,
         *,
-        role: xm.WorkUnitRole = xm.WorkUnitRole(),
+        role: xm.WorkUnitRole | None = ...,
+        dependency: dependencies.SlurmJobDependency | None = ...,
         identity: str = ...,
     ) -> asyncio.Future[SlurmWorkUnit]: ...
 
-    @typing.overload
-    def add(
-        self,
-        job: xm.Job | xm.JobGeneratorType | xm.JobConfig,
-        args: Mapping[str, Any] | None,
-        *,
-        role: xm.ExperimentUnitRole,
-        identity: str = ...,
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
-
-    @typing.overload
-    def add(
-        self,
-        job: xm.Job | xm.JobGeneratorType,
-        args: Sequence[Mapping[str, Any]],
-        *,
-        role: xm.WorkUnitRole = xm.WorkUnitRole(),
-        identity: str = ...,
-    ) -> asyncio.Future[Sequence[SlurmWorkUnit]]: ...
-
-    @typing.overload
+    @tp.overload
     def add(
         self,
         job: xm.JobType,
         *,
-        role: xm.AuxiliaryUnitRole = ...,
+        role: xm.AuxiliaryUnitRole,
+        dependency: dependencies.SlurmJobDependency | None = ...,
         identity: str = ...,
     ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
 
     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[SlurmAuxiliaryUnit]
-        | asyncio.Future[SlurmExperimentUnit]
         | 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.")
@@ -554,21 +704,76 @@ class SlurmExperiment(xm.Experiment):
             # Validate job & args
             for trial in args:
                 _validate_job(job, trial)
-            args = typing.cast(Sequence[JobArgs], args)
+            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[JobArgs],
+        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
@@ -579,7 +784,11 @@ class SlurmExperiment(xm.Experiment):
         # and collect the futures
         wu_futures = []
         for idx, trial in enumerate(args):
-            wu_futures.append(super().add(job, args=trial, role=role, identity=f"{identity}_{idx}"))
+            wu_futures.append(
+                self.add(
+                    job, args=trial, role=role, identity=f"{identity}_{idx}" if identity else ""
+                )
+            )
 
         # We'll wait until XManager has filled the queue.
         # There are two cases here, either we were given an xm.Job
@@ -589,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():
@@ -650,6 +860,78 @@ class SlurmExperiment(xm.Experiment):
             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.
@@ -666,10 +948,18 @@ class SlurmExperiment(xm.Experiment):
                         args=xm.SequentialArgs.from_collection(common_args),
                         env_vars=dict(common_env_vars),
                     ),
-                    args=sweep_args,
+                    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)
@@ -697,11 +987,11 @@ class SlurmExperiment(xm.Experiment):
 
     def _create_experiment_unit(  # type: ignore
         self,
-        args: JobArgs,
+        args: JobArgs | tp.Mapping[str, JobArgs] | None,
         role: xm.ExperimentUnitRole,
         identity: str,
-    ) -> Awaitable[SlurmWorkUnit | SlurmAuxiliaryUnit]:
-        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,
@@ -726,7 +1016,7 @@ class SlurmExperiment(xm.Experiment):
             future.set_result(work_unit)
             return future
 
-        def _create_auxiliary_unit(role: xm.AuxiliaryUnitRole) -> Awaitable[SlurmAuxiliaryUnit]:
+        def _create_auxiliary_unit(role: xm.AuxiliaryUnitRole) -> tp.Awaitable[SlurmAuxiliaryUnit]:
             auxiliary_unit = SlurmAuxiliaryUnit(
                 self,
                 self._create_task,
@@ -756,8 +1046,8 @@ class SlurmExperiment(xm.Experiment):
         experiment_id: int,
         identity: str,
         role: xm.ExperimentUnitRole,
-        args: JobArgs | None = None,
-    ) -> Awaitable[xm.ExperimentUnit]:
+        args: JobArgs | tp.Mapping[str, JobArgs] | None = None,
+    ) -> tp.Awaitable[SlurmExperimentUnit]:
         del experiment_id, identity, role, args
         raise NotImplementedError
 
@@ -797,7 +1087,7 @@ class SlurmExperiment(xm.Experiment):
     def work_unit_count(self) -> int:
         return self._work_unit_count
 
-    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)
@@ -822,8 +1112,8 @@ def get_experiment(experiment_id: int) -> SlurmExperiment:
     experiment._work_unit_id_predictor = id_predictor.Predictor(1)
 
     # Populate annotations
-    experiment.context.annotations.description = experiment_model.description
-    experiment.context.annotations.note = experiment_model.note
+    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
@@ -846,8 +1136,9 @@ def get_experiment(experiment_id: int) -> SlurmExperiment:
         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,
-                job_id=str(job_model.slurm_job_id),
+                slurm_job=str(job_model.slurm_job_id),
                 job_name=job_model.name,
             )
             work_unit._execution_handles.append(handle)