From d6a7850d9c02f2790f2da7d870e6d4a5b390fc21 Mon Sep 17 00:00:00 2001 From: Endri Bezati Date: Thu, 1 Oct 2026 19:16:27 +0200 Subject: [PATCH] runner: a run is resumed at any step of its queue trace, by replaying a firing journal `wfpy run flow.py --resume-from wf-out/ --at-step N` carries a run on from the step the IDE's stepper shows, not only from where it failed. Every firing that completes now appends to `run.wf-journal.jsonl` what it took (by digest), what it put and the actor's state after; a resume at step N starts the run again from its inputs and replays each actor's firings up to the step instead of running them, then carries on live. Replaying per actor rather than restoring the state at step N is what makes it sound with parallel workers: a step is recorded while other firings are in flight, so the state "after step N" can hold half of one, but each actor sees the same sequence of tokens however firings interleave, so its k-th firing can be replayed whenever its inputs come. Which firings a step covers is chosen per actor too -- one journal cut-off races the same way. A replayed firing checks its inputs against the original's; when they differ, replay stops (`replayStopped` in the run record) and the run carries on live. The queue trace is written when a run fails too, and its steps carry `journalSeq` and `replayed`. Phase 3 of docs/proposals/resume.md, whose section is rewritten to the design built. Claude-Session: https://claude.ai/code/session_015VK7fH1c4aKbexnq2QcuKU --- CHANGELOG.md | 10 + CLAUDE.md | 5 +- docs/proposals/resume.md | 170 +++++++------ src/wfpy/_checkpoint_runtime.py | 52 ++-- src/wfpy/_journal_runtime.py | 298 +++++++++++++++++++++++ src/wfpy/_run_artifacts.py | 8 + src/wfpy/_run_finalization.py | 27 ++- src/wfpy/cli.py | 9 + src/wfpy/runner.py | 145 ++++++++++- tests/test_resume_at_step.py | 409 ++++++++++++++++++++++++++++++++ 10 files changed, 1017 insertions(+), 116 deletions(-) create mode 100644 src/wfpy/_journal_runtime.py create mode 100644 tests/test_resume_at_step.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a423522..a17dbb8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,16 @@ All notable changes to wfpy are documented here. The format follows describe. A value neither can write (a lock, an open file, a local class) makes the checkpoint not resumable, naming it. The resumed run records `resumedFrom`. +- Resuming a run at any step of its queue trace: + `wfpy run flow.py --resume-from wf-out/ --at-step N`. Every firing is + journalled in `run.wf-journal.jsonl` (what it took, by digest, what it put, + and the actor's state after). A resume at step N starts again from the run's + inputs and replays each actor's firings up to the step instead of running + them, then carries on live. Replaying per actor is sound with parallel + workers, where restoring a snapshot taken at a step is not. A replayed firing + whose inputs differ from the original's stops the replay (`replayStopped` in + the run record). Trace steps carry `journalSeq` and `replayed`, and the queue + trace is now written when a run fails too. ### Fixed - A run whose tokens are arbitrary objects no longer fails writing its overlay diff --git a/CLAUDE.md b/CLAUDE.md index 81337c3..d27da6a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -63,6 +63,8 @@ src/wfpy/ ├── _run_finalization.py Run finalization (overlays, artifacts, run record) ├── _checkpoint_runtime.py Run checkpoints: capture/restore a plan's state, │ run.wf-checkpoint.json, resume_from +├── _journal_runtime.py Firing journal (run.wf-journal.jsonl) + per-actor +│ replay for resume_from + at_step ├── _workflow_instance.py Workflow instances (child(instance=...)) + run discovery ├── _context_runtime.py Shared-context (@context) runtime ├── _lsp_client.py LSP client for agent LSP integration @@ -90,7 +92,8 @@ Public API is re-exported from `src/wfpy/__init__.py`. - `cli._load_module()` prepends the workflow module dir, its package parents, and the nearest project root to `sys.path`. Prefer fixing that logic over adding per-example `sys.path` hacks. - `RewriteEngine.export_workflow_graph()` resolves nested child workflows recursively, including `from ... import ...` imports, with recursion guards. - `run()` defaults `queue_trace=True`; writes `run.wf-queues.json`, `run.wf-run.json`, and appends `run-log.jsonl`. `keep_intermediates=True` copies `work/` into the run output. -- A run that fails writes `run.wf-checkpoint.json` (queues, actor/control state, sub-plans, context); `run(resume_from=...)` / `wfpy run --resume-from` restores it into a freshly built plan. A firing that raises gives its tokens back (`_atomic_firing`); nested workflows and if/loop are composite, not wrapped, and mark an interrupted firing in `RuntimeActor.pending`. New per-actor runtime state must be added to `_checkpoint_runtime.capture/restore` or a resume silently loses it. +- A run that fails writes `run.wf-checkpoint.json` (queues, actor/control state, sub-plans, context); `run(resume_from=...)` / `wfpy run --resume-from` restores it into a freshly built plan. A firing that raises gives its tokens back (`_atomic_firing`); nested workflows and if/loop are composite, not wrapped, and mark an interrupted firing in `RuntimeActor.pending`. New per-actor runtime state must be added to `_checkpoint_runtime.actor_state/restore_actor_state` (shared by the checkpoint and the journal) or a resume silently loses it. +- Every wrapped firing is journalled (`run.wf-journal.jsonl`, on whenever `queue_trace` is); `run(resume_from=..., at_step=N)` replays each actor's first firings from it (`Journal.replay_firing`, called from `_step_actor`), chosen per actor by `_journal_runtime.entries_at_step` — not by one journal cut-off, which races with parallel workers. Token writes during a firing must go through `Queue.enqueue`/`dequeue` so the journal sees them. - CLI `--agent-tools` auto-discovers `agent-tools.json` beside the workflow file. Without the flag, the registry is not auto-loaded. - `wfpy plan --format graph --best-effort` falls back to the AST partial-graph builder in `partial.py`; plain `plan` does not. - Do not use system `/tmp` for repo work or scratch — use the repo-local `.tmp/` directory. diff --git a/docs/proposals/resume.md b/docs/proposals/resume.md index ed2f4ea..f8773ea 100644 --- a/docs/proposals/resume.md +++ b/docs/proposals/resume.md @@ -1,6 +1,6 @@ # Proposal: resuming a run -**Status:** proposal — phases 1 (atomic firings) and 2 (checkpoint on failure, `--resume-from`) are implemented; the queue trace and the stepper are not. +**Status:** proposal — phases 1 to 3 (atomic firings, checkpoint on failure, the firing journal and `--at-step`) are implemented; the stepper's resume button is not. **Affects:** wfpy (the runtime and the queue trace); dialogram (the queue-trace stepper gets a resume button); wfpy-ide (one command id). @@ -236,71 +236,85 @@ File tokens keep pointing into the earlier run's `work/` directory. That run has to stay where it is, which it does: nothing deletes it. `--copy-work` copies it into the new run when it should not have to. -## Resuming from a completed step, through the queue trace +## Resuming from a completed step: replaying a firing journal -Restoring the state **after** firing N needs that state, and the failure -checkpoint has only the last one. Writing a full checkpoint after every firing -is too expensive for a run with large tokens. The queue trace is already -written once per firing; it lacks only what changed **in** the firing. +**Implemented** (phase 3), in `_journal_runtime.py`, differently from the +first version of this section. That version extended the queue trace with +each step's consumed and produced tokens, and rebuilt the state after step N +by applying steps 1 to N. That is unsound with parallel workers. A step is +recorded when a firing completes, while other firings are in flight, so the +state "after step N" can hold half of a firing that completes later. A +consumer can even complete, and be recorded, before the firing that produced +its input. Applying steps 1 to N then takes a token that was never put. -### The trace becomes a journal (version 2) +### Replay per actor, not per state -Each step keeps what it has now and gains what the firing did: +Every firing that completes appends one line to `run.wf-journal.jsonl`: ```json -{"step": 12, "actorInstanceName": "review", "actorKind": "agent", "actorFireCount": 3, - "queueSizes": [ "…as today…" ], - "consumed": {"q:draft.Out->review.In": 1}, - "produced": {"q:review.Out->wf:output:Report": [{"$wfType": "file", "path": "work/review__Out__3.md"}]}, - "state": {"_round": 4}, - "runtime": {"scheduleState": "done", "chatHistory": "…appended turns only…"}} +{"seq": 12, "actor": "child/review", "fireCount": 3, + "consumed": [["draft.Out-->review.In", 1, "sha256 of the tokens"]], + "produced": [["review.Out-->WF.Report", [{"$wfType": "file", "path": "…"}]]], + "state": {"_round": 4}, "chatHistory": ["…"]} ``` -- `consumed` is a count per queue. The values are already in an earlier - step's `produced`, or in the workflow inputs, so they are not written twice. -- `produced` is the tokens the firing enqueued, encoded as checkpoint tokens - are. -- `state` and `runtime` are the actor's state after the firing, written only - when it changed. - -A step is recorded under the scheduler's lock, in completion order. The trace -is therefore a serialization of the run, even with parallel workers. - -The state after step N is computed without running anything. Start from the -inputs, then for each step up to N: dequeue `consumed`, enqueue `produced`, -and replace the actor's state. `queueSizes` is kept as it is. The stepper -still reads it, and on resume it is a check: replayed sizes that differ from -the recorded ones mean a corrupt or edited trace, and the resume is refused. - -Two changes to when the trace is written: - -- **On failure too.** Today it is built only in `_finalize_run`. A failed run - is the run most worth stepping through. Its last step is the last firing - that completed, and the error overlay names the one that did not. -- **Appended as the run goes**, to `run.wf-queues.jsonl`, one step per line, - then assembled into `run.wf-queues.json` at the end as now. A run killed - outright still leaves every completed step on disk, and the IDE can step - through a run that is still going. - -Values make the trace larger. Version 2 is on by default only when it stays -under a size (see Open questions). `--queue-trace=sizes` keeps version 1. -Version 1 is still read: a run with a version 1 trace can be stepped through, -but resumed only from its failure checkpoint. - -### Resume - -```bash -wfpy run flows/pipeline.py --resume-from wf-out/ --at-step 12 -``` - -`--at-step` uses the stepper's numbering, so the number a person reads in the -IDE is the number they pass. - -Step N is a firing of the top-level plan. A nested workflow's firing is one -step, as the stepper shows it today: resuming lands before or after a child, -not inside it. Resuming inside a child needs the child's trace too. That is -left for later, since the failure checkpoint already covers a failure inside -a child. +A header line holds the run's inputs and graph fingerprint. Tokens and state +are encoded as checkpoints encode them. + +`wfpy run flow.py --resume-from wf-out/ --at-step N` starts the run +again from that run's inputs. Each actor's first firings, those the run had +completed by step N, are **replayed** instead of run: the recorded tokens are +taken and put, and the recorded state is set. Then the run carries on live. + +This is sound under any scheduling, because in a dataflow network each actor +sees the same sequence of tokens however its firings interleave with others'. +An actor's k-th firing can therefore be replayed whenever its inputs arrive. +It does not matter which other firings happen to be in flight, or which +completed first. + +### What "by step N" means + +Which firings a step covers is chosen per actor, not by one journal cut-off +(that is the same race again: another actor's firing can be journalled +between a firing's entry and its step): + +- an actor the trace shows replays its entries up to its own last step at or + before N; +- an actor inside a nested workflow (path `child/leaf`) follows that + workflow's steps; +- an actor the trace never shows fires only inside an if or a loop, which run + with nothing else in flight, so the cut-off of step N is exact for it. + +Each trace step carries `journalSeq`, the journal's length when it was +recorded, and `replayed: true` when the firing was replayed. + +### When the run goes another way + +A replayed firing checks that the tokens it takes are the ones the original +took (by digest). The firings feeding a replayed one are replayed too, so in +a straight resume they always match. They do not when something upstream runs +live and behaves differently: a firing of the race above, or a journal that +was cut short. Then replay stops, the reason goes into the run record +(`replayStopped`), and the run carries on live from there. Replay also stops if +the run goes quiet with firings left to replay, and at a firing whose values +could not be saved (`notReplayable`). + +What a replayed firing returns is what it returned then. Resuming at step N +keeps everything before N as it happened, including what an agent answered. +A fix to an actor's code applies to its firings after N. + +### Also + +- The queue trace is now written when a run fails too: the last step is the + last firing that completed. +- A resumed run keeps its own journal, replayed firings included, so it can be + resumed at a step in turn. A run resumed from a checkpoint (phase 2) cannot, + because it did not start from inputs. Resume at a step of the run it came + from instead. +- `queue_trace=False` keeps no journal either. +- Not done: the trace is not appended as the run goes (`run.wf-queues.jsonl`). + The journal is, so a run killed outright can still be resumed at a step if + its trace survives, which today it does not. ## The stepper resumes @@ -312,9 +326,11 @@ cluster: ``` **⟲ Resume from here** runs `--resume-from --at-step `, through the profile's run driver as ▶ Run does. It is enabled when the -run's trace is version 2. On a failed run's last step it resumes from the -failure checkpoint, so "fix, then ⟲" is the whole loop. +step>`, through the profile's run driver as ▶ Run does. The number shown is +the trace's `step`, which is the number `--at-step` takes. It is enabled when +the run's trace steps carry `journalSeq`. On a failed run's last step it +replays every firing that completed and runs the failed one live, so "fix, +then ⟲" is the whole loop. On a node, **Rerun from here** finds the step of that node's last firing and resumes from the step before it, so the node fires again. @@ -338,20 +354,16 @@ flags. ## Phasing -1. **Atomic firings.** Peek, run, then commit, for every actor kind. This is - useful alone: a failure no longer loses a token. Tests: each kind fails - once, and its inputs are still queued afterwards. -2. **Checkpoint on failure, and `--resume-from`.** Serialization, graph - fingerprint, restore. Tests: a workflow that fails at a known firing is - fixed and resumed. Firings before the failure do not run again (counted), - and the outputs equal those of a clean run. Also a failure inside a nested - workflow, inside a loop, and in an agent with chat history. -3. **Queue trace version 2.** Written on failure, appended as it goes, - `consumed`, `produced` and state per step, and `--at-step`. Tests: for - every step of a run, the replayed state equals the state captured live at - that step. The stepper keeps working on version 1 and version 2 traces. +1. **Atomic firings.** *Done.* +2. **Checkpoint on failure, and `--resume-from`.** *Done.* +3. **The firing journal, and `--at-step`.** *Done.* Tests resume at every + step of a run, with one worker and with parallel ones. They check that the + firings before the step are replayed and not run, that those after it run, + and that the outputs equal a clean run's. Also covered: a failed run at its + last step, a loop over a child workflow, a resumed run resumed again, + replay stopping when inputs differ, and the CLI. 4. **The stepper's ⟲ button**, and "Rerun from here" on a node (dialogram + - wfpy-ide). + wfpy-ide). The button passes the step it shows to `--at-step`. ## Open questions @@ -362,7 +374,7 @@ flags. stateless agent asked again after a resume may answer differently than it would have the first time. Is that acceptable, or should a resume replay the recorded answer for a firing that is in the journal? -- **Trace size.** Version 2 writes every produced token, so a large non-file - token (a big string, a list) appears inline in the step that produced it. - What threshold sends it to a side file instead, as File outputs already - are? And which default: version 2 always, or only when asked for? +- **Journal size.** The journal writes every produced token, so a large + non-file token (a big string, a list) appears inline in the firing that + produced it. What threshold sends it to a side file instead, as File outputs + already are? It is on whenever the queue trace is (the default). diff --git a/src/wfpy/_checkpoint_runtime.py b/src/wfpy/_checkpoint_runtime.py index f21de88..b0f8e66 100644 --- a/src/wfpy/_checkpoint_runtime.py +++ b/src/wfpy/_checkpoint_runtime.py @@ -299,6 +299,39 @@ def _instance_state(actor: Any) -> dict[str, Any]: return {name: getattr(actor.instance, name) for name in names if hasattr(actor.instance, name)} +def actor_state(actor: Any, enc: Any, at: str) -> dict[str, Any]: + """An actor's own state -- not a control node's -- written with *enc*. + + Its task fields and the runtime attributes wfpy keeps on the instance, and + an agent's budget, conversation and CLI sessions. Shared by the checkpoint + and the firing journal, so both carry the same state. + """ + entry: dict[str, Any] = { + "state": { + name: enc(value, f"{at}.{name}") for name, value in _instance_state(actor).items() + } + } + if actor.agent_fire_budget is not None: + entry["agentFireBudget"] = actor.agent_fire_budget + if actor.chat_history: + entry["chatHistory"] = enc(actor.chat_history, f"{at}.chatHistory") + if actor.agent_cli_session_ids: + entry["agentCliSessionIds"] = dict(actor.agent_cli_session_ids) + return entry + + +def restore_actor_state(actor: Any, entry: dict[str, Any]) -> None: + """Put back what :func:`actor_state` wrote.""" + for field, value in (entry.get("state") or {}).items(): + setattr(actor.instance, field, decode(value)) + if "agentFireBudget" in entry: + actor.agent_fire_budget = entry["agentFireBudget"] + if "chatHistory" in entry: + actor.chat_history = decode(entry["chatHistory"]) + if "agentCliSessionIds" in entry: + actor.agent_cli_session_ids = dict(entry["agentCliSessionIds"]) + + def capture(plan: Any, problems: list[str], where: str = "") -> dict[str, Any]: """The state of *plan* and its sub-plans; what cannot be saved goes to *problems*.""" @@ -333,15 +366,7 @@ def enc(value: Any, at: str) -> Any: control["loopIndex"] = actor._loop_index entry["control"] = control else: - entry["state"] = { - name: enc(value, f"{at}.{name}") for name, value in _instance_state(actor).items() - } - if actor.agent_fire_budget is not None: - entry["agentFireBudget"] = actor.agent_fire_budget - if actor.chat_history: - entry["chatHistory"] = enc(actor.chat_history, f"{at}.chatHistory") - if actor.agent_cli_session_ids: - entry["agentCliSessionIds"] = dict(actor.agent_cli_session_ids) + entry.update(actor_state(actor, enc, at)) if actor.sub_plan is not None: entry["subPlan"] = capture(actor.sub_plan, problems, f"{at}/") actors[actor.name] = entry @@ -379,14 +404,7 @@ def restore(plan: Any, saved: dict[str, Any]) -> None: actor._loop_source = source actor._loop_index = control["loopIndex"] actor._loop_iter = iterator - for field, value in (entry.get("state") or {}).items(): - setattr(actor.instance, field, decode(value)) - if "agentFireBudget" in entry: - actor.agent_fire_budget = entry["agentFireBudget"] - if "chatHistory" in entry: - actor.chat_history = decode(entry["chatHistory"]) - if "agentCliSessionIds" in entry: - actor.agent_cli_session_ids = dict(entry["agentCliSessionIds"]) + restore_actor_state(actor, entry) if "subPlan" in entry and actor.sub_plan is not None: restore(actor.sub_plan, entry["subPlan"]) diff --git a/src/wfpy/_journal_runtime.py b/src/wfpy/_journal_runtime.py new file mode 100644 index 0000000..e7668d2 --- /dev/null +++ b/src/wfpy/_journal_runtime.py @@ -0,0 +1,298 @@ +"""The firing journal: what each firing did, and replaying it. + +Every firing that completes appends one line to ``run.wf-journal.jsonl``: +which actor fired, the tokens it took from each queue (as a digest), the +tokens it put on each queue, and the actor's state afterwards. A run is then +resumable from any step its queue trace shows (``run(resume_from=..., +at_step=N)``): the run starts again from its inputs, and each actor's first +firings -- those the original run had completed by step N -- are replayed from +the journal instead of run. + +Replaying per actor rather than restoring a snapshot is what makes this sound +with parallel workers. A snapshot taken at step N can hold half of a firing +still in flight; but in a dataflow network each actor sees the same sequence of +tokens however its firings interleave with others', so its k-th firing can be +replayed whenever its inputs arrive. A replayed firing checks that the tokens +it takes are the ones the original took. When they are not -- an agent upstream +answered differently -- replay stops there, and the run carries on live. + +See ``docs/proposals/resume.md``. +""" + +from __future__ import annotations + +import collections +import hashlib +import json +import threading +from pathlib import Path +from typing import IO, Any + +from wfpy._checkpoint_runtime import ( + NotResumable, + actor_state, + decode, + encode, + restore_actor_state, +) + +JOURNAL = "run.wf-journal.jsonl" +VERSION = 1 + + +def actor_path(plan: Any, actor: Any) -> str: + """The actor's name, under the nested workflows it is in: ``child/leaf``.""" + return "/".join([*plan.overlay_path_prefix, actor.name]) + + +def _digest(encoded: list[Any]) -> str: + return hashlib.sha256(json.dumps(encoded, sort_keys=True).encode("utf-8")).hexdigest() + + +def _grouped(log: list[tuple[Any, Any]]) -> list[tuple[Any, list[Any]]]: + """The values in *log* per queue, queues in the order first touched.""" + by_queue: dict[int, tuple[Any, list[Any]]] = {} + for queue, value in log: + by_queue.setdefault(id(queue), (queue, []))[1].append(value) + return list(by_queue.values()) + + +class Journal: + """Appends the run's firings, and replays another run's.""" + + def __init__(self, path: Path) -> None: + self.path = path + self.seq = 0 + self._lock = threading.Lock() + self._file: IO[str] | None = path.open("a", encoding="utf-8") + # Replay: actor path -> that actor's firings still to replay, in order. + self._replay: dict[str, collections.deque[dict[str, Any]]] = {} + self.replaying = False + self.replay_stopped: str | None = None + + # ── Writing ─────────────────────────────────────────────────────────── + + def _write(self, line: dict[str, Any]) -> None: + if self._file is None: + return + self._file.write(json.dumps(line) + "\n") + self._file.flush() + + def header( + self, + *, + workflow_name: str, + graph: str, + inputs: dict[str, Any] | None, + resumed_from: str | None = None, + ) -> None: + head: dict[str, Any] = {"version": VERSION, "workflowName": workflow_name, "graph": graph} + if resumed_from is not None: + head["resumedFrom"] = resumed_from + else: + try: + head["inputs"] = encode(inputs or {}, "inputs") + except NotResumable as exc: + head["inputsNotSaved"] = str(exc) + with self._lock: + self._write({"header": head}) + + def record( + self, + plan: Any, + actor: Any, + taken: list[tuple[Any, Any]], + put: list[tuple[Any, Any]], + context_version: int, + *, + replayed: bool = False, + ) -> int: + """Append one completed firing; returns its sequence number.""" + path = actor_path(plan, actor) + entry: dict[str, Any] = {"actor": path, "fireCount": actor.fire_count} + if replayed: + entry["replayed"] = True + try: + entry["consumed"] = [ + [queue.id, len(values), _digest([encode(v, path) for v in values])] + for queue, values in _grouped(taken) + ] + entry["produced"] = [ + [queue.id, [encode(v, f"{path}->{queue.id}") for v in values]] + for queue, values in _grouped(put) + ] + entry.update(actor_state(actor, encode, path)) + if plan.context_version != context_version: + entry["context"] = { + "store": encode(plan.context_store, "context"), + "version": plan.context_version, + "commitSeq": plan.context_commit_seq, + } + except NotResumable as exc: + # Recorded all the same, so the steps keep their numbers; a resume + # replays up to this firing and runs it live. + for key in ("consumed", "produced", "state", "chatHistory", "context"): + entry.pop(key, None) + entry["notReplayable"] = str(exc) + with self._lock: + self.seq += 1 + entry["seq"] = self.seq + self._write(entry) + return self.seq + + def close(self) -> None: + with self._lock: + if self._file is not None: + self._file.close() + self._file = None + + # ── Replaying ───────────────────────────────────────────────────────── + + def load_replay(self, entries: list[dict[str, Any]]) -> None: + """Replay *entries* (another run's firings, in order) as they come due.""" + for entry in entries: + self._replay.setdefault(entry["actor"], collections.deque()).append(entry) + self.replaying = bool(self._replay) + + def pending(self) -> bool: + return any(self._replay.values()) + + def stop_replay(self, reason: str) -> None: + """Run everything live from here; the first reason is kept.""" + if self.replaying: + self.replay_stopped = reason + self.replaying = False + self._replay.clear() + + def replay_firing(self, plan: Any, actor: Any) -> bool | None: + """Replay *actor*'s next recorded firing. + + ``None``: nothing to replay for it, so fire it live. ``False``: its + recorded inputs have not all arrived yet. ``True``: replayed. + """ + due = self._replay.get(actor_path(plan, actor)) + if not due: + return None + entry = due[0] + if "notReplayable" in entry: + self.stop_replay(f"{entry['actor']}: {entry['notReplayable']}") + return None + + in_queues = {q.id: q for queues in actor.in_queues.values() for q in queues} + taking: list[tuple[Any, int]] = [] + for queue_id, count, digest in entry["consumed"]: + queue = in_queues.get(queue_id) + if queue is None: + self.stop_replay(f"{entry['actor']}: no input queue {queue_id}") + return None + if queue.size() < count: + return False + try: + arrived = [encode(v, entry["actor"]) for v in queue.peek(count)] + except NotResumable: + arrived = None + if arrived is None or _digest(arrived) != digest: + self.stop_replay( + f"{entry['actor']} (firing {entry['fireCount']}): its inputs differ " + "from the original run's" + ) + return None + taking.append((queue, count)) + + due.popleft() + taken: list[tuple[Any, Any]] = [] + for queue, count in taking: + for _ in range(count): + taken.append((queue, queue.dequeue())) + + out_queues = {q.id: q for queues in actor.out_queues.values() for q in queues} + put: list[tuple[Any, Any]] = [] + from wfpy._run_artifacts import _record_edge_token + + for queue_id, tokens in entry["produced"]: + queue = out_queues[queue_id] + for token in tokens: + value = decode(token) + queue.enqueue(value) + _record_edge_token(plan, queue, value) + put.append((queue, value)) + + actor.fire_count = entry["fireCount"] + restore_actor_state(actor, entry) + context_version = plan.context_version + context = entry.get("context") + if context is not None: + plan.context_store = decode(context["store"]) + plan.context_version = context["version"] + plan.context_commit_seq = context["commitSeq"] + self.record(plan, actor, taken, put, context_version, replayed=True) + return True + + +def load_journal(run_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """The header and the firings of the journal in *run_dir*.""" + path = run_dir / JOURNAL + if not path.is_file(): + raise FileNotFoundError( + f"no {JOURNAL} in {run_dir}: the run was made without a queue trace, or " + "by a wfpy that did not keep one" + ) + header: dict[str, Any] | None = None + entries: list[dict[str, Any]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + record = json.loads(line) + except ValueError: + break # a line cut short by a run that was killed: the rest is gone + if "header" in record: + header = record["header"] + else: + entries.append(record) + if header is None: + raise ValueError(f"{path} has no header line") + return header, entries + + +def entries_at_step( + run_dir: Path, step: int, entries: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """The firings of *entries* that the run had completed by its step *step*. + + Chosen per actor, not by one cut-off. With parallel workers another + actor's firing can be journalled between a firing's own journal entry and + its step in the trace, so "everything journalled by step N" holds a firing + the trace puts after N. So an actor the trace shows replays its entries up + to its own last step at or before N; an actor inside a nested workflow + follows that workflow's steps (its path starts with the workflow's name). + An actor the trace never shows fires only inside an if or a loop, whose + firings run with nothing else in flight, so the cut-off of step N itself is + exact for it. + """ + if step == 0: + return [] + trace_path = run_dir / "run.wf-queues.json" + if not trace_path.is_file(): + raise FileNotFoundError(f"no run.wf-queues.json in {run_dir}: the run kept no queue trace") + steps = json.loads(trace_path.read_text(encoding="utf-8")).get("steps") or [] + if not 1 <= step <= len(steps): + raise ValueError(f"{run_dir} has no step {step}; its steps are 1 to {len(steps)}") + if any("journalSeq" not in entry for entry in steps[:step]): + raise ValueError( + f"{run_dir}'s trace carries no journal positions: the run was made by a wfpy " + "that did not keep a journal" + ) + + traced = {entry["actorInstanceName"] for entry in steps} + upto: dict[str, int] = {} + for entry in steps[:step]: + upto[entry["actorInstanceName"]] = int(entry["journalSeq"]) + at_step = int(steps[step - 1]["journalSeq"]) + + def due(entry: dict[str, Any]) -> bool: + top = entry["actor"].split("/", 1)[0] + limit = upto.get(top, 0) if top in traced else at_step + return int(entry["seq"]) <= limit + + return [entry for entry in entries if due(entry)] diff --git a/src/wfpy/_run_artifacts.py b/src/wfpy/_run_artifacts.py index d464d7c..18a8649 100644 --- a/src/wfpy/_run_artifacts.py +++ b/src/wfpy/_run_artifacts.py @@ -530,6 +530,10 @@ class QueueTraceStep: actor_kind: str actor_fire_count: int queue_sizes: list[dict[str, Any]] + # How many firings the journal held when this step was recorded: the + # firings a resume at this step replays. None without a journal. + journal_seq: int | None = None + replayed: bool = False @dataclasses.dataclass @@ -554,6 +558,8 @@ def record_fire(self, plan: Any, actor: Any) -> None: actor_kind=actor.kind, actor_fire_count=actor.fire_count, queue_sizes=sizes, + journal_seq=plan.journal.seq if plan.journal is not None else None, + replayed=bool(getattr(actor, "replayed", False)), ) ) @@ -583,6 +589,8 @@ def build(self, finished_at: str, plan: Any) -> dict[str, Any]: "actorKind": step.actor_kind, "actorFireCount": step.actor_fire_count, "queueSizes": step.queue_sizes, + **({"journalSeq": step.journal_seq} if step.journal_seq is not None else {}), + **({"replayed": True} if step.replayed else {}), } for step in self.steps ], diff --git a/src/wfpy/_run_finalization.py b/src/wfpy/_run_finalization.py index 0fc1519..86c980b 100644 --- a/src/wfpy/_run_finalization.py +++ b/src/wfpy/_run_finalization.py @@ -173,6 +173,7 @@ def _persist_run_record( wdir: Path | None = None, error: dict[str, str] | None = None, resumed_from: str | None = None, + extra: dict[str, Any] | None = None, ) -> None: """Write run.wf-run.json and append to run-log.jsonl.""" run_record = { @@ -215,6 +216,7 @@ def _persist_run_record( run_record["error"] = error if resumed_from is not None: run_record["resumedFrom"] = resumed_from + run_record.update(extra or {}) run_record_path = run_out_dir / "run.wf-run.json" try: @@ -238,6 +240,7 @@ def _persist_run_record( log_entry["error"] = error if resumed_from is not None: log_entry["resumedFrom"] = resumed_from + log_entry.update(extra or {}) run_log_path = base_out_dir / "run-log.jsonl" try: @@ -271,6 +274,22 @@ def _write_error_overlay( pass +def _write_queue_trace( + plan: Any, run_out_dir: Path, finished_at: datetime, *, verbose: bool = False +) -> None: + """Write run.wf-queues.json -- for a run that failed as for one that did not. + + A failed run is the run most worth stepping through, and the trace is what + the IDE steps through and what a resume at a step reads. + """ + if plan.queue_trace and plan.queue_trace.steps: + trace_json = plan.queue_trace.build(finished_at.isoformat(), plan) + qt_path = run_out_dir / "run.wf-queues.json" + qt_path.write_text(_runtime_json_dumps(trace_json)) + if verbose: + logger.info("Queue trace: %d steps → %s", len(plan.queue_trace.steps), qt_path) + + def _finalize_run( plan: Any, rid: str, @@ -338,13 +357,7 @@ def _finalize_run( except OSError: pass - # Write queue trace (run.wf-queues.json) - if plan.queue_trace and plan.queue_trace.steps: - trace_json = plan.queue_trace.build(finished_at.isoformat(), plan) - qt_path = run_out_dir / "run.wf-queues.json" - qt_path.write_text(_runtime_json_dumps(trace_json)) - if verbose: - logger.info("Queue trace: %d steps → %s", len(plan.queue_trace.steps), qt_path) + _write_queue_trace(plan, run_out_dir, finished_at, verbose=verbose) if verbose: logger.info("Run %s finished. Output: %s", rid, run_out_dir) diff --git a/src/wfpy/cli.py b/src/wfpy/cli.py index 718ade7..6e45141 100644 --- a/src/wfpy/cli.py +++ b/src/wfpy/cli.py @@ -340,6 +340,7 @@ def cmd_run(args: argparse.Namespace) -> None: context_summarize=getattr(args, "context_summarize", None), resume_context_from=getattr(args, "resume_context_from", None), resume_from=getattr(args, "resume_from", None), + at_step=getattr(args, "at_step", None), ) # Print outputs @@ -711,6 +712,14 @@ def main() -> None: help="Carry on a run that failed or was stopped: its run directory, or " "its run.wf-checkpoint.json. Takes no --input.", ) + run_parser.add_argument( + "--at-step", + type=int, + default=None, + help="With --resume-from: carry on from this step of that run's queue trace " + "(the step the IDE's stepper shows) instead of from where it stopped. " + "The firings up to it are replayed, not run.", + ) run_parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging") run_parser.set_defaults(func=cmd_run) diff --git a/src/wfpy/runner.py b/src/wfpy/runner.py index 93d7a51..4814886 100644 --- a/src/wfpy/runner.py +++ b/src/wfpy/runner.py @@ -44,6 +44,7 @@ import wfpy._validation_runtime as _val_runtime import wfpy._agent_staging_runtime as _staging_runtime import wfpy._checkpoint_runtime as _checkpoint_runtime +import wfpy._journal_runtime as _journal_runtime import wfpy._step_external_runtime as _ext_runtime from wfpy._step_streamblocks_runtime import _step_streamblocks_instance import wfpy._action_runtime as _action_runtime @@ -266,6 +267,12 @@ def _log_taken(queue: "Queue", values: list[Any]) -> None: log.extend((queue, value) for value in values) +def _log_put(queue: "Queue", value: Any) -> None: + log: list[tuple[Queue, Any]] | None = getattr(_firing_log, "put", None) + if log is not None: + log.append((queue, value)) + + class _atomic_firing: """Give a firing's inputs back when it fails. @@ -282,17 +289,20 @@ class _atomic_firing: """ def __enter__(self) -> "_atomic_firing": - self._outer = getattr(_firing_log, "taken", None) + self._outer = (getattr(_firing_log, "taken", None), getattr(_firing_log, "put", None)) _firing_log.taken = [] + _firing_log.put = [] return self def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: - taken: list[tuple[Queue, Any]] = _firing_log.taken - _firing_log.taken = self._outer + # What the firing took and put, in order: the journal records both. + self.taken: list[tuple[Queue, Any]] = _firing_log.taken + self.put: list[tuple[Queue, Any]] = _firing_log.put + _firing_log.taken, _firing_log.put = self._outer # A firing that succeeded keeps what it took, whatever fails after it. if exc_type is None: return - for queue, value in reversed(taken): + for queue, value in reversed(self.taken): queue.give_back(value) @@ -331,6 +341,7 @@ def __init__( def enqueue(self, value: Any) -> None: with self._lock: self.items.append(value) + _log_put(self, value) def dequeue(self) -> Any: with self._lock: @@ -414,6 +425,9 @@ def __init__( # interrupted is still marked, and a resumed run finishes it. self.pending: dict[str, Any] | None = None + # Whether its last firing was replayed from another run's journal. + self.replayed = False + # Conversation history used for UI/debug overlays; only stateful agents # feed it back into subsequent requests. self.chat_history: list[dict[str, str]] = [] @@ -446,6 +460,9 @@ def __init__(self, name: str = "") -> None: self.name = name self.actors: list[RuntimeActor] = [] self.all_queues: list[Queue] = [] + # The run's firing journal, shared with every sub-plan (see + # _journal_runtime). None when the run keeps none. + self.journal: _journal_runtime.Journal | None = None self.scopes: dict[str, ScopeRecord] = {} self.control_nodes: dict[str, ControlNodeRecord] = {} @@ -1106,6 +1123,18 @@ def _submit_ready_actors() -> bool: if not in_flight: if not control_fired and not submitted and not streamed: + # Firings still waiting to be replayed can never come + # due now: the run has gone another way. Carry on live. + journal = plan.journal + if ( + journal is not None + and journal.replaying + and not plan.overlay_path_prefix + and journal.pending() + ): + journal.stop_replay("the run went quiet with firings left to replay") + blocked_until_progress.clear() + continue break # quiescence continue @@ -1198,8 +1227,21 @@ def _step_actor( if actor.pending is not None: return _finish_control(actor, plan, out_dir, verbose, active_scopes) return _dispatch_step(actor, plan, out_dir, verbose, active_scopes) - with _atomic_firing(): - return _dispatch_step(actor, plan, out_dir, verbose, active_scopes) + + journal = plan.journal + actor.replayed = False + if journal is not None and journal.replaying: + replayed = journal.replay_firing(plan, actor) + if replayed is not None: + actor.replayed = replayed + return replayed + + context_version = plan.context_version + with _atomic_firing() as firing: + fired = _dispatch_step(actor, plan, out_dir, verbose, active_scopes) + if fired and journal is not None: + journal.record(plan, actor, firing.taken, firing.put, context_version) + return fired def _dispatch_step( @@ -2547,6 +2589,7 @@ def _run_sub_plan( sub_plan.output_name_suffix = f"__{actor.fire_count}" sub_plan.overlay_writer = parent_plan.overlay_writer + sub_plan.journal = parent_plan.journal sub_plan.overlay_path_prefix = [*parent_plan.overlay_path_prefix, actor.name] sub_plan.agent_context_writer = parent_plan.agent_context_writer sub_plan.context_store = parent_plan.context_store @@ -2645,6 +2688,41 @@ def _step_viewer(actor: RuntimeActor, plan: FifoPlan, verbose: bool) -> bool: # ═══════════════════════════════════════════════════════════════════════════ +def _replay_from_step(plan: FifoPlan, workflow_name: str, source: Path, step: int) -> Any: + """Load the firings *source*'s run had completed by *step* for replay. + + Returns the inputs that run started from: this one starts from them too. + """ + run_dir = source.expanduser().resolve() + if run_dir.is_file(): + run_dir = run_dir.parent + if step < 0: + raise ValueError(f"at_step must be 0 or more, not {step}") + header, entries = _journal_runtime.load_journal(run_dir) + if header.get("workflowName") != workflow_name: + raise ValueError( + f"{run_dir} is a run of {header.get('workflowName')!r}, not {workflow_name!r}" + ) + if header.get("graph") != _checkpoint_runtime.graph_fingerprint(plan): + raise ValueError( + f"{run_dir}: the workflow's graph changed since the run (its actors, their " + "kinds or its connections). An action's code can change before a resume; " + "its wiring cannot." + ) + if "inputs" not in header: + reason = header.get("inputsNotSaved") or ( + f"it was itself resumed from {header.get('resumedFrom')!r}; resume at a step " + "of that run instead" + if header.get("resumedFrom") + else "its inputs were not recorded" + ) + raise ValueError(f"{run_dir} cannot be resumed at a step: {reason}") + if plan.journal is None: + raise ValueError("run(at_step=...) needs the queue trace (queue_trace=True)") + plan.journal.load_replay(_journal_runtime.entries_at_step(run_dir, step, entries)) + return _checkpoint_runtime.decode(header["inputs"]) or None + + def run( target: Any, inputs: dict[str, Any] | None = None, @@ -2690,6 +2768,7 @@ def run( context_summarize: str | bool | None = None, resume_context_from: str | None = None, resume_from: str | None = None, + at_step: int | None = None, context_seed: dict[str, Any] | None = None, elicitation_handler: Any = None, interactive: bool | None = None, @@ -2747,6 +2826,11 @@ def run( ``run.wf-checkpoint.json``). The plan is built from the current source, the run's state is restored into it, and the run carries on from where it stopped; takes no *inputs*. + at_step: With *resume_from*: carry on from this step of that run's queue + trace (the step the IDE's stepper shows), rather than from where + it stopped. The run starts again from that run's inputs, and the + firings the run had completed by the step are replayed from its + journal instead of run. context_seed: Optional initial context seed object. Returns: @@ -2909,21 +2993,47 @@ def run( if verbose: logger.warning("Failed to restore chat histories: %s", e) + # Every firing is journalled, so this run can be resumed from any step. + if queue_trace: + plan.journal = _journal_runtime.Journal(run_out_dir / _journal_runtime.JOURNAL) + # Carry on from where a run stopped: its queues, its actors' state, its - # control nodes and its sub-plans, restored into the plan just built. + # control nodes and its sub-plans, restored into the plan just built. Or, + # at a step, replay the firings it had completed by then. resumed_from: str | None = None + resume_extra: dict[str, Any] = {} + if at_step is not None and not resume_from: + raise ValueError("run(at_step=...) needs resume_from=: the run whose step it is") if resume_from: if inputs: raise ValueError( "run(resume_from=...) takes no inputs: the run's inputs are already " "in its checkpoint, consumed or still queued" ) - checkpoint_path, checkpoint = _checkpoint_runtime.load_checkpoint(resume_from) - _checkpoint_runtime.restore_checkpoint( - plan, wf_def.name or "", checkpoint_path, checkpoint + if at_step is not None: + inputs = _replay_from_step(plan, wf_def.name or "", Path(resume_from), at_step) + resumed_from = Path(resume_from).expanduser().resolve().name + if Path(resume_from).is_file(): + resumed_from = Path(resume_from).expanduser().resolve().parent.name + resume_extra["resumedAtStep"] = at_step + logger.info("Resuming run %s at its step %d", resumed_from, at_step) + else: + checkpoint_path, checkpoint = _checkpoint_runtime.load_checkpoint(resume_from) + _checkpoint_runtime.restore_checkpoint( + plan, wf_def.name or "", checkpoint_path, checkpoint + ) + resumed_from = str(checkpoint.get("runId") or checkpoint_path.parent.name) + logger.info("Resuming run %s from %s", resumed_from, checkpoint_path) + + if plan.journal is not None: + plan.journal.header( + workflow_name=wf_def.name or "", + graph=_checkpoint_runtime.graph_fingerprint(plan), + # A run resumed from a checkpoint starts mid-way: there are no + # inputs to start it again from, so it cannot be resumed at a step. + inputs=None if (resume_from and at_step is None) else inputs, + resumed_from=resumed_from if at_step is None else None, ) - resumed_from = str(checkpoint.get("runId") or checkpoint_path.parent.name) - logger.info("Resuming run %s from %s", resumed_from, checkpoint_path) if verbose: logger.info("Run %s out=%s work=%s", rid, run_out_dir, wdir or "(none)") @@ -2951,6 +3061,14 @@ def _persist_run_record( started_at, finished_at_value, outputs_payload, inputs=inputs, has_external=has_external, wdir=wdir, error=error, resumed_from=resumed_from, + extra={ + **resume_extra, + **( + {"replayStopped": plan.journal.replay_stopped} + if plan.journal is not None and plan.journal.replay_stopped + else {} + ), + }, ) # Resolve max_workers: run() arg → @workflow(...) → executor default @@ -2983,6 +3101,7 @@ def _persist_run_record( _run_finalization._write_error_overlay( plan, writers.overlay_base, overlay_writer, run_out_dir, error_info, finished_at_err, ) + _run_finalization._write_queue_trace(plan, run_out_dir, finished_at_err) _persist_run_record(finished_at_err, {}, error=error_info) # Every firing still in flight has finished (the executor waited for # them), and the one that failed gave its tokens back: the plan is a @@ -3000,6 +3119,8 @@ def _persist_run_record( logger.warning("Could not write the run's checkpoint", exc_info=True) raise finally: + if plan.journal is not None: + plan.journal.close() if run_out_token is None: os.environ.pop("WF_RUN_OUT_DIR", None) else: diff --git a/tests/test_resume_at_step.py b/tests/test_resume_at_step.py new file mode 100644 index 0000000..1640b74 --- /dev/null +++ b/tests/test_resume_at_step.py @@ -0,0 +1,409 @@ +"""Resuming a run at any step of its queue trace (docs/proposals/resume.md, phase 3). + +The run starts again from its inputs, and the firings the original run had +completed by the step are replayed from its journal rather than run. What +these check is that: replayed firings do not run, the ones after the step do, +and the outputs are those a clean run gives -- with parallel workers too. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from wfpy import Port, action, connect, guard, loop, run, task, workflow + + +class Calls: + def __init__(self) -> None: + self.broken = False + self.by_actor: dict[str, int] = {} + self.salt = 0 + + def fired(self, name: str) -> None: + self.by_actor[name] = self.by_actor.get(name, 0) + 1 + + +def _pipeline(calls: Calls) -> Any: + @task + class Count: + _n: int = 0 + + class Ports: + Out = Port[int](direction="out") + + @action(consumes={}, produces={"Out": 1}) + @guard(lambda self: self._n < 4) + def emit(self) -> int: + calls.fired("c") + self._n += 1 + return self._n + + @task + class Pass: + class Ports: + In = Port[int](direction="in") + Out = Port[int](direction="out") + + @action(consumes={"In": 1}, produces={"Out": 1}) + def go(self, x: int) -> int: + calls.fired("p") + return x + calls.salt + + @task + class Double: + _seen: int = 0 + + class Ports: + In = Port[int](direction="in") + Out = Port[int](direction="out") + + @action(consumes={"In": 1}, produces={"Out": 1}) + def go(self, x: int) -> int: + if calls.broken and x == 3: + raise ValueError("broken on 3") + calls.fired("d") + self._seen += 1 + return x * 2 + + @workflow(outputs={"Out": int}) + def flow(): + c = Count() + p = Pass() + d = Double() + connect(c.Out, p.In) + connect(p.Out, d.In) + connect(d.Out, "Out") + + return flow + + +def _trace(run_dir: Path) -> list[dict[str, Any]]: + return json.loads((run_dir / "run.wf-queues.json").read_text())["steps"] + + +def _fired_by(steps: list[dict[str, Any]], step: int) -> dict[str, int]: + fired: dict[str, int] = {} + for entry in steps[:step]: + name = entry["actorInstanceName"] + fired[name] = fired.get(name, 0) + 1 + return fired + + +@pytest.mark.parametrize("workers", [1, None]) +def test_resuming_at_each_step_replays_what_came_before( + tmp_path: Path, workers: int | None +) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + assert run(flow, out_dir=str(out), run_id="first", max_workers=workers)["Out"] == [2, 4, 6, 8] + steps = _trace(out / "first") + totals = dict(calls.by_actor) + + for step in range(len(steps) + 1): + calls.by_actor.clear() + outputs = run( + flow, + out_dir=str(out), + run_id=f"at{step}", + resume_from=str(out / "first"), + at_step=step, + max_workers=workers, + ) + assert outputs["Out"] == [2, 4, 6, 8], step + before = _fired_by(steps, step) + # What fired by the step is replayed; only what came after runs. + assert calls.by_actor == { + name: totals[name] - before.get(name, 0) + for name in totals + if totals[name] - before.get(name, 0) + }, step + record = json.loads((out / f"at{step}" / "run.wf-run.json").read_text()) + assert record["resumedFrom"] == "first" + assert record["resumedAtStep"] == step + assert "replayStopped" not in record + + +def test_a_failed_run_is_resumed_at_its_last_step(tmp_path: Path) -> None: + calls = Calls() + calls.broken = True + flow = _pipeline(calls) + out = tmp_path / "wf-out" + with pytest.raises(ValueError, match="broken on 3"): + run(flow, out_dir=str(out), run_id="first", max_workers=1) + steps = _trace(out / "first") # written although the run failed + first_calls = dict(calls.by_actor) + + calls.broken = False # the fix + calls.by_actor.clear() + outputs = run( + flow, out_dir=str(out), run_id="second", resume_from=str(out / "first"), at_step=len(steps) + ) + + assert outputs["Out"] == [2, 4, 6, 8] + total = { + name: first_calls.get(name, 0) + calls.by_actor.get(name, 0) for name in ("c", "p", "d") + } + assert total == {"c": 4, "p": 4, "d": 4} # none twice + + +def test_replayed_steps_are_marked_in_the_new_trace(tmp_path: Path) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + run(flow, out_dir=str(out), run_id="first", max_workers=1) + + run( + flow, + out_dir=str(out), + run_id="second", + resume_from=str(out / "first"), + at_step=3, + max_workers=1, + ) + + marks = [bool(step.get("replayed")) for step in _trace(out / "second")] + assert marks[:3] == [True, True, True] + assert not any(marks[3:]) + + +def test_a_resumed_run_can_itself_be_resumed_at_a_step(tmp_path: Path) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + run(flow, out_dir=str(out), run_id="first", max_workers=1) + run( + flow, + out_dir=str(out), + run_id="second", + resume_from=str(out / "first"), + at_step=3, + max_workers=1, + ) + + calls.by_actor.clear() + outputs = run( + flow, + out_dir=str(out), + run_id="third", + resume_from=str(out / "second"), + at_step=5, + max_workers=1, + ) + + assert outputs["Out"] == [2, 4, 6, 8] + assert sum(calls.by_actor.values()) == len(_trace(out / "first")) - 5 + + +def test_replayed_firings_keep_their_recorded_answers(tmp_path: Path) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + run(flow, out_dir=str(out), run_id="first", max_workers=1) + last = len(_trace(out / "first")) + + calls.salt = 100 # Pass now answers differently, as an agent might + calls.by_actor.clear() + outputs = run( + flow, out_dir=str(out), run_id="second", resume_from=str(out / "first"), at_step=last + ) + + # Every firing is before the step, so every one is replayed with the + # answer it gave then: resuming at a step keeps what happened before it. + assert outputs["Out"] == [2, 4, 6, 8] + assert calls.by_actor == {} + record = json.loads((out / "second" / "run.wf-run.json").read_text()) + assert "replayStopped" not in record + + # At step 0 nothing is replayed: everything runs, and takes the new answers. + outputs = run(flow, out_dir=str(out), run_id="third", resume_from=str(out / "first"), at_step=0) + assert outputs["Out"] == [202, 204, 206, 208] + + +def test_replay_stops_when_inputs_differ(tmp_path: Path) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + run(flow, out_dir=str(out), run_id="first", max_workers=1) + steps = _trace(out / "first") + + # Pass's firings after its first are not in the journal, so they run live, + # and now answer differently (as an agent might). Double's second recorded + # firing took Pass's old answer: its inputs no longer match, and replay + # stops there rather than replaying a firing on inputs it never had. + journal = out / "first" / "run.wf-journal.jsonl" + lines = journal.read_text().splitlines() + seen_pass = 0 + kept = [] + for line in lines: + entry = json.loads(line) + if entry.get("actor") == "p": + seen_pass += 1 + if seen_pass > 1: + continue + kept.append(line) + journal.write_text("\n".join(kept) + "\n") + + calls.salt = 100 + outputs = run( + flow, + out_dir=str(out), + run_id="second", + resume_from=str(out / "first"), + at_step=len(steps), + max_workers=1, + ) + + record = json.loads((out / "second" / "run.wf-run.json").read_text()) + assert "inputs differ" in record["replayStopped"] + assert record["replayStopped"].startswith("d ") + assert outputs["Out"] == [2, 204, 206, 208] + + +def test_a_loop_and_a_child_workflow_are_resumed_at_a_step(tmp_path: Path) -> None: + calls = Calls() + + @task + class Pass: + class Ports: + In = Port[int](direction="in") + Out = Port[int](direction="out") + + @action(consumes={"In": 1}, produces={"Out": 1}) + def go(self, x: int) -> int: + calls.fired("p") + return x * 10 + + @workflow(inputs={"In": int}, outputs={"Out": int}) + def child(): + p = Pass() + connect("In", p.In) + connect(p.Out, "Out") + + @workflow(outputs={"Out": int}) + def flow(): + lp = loop([1, 2, 3]) + with lp: + ch = child() + connect(lp.item, ch.In) + connect(ch.Out, "Out") + + out = tmp_path / "wf-out" + assert run(flow, out_dir=str(out), run_id="first")["Out"] == [10, 20, 30] + steps = _trace(out / "first") + + for step in range(len(steps) + 1): + calls.by_actor.clear() + outputs = run( + flow, out_dir=str(out), run_id=f"at{step}", resume_from=str(out / "first"), at_step=step + ) + assert outputs["Out"] == [10, 20, 30], step + calls.by_actor.clear() + run(flow, out_dir=str(out), run_id="last", resume_from=str(out / "first"), at_step=len(steps)) + assert calls.by_actor == {} # every firing replayed + + +def test_a_run_resumed_from_its_checkpoint_is_not_resumed_at_a_step(tmp_path: Path) -> None: + calls = Calls() + calls.broken = True + flow = _pipeline(calls) + out = tmp_path / "wf-out" + with pytest.raises(ValueError): + run(flow, out_dir=str(out), run_id="first") + calls.broken = False + run(flow, out_dir=str(out), run_id="second", resume_from=str(out / "first")) + + with pytest.raises(ValueError, match="resume at a step of that run instead"): + run(flow, out_dir=str(out), resume_from=str(out / "second"), at_step=1) + + +def test_refusals(tmp_path: Path) -> None: + calls = Calls() + flow = _pipeline(calls) + out = tmp_path / "wf-out" + run(flow, out_dir=str(out), run_id="first") + + with pytest.raises(ValueError, match="needs resume_from"): + run(flow, out_dir=str(out), at_step=1) + with pytest.raises(ValueError, match="has no step 99"): + run(flow, out_dir=str(out), resume_from=str(out / "first"), at_step=99) + with pytest.raises(ValueError, match="takes no inputs"): + run(flow, {"In": 1}, out_dir=str(out), resume_from=str(out / "first"), at_step=1) + + @workflow(outputs={"Out": int}) + def rewired(): + connect("Out", "Out") + + rewired._wfpy_workflow.name = "flow" # the same workflow, wired differently + with pytest.raises(ValueError, match="graph changed"): + run(rewired, out_dir=str(out), resume_from=str(out / "first"), at_step=1) + + +def test_the_cli_resumes_at_a_step(tmp_path: Path) -> None: + import subprocess + import sys + + flow = tmp_path / "flow.py" + flow.write_text( + """ +from pathlib import Path + +from wfpy import Port, action, connect, guard, task, workflow + +LOG = Path(__file__).parent / "fired.log" + + +@task +class Count: + _n: int = 0 + + class Ports: + Out = Port[int](direction="out") + + @action(consumes={}, produces={"Out": 1}) + @guard(lambda self: self._n < 3) + def emit(self) -> int: + with open(LOG, "a") as log: + log.write("count\\n") + self._n += 1 + return self._n + + +@workflow(outputs={"Out": int}) +def flow(): + c = Count() + connect(c.Out, "Out") +""" + ) + out = tmp_path / "wf-out" + + def wfpy(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, "-c", "from wfpy.cli import main; main()", "run", str(flow), *args], + capture_output=True, + text=True, + cwd=tmp_path, + ) + + assert wfpy("--out-dir", str(out), "--run-id", "first").returncode == 0 + resumed = wfpy( + "--out-dir", + str(out), + "--run-id", + "second", + "--resume-from", + str(out / "first"), + "--at-step", + "2", + ) + assert resumed.returncode == 0, resumed.stderr + + record = json.loads((out / "second" / "run.wf-run.json").read_text()) + assert record["outputs"]["Out"] == [1, 2, 3] + assert record["resumedAtStep"] == 2 + # Three in the first run; in the second only the third firing ran. + assert (tmp_path / "fired.log").read_text().count("count") == 4