earthkit-workflows 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

cascade/benchmarks/anemoi.py

@@ -10,7 +10,7 @@ from earthkit.workflows import Cascade
 
 
 def get_graph(lead_time, ensemble_members, CKPT=None, date="2024-12-02T00:00"):
-    import anemoicascade as ac
+    import anemoicascade as ac  # ty: ignore[unresolved-import]
 
     CKPT = (
         CKPT

cascade/benchmarks/dask.py

@@ -5,9 +5,9 @@ from cascade.low.core import JobInstance
 from cascade.low.dask import graph2job
 
 
-def get_job(job: str) -> JobInstance:
+def get_job(job_name: str) -> JobInstance:
 
-    if job == "add":
+    if job_name == "add":
 
         def add(x, y):
             result = x + y
@@ -21,7 +21,7 @@ def get_job(job: str) -> JobInstance:
             dataset for task in job.tasks for dataset in job.outputs_of(task)
         ]
         return job
-    elif job == "groupby":
+    elif job_name == "groupby":
         df = dd.DataFrame.from_dict({"x": [0, 0, 1, 1], "y": [1, 2, 3, 4]})
         df = df.groupby("x").sum()
         job = graph2job(df.__dask_graph__())
@@ -30,4 +30,4 @@ def get_job(job: str) -> JobInstance:
         ]
         return job
     else:
-        raise NotImplementedError(job)
+        raise NotImplementedError(job_name)

cascade/benchmarks/dist.py

@@ -26,7 +26,7 @@ def dist_func_torch(a: int) -> int:
     import datetime as dt
 
     import numpy as np
-    import torch.distributed as dist
+    import torch.distributed as dist  # ty: ignore[unresolved-import]
 
     world_size = int(os.environ["CASCADE_GANG_WORLD_SIZE"])
     rank = int(os.environ["CASCADE_GANG_RANK"])
@@ -61,8 +61,8 @@ def dist_func_jax(a: int) -> int:
     os.environ["JAX_NUM_CPU_DEVICES"] = "1"
     os.environ["JAX_PLATFORM_NAME"] = "cpu"
     os.environ["JAX_PLATFORMS"] = "cpu"
-    import jax
-    import jax.numpy as jp
+    import jax  # ty: ignore[unresolved-import]
+    import jax.numpy as jp  # ty: ignore[unresolved-import]
 
     jax.config.update("jax_platforms", "cpu")
     jax.config.update("jax_platform_name", "cpu")

cascade/benchmarks/job1.py

@@ -16,10 +16,9 @@ Controlled by env var params: JOB1_{DATA_ROOT, GRID, ...}, see below
 import os
 
 import earthkit.data
-
 from earthkit.workflows.fluent import Payload
-from earthkit.workflows.plugins.pproc.fluent import from_source
-from earthkit.workflows.plugins.pproc.utils.window import Range
+from earthkit.workflows.plugins.pproc.fluent import from_source  # ty: ignore
+from earthkit.workflows.plugins.pproc.utils.window import Range  # ty: ignore
 
 # *** PARAMS ***
 
@@ -137,7 +136,7 @@ def download_inputs():
             }
             data = earthkit.data.from_source("mars", **ekp)
             with open(f"{data_root}/data_{number}_{step}.grib", "wb") as f:
-                data.write(f)
+                data.write(f) # ty: ignore
 
 
 def download_climatology():
@@ -157,7 +156,7 @@ def download_climatology():
         }
         data = earthkit.data.from_source("mars", **ekp)
         with open(f"{data_root}/data_clim_{step}.grib", "wb") as f:
-            data.write(f)
+            data.write(f) # ty: ignore
 
 
 if __name__ == "__main__":

cascade/benchmarks/matmul.py

@@ -1,9 +1,9 @@
 import os
 from typing import Any
 
-import jax
-import jax.numpy as jp
-import jax.random as jr
+import jax  # ty: ignore[unresolved-import]
+import jax.numpy as jp  # ty: ignore[unresolved-import]
+import jax.random as jr  # ty: ignore[unresolved-import]
 
 from cascade.low.builders import JobBuilder, TaskBuilder
 from cascade.low.core import JobInstance
