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

xm_slurm/experiment.py

@@ -10,11 +10,13 @@ import typing
 from concurrent import futures
 from typing import Any, Awaitable, Callable, Mapping, MutableSet, Sequence
 
+import more_itertools as mit
 from xmanager import xm
 from xmanager.xm import async_packager, id_predictor
 
-from xm_slurm import api, execution, executors
+from xm_slurm import api, config, 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
@@ -26,7 +28,7 @@ _current_job_array_queue = contextvars.ContextVar[
 
 def _validate_job(
     job: xm.JobType,
-    args_view: Mapping[str, Any],
+    args_view: JobArgs | Mapping[str, JobArgs],
 ) -> None:
     if not args_view:
         return
@@ -49,7 +51,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], typing.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}.")
 
@@ -117,28 +119,40 @@ 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,
+        args: 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=[]),
         )
 
+    @typing.overload
+    async def _submit_jobs_for_execution(
+        self,
+        job: xm.Job,
+        args_view: JobArgs,
+        identity: str | None = ...,
+    ) -> execution.SlurmHandle: ...
+
+    @typing.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.JobGroup,
+        args_view: Mapping[str, JobArgs],
+        identity: str | None = ...,
+    ) -> execution.SlurmHandle: ...
+
+    async def _submit_jobs_for_execution(self, job, args_view, identity=None):
         return await execution.launch(
             job=job,
             args=args_view,
@@ -147,9 +161,28 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
         )
 
     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.job_id,  # type: ignore
+                    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
@@ -157,6 +190,7 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
                         )
                     )
             case xm.Job():
+                _ingest_job(job)
                 self._launched_jobs.append(
                     xm.LaunchedJob(
                         name=handle.job.name,  # type: ignore
@@ -191,11 +225,12 @@ class SlurmExperimentUnit(xm.ExperimentUnit):
         states = await asyncio.gather(*[handle.get_state() for handle in self._execution_handles])
         return SlurmWorkUnitStatus.aggregate(states)
 
+    @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
 
 
@@ -204,68 +239,65 @@ class SlurmWorkUnit(xm.WorkUnit, SlurmExperimentUnit):
         self,
         experiment: "SlurmExperiment",
         create_task: Callable[[Awaitable[Any]], futures.Future],
-        args: Mapping[str, Any],
+        args: JobArgs,
         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: Mapping[str, JobArgs],
         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, 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 +314,15 @@ 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: Mapping[str, JobArgs],
         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)
         self._ingest_launched_jobs(job, slurm_handle)
 
     @property
@@ -438,37 +465,47 @@ class SlurmExperiment(xm.Experiment):
         job: xm.AuxiliaryUnitJob,
         args: Mapping[str, Any] | None = ...,
         *,
-        identity: str = "",
-    ) -> asyncio.Future[SlurmExperimentUnit]: ...
+        identity: str = ...,
+    ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
 
     @typing.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None = ...,
+        job: xm.JobGroup,
+        args: Mapping[str, Mapping[str, Any]] | None,
         *,
-        role: xm.WorkUnitRole = ...,
-        identity: str = "",
+        role: xm.WorkUnitRole = xm.WorkUnitRole(),
+        identity: str = ...,
     ) -> asyncio.Future[SlurmWorkUnit]: ...
 
     @typing.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None,
+        job: xm.JobGroup,
+        args: Mapping[str, Mapping[str, Any]] | None,
         *,
         role: xm.ExperimentUnitRole,
-        identity: str = "",
+        identity: str = ...,
     ) -> asyncio.Future[SlurmExperimentUnit]: ...
 
     @typing.overload
     def add(
         self,
-        job: xm.JobType,
-        args: Mapping[str, Any] | None = ...,
+        job: xm.Job | xm.JobGeneratorType | xm.JobConfig,
+        args: Mapping[str, Any] | None,
+        *,
+        role: xm.WorkUnitRole = xm.WorkUnitRole(),
+        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 = "",
+        identity: str = ...,
     ) -> asyncio.Future[SlurmExperimentUnit]: ...
 
     @typing.overload
@@ -477,11 +514,20 @@ class SlurmExperiment(xm.Experiment):
         job: xm.Job | xm.JobGeneratorType,
         args: Sequence[Mapping[str, Any]],
         *,
-        role: xm.WorkUnitRole = ...,
-        identity: str = "",
+        role: xm.WorkUnitRole = xm.WorkUnitRole(),
+        identity: str = ...,
     ) -> asyncio.Future[Sequence[SlurmWorkUnit]]: ...
 
+    @typing.overload
     def add(
+        self,
+        job: xm.JobType,
+        *,
+        role: xm.AuxiliaryUnitRole = ...,
+        identity: str = ...,
+    ) -> asyncio.Future[SlurmAuxiliaryUnit]: ...
+
+    def add(  # type: ignore
         self,
         job: xm.JobType,
         args: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None,
@@ -489,18 +535,14 @@ class SlurmExperiment(xm.Experiment):
         role: xm.ExperimentUnitRole = xm.WorkUnitRole(),
         identity: str = "",
     ) -> (
-        asyncio.Future[SlurmExperimentUnit]
+        asyncio.Future[SlurmAuxiliaryUnit]
+        | asyncio.Future[SlurmExperimentUnit]
         | asyncio.Future[SlurmWorkUnit]
         | asyncio.Future[Sequence[SlurmWorkUnit]]
     ):
         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,6 +551,11 @@ 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 = typing.cast(Sequence[JobArgs], args)
+
             return asyncio.wrap_future(
                 self._create_task(self._launch_job_array(job, args, role, identity))
             )
@@ -518,7 +565,7 @@ class SlurmExperiment(xm.Experiment):
     async def _launch_job_array(
         self,
         job: xm.Job | xm.JobGeneratorType,
-        args: Sequence[Mapping[str, Any]],
+        args: Sequence[JobArgs],
         role: xm.WorkUnitRole,
         identity: str = "",
     ) -> Sequence[SlurmWorkUnit]:
@@ -531,10 +578,9 @@ 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(super().add(job, args=trial, role=role, identity=f"{identity}_{idx}"))
 
-        # 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.
@@ -551,9 +597,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,64 +621,86 @@ 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)
         ]
 
         # 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),
+                    ),
+                    args=sweep_args,
+                    experiment_id=self.experiment_id,
+                    identity=identity,
+                )
+            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,
         role: xm.ExperimentUnitRole,
         identity: str,
-    ) -> Awaitable[SlurmWorkUnit]:
-        del identity
-
+    ) -> Awaitable[SlurmWorkUnit | SlurmAuxiliaryUnit]:
         def _create_work_unit(role: xm.WorkUnitRole) -> Awaitable[SlurmWorkUnit]:
             work_unit = SlurmWorkUnit(
                 self,
@@ -636,26 +708,55 @@ class SlurmExperiment(xm.Experiment):
                 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) -> 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,
+        args: JobArgs | None = None,
     ) -> Awaitable[xm.ExperimentUnit]:
         del experiment_id, identity, role, args
         raise NotImplementedError
@@ -672,7 +773,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,14 +790,13 @@ 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]:
         """Gets work units created via self.add()."""
         return {
@@ -716,8 +816,50 @@ 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
+    experiment.context.annotations.note = experiment_model.note
+    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(
+                ssh=slurm_ssh_config,
+                job_id=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 +875,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