hi @ustulation
great analysis @ustulation - I traced it through the runtime source (pregel/_loop.py, _algo.py, cache/base) and your model is mostly right, but the missing piece you want turns out to be structural. Point by point:
GOOD: Your 2-level model is correct, and level-1 wins over level-2. Replay re-attaches a completed task’s writes and the runner only executes tasks where not t.writes; the node cache (match_cached_writes) is consulted only for tasks still empty after replay:
# pregel/_loop.py
cached = {(t.cache_key.ns, t.cache_key.key): t
for t in self.tasks.values()
if t.cache_key and not t.writes} # ← cache only for non-replayed tasks
GOOD: The new-checkpoint trick works. Passing non-None input takes the loop’s “fresh input” branch (discard any unfinished tasks from previous checkpoint), so the new checkpoint has no pending writes → every task is “considered” → the cache decides. Sound.
CAUTION: But key_func only sees the task’s input args - no thread_id/run_id:
# pregel/_algo.py
args_key = cache_policy.key_func(*call.input[0], call.input[1])
So {thread-id + cache-run-id + task-name} has to be passed in as task arguments, and “flip the run-id at task-x” is a per-task decision your entrypoint makes - i.e. still app-level control flow.
WARNING: The “also cache under another key” feature isn’t expressible today - and not just unimplemented. The loop computes one cache_key per task and uses it for both cache.get and cache.set (write happens only on a real miss - put_writes), so there’s no copy-forward. Worse, the stored key is an opaque hash (xxh3_128_hexdigest(key_func(input))), and BaseCache.clear() deletes by namespace = function identity only - so even a custom BaseCache can’t decompose keys to alias old→new run-ids, and you can’t surgically evict one (thread, run, task) entry. The node cache is a pure input-keyed memoizer, not a run-versioned store - your messiness on the 3rd restart is fundamental, not a bug. Worth filing as a feature request, but it’d need a core change (decouple read-key from write-key).
WARNING: The rename hack: task_id = hash(checkpoint_id, ns, step, node, PUSH, idx) - name is a component so renaming does invalidate, but so do step/idx/checkpoint_id. It’s an internal id contract; your “might change in future” instinct is right. Skip it.
GOOD: What actually hits your PS goal (auto-capture, no manual per-task state, composes over N restarts): use entrypoint.final + previous to persist a {task_name: output} map inside the checkpoint itself:
@entrypoint(checkpointer=checkpointer)
def workflow(inp, *, previous: dict | None = None):
done = dict(previous or {})
def step(name, fn, *a):
if inp.restart_from and _at_or_after(name, inp.restart_from):
done.pop(name, None)
if name not in done:
done[name] = fn(*a).result()
return done[name]
r0 = step("task0", task0, inp.x)
r1 = step("task1", task1, r0)
r2 = step("task2", task2, r1)
return entrypoint.final(value=step("task3", task3, r2), save=done)
One logical slot per task → no run-id sprawl, no external store, no impl-detail reliance, and it stays type-checkable. If you’d rather use cache_policy, use it as a plain memoizer with a stable key and task_x.clear_cache(cache) to force re-runs (works cleanest single-thread, since clear is per-function-namespace).
(Refs: docs - Functional API for entrypoint.final/previous/determinism; source - pregel/_loop.py, pregel/_algo.py, cache/base.)
