← back to Exo
Keep image cache fresh (temporary fix) (#1961)
290e3fd927e2264c4ea2b15cf62baee9f9cf584b · 2026-04-23 11:39:36 +0100 · rltakashige
## Motivation
When a new node joins, it might not have the cache.
Caveat:
This is potentially fallible if a new node joins and updates real
topology, but the API topology hasn't caught up with this fact and the
user queues up a new text generation. In practice, there is only a split
second where this is the case, and this is only for users of the
dashboard interface. We should fix this properly after the release.
Files touched
M src/exo/api/main.pyM src/exo/shared/types/text_generation.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/tests/unittests/test_plan/test_download_and_loading.pyM src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.pyM src/exo/worker/tests/unittests/test_plan/test_task_forwarding.pyM src/exo/worker/tests/unittests/test_plan/test_warmup.py
Diff
commit 290e3fd927e2264c4ea2b15cf62baee9f9cf584b
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Apr 23 11:39:36 2026 +0100
Keep image cache fresh (temporary fix) (#1961)
## Motivation
When a new node joins, it might not have the cache.
Caveat:
This is potentially fallible if a new node joins and updates real
topology, but the API topology hasn't caught up with this fact and the
user queues up a new text generation. In practice, there is only a split
second where this is the case, and this is only for users of the
dashboard interface. We should fix this properly after the release.
---
src/exo/api/main.py | 35 +++++--------
src/exo/shared/types/text_generation.py | 2 -
src/exo/worker/main.py | 58 +++++++++-------------
src/exo/worker/plan.py | 25 ++++++----
.../test_plan/test_download_and_loading.py | 5 ++
.../unittests/test_plan/test_runner_lifecycle.py | 5 ++
.../unittests/test_plan/test_task_forwarding.py | 5 ++
.../tests/unittests/test_plan/test_warmup.py | 8 +++
8 files changed, 75 insertions(+), 68 deletions(-)
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 400ff9ab..37557f5b 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -185,7 +185,10 @@ from exo.shared.types.tasks import (
from exo.shared.types.tasks import (
TextGeneration as TextGenerationTask,
)
-from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+ Base64ImageHash,
+ TextGenerationTaskParams,
+)
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
@@ -234,6 +237,7 @@ 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()
@@ -283,6 +287,7 @@ 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")
@@ -737,8 +742,6 @@ class API:
"TODO: we should send a notification to the user to download the model"
)
- _sent_image_hashes: set[str] = set()
-
async def _send_text_generation_with_images(
self, task_params: TextGenerationTaskParams
) -> TextGeneration:
@@ -750,23 +753,19 @@ class API:
return command
hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
+ all_hashes = {idx: Base64ImageHash(h) for idx, h in enumerate(hashes)}
+ task_params = task_params.model_copy(
+ update={"images": [], "image_hashes": all_hashes}
+ )
+ command = TextGeneration(task_params=task_params)
- cached_hashes: dict[int, str] = {}
new_images: list[tuple[int, str]] = []
for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
- if h in self._sent_image_hashes:
- cached_hashes[idx] = h
- else:
+ if h not in self._sent_image_hashes:
self._sent_image_hashes.add(h)
new_images.append((idx, img))
- wrapped_hashes = {idx: Base64Image(h) for idx, h in cached_hashes.items()}
-
if not new_images:
- task_params = task_params.model_copy(
- update={"images": [], "image_hashes": wrapped_hashes}
- )
- command = TextGeneration(task_params=task_params)
await self._send(command)
return command
@@ -775,16 +774,6 @@ class API:
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]))
- task_params = task_params.model_copy(
- update={
- "images": [],
- "image_hashes": wrapped_hashes,
- "total_input_chunks": len(all_chunks),
- "image_count": len(new_images),
- }
- )
- command = TextGeneration(task_params=task_params)
-
for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
await self._send(
SendInputChunk(
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index dc763248..1d3fd526 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -114,8 +114,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
frequency_penalty: float | None = None
images: list[Base64Image] = Field(default_factory=list)
image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
- total_input_chunks: int = 0
- image_count: int = 0
def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
from exo.shared.models.model_cards import get_card
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index f983b9ca..c676c85c 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -152,6 +152,26 @@ class Worker:
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)
@@ -170,6 +190,7 @@ class Worker:
self.state.runners,
self.state.tasks,
self.input_chunk_buffer,
+ self.image_cache,
self._instance_backoff,
self._download_backoff,
)
@@ -307,42 +328,11 @@ class Worker:
del self.input_chunk_counts[cmd_id]
await self._start_runner_task(modified_task)
- case TextGeneration() if (
- task.task_params.image_hashes
- or task.task_params.total_input_chunks > 0
- ):
+ case TextGeneration() if task.task_params.image_hashes:
cmd_id = task.command_id
- by_index: dict[int, Base64Image] = {}
-
- for idx, h in task.task_params.image_hashes.items():
- assert h in self.image_cache
- by_index[idx] = self.image_cache[h]
-
- if task.task_params.total_input_chunks > 0:
- chunk_buffer = self.input_chunk_buffer.get(cmd_id, {})
- per_image: defaultdict[int, list[InputImageChunk]] = (
- defaultdict(list)
- )
- for chunk in chunk_buffer.values():
- per_image[chunk.image_index].append(chunk)
- for img_idx in sorted(per_image):
- sorted_chunks = sorted(
- per_image[img_idx], 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
- by_index[img_idx] = img
- logger.info(
- f"Assembled {len(per_image)} VLM image(s) "
- f"from {len(chunk_buffer)} chunks"
- )
-
resolved_images = [
- Base64Image(by_index[i]) for i in sorted(by_index)
+ self.image_cache[h]
+ for _, h in sorted(task.task_params.image_hashes.items())
]
modified_task = task.model_copy(
update={
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 07aeeab8..1ca2cfaf 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -19,6 +19,7 @@ from exo.shared.types.tasks import (
TaskStatus,
TextGeneration,
)
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
from exo.shared.types.worker.downloads import (
DownloadCompleted,
DownloadFailed,
@@ -52,6 +53,7 @@ def plan(
all_runners: Mapping[RunnerId, RunnerStatus], # all global
tasks: Mapping[TaskId, Task],
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+ image_cache: Mapping[Base64ImageHash, Base64Image],
instance_backoff: KeyedBackoff[InstanceId],
download_backoff: KeyedBackoff[ModelId],
) -> Task | None:
@@ -66,7 +68,7 @@ def plan(
or _init_distributed_backend(runners, all_runners)
or _load_model(runners, all_runners, global_download_status)
or _ready_to_warmup(runners, all_runners)
- or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
+ or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
)
@@ -300,6 +302,7 @@ def _pending_tasks(
tasks: Mapping[TaskId, Task],
all_runners: Mapping[RunnerId, RunnerStatus],
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+ image_cache: Mapping[Base64ImageHash, Base64Image],
) -> Task | None:
for task in tasks.values():
# for now, just forward chat completions
@@ -309,16 +312,20 @@ def _pending_tasks(
if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
continue
- # For tasks with images, verify all input chunks have been received
- expected_image_chunks = 0
- if isinstance(task, (ImageEdits, TextGeneration)):
- expected_image_chunks = task.task_params.total_input_chunks
- if expected_image_chunks > 0:
- cmd_id = task.command_id
- received = len(input_chunk_buffer.get(cmd_id, {}))
- if received < expected_image_chunks:
+ if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
+ received = len(input_chunk_buffer.get(task.command_id, {}))
+ if received < task.task_params.total_input_chunks:
continue # Wait for all chunks to arrive
+ if (
+ isinstance(task, TextGeneration)
+ and task.task_params.image_hashes
+ and not all(
+ h in image_cache for h in task.task_params.image_hashes.values()
+ )
+ ):
+ continue # Wait for all images to be assembled into the cache
+
for runner in runners.values():
if task.instance_id != runner.bound_instance.instance.instance_id:
continue
diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
index 461d8f12..b574a9b2 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
@@ -54,6 +54,7 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -109,6 +110,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -154,6 +156,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -204,6 +207,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -227,6 +231,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
diff --git a/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py b/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py
index 1ac9dee0..bfde8a1d 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py
@@ -54,6 +54,7 @@ def test_plan_kills_runner_when_instance_missing():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -96,6 +97,7 @@ def test_plan_kills_runner_when_sibling_failed():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -130,6 +132,7 @@ def test_plan_creates_runner_when_missing_for_node():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -171,6 +174,7 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -203,6 +207,7 @@ def test_plan_does_not_create_runner_for_unassigned_node():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
diff --git a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py b/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
index bb87268f..152fc6e1 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py
@@ -78,6 +78,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
all_runners=all_runners,
tasks={TASK_1_ID: task},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -131,6 +132,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
all_runners=all_runners,
tasks={TASK_1_ID: task},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -181,6 +183,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
all_runners=all_runners,
tasks={foreign_task.task_id: foreign_task},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -249,6 +252,7 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
all_runners=all_runners,
tasks={TASK_1_ID: completed_task, other_task_id: other_task},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -291,6 +295,7 @@ def test_plan_returns_none_when_nothing_to_do():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
diff --git a/src/exo/worker/tests/unittests/test_plan/test_warmup.py b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
index 612868e7..46e372f6 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
@@ -63,6 +63,7 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -107,6 +108,7 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -150,6 +152,7 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -197,6 +200,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -216,6 +220,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -262,6 +267,7 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -307,6 +313,7 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
@@ -351,6 +358,7 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
+ image_cache={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
← 3894cf13 Fix Gemma 4 E2B TP + DeepSeek V32 thinking parsing (#1967)
·
back to Exo
·
chore(app): hardcode bug report presigned-URL endpoint (#197 45248c5c →