diff --git a/docs/WORKFLOW_GUIDE.md b/docs/WORKFLOW_GUIDE.md index 561f61f5..ad811e91 100644 --- a/docs/WORKFLOW_GUIDE.md +++ b/docs/WORKFLOW_GUIDE.md @@ -818,7 +818,28 @@ a pipeline's own `seed` overrides its step's, which overrides the workflow's: } ``` -Omit `seed` entirely to let the workflow draw a random one at run time. +Omit `seed` entirely to let the workflow draw a random one at run time. The seed a run +actually used - drawn or named - is recorded in its `manifest.json`, so a run you liked +can be reproduced after the fact. + +Any of the three levels accepts a `variable:` reference, which is how a seed becomes +settable per run without editing the file: + +```json +{ + "variables": { "seed": 42 }, + "seed": "variable:seed", + "steps": [ ... ] +} +``` + +```bash +python -m dw.run workflows/ZImage.json seed=1234 +``` + +Declare the variable with an integer default, as above: the value from the command line +arrives as a string and is converted to the declared type. A string that is not a +`variable:` reference is rejected by the schema. The seed also reaches sub-workflows: a delegated `workflow` step runs the child under the parent's seed unless the child names its own. Without that a child draws its own diff --git a/docs/proposals/resume.md b/docs/proposals/resume.md new file mode 100644 index 00000000..7690beeb --- /dev/null +++ b/docs/proposals/resume.md @@ -0,0 +1,161 @@ +# Proposal: resuming a run — making the step cache survive the process + +Status: draft / scoping only — no implementation. + +## Problem + +A multi-step run that dies part way through — the OOM killer taking the +worker, a cancelled job, a crashed box, a mistake in step 6 of 8 — starts again +at step 1 the next time it is asked for. It regenerates work that is already +sitting on disk, and on a video workflow that is measured in GPU-hours rather +than seconds. Reported from an MCP session as roughly an hour of GPU time lost +to a single crash. + +The engine already has the two halves this needs and has never connected them: + +- **The step cache** (`dw/step_cache.py`) knows precisely when a step may be + skipped, and its criteria are already the strict ones. But it is a + process-wide singleton, so it dies with the worker that crashed. +- **Run directories** (`dw/runs.py`) already write every step's files under + `///` with a `manifest.json` beside them, and + `output:` already resolves a path *out of* that tree from another run. But + nothing reads a manifest back as an input to execution. + +So the durable record of "step 4 produced these files" exists on disk and is +thrown away, while the component that would act on it lives only in RAM. + +## Current behavior + +| Piece | Where | Lifetime | +|---|---|---| +| Step skip decision | `StepCache.get`, keyed `(workflow_id, step_name)` | Process | +| Cached entry payload | `step_data`, `step_seed`, `result`, `output_dir`, `generation`, `upstream_generations`, `retained` | Process | +| Files a step wrote | `manifest["steps"][i]["files"]`, relative to the run dir | Disk, forever | +| Seed the run used | `manifest["seed"]` | Disk, forever (as of the seed fix) | +| Run identity | `new_run_id({"workflow": workflow_def, "arguments": arguments})` | Disk, in the directory name | + +`Workflow.run` consults the cache for every step, server jobs included, and a +hit reports the earlier run's files with `reused: true` and writes nothing new. +That behavior — the thing a resume wants — is fully built. It just cannot be +reached across a process boundary. + +Note also that the cache is disabled outright for a workflow that sets no seed +(`cache_enabled_this_run`), because a fresh random seed each run means no entry +could ever match. Any resume inherits that constraint: **a run is resumable only +if it is reproducible.** The companion seed changes (variable `seed`, and +recording the drawn seed in the manifest) are a prerequisite, not a nicety. + +## Proposed shape + +**Rehydrate the existing step cache from manifests, rather than build a second +resume path.** + +The alternative — an explicit `--resume-from ` that replays a run and +skips steps by name — means a second set of rules about when skipping is safe, +sitting beside the six criteria `step_cache.py` already documents and enforces. +Two mechanisms answering the same question is how they drift. Rehydration adds +one loader and reuses every existing invariant. + +### Mechanism + +1. On `Workflow.run` start, when resume is enabled, read the manifests under + this workflow's identity in the output root, newest first. +2. For each manifest whose recorded seed and step definitions still match this + run, synthesize `StepCache` entries for its steps and `put` them before the + step loop begins. +3. Everything after that is unchanged. `StepCache.get` decides. Its criterion 5 + (every named file must still exist) already handles a partially deleted run + directory; criterion 4 (output root must match) already handles a resume + pointed somewhere else. + +For entry validity across processes, the manifest must carry enough to +reconstruct the comparison. Today it holds the run's `arguments` and seed but +not each step's *resolved* definition. Two options: + +- **(a)** Store a hash of each step's resolved `step_data` in the manifest entry + and compare hashes rather than values. Cheap, small, and matches how + `pipeline_cache_key` and `new_run_id` already establish identity. +- **(b)** Store the resolved `step_data` itself. Larger and leaks realized + values (a resolved argument can hold a PIL image), so (a) is preferred. + +With (a), `StepCache` needs to accept an entry whose `step_data` is known only +by digest. That is a real change to its shape and the main design decision here. + +### The hard boundary: results vs. files + +A cache entry may hold a live `Result` (`retained`) or only `saved_files`. Only +the second survives a process. So: + +- A step whose downstream consumers need only its **files** — which is most + image and video work, and everything the `output:` reference already supports + — can be resumed. +- A step whose result is consumed **in memory** by a later step (latents, an + embedding, a mask handed to an inpaint step via `previous_result:`) cannot be + rehydrated from disk, because those bytes were never written. + +The proposal is to let such a step **miss** rather than to invent a +serialization format for arbitrary pipeline outputs. `StepCache.get` already +takes `needs_result` and already returns None when an entry lacks one, so this +falls out of the existing code with no new branch. Consequence worth stating +plainly: a workflow whose steps chain in memory resumes at the first such step +and re-runs from there. A workflow whose steps chain through saved files — +the multi-stage pattern the GYRE and marmot workflows use — resumes at the true +failure point. + +A later stage could widen this by reloading saved media into a `Result` (the +loaders already exist for `asset:` and `output:`), turning "files only" entries +into serviceable ones for the common image/video case. Deliberately out of +scope for stage one. + +### Surface + +Resume should be **off by default and explicit**, at least initially. A run +that silently reuses a previous run's files is the kind of behavior that is +delightful until the one time it is wrong, and the step cache's in-process +version is already scoped to a session where the user can see what happened. + +- CLI: `--resume` (reuse whatever matches from previous runs of this workflow), + and `--resume-from ` to name one explicitly. +- Server/MCP: a `resume` flag on job submission, defaulting off. Note that a + job carries its own `output_dir`, so resume is naturally confined to its + workspace with no extra work. +- Reporting: reused steps already surface as `reused: true` in the manifest and + in `step_end` events, so the gallery and job detail need no changes to show + what was skipped. Worth confirming the UI actually distinguishes it. + +## Staging + +1. **Prerequisite (done): seed.** `seed` accepts a `variable:` reference, and + the manifest records the seed actually used. Without both, a seedless + workflow has nothing to resume against. +2. **Manifest carries step identity.** Add a per-step digest of the resolved + `step_data` to the manifest. Inert on its own — nothing reads it yet — and + independently useful for answering "did this run actually run the same thing". +3. **Rehydration.** A loader that turns manifests into `StepCache` entries, plus + whatever `StepCache` needs to compare by digest. Behind `--resume`. +4. **Reload saved media into results.** Widen resumability past files-only + steps, if stage 3 shows it is the common blocker. + +Stages 1–3 are the useful unit; 4 is an optimization. + +## Open questions + +- **Scope of a match.** Should resume consider only the most recent matching + run, or any run of that identity? Any-run is strictly more useful and makes + the "fix one still" workflow work across days, but it also means a step can + be served from a run the user has forgotten about. Leaning: most recent + matching run in stage 3, widened later if it proves too narrow. +- **Interaction with pruning.** If output pruning ever lands, a resumable run + is one whose files must not be pruned. Criterion 5 makes a pruned run fail + safe (it misses and re-runs), so this is a quality-of-life question, not a + correctness one. +- **Sub-workflows.** A sub-workflow inherits the parent's run directory and + writes no manifest of its own, so its steps are invisible to a manifest-based + resume. Its steps *do* appear in the parent's rolled-up manifest, but keyed by + the child's `workflow_id`. Needs a decision before stage 3. +- **Flat layout.** `--output-layout flat` writes no manifests at all, so resume + is simply unavailable there. Should it say so, or stay silent? +- **Confinement.** Reading manifests means reading paths written by an earlier + run and republishing them. They are already confined to the output root and + already validated on the `output:` path, but the resume loader is new code + reading path-shaped data off disk and should go through the same validators. diff --git a/dw/repl_commands.py b/dw/repl_commands.py index 48b5c5b1..31152fb4 100644 --- a/dw/repl_commands.py +++ b/dw/repl_commands.py @@ -741,8 +741,11 @@ def end_inline(): break except Exception as e: end_inline() + # A worker killed by a signal sends nothing and leaves no + # traceback, so its exit code is the only diagnosis there is + detail = self.repl.worker_manager.crash_details() print("\n" + "=" * 80) - print(f"ERROR receiving results: {e}") + print(f"ERROR receiving results: {detail or e}") print("=" * 80) print("Worker communication failed. Shutting down worker.\n") self.repl.worker_manager.shutdown_worker() diff --git a/dw/repl_worker.py b/dw/repl_worker.py index bc33a350..19e6d7d3 100644 --- a/dw/repl_worker.py +++ b/dw/repl_worker.py @@ -8,6 +8,7 @@ import multiprocessing import queue as queue_module import logging +import signal from typing import Optional from .worker import worker_main @@ -122,6 +123,35 @@ def cancel(self): """Ask the worker to cancel the workflow it is running.""" self.send_command({"type": "cancel"}) + def crash_details(self): + """Why the worker process is gone, as far as the OS will say. + + A worker killed by a signal - the OOM killer's SIGKILL above all - + never reaches worker_main's except clause, so there is no traceback + to report and the exit code is the whole diagnosis. Returns None + while the process is alive or was never started. + """ + process = self.worker_process + if process is None or process.is_alive(): + return None + exitcode = process.exitcode + if exitcode is None: + return None + if exitcode < 0: + signal_number = -exitcode + try: + name = signal.Signals(signal_number).name + except ValueError: + name = f"signal {signal_number}" + detail = f"killed by {name}" + if signal_number == signal.SIGKILL: + # By far the likeliest cause on a box that loads models + # measured in tens of gigabytes, and the one thing the user + # can act on - no Python-level error will have been logged + detail += " (typically the out-of-memory killer)" + return detail + return f"exited with code {exitcode}" + def mark_crashed(self): """Record that the worker process died on its own - no shutdown handshake to attempt, just clear the tracking state.""" diff --git a/dw/server/jobs.py b/dw/server/jobs.py index f6c12c18..7a33995f 100644 --- a/dw/server/jobs.py +++ b/dw/server/jobs.py @@ -9,6 +9,7 @@ import os import copy import json +import queue import sqlite3 import time import uuid @@ -724,7 +725,22 @@ def _consume_results(self, job): (status, error, traceback) for _run_job to apply once the manager no longer counts the job as current.""" while True: - message = self.worker_manager.get_result() + try: + message = self.worker_manager.get_result() + except RuntimeError as e: + # The worker died without managing to send anything - a + # signal, not an exception, so worker_main's handler never + # ran and there is no traceback to be had. The exit code is + # the only diagnosis available, and marking the crash here + # matters beyond this job: the manager would otherwise go on + # believing a dead process is active, and every later call + # that talks to it (memory_status above all) would fail + # against a queue nobody is reading + detail = self.worker_manager.crash_details() + self.worker_manager.mark_crashed() + reason = f"Worker process died: {detail or e}" + logger.error(f"Job {job.id}: {reason}") + return (FAILED, reason, None) message_type = message.get("type") if message_type == "progress": @@ -812,6 +828,19 @@ def memory_status(self, timeout=5): try: self.worker_manager.send_command({"type": "memory_status"}) result = self.worker_manager.get_result(timeout=timeout) + except (RuntimeError, queue.Empty) as e: + # A dead worker is exactly when the last reading taken before it + # died is worth the most, so report that rather than failing the + # request. Without this the caller gets a 503 at the one moment + # it most wants a number + detail = self.worker_manager.crash_details() + logger.warning(f"Worker unavailable for memory status: {detail or e}") + if detail is not None: + # crash_details only answers for a process the OS has reaped, + # so a worker that is merely slow to reply keeps its state - + # a timeout is not evidence of death + self.worker_manager.mark_crashed() + return {"live": False, "info": self.last_memory} finally: self._worker_lock.release() if result.get("type") == "memory_status": diff --git a/dw/workflow.py b/dw/workflow.py index 8577a159..28555b04 100644 --- a/dw/workflow.py +++ b/dw/workflow.py @@ -299,6 +299,12 @@ def run( # this alone died on an exception the manifest should say so about status = "failed" run_id = None + # The seed this run actually used, which is not self.workflow_definition's: + # run() works on a deep copy, and a workflow naming no seed draws a random + # one into that copy. Recording the original would write null into the + # manifest of every seedless run and lose the only record of what produced + # its files - the seed is what makes a run repeatable + resolved_seed = None started_at = datetime.now(timezone.utc).isoformat() try: # CRITICAL: Work on a copy to avoid mutating the original workflow definition @@ -336,6 +342,20 @@ def run( # dict.get default, torch.seed() would run on every call and reseed # the global RNG even when the workflow names an explicit seed default_seed = workflow_def.get("seed") + # The schema lets 'seed' be a string so it can hold a 'variable:' + # reference, which the substitution above has already resolved - + # but a variable overridden from the command line arrives as a + # string whenever the workflow declared no integer default to + # coerce against, and manual_seed would fail deep inside the run + if isinstance(default_seed, str): + try: + default_seed = int(default_seed) + except ValueError: + raise ValueError( + f"Workflow {workflow_id} seed must be an integer, " + f"got {default_seed!r}" + ) + workflow_def["seed"] = default_seed # A workflow that names no seed gets a fresh one every run, so no # step's cache entry can ever match again - skip the cache # wholesale rather than deep-copying every step's realized images @@ -353,6 +373,7 @@ def run( # global RNG the process may have seeded for reproducibility default_seed = torch.Generator().seed() workflow_def["seed"] = default_seed + resolved_seed = default_seed # One execution, one directory - opened here, after variable # substitution and the seed have settled, so the run's identity @@ -609,11 +630,13 @@ def run( # write are on disk either way, and what produced them is exactly # what a failed run needs to explain itself if self._run_dir and not self._run_dir_inherited: - self._write_run_manifest(run_id, status, started_at, arguments) + self._write_run_manifest( + run_id, status, started_at, arguments, resolved_seed + ) deactivate_output_root(output_root_token) deactivate_context(context_token) - def _write_run_manifest(self, run_id, status, started_at, arguments): + def _write_run_manifest(self, run_id, status, started_at, arguments, seed): """Leave a record of the run beside the files it wrote. A server run is in jobs.sqlite as well, but a CLI run has never been @@ -637,7 +660,7 @@ def _write_run_manifest(self, run_id, status, started_at, arguments): "file": self.file_spec, "identity": workflow_identity(self.file_spec, self.name), }, - "seed": self.workflow_definition.get("seed"), + "seed": seed, "arguments": arguments or {}, "steps": [ { diff --git a/dw/workflow_schema.json b/dw/workflow_schema.json index 9b4a6f18..c29b9408 100644 --- a/dw/workflow_schema.json +++ b/dw/workflow_schema.json @@ -14,8 +14,12 @@ "$ref": "#/$defs/arguments" }, "seed": { - "description": "Default seed for the entire workflow", - "type": "integer", + "description": "Default seed for the entire workflow. Accepts a 'variable:' reference.", + "type": [ + "integer", + "string" + ], + "pattern": "^variable:", "format": "int64" }, "steps": { @@ -63,8 +67,12 @@ "type": "string" }, "seed": { - "description": "Default seed for the entire step", - "type": "integer", + "description": "Default seed for the entire step. Accepts a 'variable:' reference.", + "type": [ + "integer", + "string" + ], + "pattern": "^variable:", "format": "int64" }, "release_pipeline": { @@ -369,8 +377,12 @@ ] }, "seed": { - "description": "The seed for this pipeline", - "type": "integer", + "description": "The seed for this pipeline. Accepts a 'variable:' reference.", + "type": [ + "integer", + "string" + ], + "pattern": "^variable:", "format": "int64" }, "chain": { diff --git a/tests/test_repl_worker.py b/tests/test_repl_worker.py index cded3aa7..01ad44eb 100644 --- a/tests/test_repl_worker.py +++ b/tests/test_repl_worker.py @@ -12,8 +12,9 @@ class FakeProcess: - def __init__(self, alive=True): + def __init__(self, alive=True, exitcode=None): self.alive = alive + self.exitcode = exitcode def is_alive(self): return self.alive @@ -57,6 +58,40 @@ def test_an_explicit_timeout_raises_empty_not_runtime_error(self, manager): manager.get_result(timeout=0.01) +class TestCrashDetails: + """A worker killed by a signal never reaches worker_main's handler, so + no traceback exists and the exit code is the whole diagnosis.""" + + def test_a_live_worker_has_nothing_to_report(self, manager): + assert manager.crash_details() is None + + def test_a_worker_not_yet_reaped_has_nothing_to_report(self, manager): + manager.worker_process.alive = False + assert manager.crash_details() is None + + def test_sigkill_names_the_out_of_memory_killer(self, manager): + manager.worker_process.alive = False + manager.worker_process.exitcode = -9 + detail = manager.crash_details() + assert "SIGKILL" in detail + assert "out-of-memory" in detail + + def test_another_signal_is_named_without_the_oom_guess(self, manager): + manager.worker_process.alive = False + manager.worker_process.exitcode = -11 + detail = manager.crash_details() + assert "SIGSEGV" in detail + assert "out-of-memory" not in detail + + def test_a_nonzero_exit_reports_its_code(self, manager): + manager.worker_process.alive = False + manager.worker_process.exitcode = 1 + assert manager.crash_details() == "exited with code 1" + + def test_it_answers_for_a_manager_that_never_started_one(self): + assert WorkerManager().crash_details() is None + + class TestInactiveWorker: def test_send_and_receive_refuse_an_inactive_worker(self): manager = WorkerManager() diff --git a/tests/test_runs.py b/tests/test_runs.py index 4037b2d2..4baf4b31 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -253,6 +253,42 @@ def test_the_manifest_describes_the_run(self, tmp_path, fake_pipeline): assert not os.path.isabs(name) assert (run_dir / name).exists() + def test_the_manifest_records_the_seed_a_seedless_run_drew( + self, tmp_path, fake_pipeline + ): + """A workflow naming no seed gets a random one, and the manifest is + the only place it is ever written down - reporting null there loses + the one value needed to reproduce the files sitting beside it.""" + from dw.workflow import Workflow + + definition = _workflow_definition() + del definition["seed"] + workflow = Workflow(definition, str(tmp_path), "/w/workflows/ltx2/Gyre.json") + workflow.run({"prompt": "a cat"}) + + run_dir = next((tmp_path / "ltx2" / "Gyre").iterdir()) + manifest = json.loads((run_dir / "manifest.json").read_text()) + assert isinstance(manifest["seed"], int) + # and the definition it was built from is untouched, so the next run + # draws its own seed rather than inheriting this one + assert "seed" not in definition + + def test_a_seed_can_come_from_a_variable(self, tmp_path, fake_pipeline): + """'seed' accepts a 'variable:' reference, so a caller can re-run a + workflow at the seed a previous run's manifest reported.""" + from dw.workflow import Workflow + + definition = _workflow_definition() + definition["variables"] = {"seed": 7} + definition["seed"] = "variable:seed" + workflow = Workflow(definition, str(tmp_path), "/w/workflows/ltx2/Gyre.json") + workflow.validate() + workflow.run({"seed": "1234"}) # as it arrives from the command line + + run_dir = next((tmp_path / "ltx2" / "Gyre").iterdir()) + manifest = json.loads((run_dir / "manifest.json").read_text()) + assert manifest["seed"] == 1234 + def test_a_failed_run_still_records_what_it_wrote(self, tmp_path, fake_pipeline): from dw.workflow import Workflow diff --git a/tests/test_server.py b/tests/test_server.py index 5d2f04cb..07c38867 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -70,6 +70,24 @@ def mark_crashed(self): self.crashed = True +class DyingWorkerManager(ScriptedWorkerManager): + """A worker killed by a signal: it sends nothing, and get_result raises + the way the real WorkerManager's liveness poll does.""" + + def __init__( + self, detail="killed by SIGKILL (typically the " "out-of-memory killer)" + ): + super().__init__(script=lambda command: []) + self.detail = detail + self.crashed = False + + def get_result(self, timeout=None): + raise RuntimeError("Worker process died while waiting for results") + + def crash_details(self): + return self.detail + + def success_script(command): yield { "type": "progress", @@ -444,6 +462,74 @@ def test_workflow_browsing_and_confinement(server): assert client.get("/api/workflows/nope").status_code == 404 +def test_a_silently_killed_worker_reports_its_exit_and_frees_the_manager(tmp_path): + """The OOM killer leaves no traceback, so the exit code is the whole + diagnosis - and the manager must stop believing the dead process is + alive, or every later call that talks to it fails against a queue + nobody reads.""" + workflow_dir = tmp_path / "workflows" + workflow_dir.mkdir() + manager = JobManager( + str(tmp_path / "outputs"), + worker_manager=DyingWorkerManager(), + history_path=str(tmp_path / "jobs.sqlite"), + ) + app = create_app( + workflow_dir=str(workflow_dir), + output_dir=str(tmp_path / "outputs"), + job_manager=manager, + ) + with TestClient(app, base_url="http://localhost") as client: + # the last reading taken while the worker was still alive + manager.last_memory = {"gpu_available": True, "used": 42} + + job = client.post("/api/jobs", json={"workflow": valid_workflow()}).json() + detail = wait_for_status(client, job["id"], TERMINAL_STATES) + assert detail["status"] == "failed" + assert "SIGKILL" in detail["error"] + assert "out-of-memory" in detail["error"] + assert detail["traceback"] is None + + assert manager.worker_manager.crashed is True + assert client.get("/api/health").json()["current_job"] is None + + # the one moment the last reading is worth the most is the one that + # used to answer 503 + response = client.get("/api/memory") + assert response.status_code == 200 + assert response.json() == { + "live": False, + "info": {"gpu_available": True, "used": 42}, + } + + +def test_a_slow_worker_is_not_declared_dead(tmp_path): + """A memory poll that times out against a live worker must fall back to + the cached reading without marking the process crashed.""" + + class SilentWorkerManager(ScriptedWorkerManager): + def get_result(self, timeout=None): + raise queue.Empty() + + def crash_details(self): + return None # the process is alive, just slow + + manager = JobManager( + str(tmp_path / "outputs"), + worker_manager=SilentWorkerManager(success_script), + history_path=str(tmp_path / "jobs.sqlite"), + ) + manager.worker_manager.worker_active = True + manager.last_memory = {"gpu_available": True} + + assert manager.memory_status(timeout=0.01) == { + "live": False, + "info": {"gpu_available": True}, + } + assert manager.worker_manager.worker_active is True + assert getattr(manager.worker_manager, "crashed", False) is False + + def test_health_and_memory(server): import socket diff --git a/workflows/ZImage.json b/workflows/ZImage.json index bdbf33c9..92767920 100644 --- a/workflows/ZImage.json +++ b/workflows/ZImage.json @@ -6,8 +6,10 @@ "num_inference_steps": 9, "guidance_scale": 0.0, "width": 960, - "height": 544 + "height": 544, + "seed": 42 }, + "seed": "variable:seed", "id": "ZImage", "description": "Text-to-image with Z-Image Turbo - photorealistic output in nine steps.", "steps": [ @@ -39,4 +41,4 @@ } } ] -} \ No newline at end of file +}