@@ -65,7 +65,7 @@ def execute_locally():
 
     from multiprocessing.shared_memory import SharedMemory
 
-    mem = SharedMemory("benchmark_tmp", create=True, size=m0.nbytes)
+    mem = SharedMemory("benchmark_tmp", create=True, size=m0.nbytes); assert mem.buf is not None
     mem.buf[:] = m0.tobytes()
 
 

cascade/benchmarks/tests.py

@@ -32,7 +32,7 @@ from cascade.executor.runner.memory import Memory, ds2shmid
 from cascade.executor.runner.packages import PackagesEnv
 from cascade.executor.runner.runner import ExecutionContext, run
 from cascade.low.builders import TaskBuilder
-from cascade.low.core import DatasetId
+from cascade.low.core import DatasetId, WorkerId
 from cascade.shm.server import entrypoint as shm_server
 
 logger = logging.getLogger(__name__)
@@ -75,7 +75,7 @@ def simple_runner(callback: BackboneAddress, executionContext: ExecutionContext)
         raise ValueError(f"expected 1 task, gotten {len(tasks)}")
     taskId = tasks[0]
     taskInstance = executionContext.tasks[taskId]
-    with Memory(callback, "testWorker") as memory, PackagesEnv() as pckg:
+    with Memory(callback, WorkerId(host="testHost", worker="testWorker")) as memory, PackagesEnv() as pckg:
         # for key, value in taskSequence.extra_env.items():
         #    os.environ[key] = value
 
@@ -142,7 +142,7 @@ def run_test(
         while perf_counter_ns() < end:
             mess = listener.recv_messages()
             if mess == [
-                DatasetPublished(origin="testWorker", ds=output, transmit_idx=None)
+                DatasetPublished(origin=WorkerId(host="testHost", worker="testWorker"), ds=output, transmit_idx=None)
             ]:
                 break
             elif not mess:

cascade/benchmarks/util.py

@@ -29,7 +29,7 @@ from cascade.executor.comms import callback
 from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.executor.executor import Executor
 from cascade.executor.msg import BackboneAddress, ExecutorShutdown
-from cascade.low.core import DatasetId, JobInstance
+from cascade.low.core import DatasetId, JobInstance, JobInstanceRich
 from cascade.low.func import msum
 from cascade.scheduler.precompute import precompute
 from earthkit.workflows.graph import Graph, deduplicate_nodes
@@ -37,15 +37,16 @@ from earthkit.workflows.graph import Graph, deduplicate_nodes
 logger = logging.getLogger("cascade.benchmarks")
 
 
-def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
+def get_job(benchmark: str | None, instance_path: str | None) -> JobInstanceRich:
     # NOTE because of os.environ, we don't import all... ideally we'd have some file-based init/config mech instead
     if benchmark is not None and instance_path is not None:
         raise TypeError("specified both benchmark name and job instance")
     elif instance_path is not None:
         with open(instance_path, "rb") as f:
             d = orjson.loads(f.read())
-            return JobInstance(**d)
+            return JobInstanceRich(**d)
     elif benchmark is not None:
+        instance: JobInstance
         if benchmark.startswith("j1"):
             import cascade.benchmarks.job1 as job1
 
@@ -58,25 +59,26 @@ def get_job(benchmark: str | None, instance_path: str | None) -> JobInstance:
                 msum((v for k, v in graphs.items() if k.startswith(prefix)), Graph)
             )
             graphs["j1.all"] = union("j1.")
-            return cascade.low.into.graph2job(graphs[benchmark])
+            instance = cascade.low.into.graph2job(graphs[benchmark])
         elif benchmark.startswith("generators"):
             import cascade.benchmarks.generators as generators
 
-            return generators.get_job()
+            instance = generators.get_job()
         elif benchmark.startswith("matmul"):
             import cascade.benchmarks.matmul as matmul
 
-            return matmul.get_job()
+            instance = matmul.get_job()
         elif benchmark.startswith("dist"):
             import cascade.benchmarks.dist as dist
 
-            return dist.get_job()
+            instance = dist.get_job()
         elif benchmark.startswith("dask"):
             import cascade.benchmarks.dask as dask
 
-            return dask.get_job(benchmark[len("dask.") :])
+            instance = dask.get_job(benchmark[len("dask.") :])
         else:
             raise NotImplementedError(benchmark)
+        return JobInstanceRich(jobInstance=instance, checkpointSpec=None)
     else:
         raise TypeError("specified neither benchmark name nor job instance")
 
@@ -116,7 +118,7 @@ def get_gpu_count(host_idx: int, worker_count: int) -> int:
 
 
 def launch_executor(
-    job_instance: JobInstance,
+    job: JobInstanceRich,
     controller_address: BackboneAddress,
     workers_per_host: int,
     portBase: int,
@@ -136,7 +138,7 @@ def launch_executor(
         logger.info(f"will set {gpu_count} gpus on host {i}")
         os.environ["CASCADE_GPU_COUNT"] = str(gpu_count)
         executor = Executor(
-            job_instance,
+            job.jobInstance,
             controller_address,
             workers_per_host,
             f"h{i}",
@@ -154,7 +156,7 @@ def launch_executor(
 
 
 def run_locally(
-    job: JobInstance,
+    job: JobInstanceRich,
     hosts: int,
     workers: int,
     portBase: int = 12345,
@@ -195,7 +197,7 @@ def run_locally(
             ps.append(p)
 
         # compute preschedule
-        preschedule = precompute(job)
+        preschedule = precompute(job.jobInstance)
 
         # check processes started healthy
         for i, p in enumerate(ps):
@@ -240,9 +242,9 @@ def main_local(
     port_base: int = 12345,
     log_base: str | None = None,
 ) -> None:
-    jobInstance = get_job(job, instance)
+    jobInstanceRich = get_job(job, instance)
     run_locally(
-        jobInstance,
+        jobInstanceRich,
         hosts,
         workers_per_host,
         report_address=report_address,
@@ -266,17 +268,17 @@ def main_dist(
     """
     launch = perf_counter_ns()
 
-    jobInstance = get_job(job, instance)
+    jobInstanceRich = get_job(job, instance)
 
     if idx == 0:
         logging.config.dictConfig(logging_config)
         tp = ThreadPoolExecutor(max_workers=1)
-        preschedule_fut = tp.submit(precompute, jobInstance)
+        preschedule_fut = tp.submit(precompute, jobInstanceRich.jobInstance)
         b = Bridge(controller_url, hosts)
         preschedule = preschedule_fut.result()
         tp.shutdown()
         start = perf_counter_ns()
-        run(jobInstance, b, preschedule, report_address=report_address)
+        run(jobInstanceRich, b, preschedule, report_address=report_address)
         end = perf_counter_ns()
         print(
             f"compute took {(end-start)/1e9:.3f}s, including startup {(end-launch)/1e9:.3f}s"
@@ -284,12 +286,13 @@ def main_dist(
     else:
         gpu_count = get_gpu_count(0, workers_per_host)
         launch_executor(
-            jobInstance,
+            jobInstanceRich,
             controller_url,
             workers_per_host,
             12345,
             idx,
             shm_vol_gb,
             gpu_count,
-            f"tcp://{platform.get_bindabble_self()}",
+            log_base = None, # TODO handle log collection for dist scenario
+            url_base = f"tcp://{platform.get_bindabble_self()}",
         )

cascade/controller/act.py

@@ -10,6 +10,7 @@
 
 import logging
 
+import cascade.executor.checkpoints as checkpoints
 from cascade.controller.core import State
 from cascade.executor.bridge import Bridge
 from cascade.executor.msg import TaskSequence
@@ -76,6 +77,12 @@ def flush_queues(bridge: Bridge, state: State, context: JobExecutionContext):
     for dataset, host in state.drain_fetching_queue():
         bridge.fetch(dataset, host)
 
+    for dataset, host in state.drain_persist_queue():
+        if context.checkpoint_spec is None:
+            raise TypeError(f"unexpected persist need when checkpoint storage not configured")
+        persist_params = checkpoints.serialize_persist_params(context.checkpoint_spec)
+        bridge.persist(dataset, host, context.checkpoint_spec.storage_type, persist_params)
+
     for ds in state.drain_purging_queue():
         for host in context.purge_dataset(ds):
             logger.debug(f"issuing purge of {ds=} to {host=}")

cascade/controller/core.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from typing import Any, Iterator
 
 import cascade.executor.serde as serde
-from cascade.executor.msg import DatasetTransmitPayload
+from cascade.executor.msg import DatasetPersistSuccess, DatasetTransmitPayload
 from cascade.low.core import DatasetId, HostId, TaskId
 
 logger = logging.getLogger(__name__)
@@ -16,10 +16,14 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class State:
-    # key add by core.initialize, value add by notify.notify
+    # key add by core.init_state, value add by notify.notify
     outputs: dict[DatasetId, Any]
+    # key add by core.init_state, value add by notify.notify
+    to_persist: set[DatasetId]
     # add by notify.notify, remove by act.flush_queues
     fetching_queue: dict[DatasetId, HostId]
+    # add by notify.notify, remove by act.flush_queues
+    persist_queue: dict[DatasetId, HostId]
     # add by notify.notify, removed by act.flush_queues
     purging_queue: list[DatasetId]
     # add by core.init_state, remove by notify.notify
@@ -31,13 +35,16 @@ class State:
         for e in self.outputs.values():
             if e is None:
                 return True
+        if self.to_persist:
+            return True
         return False
 
     def _consider_purge(self, dataset: DatasetId) -> None:
         """If dataset not required anymore, add to purging_queue"""
         no_dependants = not self.purging_tracker.get(dataset, None)
         not_required_output = self.outputs.get(dataset, 1) is not None
-        if no_dependants and not_required_output:
+        not_required_persist = not dataset in self.to_persist
+        if all((no_dependants, not_required_output, not_required_persist)):
             logger.debug(f"adding {dataset=} to purging queue")
             if dataset in self.purging_tracker:
                 self.purging_tracker.pop(dataset)
@@ -52,6 +59,14 @@ class State:
         ):
             self.fetching_queue[dataset] = at
 
+    def consider_persist(self, dataset: DatasetId, at: HostId) -> None:
+        """If required as persist and not yet acknowledged, add to persist queue"""
+        if (
+            dataset in self.to_persist
+            and dataset not in self.persist_queue
+        ):
+            self.persist_queue[dataset] = at
+
     def receive_payload(self, payload: DatasetTransmitPayload) -> None:
         """Stores deserialized value into outputs, considers purge"""
         # NOTE ifneedbe get annotation from job.tasks[event.ds.task].definition.output_schema[event.ds.output]
@@ -60,6 +75,11 @@ class State:
         )
         self._consider_purge(payload.header.ds)
 
+    def acknowledge_persist(self, payload: DatasetPersistSuccess) -> None:
+        """Marks acknowledged, considers purge"""
+        self.to_persist.discard(payload.ds)
+        self._consider_purge(payload.ds)
+
     def task_done(self, task: TaskId, inputs: set[DatasetId]) -> None:
         """Marks that the inputs are not needed for this task anymore, considers purge of each"""
         for sourceDataset in inputs:
@@ -76,15 +96,22 @@ class State:
             yield dataset, host
         self.fetching_queue = {}
 
+    def drain_persist_queue(self) -> Iterator[tuple[DatasetId, HostId]]:
+        for dataset, host in self.persist_queue.items():
+            yield dataset, host
+        self.persist_queue = {}
+
 
-def init_state(outputs: set[DatasetId], edge_o: dict[DatasetId, set[TaskId]]) -> State:
+def init_state(outputs: set[DatasetId], to_persist: set[DatasetId], edge_o: dict[DatasetId, set[TaskId]]) -> State:
     purging_tracker = {
         ds: {task for task in dependants} for ds, dependants in edge_o.items()
     }
 
     return State(
         outputs={e: None for e in outputs},
+        to_persist={e for e in to_persist},
         fetching_queue={},
         purging_queue=[],
         purging_tracker=purging_tracker,
+        persist_queue={},
     )

cascade/controller/impl.py

@@ -14,7 +14,7 @@ from cascade.controller.core import State, init_state
 from cascade.controller.notify import notify
 from cascade.controller.report import Reporter
 from cascade.executor.bridge import Bridge, Event
-from cascade.low.core import JobInstance, type_dec
+from cascade.low.core import JobInstance, JobInstanceRich, type_dec
 from cascade.low.execution_context import init_context
 from cascade.low.tracing import ControllerPhases, Microtrace, label, mark, timer
 from cascade.scheduler.api import assign, init_schedule, plan
@@ -24,7 +24,7 @@ logger = logging.getLogger(__name__)
 
 
 def run(
-    job: JobInstance,
+    job: JobInstanceRich,
     bridge: Bridge,
     preschedule: Preschedule,
     report_address: str | None = None,
@@ -34,7 +34,8 @@ def run(
     outputs = set(context.job_instance.ext_outputs)
     logger.debug(f"starting with {env=} and {report_address=}")
     schedule = timer(init_schedule, Microtrace.ctrl_init)(preschedule, context)
-    state = init_state(outputs, context.edge_o)
+    to_persist = set(job.checkpointSpec.to_persist) if job.checkpointSpec is not None else set()
+    state = init_state(outputs, to_persist, context.edge_o)
 
     label("host", "controller")
     events: list[Event] = []
@@ -44,7 +45,7 @@ def run(
 
     try:
         total_gpus = sum(worker.gpu for worker in env.workers.values())
-        needs_gpus = any(task.definition.needs_gpu for task in job.tasks.values())
+        needs_gpus = any(task.definition.needs_gpu for task in job.jobInstance.tasks.values())
         if needs_gpus and total_gpus == 0:
             raise ValueError("environment contains no gpu yet job demands one")
 

cascade/controller/notify.py

@@ -17,7 +17,7 @@ from typing import Iterable
 from cascade.controller.core import State
 from cascade.controller.report import Reporter
 from cascade.executor.bridge import Event
-from cascade.executor.msg import DatasetPublished, DatasetTransmitPayload
+from cascade.executor.msg import DatasetPersistSuccess, DatasetPublished, DatasetTransmitPayload
 from cascade.low.core import DatasetId, HostId, WorkerId
 from cascade.low.execution_context import DatasetStatus, JobExecutionContext
 from cascade.low.func import assert_never
@@ -89,6 +89,7 @@ def notify(
             context.host2ds[host][event.ds] = DatasetStatus.available
             context.ds2host[event.ds][host] = DatasetStatus.available
             state.consider_fetch(event.ds, host)
+            state.consider_persist(event.ds, host)
             consider_computable(schedule, state, context, event.ds, host)
             if event.transmit_idx is not None:
                 mark(
@@ -121,5 +122,7 @@ def notify(
         elif isinstance(event, DatasetTransmitPayload):
             state.receive_payload(event)
             reporter.send_result(event.header.ds, event.value)
+        elif isinstance(event, DatasetPersistSuccess):
+            state.acknowledge_persist(event)
         else:
             assert_never(event)

cascade/executor/bridge.py

@@ -16,6 +16,9 @@ from cascade.executor.comms import default_message_resend_ms as resend_grace_ms
 from cascade.executor.executor import heartbeat_grace_ms as executor_heartbeat_grace_ms
 from cascade.executor.msg import (
     Ack,
+    DatasetPersistCommand,
+    DatasetPersistFailure,
+    DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
     DatasetTransmitCommand,
@@ -29,14 +32,15 @@ from cascade.executor.msg import (
     TaskFailure,
     TaskSequence,
 )
-from cascade.low.core import DatasetId, Environment, HostId, Worker, WorkerId
+from cascade.low.core import CheckpointStorageType, DatasetId, Environment, HostId, Worker, WorkerId
 from cascade.low.func import assert_never
 
 logger = logging.getLogger(__name__)
 
-Event = DatasetPublished | DatasetTransmitPayload
-ToShutdown = TaskFailure | ExecutorFailure | DatasetTransmitFailure | ExecutorExit
-Unsupported = TaskSequence | DatasetPurge | DatasetTransmitCommand | ExecutorShutdown
+Event = DatasetPublished | DatasetTransmitPayload | DatasetPersistSuccess
+# TODO consider retries here, esp on the PersistFailure
+ToShutdown = TaskFailure | ExecutorFailure | DatasetTransmitFailure | DatasetPersistFailure | ExecutorExit
+Unsupported = TaskSequence | DatasetPurge | DatasetTransmitCommand | DatasetPersistCommand | ExecutorShutdown
 
 
 class Bridge:
@@ -158,6 +162,15 @@ class Bridge:
         self.transmit_idx_counter += 1
         self.sender.send("data." + source, m)
 
+    def persist(self, ds: DatasetId, source: HostId, storage_type: CheckpointStorageType, persist_params: str) -> None:
+        m = DatasetPersistCommand(
+            source=source,
+            ds=ds,
+            storage_type=storage_type,
+            persist_params=persist_params,
+        )
+        self.sender.send("data." + source, m)
+
     def fetch(self, ds: DatasetId, source: HostId) -> None:
         m = DatasetTransmitCommand(
             source=source,

cascade/executor/checkpoints.py

@@ -0,0 +1,42 @@
+# (C) Copyright 2025- ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+"""Handles the checkpoint management: storage, retrieval"""
+
+import pathlib
+
+from cascade.executor.msg import DatasetPersistCommand
+from cascade.low.core import CheckpointSpec
+from cascade.low.func import assert_never
+from cascade.shm.client import AllocatedBuffer
+
+
+def persist_dataset(command: DatasetPersistCommand, buf: AllocatedBuffer) -> None:
+    match command.storage_type:
+        case "fs":
+            root = pathlib.Path(command.persist_params)
+            root.mkdir(parents=True, exist_ok=True)
+            file = root / repr(command.ds)
+            # TODO what about overwrites / concurrent writes? Append uuid?
+            file.write_bytes(buf.view())
+        case s:
+            assert_never(s)
+
+def serialize_persist_params(spec: CheckpointSpec) -> str:
+    # NOTE we call this every time we store, ideally call this once when building `low.execution_context`
+    match spec.storage_type:
+        case "fs":
+            if not isinstance(spec.storage_params, str):
+                raise TypeError(f"expected checkpoint storage params to be str, gotten {spec.storage_params.__class__}")
+            if spec.persist_id is None:
+                raise TypeError(f"serialize_persist_params called, but persist_id is None")
+            root = pathlib.Path(spec.storage_params)
+            return str(root / spec.persist_id)
+        case s:
+            assert_never(s)
+

cascade/executor/data_server.py

@@ -15,16 +15,20 @@ large data object.
 
 import logging
 import logging.config
-from concurrent.futures import ALL_COMPLETED, FIRST_COMPLETED
+from concurrent.futures import ALL_COMPLETED, FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
 from concurrent.futures import Executor as PythonExecutor
-from concurrent.futures import Future, ThreadPoolExecutor, wait
 from time import time_ns
+from typing import cast
 
 import cascade.shm.client as shm_client
+from cascade.executor.checkpoints import persist_dataset
 from cascade.executor.comms import Listener, callback, send_data
 from cascade.executor.msg import (
     Ack,
     BackboneAddress,
+    DatasetPersistCommand,
+    DatasetPersistFailure,
+    DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
     DatasetTransmitCommand,
@@ -59,7 +63,7 @@ class DataServer:
         self.cap = 2
         self.ds_proc_tp: PythonExecutor = ThreadPoolExecutor(max_workers=self.cap)
         self.futs_in_progress: dict[
-            DatasetTransmitCommand | DatasetTransmitPayload, Future
+            DatasetTransmitCommand | DatasetTransmitPayload | DatasetPersistCommand, Future
         ] = {}
         self.awaiting_confirmation: dict[int, tuple[DatasetTransmitCommand, int]] = {}
         self.invalid: set[DatasetId] = (
@@ -149,6 +153,26 @@ class DataServer:
             time_ns()
         )  # not actually consumed but uniform signature with send_payload simplifies typing
 
+    def persist_payload(self, command: DatasetPersistCommand) -> int:
+        buf: None | shm_client.AllocatedBuffer = None
+        try:
+            if command.source != self.host:
+                raise ValueError(f"invalid {command=}")
+            buf = shm_client.get(key=ds2shmid(command.ds))
+            persist_dataset(command, buf)
+            logger.debug(f"dataset for {command} persisted")
+            callback(self.maddress, DatasetPersistSuccess(host=self.host, ds=command.ds))
+        except Exception as e:
+            logger.exception(f"failed to persist dataset for {command}, reporting up")
+            callback(
+                self.maddress,
+                DatasetPersistFailure(host=self.host, detail=f"{repr(command)} -> {repr(e)}"),
+            )
+        finally:
+            if buf is not None:
+                buf.close()
+        return time_ns()
+
     def send_payload(self, command: DatasetTransmitCommand) -> int:
         buf: None | shm_client.AllocatedBuffer = None
         payload: None | DatasetTransmitPayload = None
@@ -171,7 +195,7 @@ class DataServer:
                 ds=command.ds,
                 deser_fun=buf.deser_fun,
             )
-            payload = DatasetTransmitPayload(header, value=buf.view())
+            payload = DatasetTransmitPayload(header, value=cast(bytes, buf.view()))
             syn = Syn(command.idx, self.dlistener.address)
             send_data(command.daddress, payload, syn)
             logger.debug(f"payload for {command} sent")
@@ -218,6 +242,14 @@ class DataServer:
                         self.awaiting_confirmation[m.idx] = (m, -1)
                         fut = self.ds_proc_tp.submit(self.send_payload, m)
                         self.futs_in_progress[m] = fut
+                    elif isinstance(m, DatasetPersistCommand):
+                        if m.ds in self.invalid:
+                            raise ValueError(
+                                f"unexpected persist command {m} as the dataset was already purged"
+                            )
+                        # TODO mark?
+                        fut = self.ds_proc_tp.submit(self.persist_payload, m)
+                        self.futs_in_progress[m] = fut
                     elif isinstance(m, DatasetTransmitPayload):
                         if m.header.ds in self.invalid:
                             logger.warning(
@@ -238,9 +270,10 @@ class DataServer:
                         self.acks.add(m.idx)
                     elif isinstance(m, DatasetPurge):
                         # we need to handle potential commands transmitting this dataset, as otherwise they'd fail
+                        # TODO submit this as a future? This actively blocks the whole server
                         to_wait = []
                         for commandProg, fut in self.futs_in_progress.items():
-                            if isinstance(commandProg, DatasetTransmitCommand):
+                            if isinstance(commandProg, DatasetTransmitCommand|DatasetPersistCommand):
                                 val = commandProg.ds
                             elif isinstance(commandProg, DatasetTransmitPayload):
                                 val = commandProg.header.ds

cascade/executor/executor.py

@@ -35,6 +35,8 @@ from cascade.executor.data_server import start_data_server
 from cascade.executor.msg import (
     Ack,
     BackboneAddress,
+    DatasetPersistFailure,
+    DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
     DatasetTransmitFailure,
@@ -306,7 +308,7 @@ class Executor:
                             callback(worker_address(worker), m)
                         self.datasets.add(m.ds)
                         self.to_controller(m)
-                    elif isinstance(m, DatasetTransmitFailure):
+                    elif isinstance(m, DatasetTransmitFailure|DatasetPersistSuccess|DatasetPersistFailure):
                         self.to_controller(m)
                     else:
                         # NOTE transmit and store are handled in DataServer (which has its own socket)