feat: store input chunks in state

src/exo/api/main.py

@@ -254,7 +254,6 @@ class API:
         self.node_id: NodeId = node_id
         self.last_completed_election: int = 0
         self.port = port
-        self._sent_image_hashes: set[str] = set()
 
         self.paused: bool = False
         self.paused_ev: anyio.Event = anyio.Event()
@@ -304,7 +303,6 @@ class API:
         self.event_receiver.close()
         self.event_receiver = event_receiver
         self._tg.start_soon(self._apply_state)
-        self._sent_image_hashes = set()
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -826,18 +824,8 @@ class API:
         )
         command = TextGeneration(task_params=task_params)
 
-        new_images: list[tuple[int, str]] = []
-        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
-            if h not in self._sent_image_hashes:
-                self._sent_image_hashes.add(h)
-                new_images.append((idx, img))
-
-        if not new_images:
-            await self._send(command)
-            return command
-
         all_chunks: list[tuple[int, str]] = []
-        for img_idx, img_data in new_images:
+        for img_idx, img_data in enumerate(images):
             for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                 all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
 

src/exo/shared/apply.py

@@ -40,7 +40,14 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import (
+    ImageEdits,
+    ImageGeneration,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -72,7 +79,6 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
-            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
             | CustomModelCardAdded()
@@ -93,6 +99,8 @@ def event_apply(event: Event, state: State) -> State:
             return apply_runner_status_updated(event, state)
         case TaskCreated():
             return apply_task_created(event, state)
+        case InputChunkReceived():
+            return apply_input_chunk_received(event, state)
         case TaskDeleted():
             return apply_task_deleted(event, state)
         case TaskFailed():
@@ -157,10 +165,32 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
     return state.model_copy(update={"tasks": new_tasks})
 
 
+def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
+    command_chunks = {
+        **state.input_chunks.get(event.command_id, {}),
+        event.chunk.chunk_index: event.chunk,
+    }
+    return state.model_copy(
+        update={
+            "input_chunks": {**state.input_chunks, event.command_id: command_chunks}
+        }
+    )
+
+
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
+    task = state.tasks.get(event.task_id)
     new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
+    if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+        new_input_chunks = {
+            command_id: chunks
+            for command_id, chunks in state.input_chunks.items()
+            if command_id != task.command_id
+        }
+        return state.model_copy(
+            update={"tasks": new_tasks, "input_chunks": new_input_chunks}
+        )
     return state.model_copy(update={"tasks": new_tasks})
 
 

src/exo/shared/tests/test_apply/test_apply_input_chunks.py

@@ -0,0 +1,85 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.chunks import InputImageChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import (
+    IndexedEvent,
+    InputChunkReceived,
+    TaskCreated,
+    TaskDeleted,
+)
+from exo.shared.types.state import State
+from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
+from exo.shared.types.text_generation import (
+    InputMessage,
+    InputMessageContent,
+    TextGenerationTaskParams,
+)
+from exo.shared.types.worker.instances import InstanceId
+
+
+def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
+    command_id = CommandId("command-1")
+    chunk = InputImageChunk(
+        model=ModelId("mlx-community/test-model"),
+        command_id=command_id,
+        data="abc",
+        chunk_index=0,
+        total_chunks=1,
+        image_index=0,
+    )
+
+    state = apply(
+        State(),
+        IndexedEvent(
+            idx=0,
+            event=InputChunkReceived(command_id=command_id, chunk=chunk),
+        ),
+    )
+
+    assert state.input_chunks == {command_id: {0: chunk}}
+
+
+def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
+    command_id = CommandId("command-1")
+    task_id = TaskId("task-1")
+    chunk = InputImageChunk(
+        model=ModelId("mlx-community/test-model"),
+        command_id=command_id,
+        data="abc",
+        chunk_index=0,
+        total_chunks=1,
+        image_index=0,
+    )
+    task = TextGeneration(
+        task_id=task_id,
+        instance_id=InstanceId("instance-1"),
+        task_status=TaskStatus.Pending,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=ModelId("mlx-community/test-model"),
+            input=[
+                InputMessage(role="user", content=InputMessageContent("hello")),
+            ],
+        ),
+    )
+
+    state = State()
+    state = apply(
+        state,
+        IndexedEvent(
+            idx=0,
+            event=InputChunkReceived(command_id=command_id, chunk=chunk),
+        ),
+    )
+    state = apply(
+        state,
+        IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
+    )
+    state = apply(
+        state,
+        IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
+    )
+
+    assert state.tasks == {}
+    assert state.input_chunks == {}

src/exo/shared/types/state.py

