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

cascade/benchmarks/util.py

@@ -209,7 +209,7 @@ def run_locally(
 
         # start bridge itself
         logger.debug("starting bridge")
-        b = Bridge(c, hosts)
+        b = Bridge(c, hosts, job.checkpointSpec)
         start = perf_counter_ns()
         result = run(job, b, preschedule, report_address=report_address)
         end = perf_counter_ns()
@@ -274,7 +274,7 @@ def main_dist(
         logging.config.dictConfig(logging_config)
         tp = ThreadPoolExecutor(max_workers=1)
         preschedule_fut = tp.submit(precompute, jobInstanceRich.jobInstance)
-        b = Bridge(controller_url, hosts)
+        b = Bridge(controller_url, hosts, jobInstanceRich.checkpointSpec)
         preschedule = preschedule_fut.result()
         tp.shutdown()
         start = perf_counter_ns()

cascade/controller/act.py

@@ -9,12 +9,14 @@
 """Implements the invocation of Bridge/Executor methods given a sequence of Actions"""
 
 import logging
+from typing import Iterable, Iterator, cast
 
-import cascade.executor.checkpoints as checkpoints
 from cascade.controller.core import State
 from cascade.executor.bridge import Bridge
-from cascade.executor.msg import TaskSequence
-from cascade.low.execution_context import JobExecutionContext
+from cascade.executor.checkpoints import build_retrieve_command, possible_repersist, retrieve_dataset
+from cascade.executor.msg import DatasetPublished, TaskSequence
+from cascade.low.core import DatasetId
+from cascade.low.execution_context import JobExecutionContext, VirtualCheckpointHost
 from cascade.low.tracing import TaskLifecycle, TransmitLifecycle, mark
 from cascade.scheduler.core import Assignment
 
@@ -75,17 +77,46 @@ def flush_queues(bridge: Bridge, state: State, context: JobExecutionContext):
     """
 
     for dataset, host in state.drain_fetching_queue():
-        bridge.fetch(dataset, host)
+        if host != VirtualCheckpointHost:
+            bridge.fetch(dataset, host)
+        else:
+            # NOTE we would rather not be here, but we dont generally expect
+            # checkpointed datasets to be outputs. If needbe, send a command
+            # to any worker, or spawn a thread with this
+            logger.warning(f"execute checkpoint retrieve on controller")
+            # NOTE the host is the virtual one so the message is not really valid, but no big deal
+            virtual_command = build_retrieve_command(bridge.checkpoint_spec, dataset, host)
+            buffer = retrieve_dataset(virtual_command)
+            try:
+                # the cast is wrong but ty is bit confused about memoryview anyway
+                state.receive_payload(dataset, cast(bytes, buffer.view()), buffer.deser_fun)
+            finally:
+                buffer.close()
 
     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)
+        if host != VirtualCheckpointHost:
+            bridge.persist(dataset, host)
+        else:
+            possible_repersist(dataset, bridge.checkpoint_spec)
+            state.acknowledge_persist(dataset)
 
     for ds in state.drain_purging_queue():
         for host in context.purge_dataset(ds):
-            logger.debug(f"issuing purge of {ds=} to {host=}")
-            bridge.purge(host, ds)
+            if host != VirtualCheckpointHost:
+                logger.debug(f"issuing purge of {ds=} to {host=}")
+                bridge.purge(host, ds)
 
     return state
+
+def virtual_checkpoint_publish(datasets: Iterable[DatasetId]) -> Iterator[DatasetPublished]:
+    """Virtual in the sense of not actually sending any message, but instead simulating
+    a response so that controller.notify can bring the contexts into the right state.
+    Invoked once, at the job start, after the checkpoint has been listed"""
+    return (
+        DatasetPublished(
+            origin=VirtualCheckpointHost,
+            ds=dataset,
+            transmit_idx=None,
+        )
+        for dataset in datasets
+    )

cascade/controller/core.py

@@ -10,6 +10,7 @@ from typing import Any, Iterator
 import cascade.executor.serde as serde
 from cascade.executor.msg import DatasetPersistSuccess, DatasetTransmitPayload
 from cascade.low.core import DatasetId, HostId, TaskId
+from cascade.low.execution_context import VirtualCheckpointHost
 
 logger = logging.getLogger(__name__)
 
@@ -67,18 +68,16 @@ class State:
         ):
             self.persist_queue[dataset] = at
 
-    def receive_payload(self, payload: DatasetTransmitPayload) -> None:
+    def receive_payload(self, ds: DatasetId, payload: bytes, deser_fun: str) -> None:
         """Stores deserialized value into outputs, considers purge"""
         # NOTE ifneedbe get annotation from job.tasks[event.ds.task].definition.output_schema[event.ds.output]
-        self.outputs[payload.header.ds] = serde.des_output(
-            payload.value, "Any", payload.header.deser_fun
-        )
-        self._consider_purge(payload.header.ds)
+        self.outputs[ds] = serde.des_output(payload, "Any", deser_fun)
+        self._consider_purge(ds)
 
-    def acknowledge_persist(self, payload: DatasetPersistSuccess) -> None:
+    def acknowledge_persist(self, ds: DatasetId) -> None:
         """Marks acknowledged, considers purge"""
-        self.to_persist.discard(payload.ds)
-        self._consider_purge(payload.ds)
+        self.to_persist.discard(ds)
+        self._consider_purge(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"""

cascade/controller/impl.py

@@ -9,15 +9,17 @@
 import logging
 
 import cascade.executor.serde as serde
-from cascade.controller.act import act, flush_queues
+from cascade.controller.act import act, flush_queues, virtual_checkpoint_publish
 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.executor.checkpoints import list_persisted_datasets
 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
+from cascade.scheduler.checkpoints import trim_with_persisted, virtual_update_schedule
 from cascade.scheduler.core import Preschedule
 
 logger = logging.getLogger(__name__)
@@ -30,6 +32,9 @@ def run(
     report_address: str | None = None,
 ) -> State:
     env = bridge.get_environment()
+    persisted = list_persisted_datasets(job.checkpointSpec) if job.checkpointSpec is not None else []
+    jobInstance, preschedule, persisted_valid = trim_with_persisted(job, preschedule, set(persisted))
+    job.jobInstance = jobInstance
     context = init_context(env, job, preschedule.edge_o, preschedule.edge_i)
     outputs = set(context.job_instance.ext_outputs)
     logger.debug(f"starting with {env=} and {report_address=}")
@@ -42,6 +47,7 @@ def run(
     for serdeTypeEnc, (serdeSer, serdeDes) in context.job_instance.serdes.items():
         serde.SerdeRegistry.register(type_dec(serdeTypeEnc), serdeSer, serdeDes)
     reporter = Reporter(report_address)
+    notify_wrapper = lambda events: notify(state, schedule, context, events, reporter)
 
     try:
         total_gpus = sum(worker.gpu for worker in env.workers.values())
@@ -49,6 +55,10 @@ def run(
         if needs_gpus and total_gpus == 0:
             raise ValueError("environment contains no gpu yet job demands one")
 
+        virtual_update_schedule(persisted_valid, schedule, context)
+        virtual_events = virtual_checkpoint_publish(persisted_valid)
+        timer(notify_wrapper, Microtrace.ctrl_notify)(virtual_events)
+
         while (
             state.has_awaitable()
             or context.has_awaitable()
@@ -68,11 +78,9 @@ def run(
 
             mark({"action": ControllerPhases.wait})
             if state.has_awaitable() or context.has_awaitable():
-                logger.debug(f"about to await bridge with {context.ongoing_total=}")
+                logger.debug(f"about to await bridge with {context.ongoing_total=}, {context.remaining=} and {state.has_awaitable()=}")
                 events = timer(bridge.recv_events, Microtrace.ctrl_wait)()
-                timer(notify, Microtrace.ctrl_notify)(
-                    state, schedule, context, events, reporter
-                )
+                timer(notify_wrapper, Microtrace.ctrl_notify)(events)
                 logger.debug(f"received {len(events)} events")
     except Exception as ex:
         logger.error("crash in controller, shuting down")

cascade/controller/notify.py

@@ -12,14 +12,14 @@
 # Thus the caller always *must* use the return value and cease using the input.
 
 import logging
-from typing import Iterable
+from typing import Iterable, cast
 
 from cascade.controller.core import State
 from cascade.controller.report import Reporter
 from cascade.executor.bridge import Event
-from cascade.executor.msg import DatasetPersistSuccess, DatasetPublished, DatasetTransmitPayload
+from cascade.executor.msg import DatasetPersistSuccess, DatasetPublished, DatasetRetrieveSuccess, DatasetTransmitPayload
 from cascade.low.core import DatasetId, HostId, WorkerId
-from cascade.low.execution_context import DatasetStatus, JobExecutionContext
+from cascade.low.execution_context import DatasetStatus, JobExecutionContext, VirtualCheckpointHost
 from cascade.low.func import assert_never
 from cascade.low.tracing import TaskLifecycle, TransmitLifecycle, mark
 from cascade.scheduler.api import gang_check_ready
@@ -103,26 +103,34 @@ def notify(
             elif context.is_last_output_of(event.ds):
                 worker = event.origin
                 task = event.ds.task
-                if not isinstance(worker, WorkerId):
+                isWorker = isinstance(worker, WorkerId)
+                isVirtual = worker == VirtualCheckpointHost
+                if not isWorker and not isVirtual:
                     raise ValueError(
                         f"malformed event, expected origin to be WorkerId: {event}"
                     )
                 logger.debug(f"last output of {task}, assuming completion")
-                mark(
-                    {
-                        "task": task,
-                        "action": TaskLifecycle.completed,
-                        "worker": repr(worker),
-                        "host": "controller",
-                    }
-                )
                 state.task_done(task, context.edge_i.get(event.ds.task, set()))
-                context.task_done(task, worker)
+                if isWorker: 
+                    mark(
+                        {
+                            "task": task,
+                            "action": TaskLifecycle.completed,
+                            "worker": repr(worker),
+                            "host": "controller",
+                        }
+                    )
+                    worker = cast(WorkerId, worker) # ty cant yet derive this to be true
+                    context.task_done_at(task, worker)
+                else:
+                    context.task_done(task)
                 reporter.send_progress(context)
         elif isinstance(event, DatasetTransmitPayload):
-            state.receive_payload(event)
+            state.receive_payload(event.header.ds, event.value, event.header.deser_fun)
             reporter.send_result(event.header.ds, event.value)
         elif isinstance(event, DatasetPersistSuccess):
-            state.acknowledge_persist(event)
+            state.acknowledge_persist(event.ds)
+        elif isinstance(event, DatasetRetrieveSuccess):
+            pass
         else:
             assert_never(event)

cascade/executor/bridge.py

@@ -11,6 +11,7 @@
 import logging
 import time
 
+from cascade.executor.checkpoints import build_persist_command, build_retrieve_command, serialize_params
 from cascade.executor.comms import GraceWatcher, Listener, ReliableSender
 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
@@ -21,6 +22,9 @@ from cascade.executor.msg import (
     DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
+    DatasetRetrieveCommand,
+    DatasetRetrieveFailure,
+    DatasetRetrieveSuccess,
     DatasetTransmitCommand,
     DatasetTransmitFailure,
     DatasetTransmitPayload,
@@ -32,19 +36,21 @@ from cascade.executor.msg import (
     TaskFailure,
     TaskSequence,
 )
-from cascade.low.core import CheckpointStorageType, DatasetId, Environment, HostId, Worker, WorkerId
+from cascade.low.core import CheckpointSpec, DatasetId, Environment, HostId, Worker, WorkerId
+from cascade.low.execution_context import VirtualCheckpointHost
 from cascade.low.func import assert_never
 
 logger = logging.getLogger(__name__)
 
-Event = DatasetPublished | DatasetTransmitPayload | DatasetPersistSuccess
-# TODO consider retries here, esp on the PersistFailure
-ToShutdown = TaskFailure | ExecutorFailure | DatasetTransmitFailure | DatasetPersistFailure | ExecutorExit
+Event = DatasetPublished | DatasetTransmitPayload | DatasetPersistSuccess | DatasetRetrieveSuccess
+# TODO consider retries here, esp on the Persist/Retrieve Failures
+ToShutdown = TaskFailure | ExecutorFailure | DatasetRetrieveFailure | DatasetTransmitFailure | DatasetPersistFailure | ExecutorExit
 Unsupported = TaskSequence | DatasetPurge | DatasetTransmitCommand | DatasetPersistCommand | ExecutorShutdown
 
 
 class Bridge:
-    def __init__(self, controller_url: str, expected_executors: int) -> None:
+    def __init__(self, controller_url: str, expected_executors: int, checkpoint_spec: CheckpointSpec|None=None) -> None:
+        self.checkpoint_spec = checkpoint_spec
         self.mlistener = Listener(controller_url)
         self.heartbeat_checker: dict[HostId, GraceWatcher] = {}
         self.transmit_idx_counter = 0
@@ -152,24 +158,23 @@ class Bridge:
         self._send(host, m)
 
     def transmit(self, ds: DatasetId, source: HostId, target: HostId) -> None:
-        m = DatasetTransmitCommand(
-            source=source,
-            target=target,
-            daddress=self.sender.hosts["data." + target][1],
-            ds=ds,
-            idx=self.transmit_idx_counter,
-        )
-        self.transmit_idx_counter += 1
-        self.sender.send("data." + source, m)
+        if source == VirtualCheckpointHost:
+            command = build_retrieve_command(self.checkpoint_spec, ds, target)
+            self.sender.send("data." + target, command)
+        else:
+            m = DatasetTransmitCommand(
+                source=source,
+                target=target,
+                daddress=self.sender.hosts["data." + target][1],
+                ds=ds,
+                idx=self.transmit_idx_counter,
+            )
+            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 persist(self, ds: DatasetId, source: HostId) -> None:
+        command = build_persist_command(self.checkpoint_spec, ds, source)
+        self.sender.send("data." + source, command)
 
     def fetch(self, ds: DatasetId, source: HostId) -> None:
         m = DatasetTransmitCommand(

cascade/executor/checkpoints.py

@@ -8,35 +8,127 @@
 
 """Handles the checkpoint management: storage, retrieval"""
 
+import io
+import logging
+import os
 import pathlib
 
-from cascade.executor.msg import DatasetPersistCommand
-from cascade.low.core import CheckpointSpec
+from cascade.executor.msg import DatasetPersistCommand, DatasetRetrieveCommand
+from cascade.executor.platform import advise_seqread
+from cascade.executor.runner.memory import ds2shmid
+from cascade.executor.serde import DefaultSerde
+from cascade.low.core import CheckpointSpec, DatasetId, HostId
+from cascade.low.execution_context import VirtualCheckpointHost
 from cascade.low.func import assert_never
-from cascade.shm.client import AllocatedBuffer
+from cascade.shm.client import AllocatedBuffer, allocate
 
+logger = logging.getLogger(__name__)
+
+def serialize_params(spec: CheckpointSpec, id_: str) -> str:
+    """id_ is either the persist id or retrieve id from the spec"""
+    # 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__}")
+            root = pathlib.Path(spec.storage_params)
+            return str(root / id_)
+        case s:
+            assert_never(s)
+
+def build_persist_command(checkpoint_spec: CheckpointSpec|None, ds: DatasetId, hostId: HostId) -> DatasetPersistCommand:
+    if checkpoint_spec is None:
+        raise ValueError(f"unexpected persist need when checkpoint storage not configured")
+    id_ = checkpoint_spec.persist_id
+    if not id_:
+        raise ValueError(f"unexpected persist need when there is no persist id")
+    persist_params = serialize_params(checkpoint_spec, id_)
+    return DatasetPersistCommand(
+        source=hostId,
+        ds=ds,
+        storage_type=checkpoint_spec.storage_type,
+        persist_params=persist_params,
+    )
 
 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)
+            file = root / command.ds.ser()
             # 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`
+def list_persisted_datasets(spec: CheckpointSpec) -> list[DatasetId]:
     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)
+            if not spec.persist_id:
+                raise ValueError("unexpected list persisted when there is no persist id")
+            root = pathlib.Path(spec.storage_params) / spec.persist_id
+            if not root.exists():
+                return [] # we mkdir only at a first persist, so absence of folder is valid emptiness
+            files = (x for x in root.iterdir() if x.is_file())
+            return [DatasetId.des(file.parts[-1]) for file in files]
         case s:
             assert_never(s)
 
+def build_retrieve_command(checkpoint_spec: CheckpointSpec|None, ds: DatasetId, hostId: HostId) -> DatasetRetrieveCommand:
+    if checkpoint_spec is None:
+        raise ValueError(f"unexpected retrieve need when checkpoint storage not configured")
+    id_ = checkpoint_spec.retrieve_id
+    if not id_:
+        raise ValueError(f"unexpected retrieve when there is no retrive id")
+    retrieve_params = serialize_params(checkpoint_spec, id_)
+    return DatasetRetrieveCommand(
+        target=hostId,
+        ds=ds,
+        storage_type=checkpoint_spec.storage_type,
+        retrieve_params=retrieve_params,
+    )
+
+def retrieve_dataset(command: DatasetRetrieveCommand) -> AllocatedBuffer:
+    match command.storage_type:
+        case "fs":
+            shm_key = ds2shmid(command.ds)
+            fpath = pathlib.Path(command.retrieve_params) / command.ds.ser()
+            fd = os.open(fpath, os.O_RDONLY)
+            try:
+                advise_seqread(fd)
+                size = os.fstat(fd).st_size
+                # TODO dont use default serde, get it via the command
+                buf = allocate(shm_key, size, DefaultSerde)
+                # once on 3.14+, replace with this
+                # os.readinto(fd, buf.view())
+                with io.FileIO(fd, closefd=False) as raw_io:
+                    raw_io.readinto(buf.view())
+            finally:
+                os.close(fd)
+            return buf
+        case s:
+            assert_never(s)
+
+def possible_repersist(dataset: DatasetId, checkpointSpec: CheckpointSpec|None) -> None:
+    # NOTE blocking -> unfortunate for controller, but we dont expect this to be frequent/hot.
+    # If needbe, spawn a thread or something. In that case needs a completion callback
+    if not checkpointSpec:
+        raise ValueError(f"unexpected repersist when checkpoint storage not configured")
+    if not checkpointSpec.retrieve_id:
+        raise ValueError(f"unexpected repersist when no retrieve id")
+    if not checkpointSpec.persist_id:
+        raise ValueError(f"unexpected repersist when no persist id")
+
+    if checkpointSpec.retrieve_id == checkpointSpec.persist_id:
+        # we assume reproducibility---bold!---so we better warn about it
+        logger.warning(f"no-op for persist of {dataset} as was already persisted under the same id {checkpointSpec.retrieve_id}")
+        return
+
+    # NOTE the host is the virtual one so the message is not really valid, but no big deal
+    retrieve_command = build_retrieve_command(checkpointSpec, dataset, VirtualCheckpointHost)
+    persist_command = build_persist_command(checkpointSpec, dataset, VirtualCheckpointHost)
+    buffer = retrieve_dataset(retrieve_command)
+    try:
+        persist_dataset(persist_command, buffer)
+    finally:
+        buffer.close()

cascade/executor/data_server.py

@@ -21,7 +21,7 @@ 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.checkpoints import persist_dataset, retrieve_dataset
 from cascade.executor.comms import Listener, callback, send_data
 from cascade.executor.msg import (
     Ack,
@@ -31,6 +31,9 @@ from cascade.executor.msg import (
     DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
+    DatasetRetrieveCommand,
+    DatasetRetrieveFailure,
+    DatasetRetrieveSuccess,
     DatasetTransmitCommand,
     DatasetTransmitFailure,
     DatasetTransmitPayload,
@@ -63,7 +66,7 @@ class DataServer:
         self.cap = 2
         self.ds_proc_tp: PythonExecutor = ThreadPoolExecutor(max_workers=self.cap)
         self.futs_in_progress: dict[
-            DatasetTransmitCommand | DatasetTransmitPayload | DatasetPersistCommand, Future
+            DatasetTransmitCommand | DatasetTransmitPayload | DatasetPersistCommand | DatasetRetrieveCommand, Future
         ] = {}
         self.awaiting_confirmation: dict[int, tuple[DatasetTransmitCommand, int]] = {}
         self.invalid: set[DatasetId] = (
@@ -97,7 +100,7 @@ class DataServer:
                 return
             wait(self.futs_in_progress.values(), return_when=FIRST_COMPLETED)
 
-    def store_payload(self, payload: DatasetTransmitPayload) -> int:
+    def _store_payload(self, payload: DatasetTransmitPayload) -> int:
         try:
             l = len(payload.value)
             try:
@@ -153,11 +156,30 @@ class DataServer:
             time_ns()
         )  # not actually consumed but uniform signature with send_payload simplifies typing
 
-    def persist_payload(self, command: DatasetPersistCommand) -> int:
+    def _retrieve_dataset(self, command: DatasetRetrieveCommand) -> int:
+        buf: None | shm_client.AllocatedBuffer = None
+        try:
+            if command.target != self.host:
+                raise ValueError(f"invalid host in {command=}")
+            buf = retrieve_dataset(command)
+            logger.debug(f"dataset for {command} retrieved")
+            callback(self.maddress, DatasetRetrieveSuccess(host=self.host, ds=command.ds))
+        except Exception as e:
+            logger.exception(f"failed to retrieve dataset for {command}, reporting up")
+            callback(
+                self.maddress,
+                DatasetRetrieveFailure(host=self.host, detail=f"{repr(command)} -> {repr(e)}"),
+            )
+        finally:
+            if buf is not None:
+                buf.close()
+        return time_ns()
+
+    def _persist_dataset(self, command: DatasetPersistCommand) -> int:
         buf: None | shm_client.AllocatedBuffer = None
         try:
             if command.source != self.host:
-                raise ValueError(f"invalid {command=}")
+                raise ValueError(f"invalid host in {command=}")
             buf = shm_client.get(key=ds2shmid(command.ds))
             persist_dataset(command, buf)
             logger.debug(f"dataset for {command} persisted")
@@ -173,7 +195,7 @@ class DataServer:
                 buf.close()
         return time_ns()
 
-    def send_payload(self, command: DatasetTransmitCommand) -> int:
+    def _send_payload(self, command: DatasetTransmitCommand) -> int:
         buf: None | shm_client.AllocatedBuffer = None
         payload: None | DatasetTransmitPayload = None
         try:
@@ -240,7 +262,7 @@ class DataServer:
                             }
                         )
                         self.awaiting_confirmation[m.idx] = (m, -1)
-                        fut = self.ds_proc_tp.submit(self.send_payload, m)
+                        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:
@@ -248,7 +270,11 @@ class DataServer:
                                 f"unexpected persist command {m} as the dataset was already purged"
                             )
                         # TODO mark?
-                        fut = self.ds_proc_tp.submit(self.persist_payload, m)
+                        fut = self.ds_proc_tp.submit(self._persist_dataset, m)
+                        self.futs_in_progress[m] = fut
+                    elif isinstance(m, DatasetRetrieveCommand):
+                        # TODO mark?
+                        fut = self.ds_proc_tp.submit(self._retrieve_dataset, m)
                         self.futs_in_progress[m] = fut
                     elif isinstance(m, DatasetTransmitPayload):
                         if m.header.ds in self.invalid:
@@ -263,7 +289,7 @@ class DataServer:
                                 "target": self.host,
                             }
                         )
-                        fut = self.ds_proc_tp.submit(self.store_payload, m)
+                        fut = self.ds_proc_tp.submit(self._store_payload, m)
                         self.futs_in_progress[m] = fut
                     elif isinstance(m, Ack):
                         logger.debug(f"confirmed transmit {m.idx}")
@@ -273,7 +299,7 @@ class DataServer:
                         # 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|DatasetPersistCommand):
+                            if isinstance(commandProg, DatasetTransmitCommand|DatasetPersistCommand|DatasetRetrieveCommand):
                                 val = commandProg.ds
                             elif isinstance(commandProg, DatasetTransmitPayload):
                                 val = commandProg.header.ds
@@ -323,7 +349,7 @@ class DataServer:
                         self.awaiting_confirmation.pop(e)
                     else:
                         logger.warning(f"submitting a retry of {command}")
-                        fut = self.ds_proc_tp.submit(self.send_payload, command)
+                        fut = self.ds_proc_tp.submit(self._send_payload, command)
                         self.futs_in_progress[command] = fut
                         self.awaiting_confirmation[e] = (command, -1)
             except:

cascade/executor/executor.py

@@ -39,6 +39,8 @@ from cascade.executor.msg import (
     DatasetPersistSuccess,
     DatasetPublished,
     DatasetPurge,
+    DatasetRetrieveFailure,
+    DatasetRetrieveSuccess,
     DatasetTransmitFailure,
     ExecutorExit,
     ExecutorFailure,
@@ -60,6 +62,8 @@ from cascade.shm.server import entrypoint as shm_server
 logger = logging.getLogger(__name__)
 heartbeat_grace_ms = 2 * comms_default_timeout_ms
 
+# messages from the data server which need to go to controller, but have no additional logic here
+JustForwardToController = DatasetTransmitFailure|DatasetPersistSuccess|DatasetPersistFailure|DatasetRetrieveFailure
 
 def address_of(port: int) -> BackboneAddress:
     return f"tcp://{platform.get_bindabble_self()}:{port}"
@@ -308,7 +312,12 @@ class Executor:
                             callback(worker_address(worker), m)
                         self.datasets.add(m.ds)
                         self.to_controller(m)
-                    elif isinstance(m, DatasetTransmitFailure|DatasetPersistSuccess|DatasetPersistFailure):
+                    elif isinstance(m, DatasetRetrieveSuccess):
+                        availability_notification = DatasetPublished(ds=m.ds, origin=self.host, transmit_idx=None)
+                        for worker in self.workers:
+                            callback(worker_address(worker), availability_notification)
+                        self.to_controller(m)
+                    elif isinstance(m, JustForwardToController):
                         self.to_controller(m)
                     else:
                         # NOTE transmit and store are handled in DataServer (which has its own socket)

cascade/executor/msg.py

@@ -138,6 +138,22 @@ class DatasetPersistSuccess:
     host: HostId
     ds: DatasetId
 
+@dataclass(frozen=True)
+class DatasetRetrieveCommand:
+    target: HostId
+    ds: DatasetId
+    storage_type: CheckpointStorageType
+    retrieve_params: str # storage-type-specific serialization of params
+
+@dataclass(frozen=True)
+class DatasetRetrieveFailure:
+    host: HostId
+    detail: str
+
+@dataclass(frozen=True)
+class DatasetRetrieveSuccess:
+    host: HostId
+    ds: DatasetId
 
 @dataclass(frozen=True)
 class ExecutorFailure:
@@ -197,6 +213,9 @@ Message = (
     | DatasetPersistCommand
     | DatasetPersistFailure
     | DatasetPersistSuccess
+    | DatasetRetrieveCommand
+    | DatasetRetrieveFailure
+    | DatasetRetrieveSuccess
     | ExecutorFailure
     | ExecutorExit
     | ExecutorRegistration