@@ -6,7 +6,8 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId
+from exo.shared.types.chunks import InputImageChunk
+from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -45,6 +46,9 @@ class State(FrozenModel):
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
+    # Durable request input chunks for active image requests. Workers rebuild
+    # local image caches from this state instead of reading events directly.
+    input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)

src/exo/worker/main.py

@@ -24,7 +24,6 @@ from exo.shared.types.events import (
     CustomModelCardDeleted,
     Event,
     IndexedEvent,
-    InputChunkReceived,
     InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
@@ -141,37 +140,6 @@ class Worker:
                 if isinstance(event, InstanceDeleted):
                     self._instance_backoff.reset(event.instance_id)
 
-                # Buffer input image chunks for image editing
-                if isinstance(event, InputChunkReceived):
-                    cmd_id = event.command_id
-                    if cmd_id not in self.input_chunk_buffer:
-                        self.input_chunk_buffer[cmd_id] = {}
-                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
-
-                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
-                        event.chunk
-                    )
-
-                    if (
-                        len(self.input_chunk_buffer[cmd_id])
-                        == self.input_chunk_counts[cmd_id]
-                    ):
-                        per_image: defaultdict[int, list[InputImageChunk]] = (
-                            defaultdict(list)
-                        )
-                        for chunk in self.input_chunk_buffer[cmd_id].values():
-                            per_image[chunk.image_index].append(chunk)
-                        for chunks_for_image in per_image.values():
-                            sorted_chunks = sorted(
-                                chunks_for_image, key=lambda c: c.chunk_index
-                            )
-                            img = Base64Image("".join(c.data for c in sorted_chunks))
-                            self.image_cache[
-                                Base64ImageHash(
-                                    hashlib.sha256(img.encode("ascii")).hexdigest()
-                                )
-                            ] = img
-
                 if isinstance(event, CustomModelCardAdded):
                     await event.model_card.save_to_custom_dir()
                     add_to_card_cache(event.model_card)
@@ -179,6 +147,35 @@ class Worker:
                 if isinstance(event, CustomModelCardDeleted):
                     await delete_custom_card(event.model_id)
 
+                self._sync_input_views_from_state()
+
+    def _sync_input_views_from_state(self) -> None:
+        self.input_chunk_buffer = {
+            command_id: dict(chunks)
+            for command_id, chunks in self.state.input_chunks.items()
+        }
+        self.input_chunk_counts = {
+            command_id: next(iter(chunks.values())).total_chunks
+            for command_id, chunks in self.input_chunk_buffer.items()
+            if chunks
+        }
+
+        self.image_cache = {}
+        for command_id, chunks in self.input_chunk_buffer.items():
+            expected_chunks = self.input_chunk_counts.get(command_id)
+            if expected_chunks is None or len(chunks) != expected_chunks:
+                continue
+
+            per_image: defaultdict[int, list[InputImageChunk]] = defaultdict(list)
+            for chunk in chunks.values():
+                per_image[chunk.image_index].append(chunk)
+            for chunks_for_image in per_image.values():
+                sorted_chunks = sorted(chunks_for_image, key=lambda c: c.chunk_index)
+                image = Base64Image("".join(chunk.data for chunk in sorted_chunks))
+                self.image_cache[
+                    Base64ImageHash(hashlib.sha256(image.encode("ascii")).hexdigest())
+                ] = image
+
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
@@ -189,7 +186,7 @@ class Worker:
                 self.state.instances,
                 self.state.runners,
                 self.state.tasks,
-                self.input_chunk_buffer,
+                self.state.input_chunks,
                 self.image_cache,
                 self._instance_backoff,
                 self._download_backoff,
@@ -321,15 +318,9 @@ class Worker:
                             advanced_params=task.task_params.advanced_params,
                         ),
                     )
-                    # Cleanup buffers
-                    if cmd_id in self.input_chunk_buffer:
-                        del self.input_chunk_buffer[cmd_id]
-                    if cmd_id in self.input_chunk_counts:
-                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
 
                 case TextGeneration() if task.task_params.image_hashes:
-                    cmd_id = task.command_id
                     resolved_images = [
                         self.image_cache[h]
                         for _, h in sorted(task.task_params.image_hashes.items())
@@ -341,10 +332,6 @@ class Worker:
                             )
                         }
                     )
-                    if cmd_id in self.input_chunk_buffer:
-                        del self.input_chunk_buffer[cmd_id]
-                    if cmd_id in self.input_chunk_counts:
-                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
                 case LoadModel(instance_id=instance_id):
                     if (instance := self.state.instances.get(instance_id)) is not None: