From 97c8d266297f0fc27a5c4b19aa5df8c888d420da Mon Sep 17 00:00:00 2001 From: Don Kackman Date: Wed, 9 Sep 2026 13:07:52 -0500 Subject: [PATCH 1/2] feat: enhance rerun functionality with new seed option - Added support for rerunning jobs with a new seed, allowing users to generate different outputs from workflows that pin their seed to a variable. - Updated API to include `seed_variable` in job workflow responses, indicating if a new seed can be drawn. - Modified job rerun logic to handle new seed requests and ensure proper error handling for workflows without seed variables. - Enhanced workspace disk usage reporting, caching results for efficiency and accuracy. - Improved UI to reflect new rerun capabilities, including buttons for rerunning with or without a new seed. - Added tests to validate new seed functionality and ensure correct behavior for reruns. --- CLAUDE.md | 2 +- docs/MCP.md | 2 +- docs/SERVER.md | 10 +- docs/proposals/codeql-sanitizer-model.md | 114 -------- docs/proposals/job-record-and-export.md | 165 ----------- docs/proposals/live-latent-preview.md | 340 ----------------------- docs/proposals/resume.md | 4 +- dw/server/app.py | 33 ++- dw/server/jobs.py | 45 ++- dw/workspace.py | 122 +++++++- dw_mcp/diagnose.py | 17 +- dw_mcp/server.py | 18 +- tests/test_rerun_new_seed.py | 153 ++++++++++ tests/test_server_workspaces.py | 30 ++ ui/src/lib/api.ts | 23 +- ui/src/lib/pages/JobPage.svelte | 67 ++++- ui/src/lib/pages/ServerPage.svelte | 20 ++ ui/src/lib/results.test.ts | 23 +- ui/src/lib/results.ts | 16 +- ui/src/lib/types.ts | 3 + ui/src/lib/workspace.svelte.ts | 11 +- 21 files changed, 554 insertions(+), 664 deletions(-) delete mode 100644 docs/proposals/codeql-sanitizer-model.md delete mode 100644 docs/proposals/job-record-and-export.md delete mode 100644 docs/proposals/live-latent-preview.md create mode 100644 tests/test_rerun_new_seed.py diff --git a/CLAUDE.md b/CLAUDE.md index c66480fb..470dcea7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -231,7 +231,7 @@ same reason - default setup cannot load a pack. `/exports//` and `GET /exports/.zip` streams it. - **Step cache**: a process-wide singleton (`dw/step_cache.py`) consulted by every `Workflow.run`, including server jobs; entries are keyed by `(workflow id, step name)` and validated against the output *root*, never the per-run directory - a run directory is new every execution and would - defeat the cache; disabled entirely when the workflow sets no `seed`; a hit reports the earlier run's files with `reused: true` and writes nothing new; `memory clear` drops it + defeat the cache; disabled entirely when the workflow sets no `seed`; a hit reports the earlier run's files with `reused: true` and writes nothing new; `memory clear` drops it. This is why "Run again" on a seeded workflow finishes instantly and generates nothing - the job page says so when every step was reused, and `POST /api/jobs/{id}/rerun` with `{"new_seed": true}` (MCP `rerun_job(new_seed=True)`) draws a fresh seed into the workflow's seed variable, which is the way to get a different image ## JSON Workflow Structure diff --git a/docs/MCP.md b/docs/MCP.md index 81efcf84..8803da98 100644 --- a/docs/MCP.md +++ b/docs/MCP.md @@ -285,7 +285,7 @@ references written in the same session. | `get_job_events(job_id, after=-1, limit=200)` | `job_id`, `after`, `limit` | Get a page of a job's progress events | | `wait_for_job(job_id, timeout_seconds=20)` | `job_id`, `timeout_seconds` | Block until a job reaches a terminal status, or `timeout_seconds` elapses (capped well under a generation's real runtime). Use instead of hand-polling `get_job`/`get_job_events` in a loop; if it returns `still_running: true`, call it again | | `cancel_job(job_id)` | `job_id` | Ask a queued or running job to stop | -| `rerun_job(job_id, acknowledged_cost=False)` | `job_id`, `acknowledged_cost` | Queue a fresh job from a previous job's stored specification. Costs GPU time, so it passes the same gate as `run_workflow` | +| `rerun_job(job_id, acknowledged_cost=False, new_seed=False)` | `job_id`, `acknowledged_cost`, `new_seed` | Queue a fresh job from a previous job's stored specification. Costs GPU time, so it passes the same gate as `run_workflow`. `new_seed=true` draws a fresh seed into the workflow's seed variable — without it a seeded workflow's rerun repeats its arguments exactly and the step cache serves the whole run from the earlier one's files (`reused: true`), generating nothing. `get_job_workflow`'s `seed_variable` says whether there is one | | `move_job(job_id, direction)` | `job_id`, `direction` (`up`\|`down`\|`front`\|`back`) | Reorder a queued job | ### Models diff --git a/docs/SERVER.md b/docs/SERVER.md index 75b9808c..0a1aa6a2 100644 --- a/docs/SERVER.md +++ b/docs/SERVER.md @@ -135,13 +135,13 @@ from another machine: | `POST /api/jobs` | Queue a run: `{"workflow_path": ...}` or an inline `{"workflow": {...}, "base_dir": ...}`, plus `arguments` for variable overrides. `workflow_path` accepts a stored workflow name as listed by `/api/workflows` (with or without `.json`, nested names included), or a relative/absolute path that still resolves under `--workflow-dir` - confined the same way the `/api/workflows` CRUD routes are; a path that names a real file outside that directory is rejected with 400, not opened. Answers with argument warnings from signature checking. | | `GET /api/jobs` | Queue + history summaries | | `GET /api/jobs/{id}` | Full detail: spec, events, manifest, error. A manifest entry for a step served from the step cache carries `reused: true` | -| `GET /api/jobs/{id}/workflow` | The workflow the job ran: `{id, definition, realized}`. `realized: true` is the copy the run itself wrote (`workflow.json` in its run directory), with arguments, seed, prompts and `output:latest` pinned; `false` falls back to the submitted definition, which is what a job from before run tracking has. 404 means neither is readable - the job itself still is | +| `GET /api/jobs/{id}/workflow` | The workflow the job ran: `{id, definition, realized, seed_variable}`. `seed_variable` names the variable a `new_seed` rerun would draw into (null when the workflow has none), read from the workflow as written rather than the realized copy, whose seed is pinned. `realized: true` is the copy the run itself wrote (`workflow.json` in its run directory), with arguments, seed, prompts and `output:latest` pinned; `false` falls back to the submitted definition, which is what a job from before run tracking has. 404 means neither is readable - the job itself still is | | `POST /api/jobs/{id}/export?workspace=&overwrite=` | Gather one finished job into `/exports//`: `workflow.json`, `manifest.json`, `job.json`, `README.md`, `assets/`, `inputs/`, `outputs/`. 201 with the file list, total bytes, anything it could not find, a `zip_url`, and the three JSON files inline. 404 unknown job, 409 for a job still running or an existing export without `overwrite` | | `GET /exports/{id}.zip?workspace=` | The same tree as one archive, built on request rather than kept as a second copy. Entries are named `/`. Ungated exactly as `/outputs` is | | `GET /api/jobs/{id}/events` | Server-sent events stream; `?after=N` / `Last-Event-ID` replay missed events, so reconnects are lossless | | `GET /api/jobs/{id}/event-log?after=-1&limit=200` | The same events as the SSE stream, as one JSON page: `{id, status, events, last_seq, truncated, note}`. `after` is exclusive; page by passing back the previous `last_seq`. A job restored from history serves the bounded event tail persisted with it; a job that finished before events were retained returns an empty list and a `note` saying so. | | `POST /api/jobs/{id}/cancel` | Cooperative cancel (takes effect at the next step boundary or denoise step) | -| `POST /api/jobs/{id}/rerun` | Re-queue a finished job's spec | +| `POST /api/jobs/{id}/rerun` | Re-queue a finished job's spec. Body `{"new_seed": true}` draws a fresh seed into the workflow's seed variable instead of repeating the original arguments; 400 when the workflow pins its seed to a literal or names none. A plain rerun of a seeded workflow is served whole from the step cache — the earlier run's files, republished with `reused: true`, generating nothing | | `POST /api/jobs/{id}/move` | Reorder a queued job: `{"direction": "up"\|"down"\|"front"\|"back"}`. Job listings carry each waiting job's `queue_position`. | One job runs at a time (it is one GPU); submissions queue in order, and @@ -260,7 +260,11 @@ The editor's forms come from these; they are just as usable from scripts: `default` workspace and a named one is a subdirectory beside them, sharing the root's one prompt library. Delete answers with what it would remove and refuses until acknowledged, refuses the default, and refuses a workspace - with jobs still queued. A workspace is a namespace, **not** a security + with jobs still queued. Each listed workspace carries a `usage` + (`{files, bytes}`) — roughly how much disk its own folders hold, walked at + most once a minute per workspace and deliberately approximate; the shared + prompt library counts against the `default` workspace alone rather than + once per workspace. A workspace is a namespace, **not** a security boundary: the API token is all-or-nothing `exports/` sits beside the workspace's own folders, holding one directory per diff --git a/docs/proposals/codeql-sanitizer-model.md b/docs/proposals/codeql-sanitizer-model.md deleted file mode 100644 index 54448020..00000000 --- a/docs/proposals/codeql-sanitizer-model.md +++ /dev/null @@ -1,114 +0,0 @@ -# Proposal: making CodeQL see the path validators - -Status: draft, 2026-09-08. No implementation. - -## Problem - -Every release adds a few `py/path-injection` alerts from code scanning, and -every one so far has had the same shape: a path is joined and then confined -by `validate_path` in `dw/security.py` (or one of the wrappers around it), -and CodeQL flags the `os.*` or `shutil.*` call downstream because it does -not treat the validator as a sanitizer. The count on 2026-09-08 is 28 open -and 14 dismissed as "mitigated", all of that rule, plus the ten on #56 that -were dismissed the same way after a review found one real gap among them -(an `export_directory` job id of `.` named the exports root itself; fixed -in 7ead046 with a shape check). - -Dismissing by hand works but scales badly, and it teaches the reviewer to -skim: the one real finding on #56 was flagged in exactly the same words as -the nine that were not. - -## What already exists, and why it is not working - -`.github/codeql/extensions/dw-models/` is a CodeQL model pack, added with -the workspaces PR (#37, 2026-09-05), that declares the validators' return -values as `path-injection` barriers with the `barrierModel` extensible -predicate: - -```yaml -- ["dw.security", "Member[validate_path].ReturnValue", "path-injection"] -``` - -The approach is right. `barrierModel` and `barrierGuardModel` were added -to models-as-data for Python in CodeQL 2.25.2 (April 2026), the path -injection query's `SanitizerFromModel` reads barriers of kind -`path-injection`, code scanning default setup picks up model packs from -`.github/codeql/extensions/` without a workflow, and the repository is -scanned with CodeQL 2.26.4. Yet every alert listed above was created -*after* the pack landed: 26 to 43 on 2026-09-06 and 07, 44 to 53 on the -08th. The pack is not taking effect, and nothing in the scanning UI says -why. - -Hypotheses, most likely first: - -1. **The type string does not resolve.** Every caller imports the - validators relatively (`from ..security import validate_path`, - `from .security import ...`). The `type` column of a Python model is - matched through the API graph's module-import resolution, which is - built for absolute imports of library modules; whether a relative - import inside the analysed package resolves to `dw.security` is not - documented and is the first thing to test. -2. **The pack is not being loaded at all.** A malformed - `codeql-pack.yml`, a name collision, or a default-setup restriction - the docs do not spell out. Default setup exposes no log for this. -3. **The barrier is on the wrong node.** A barrier on the return value - stops flow that *passes through* the call. The flow CodeQL reports may - instead go around it: `validate_path` returns a value computed from - its argument, and if the analysed body contributes a second path - (the argument itself flowing to a sink inside `validate_path`, then - summarised) the return-value barrier would not cut it. Unlikely, but - the alert's path view would show it. -4. **`base_dir=None` is a real hole in the model.** `validate_path(p, - None)` resolves and pattern-checks but confines to nothing, so a - barrier on its return value claims more than the function guarantees. - That does not explain the alerts, but it is why the model should sit - on the wrappers that always pass a root (`validate_output_path`, - `validate_workflow_path`, `validate_prompt_path`) and on - `validate_path` only when the root argument is present, which the - model language cannot express. Decide whether to accept the - over-claim or narrow the model. - -## Proposed shape - -**Reproduce locally, then fix the model, then let the scanner confirm.** -Default setup gives no feedback loop; the CodeQL CLI does. - -1. Install the CodeQL CLI (`gh extension install github/gh-codeql`, which - also fetches the bundle) on the Mac. -2. `codeql database create --language=python dw-db` at the repo root, - then `codeql database analyze dw-db codeql/python-queries:codeql-suites/python-code-scanning.qls --model-packs dkackman/dw-models --format=sarif-latest --output=out.sarif` - with the pack path supplied through `--additional-packs - .github/codeql/extensions`. Count the `py/path-injection` results. -3. Iterate on `dw-security.model.yml` until the count drops to the - findings that are real: try an absolute-import form of the type, try - the wrappers rather than `validate_path`, try `barrierGuardModel` on - the `SecurityError` branch. Each try is one edit and one analyze. -4. Commit the working model. Code scanning re-runs on the next push and - the open alerts close as "fixed" on their own, which is the only - confirmation that matters. -5. Record in `dw/security.py`'s module docstring that the validators are - modelled, and where, so the next person adding a validator adds a row. - -A `scripts/codeql-local.sh` that does steps 2 and 3 in one command is -worth keeping, since the model will need the same loop each time a -validator is added. - -## What this is not - -- Not a workflow migration. Default setup honours repository model packs, - and advanced setup is only needed if step 3 shows the fix requires a - custom query, which nothing so far suggests. -- Not a reason to stop reading the alerts. Once the model holds, an alert - of this rule means a path that did not go through a validator, which is - exactly the signal the rule is for. - -## Open questions - -- Whether the fourteen "mitigated" dismissals and the ten from #56 should - be reopened once the model works, so the scanner re-evaluates them. They - would close as fixed if the model is right and stay open if it is not, - which is a useful check; but reopening is a manual click each. -- Whether `validate_path` with no root should keep returning a path at - all, or whether every caller should be made to pass one. That is a - security question independent of CodeQL, and the model's over-claim in - hypothesis 4 is the argument for looking at it. diff --git a/docs/proposals/job-record-and-export.md b/docs/proposals/job-record-and-export.md deleted file mode 100644 index 992f5dc0..00000000 --- a/docs/proposals/job-record-and-export.md +++ /dev/null @@ -1,165 +0,0 @@ -# Proposal: the realized workflow as a job's record, and exporting a job - -Status: implemented 2026-09-08 (design: -`docs/superpowers/specs/2026-09-08-job-record-and-export-design.md`, plan: -`docs/superpowers/plans/2026-09-08-job-record-and-export.md`). Supplies the -"manifest carries step identity" stage of [resume.md](resume.md), which stays -the design for resuming a run; this proposal is the record that resume reads, -plus the two ways to get it off the server. - -## Problem - -On 2026-09-08 an MCP session assembled a six-shot MiniMax H3 short as an -inline workflow, ran it for most of an hour, and lost it at the last step to -a bad argument. The retry regenerated every shot. Three things about that -run are worth fixing, and none of them is the bad argument: - -- **The workflow that produced the cut existed nowhere the user could read.** - The job store keeps an inline definition in its spec column and the job - page can draw it, but no MCP tool returns it, and the run directory holds - a manifest that says only `"file": ".../__inline__.json"`. -- **A template job records less than that.** Its spec holds a path and the - arguments. Edit the template, or a stored prompt it names, or make a newer - run that `output:latest` now picks, and the record no longer describes what - ran. The manifest pins the seed and the arguments, but not the definition. -- **Nothing packages a run.** Manifest, workflow, the assets it referenced - and the files it made are four places on one machine, and `download_output` - moves one file at a time onto that same machine. A run that is worth - keeping, sharing or committing has to be gathered by hand. - -MCP is turning out to be the primary way work is submitted, and an agent -writes inline workflows freely: an inline job is now the common case, not -the exception the job store treated it as. - -## What exists today - -| Piece | Where | What it records | -|---|---|---| -| Submitted definition | `jobs.sqlite` `spec` (`workflow` for inline, `workflow_path` for a file) | What was sent, not what ran | -| Read-only view of it | `GET /api/jobs/{id}/workflow`, `JobManager.definition` | UI only; no MCP tool | -| Rerun | `POST /api/jobs/{id}/rerun`, `rerun_job` | Resubmits the spec verbatim, arguments included | -| Run directory | `///` (`dw/runs.py`) | Every file a step wrote | -| Manifest | `manifest.json` beside them | Status, times, engine version, device, seed actually used, arguments, per-step files | -| Run id | Directory name only | Never reported to the job or the job store | -| Reuse of a file | `output://` | Resolves out of that tree | -| Keeping a file | `keep_output` | Copies one output into the assets library | - -The record of a job is therefore split between a database row that knows the -definition and a directory that knows the outcome, and neither is complete. - -## Proposed shape - -### 1. Every run writes its realized workflow beside its manifest - -`Workflow.run` writes `workflow.json` into the run directory at the start of -the run, before the first step, so a crash or a cancel still leaves it. The -manifest points at it (`"workflow": {..., "realized": "workflow.json"}`). -CLI runs and server jobs get it alike, since the manifest is written there -too. A sub-workflow inherits its parent's run directory and writes none of -its own, as with the manifest. - -*Realized* means every mutable input is pinned, so the file reproduces the -run whatever changes later: - -| Reference | In the realized file | -|---|---| -| `variables` and the job's `arguments` | Arguments folded into the variable defaults; `variable:` references stay, so the file remains runnable with overrides | -| `seed` | The seed the run actually used (drawn if the workflow named none) | -| `prompt:name` | The stored prompt's text, inlined | -| `output:/latest/` | Rewritten to the concrete run id it resolved to | -| `output:` with an explicit run id | Kept | -| `asset:name` | Kept: a name in the library, which the export bundles | -| `constant:` | Kept: a value in the Python module the manifest's `dw_version` pins | -| `builtin:` sub-workflow | Kept: packaged with the engine, pinned by `dw_version` | -| sub-workflow by local path | Kept; the manifest records the file's SHA-256 | -| `previous_result:` | Kept: it names a step in the same file | - -The realized file is a valid workflow. `validate_workflow` accepts it and -`run_workflow(inline_workflow=...)` reproduces the run, which is the -simplest possible resume-by-hand and needs no new engine path. - -### 2. The job knows its run - -The worker reports the run id and the run directory (relative to the output -root) in a `run_start` event; `Job` records them and `jobs.sqlite` gains a -`run_id` column. `JobManager.realized(job_id)` reads the run's -`workflow.json`. A job from before this feature has no run id: the manager -falls back to the submitted definition and says so, rather than deriving a -run directory from file paths. - -### 3. MCP can read it - -`get_job_workflow(job_id)` returns the realized workflow, with `realized: -true`, or the submitted definition with `realized: false` and a sentence -saying why. Small enough to return inline: a workflow is text. - -### 4. Exporting a job - -`POST /api/jobs/{id}/export` gathers one job into -`/exports//`: - -``` -exports// - README.md what this is, how it was made, how to run it again - workflow.json the realized workflow (section 1) - manifest.json the run's manifest - job.json the job row: status, times, arguments, warnings, error - assets/ every asset:name the workflow references, by name - inputs/ every file an output: reference named from another run - outputs/ every file the manifest lists, in the run's layout -``` - -The tree is git-ready as it stands: text at the top, media in folders, no -absolute paths (the README notes that large media belongs in Git LFS). -`GET /exports/.zip` streams the same tree as one archive for a -browser or a laptop, built on request from the directory rather than kept as -a second copy. `exports` joins the reserved workspace names. - -MCP `export_job(job_id)` runs the export and returns the directory, the zip -URL, and the file list with sizes; the JSON files it also returns inline. -Like `download_output`, the directory is on the machine running the server, -and the tool says so. - -### 5. Recoverability - -This proposal writes the record; [resume.md](resume.md) is the design for -acting on it. Its stage 2, a per-step identity in the manifest, is served by -the realized file: each step's definition is there to compare, and the -manifest's per-step files say what it made. Its stage 3, rehydrating the -step cache from a previous run, is what turns the lighthouse retry into a -one-step rerun, and is planned as the next piece after this one lands. Until -then the realized file is what a user hands back to `run_workflow` to -reproduce a run, and what an agent edits to change one step of it. - -## Staging - -1. **The realized file.** Realization rules, written at run start, manifest - pointer, tests for every row of the table above. CLI and server alike. -2. **Run id on the job, and the MCP read.** The `run_start` event, the - column, `JobManager.realized`, `get_job_workflow`. -3. **Export.** The route, the directory layout, the README, the zip, the MCP - tool, the reserved name. -4. **Resume**, as resume.md stage 3, its own proposal and plan. - -Stages 1 and 2 are the useful unit; 3 stands on them. - -## Open questions - -- **Sub-workflows by local path.** Inlining them would make the realized file - self-contained, but the schema's `workflow` step takes a path, not a - definition. Decided in the design: keep the path, record the file's - digest in the manifest, and inline only if the digest turns out to be - what people trip over. -- **Realizing `prompt:` loses the name.** Decided in the design: the - manifest lists the stored prompts the run inlined, so the realized file - stays schema-clean and the name is not lost. -- **Export size.** A six-shot H3 run is hundreds of megabytes of video. The - export copies rather than hard-links, since `exports/` is what a user moves - or deletes; the response reports the total so an agent can quote it. -- **`inputs/` and the realized file.** The export copies what an `output:` - reference named from another run, but does not rewrite the reference, - since the realized file is the immutable record of the run. The README - says where each input came from. -- **History rebuild.** A run directory with `workflow.json` and - `manifest.json` is enough to reconstruct a job row without the database. - Not proposed here; noted because this is what makes it possible. diff --git a/docs/proposals/live-latent-preview.md b/docs/proposals/live-latent-preview.md deleted file mode 100644 index f370b217..00000000 --- a/docs/proposals/live-latent-preview.md +++ /dev/null @@ -1,340 +0,0 @@ -# Proposal: live latent preview during generation - -Status: design only, no code changes. Responds to audit finding **U1 — no live -preview during generation**: "the single biggest 'where's the ComfyUI thing' -moment" for users coming from node-based tools. - -## Where things stand today - -The step callback already exists and already receives the latents. In -`dw/pipeline_processors/pipeline.py`, `_with_step_callback` (around -`pipeline.py:503-532`) injects a `callback_on_step_end` into any pipeline call -whose signature accepts one: - -```python -def on_step_end(pipe, step_index, timestep, callback_kwargs): - total = getattr(pipe, "_num_timesteps", None) or num_steps - run_context.emit("pipeline_step", step=step_index + 1, total_steps=total) - if total is not None and step_index + 1 >= total: - emit_phase("decoding") - run_context.check_cancelled() - return callback_kwargs -``` - -`callback_kwargs` is the dict diffusers passes to every `callback_on_step_end` -implementation, and for essentially every pipeline in this family it carries -`callback_kwargs["latents"]` — the current denoised-so-far latent tensor, -still on the accelerator, still in the pipeline's native latent space. Today -the callback reads nothing from it and returns it unchanged. So the audit's -premise is accurate: the raw material for a preview is sitting in the -callback's hand every step, and only `step_index`/`total` ever leave the -callback as `pipeline_step`, plus the coarse `phase` string -(`loading` / `cached` / `generating` / `decoding` / `saving` / `task`, defined -in `dw/events.py`). - -**Event path today**: `RunContext.emit()` (`dw/events.py:34`) → the worker -subprocess's `on_event` callback (`dw/worker.py:207`, wired as -`on_event=lambda event: self.result_queue.put({"type": "progress", **event})`) -→ crosses a `multiprocessing.Queue` (pickled) into the server process → -`JobManager` appends it to `job.events` (`dw/server/jobs.py:215`, -`add_event`) → the SSE endpoint `/api/jobs/{id}/events` -(`dw/server/app.py:394-429`) streams new events to the browser as they land → -`ui/src/lib/progress.ts` (`stepProgress`) reduces the event list into a -step/total counter and a phase label → `JobPage.svelte` renders a progress -bar and phase text (`ui/src/lib/pages/JobPage.svelte:166-204`). - -Two details of that path matter a lot for the design below and are easy to -miss from the UI side alone: - -1. **The callback runs inside the worker subprocess**, not the server - process. Anything produced from `latents` — decode, resize, JPEG-encode — - either happens in the worker before the event is queued, or the raw - tensor would have to cross the `multiprocessing.Queue` itself (pickling a - GPU/CPU tensor through IPC), which is worse in every dimension than - encoding first and sending bytes. -2. **`job.events` is not just a live stream — it is what gets persisted.** - `dw/server/jobs.py` keeps the last `MAX_PERSISTED_EVENTS` events per job - and writes them into `~/.diffusers_helper/jobs.sqlite` as the job's - history tail. An event type that carries image bytes must not be treated - as a normal event for persistence purposes, or the job history database - silently grows by a JPEG per preview frame per job, forever. - -## 1. Decode approach - -Three ways to turn `callback_kwargs["latents"]` into something a browser can -show, in order of speed: - -**A. A tiny per-family approximate decoder (TAESD-style).** Madebyollin's -`taesd` and `taesdxl` give a near-instant, low-quality RGB approximation -directly from SD/SDXL latents — a few ms on GPU, cheap enough to run every -step. A Flux-compatible tiny decoder (`taef1`) exists too. These are small -(~5MB) separate weight files, loaded once and cached like any other -component. - -- *Pro*: fast enough to not matter, works well within the resolution-vs-speed - trade-off, this is what ComfyUI and most preview implementations actually - use. -- *Con*: it is a new dependency (or vendored weights) per pipeline **family**, - not universal. The latent space a tiny decoder was trained against is - architecture-specific — an SD1.5 TAESD does not decode SDXL latents - correctly, let alone a DiT-family model's latents. Coverage has to be - built and verified per family, and it silently produces garbage (not an - error) on a mismatched family if someone lists the wrong one. - -**B. The real VAE at reduced resolution / reduced precision.** Since the -pipeline's own VAE is already loaded for the final decode, the callback -could call `pipeline.vae.decode()` on a downscaled crop or a -`nearest`-interpolated shrink of the latent tensor, then downsize the -resulting image further before encoding. - -- *Pro*: no new dependency, no new weights, no per-family coverage question — - if the pipeline runs at all, this works, because it reuses the exact - decoder the final image goes through. Simplest to reason about and to - ship first. -- *Con*: slower. A full VAE decode is the same operation that already shows - up as the `"decoding"` phase after the denoise loop finishes, and on some - pipelines (video, high-res image) that step alone is seconds — the - callback comment even flags it as long enough to leave "the bar sitting at - 100%" for video. Doing that every N steps multiplies that cost by - `num_steps / N`. Rough order of magnitude on a modern CUDA GPU, a single - SDXL-size VAE decode is ~100-300ms; on MPS it is commonly 3-6x slower; on - CPU it can be many seconds. Even decoding a downsampled latent doesn't - avoid the VAE's own architecture cost (it's convolutional and scales with - input size, so a 4x-smaller latent decode is meaningfully cheaper, roughly - proportional to pixel count — call it 4-16x faster than a full decode, but - still not free, and still on the pipeline's own device, i.e. still - competing with the denoise loop for the same accelerator). - -**C. No preview for pipelines without a fast path; degrade to today's -counter.** Not every pipeline in this repo has a matching tiny decoder or an -easily-introspectable single-frame VAE. `workflows/*.json` and -`dw/workflows/*.json` in this repo currently exercise `StableDiffusionPipeline`, -`StableDiffusion3Pipeline`, `StableDiffusionControlNetPipeline`, -`ZImagePipeline`, `Krea2Pipeline`, and `MochiPipeline` — i.e. mostly -SD-family and DiT-family image pipelines, plus one video pipeline (Mochi). -The engine additionally documents support for exotic families the audit -calls out by name — LTX-2 (video+audio, muxed via PyAV per CLAUDE.md), -MiniMax H3 (dual video/audio latent schedules, on-demand component -residency) — for which no public TAESD-equivalent exists and whose VAE -decode is itself expensive and multi-stage (video VAEs decode a temporal -window, not a single frame; audio has no "frame" to preview at all). - -- *Pro*: honest, and it's the correct fallback regardless of which of A/B is - chosen for the covered families — a `NotImplementedError`-shaped gap is - fine as long as the UI treats "no preview available" as a normal, - expected state rather than an error. -- *Con*: none really — this option isn't a real alternative to A/B so much as - the required behavior at the edges of whichever one is picked. - -**Recommendation**: start with **B** (real VAE, reduced resolution) as the -first-version decode path — zero new dependencies, works for every pipeline -that has a VAE attribute and a `callback_on_step_end` parameter, and the -performance question is answered by throttling frequency (§2) rather than by -the decode implementation. Treat **A** as a fast-follow per family, gated on -someone actually measuring the B-path cost as too high for that family -in practice; it's the better long-term answer but shouldn't block v1. **C** -is not optional — it is what every pipeline falls back to when neither A nor -B applies, from day one. - -## 2. Performance cost: every step vs. every N steps - -Decoding every step maximizes preview smoothness but means the preview cost -is paid `num_inference_steps` times per run, on the same accelerator that is -doing the actual denoising work — there is no separate "preview GPU." That -cost is not symmetric across backends, and CLAUDE.md's device-support notes -explain why: - -- **CUDA**: has spare headroom for this most of the time — TF32 matmul, - cuDNN benchmark mode, and enough throughput that a partial VAE decode - every few steps is close to free relative to a 20-50 step denoise loop. -- **MPS**: explicitly the platform where the codebase already treats - "sequential" offload as too expensive and downgrades it to "model" - offload with a warning, because per-submodule streaming on unified memory - has no separate CPU/accelerator pools to arbitrage. A preview decode - competing for the same unified memory and the same execution unit as the - denoise loop is a strictly worse deal on MPS than on CUDA — no autocast, - attention slicing already on by default (i.e. already trading speed for - memory), no torch.compile. Every-step preview here risks visibly slowing - down the generation it's supposed to be a nicety alongside. -- **CPU**: the existing "CPU is slow, expect a warning" path. A preview - decode every step here is not a nicety, it's a tax nobody asked for; N - should default much larger (or previews default off) on CPU. - -**Recommendation**: decode on an interval (every N steps, N configurable, -default something like every 3-5 steps or ~every 10% of the run, whichever -is coarser) rather than every step, and pick N adaptively by device type -using the same `get_device_type()` the rest of the codebase already uses to -branch CUDA/MPS/CPU behavior (never `== "cuda"`, per CLAUDE.md). Also always -skip the last one or two steps' preview in favor of just waiting for the -real final decode — a preview that lands 200ms before the real image does -is wasted work. This is a pure frequency knob, not a decode-quality knob: -whichever of A/B from §1 is used, N-step throttling is what actually -protects wall-clock time, and it composes with either. - -## 3. Transport: how the preview image reaches the browser - -**Option 1: base64 inline in the existing SSE event stream.** Add a new -`preview` event type alongside `pipeline_step`/`phase`, with `data` as a -base64 JPEG, emitted through the same `RunContext.emit()` → worker queue → -`JobManager.add_event()` → SSE path everything else uses. - -- *Pro*: reuses 100% of the existing plumbing — no new endpoint, no new - polling loop in the UI, ordering falls out for free (the SSE stream is - already ordered and resumable via `seq`/`Last-Event-ID`). -- *Con*: event size becomes proportional to preview frequency × image size, - and — the detail that's easy to miss from the UI side — **every event - appended via `job.add_event()` is a candidate for persistence**. `dw/server/jobs.py` - keeps the trailing `MAX_PERSISTED_EVENTS` per job and writes that tail - into `jobs.sqlite` as JSON. A `preview` event carrying kilobytes of base64 - would blow that budget out compared to today's few-hundred-byte JSON - events, and would write image bytes into a SQLite history table that was - designed for a text/number event tail. This is fixable (see recommendation) - but is not free by construction the way it looks at first glance. - -**Option 2: a separate polling endpoint** (`GET -/api/jobs/{id}/preview` returning the latest frame, or a 204 if none yet), -polled by the UI on an interval (e.g. every 500ms while `running`). - -- *Pro*: completely decouples preview traffic from the event/history system — - nothing about it touches `job.events` or `jobs.sqlite`, so no persistence - concern at all. Simple to reason about: it's just "what's the latest - frame," no ordering or replay semantics needed. -- *Con*: a second connection concept alongside SSE (poll timers, not just an - `EventSource`), and it either always shows the "latest" frame (fine, since - older previews are worthless anyway) or needs its own tiny sequence number - if the UI wants to avoid redundant re-renders of the same frame. - -**Option 3: write preview frames to a temp file, UI polls a static path.** -Worker writes `~/.diffusers_helper/previews/{job_id}.jpg` (or similar) each -N steps; the UI does `` -on a timer. - -- *Pro*: avoids putting image bytes through IPC as event payloads at all, - and disk I/O for a JPEG is cheap. -- *Con*: adds filesystem lifecycle management that doesn't otherwise exist - for a job — cleanup on completion/cancellation/crash, a new place path - traversal / naming needs `validate_path()` treatment per the security - rules, and multi-worker or multi-job-concurrency considerations (this - doesn't apply today since there's one worker subprocess, but it's a - needless new constraint to bake in). It's strictly worse than Option 2 for - no offsetting benefit here — Option 2 gets the same "just fetch the - latest thing" simplicity via HTTP response body instead of a file, without - a new directory to manage or secure. - -**Recommendation: Option 2**, a small dedicated polling endpoint, **not** -threading preview frames through the SSE/event-log system. The event stream -is the right place for state that participates in job history (`phase`, -`pipeline_step`, `log`, `job_status` all make sense to see when you reload a -job's page later, or as the persisted tail in `jobs.sqlite`); a preview -frame does not — nobody wants a 10-year-old job history page bringing back -a base64 image of a mostly-noisy step 8/30. Concretely: the worker holds the -latest encoded preview frame (as bytes, in memory, keyed by job id — not -routed through `RunContext.emit`/`job.events` at all, so it never touches -persistence), and a new endpoint on the job manager exposes "give me the -latest frame for this job or 404/204 if there is none yet." This sidesteps -the base64-in-SQLite problem by construction rather than by having to -special-case one event type's persistence behavior. Option 1 is the one to -avoid specifically because of the JobManager/sqlite coupling discovered -above, not because SSE itself is a bad transport for images in general. - -## 4. UI surface and the cancellation connection - -The natural place is `JobPage.svelte`, next to the progress bar it already -renders (`ui/src/lib/pages/JobPage.svelte:166-204`, the `{#if denoise}` -block with the fill bar and step counter). A preview thumbnail — modest -size, maybe 256-384px on the long edge — sitting above or beside that bar, -updated by the Option-2 poll while `running` is true, is a small, additive -change to a component that already owns the "this job is actively -generating" rendering branch. - -This is exactly where the audit's framing connects preview to cancellation -(S6, per the audit's own numbering): the Cancel button already sits right -there (`JobPage.svelte:135-142`, "stop this run at the next step — models -stay cached"). Today a user decides to cancel based on a step counter and an -ETA — abstract numbers. A live preview turns that into an informed decision: -"this composition is wrong, kill it now" instead of waiting out a 30-step -run to find out. The two features multiply each other's value more than -either does alone; this is a good argument for landing them in the same UI -change even though they're separable pieces of work. No cancellation -*semantics* need to change — `run_context.check_cancelled()` already fires -every step in `on_step_end` — this is purely about giving the user something -worth acting on earlier. - -## 5. Scope for a first version - -Evidence from the repo on which families are actually exercised: -`workflows/*.json` (the runnable top-level examples) and `dw/workflows/*.json` -(the packaged built-ins) together reference `StableDiffusionPipeline` (3), -`ZImagePipeline` (2), `Krea2Pipeline` (2), `StableDiffusionControlNetPipeline` -(1), `StableDiffusion3Pipeline` (1), and `MochiPipeline` (1) — nine workflow -files total, all image pipelines except Mochi, and all standard -UNet/DiT-with-a-VAE architectures with nothing exotic about their latent -space. There is no LTX-2 or MiniMax H3 example workflow in the repo despite -CLAUDE.md documenting support for them — those are the pipelines described -in §1 as needing option C (no preview) regardless of which decode approach -is chosen for the rest. - -**v1 scope (recommended)**: -- Decode approach: **B** (real VAE, reduced resolution), gated to pipelines - whose loaded `pipeline` object exposes a `.vae` with `.decode()` and where - `callback_on_step_end` is already wired in (i.e. reuse the exact - `"callback_on_step_end" not in parameters` check `_with_step_callback` - already does — no new pipeline-capability detection needed). -- Frequency: every N steps, N chosen by `get_device_type()` (§2), never on - the final 1-2 steps. -- Transport: Option 2, a polling endpoint outside the event/persistence - path (§3). -- UI: a thumbnail in `JobPage.svelte`'s existing running-job panel, polled - only while `running` (§4). -- Explicitly out of scope for v1: LTX-2, MiniMax H3, and any other - video/audio pipeline where a "frame" isn't a well-defined single-step - concept; TAESD-family fast decoders (tracked as a fast-follow per §1); - any change to the SSE event stream or `jobs.sqlite` schema. - -**Stretch goals**: TAESD/TAESDXL/TAEF1 fast decoders per family once B's -real-world cost is measured and found wanting on a specific family; a -video-pipeline preview (e.g. decode-and-show the first/most-recent frame of -a Mochi latent) once there's a concrete workflow using it in this repo to -validate against; folding preview state into `progress.ts`'s reducer if the -polling model ever needs to become event-driven for some new consumer (the -MCP server's `get_job_events` polling twin, for instance). - -## 6. Effort and risk - -**Effort**: moderate, not large, if scoped as above. - -- Backend: extend `_with_step_callback` to optionally decode-and-cache a - preview frame every N steps (new code path in `pipeline.py`, guarded by a - capability check so pipelines without a VAE are unaffected); a small new - in-memory latest-frame store keyed by job id in `dw/server/jobs.py` - (parallel to, not part of, `job.events`); one new FastAPI endpoint in - `dw/server/app.py`. No new dependency for v1 (Option B avoids the TAESD - question entirely), no schema change to `jobs.sqlite`. -- Frontend: one new poll loop and a thumbnail element in `JobPage.svelte`; - no changes needed to `progress.ts`'s event reducer since preview state - isn't an event-stream concern. -- MCP surface: `dw_mcp` wraps the REST API per CLAUDE.md's MCP section; a - preview endpoint would need a corresponding tool or explicit non-coverage - note, matching how the SSE stream itself is already excluded there in - favor of `get_job_events`. - -**Risk**: -- *Correctness*: modest — worst case a preview frame is visually wrong for - a moment (a plain "decode with the real VAE" call is intrinsically - correct for the color space; there's no new numerical logic to get - subtly wrong the way a mismatched TAESD variant could). -- *Performance regression*: the main real risk, specifically on MPS/CPU as - described in §2 — needs a device-aware default and probably a - user-visible on/off toggle (or an automatic "generation is slow, previews - auto-disabled" heuristic) so a slow backend doesn't silently get slower - because of an opt-out-only feature. -- *Concurrency*: today one worker subprocess runs one job at a time, so a - single "latest frame per job id" store is sufficient; if the worker model - ever becomes multi-job, the preview store needs the same job-id keying - discipline the rest of `JobManager` already uses — not a new problem, - just something to keep consistent. -- *Scope creep*: the biggest practical risk is reaching for TAESD coverage - or video-pipeline previews in v1 instead of treating them as the stretch - goals in §5 — B-and-throttle-and-poll is enough to deliver the "oh, it's - actually painting" moment the audit is asking for, without turning this - into a per-pipeline-family research project before anything ships. diff --git a/docs/proposals/resume.md b/docs/proposals/resume.md index 4c62f830..19dd21cd 100644 --- a/docs/proposals/resume.md +++ b/docs/proposals/resume.md @@ -129,7 +129,9 @@ version is already scoped to a session where the user can see what happened. the manifest records the seed actually used. Without both, a seedless workflow has nothing to resume against. 2. **Manifest carries step identity.** *Satisfied by the realized workflow* - ([job-record-and-export.md](job-record-and-export.md)): every run now writes + (shipped 2026-09-08; see CLAUDE.md's run-directory notes, and + `docs/superpowers/specs/2026-09-08-job-record-and-export-design.md` for the + design): every run now writes `workflow.json` beside its manifest with every mutable input pinned, so each step's definition as it actually ran is on disk to compare against, and the manifest's per-step files say what it made. A per-step digest may still be diff --git a/dw/server/app.py b/dw/server/app.py index 80856947..f6456ec6 100644 --- a/dw/server/app.py +++ b/dw/server/app.py @@ -68,9 +68,11 @@ create_workspace, delete_workspace, example_libraries, + forget_workspace_usage, named_workspace, workspace_contents, workspace_names, + workspace_usage, ) from ..workflow_sources import ( EXAMPLES_ORIGIN, @@ -863,13 +865,26 @@ def get_job_workflow(job_id: str): "id": job_id, "definition": definition, "realized": realized is not None, + # Which variable a new-seed rerun would draw into, or null when + # there is none - read from the workflow as written, since the + # realized copy above has its seed pinned to the integer it used + "seed_variable": manager.seed_variable(job_id), } + class RerunRequest(BaseModel): + new_seed: bool = Field( + default=False, + description="Draw a fresh seed into the workflow's seed variable. " + "Without it a rerun repeats the original arguments exactly, which " + "the step cache serves from the earlier run - the same seed and " + "inputs would produce the same files.", + ) + @app.post("/api/jobs/{job_id}/rerun", status_code=201) - def rerun_job(job_id: str): + def rerun_job(job_id: str, body: RerunRequest = RerunRequest()): """Queue a fresh job from a previous job's stored spec.""" try: - job = manager.rerun(job_id) + job = manager.rerun(job_id, new_seed=body.new_seed) except Exception as e: raise HTTPException(status_code=400, detail=str(e)) if job is None: @@ -1197,9 +1212,17 @@ def list_workspaces(): # first entry (always the default, see its docstring) is a named # workspace to describe individually names = workspace_names(root)[1:] if root else [] - described = [app.state.default_workspace.describe()] + listed = [app.state.default_workspace] for name in names: - described.append(named_workspace(root, name).describe()) + listed.append(named_workspace(root, name)) + described = [] + for space in listed: + entry = space.describe() + # Roughly how much disk it holds, cached for a minute inside + # workspace_usage - a listing is a glance, and a job writing + # into outputs moves the number continuously anyway + entry["usage"] = workspace_usage(space) + described.append(entry) return { "workspace_root": root.root if root else None, "default": DEFAULT_WORKSPACE_NAME, @@ -1217,6 +1240,7 @@ def add_workspace(request: WorkspaceRequest): raise HTTPException(status_code=400, detail=str(e)) except FileExistsError as e: raise HTTPException(status_code=409, detail=str(e)) + forget_workspace_usage() logger.info(f"Created workspace {request.name} at {created.root}") return created.describe() @@ -1271,6 +1295,7 @@ def remove_workspace(name: str, acknowledged: bool = False): ) except (ValueError, FileNotFoundError) as e: raise HTTPException(status_code=400, detail=str(e)) + forget_workspace_usage() logger.info(f"Deleted workspace {name}") return {"name": name, "deleted": True, "contents": contents} diff --git a/dw/server/jobs.py b/dw/server/jobs.py index 55cde2fc..22992614 100644 --- a/dw/server/jobs.py +++ b/dw/server/jobs.py @@ -10,6 +10,7 @@ import copy import json import queue +import random import sqlite3 import time import uuid @@ -26,6 +27,7 @@ validate_path, validate_workflow_path, ) +from ..realize import VARIABLE_PREFIX from ..runs import REALIZED_FILE_NAME from ..settings import resolve_path from ..workspace import DEFAULT_WORKSPACE_NAME @@ -604,13 +606,40 @@ def realized(self, job_id): logger.debug(f"No realized workflow for job {job_id}: {e}") return None - def rerun(self, job_id): + def seed_variable(self, job_id): + """The variable this job's workflow draws its seed from, or None. + + Read from the workflow as written, never from the realized copy the + run wrote: realization pins the top-level seed to the integer the run + used, so a realized workflow always looks like it names a literal. + + None means a new-seed rerun has nowhere to put one - either the seed + is a literal (an argument cannot override it) or the workflow names + no seed at all, in which case every run already draws a fresh one and + the step cache is off. + """ + definition = self.definition(job_id) + seed = (definition or {}).get("seed") + if not isinstance(seed, str) or not seed.startswith(VARIABLE_PREFIX): + return None + name = seed.removeprefix(VARIABLE_PREFIX) + return name if name in (definition.get("variables") or {}) else None + + def rerun(self, job_id, new_seed=False): """Queue a fresh job from a previous job's spec. Every root the original ran against (workflow_dir/output_dir/ asset_dir/workspace) rides along, not just the workflow identity - otherwise a rerun of a job from a named workspace would fall back to the manager's process-wide default and silently run somewhere else. + + `new_seed` draws a fresh seed into the workflow's seed variable. A + plain rerun of a seeded workflow repeats its arguments exactly, which + makes every step a step-cache hit: it republishes the earlier run's + files in a fraction of a second and generates nothing. That is the + cache doing its job - the same seed and the same inputs would produce + the same pixels - so the way to actually get another image is to + change the seed, and this is that. """ job = self.jobs.get(job_id) if job is not None: @@ -627,6 +656,20 @@ def rerun(self, job_id): } arguments = historical["arguments"] + if new_seed: + variable = self.seed_variable(job_id) + if variable is None: + raise ValueError( + "This workflow does not draw its seed from a variable, so " + "a rerun cannot change it. A workflow with no seed at all " + "already draws a fresh one every run." + ) + # Bounded to 53 bits rather than the 64 torch allows: this number + # goes out as JSON and comes back through a browser, where every + # integer is a double, and a seed that changed on the way through + # would be a seed nobody can reproduce + arguments = {**arguments, variable: random.getrandbits(53)} + workspace = spec.get("workspace") if ( workspace diff --git a/dw/workspace.py b/dw/workspace.py index 1b2dc102..3625b0fb 100644 --- a/dw/workspace.py +++ b/dw/workspace.py @@ -34,6 +34,7 @@ """ import os +import time from pathlib import Path # Set by an entry point that resolved a workspace, so a spawned worker @@ -453,26 +454,127 @@ def create_workspace(workspace, name): return named_workspace(workspace, name).ensure() +def _tree_usage(directory): + """Files and bytes under one directory, as (files, bytes). + + Walks with scandir and stats through the DirEntry, which reuses the stat + the directory read already did - the difference matters on an outputs + tree of thousands of generated files. Symlinks are counted as neither + file nor directory, so a link into a model cache cannot inflate the + number or send the walk outside the workspace. Anything unreadable is + skipped: this is a size to glance at, not an audit. + """ + files = 0 + total = 0 + stack = [directory] + while stack: + current = stack.pop() + try: + entries = list(os.scandir(current)) + except OSError: + continue + for entry in entries: + try: + if entry.is_dir(follow_symlinks=False): + stack.append(entry.path) + elif entry.is_file(follow_symlinks=False): + total += entry.stat(follow_symlinks=False).st_size + files += 1 + except OSError: + continue + return files, total + + def workspace_contents(workspace): """How much a workspace holds, for a client about to offer to delete it: file counts and total bytes per folder. Counting is the point - the number is what makes 'delete this workspace' an informed choice.""" summary = {} for folder in NAMED_SUBDIRS: - directory = os.path.join(workspace.root, folder) - files = 0 - total = 0 - for current, _dirs, names in os.walk(directory): - for entry in names: - try: - total += os.path.getsize(os.path.join(current, entry)) - except OSError: - continue - files += 1 + files, total = _tree_usage(os.path.join(workspace.root, folder)) summary[folder] = {"files": files, "bytes": total} return summary +# How long a computed workspace size stays good enough to answer with. The +# number is a glance at how much a workspace holds, not an accounting +# figure, and re-walking a large outputs tree for every listing would cost +# more than the precision is worth - a running job moves it continuously +# anyway +USAGE_CACHE_SECONDS = 60 + +# the folders counted (a tuple of paths) -> (monotonic time, usage) +_usage_cache = {} + + +def _own_folders(workspace): + """The folders whose bytes count as this workspace's own. + + A workspace's four folder properties are the candidates, deduplicated, + minus any that is not actually inside its root: a named workspace's + prompt library belongs to the root and is shared by every workspace, so + counting it here would count it once per workspace. A workspace with no + root of its own (a server configured folder by folder, where each can + point anywhere) has nothing to exclude, and every folder it names is its + own by definition. + """ + folders = [] + for path in ( + workspace.workflows, + workspace.prompts, + workspace.assets, + workspace.outputs, + ): + if not path or path in folders: + continue + if workspace.root and not _is_within(path, workspace.root): + continue + folders.append(path) + return folders + + +def _is_within(path, root): + """Whether a path is the root or sits under it.""" + try: + return os.path.commonpath([path, root]) == root + except ValueError: # different drives on Windows + return False + + +def workspace_usage(workspace, max_age=USAGE_CACHE_SECONDS): + """Roughly how much disk a workspace occupies: files and total bytes. + + Only the folders that are the workspace's own are counted, per + _own_folders - which is what keeps the shared prompt library from being + added to every workspace's total, and named workspaces from being + counted inside the default one (they sit beside its folders, not in + them). + """ + folders = _own_folders(workspace) + key = tuple(folders) + now = time.monotonic() + cached = _usage_cache.get(key) + if cached and now - cached[0] < max_age: + return cached[1] + files = 0 + total = 0 + for directory in folders: + if not os.path.isdir(directory): + continue + count, size = _tree_usage(directory) + files += count + total += size + usage = {"files": files, "bytes": total} + _usage_cache[key] = (now, usage) + return usage + + +def forget_workspace_usage(): + """Drop every cached size - after creating or deleting a workspace, + where a stale answer would be visibly wrong rather than merely old.""" + _usage_cache.clear() + + def delete_workspace(workspace, name): """Remove a named workspace and everything in it. diff --git a/dw_mcp/diagnose.py b/dw_mcp/diagnose.py index 83973085..8ae645fd 100644 --- a/dw_mcp/diagnose.py +++ b/dw_mcp/diagnose.py @@ -82,6 +82,9 @@ def get_job_workflow(client, job_id): "job_id": job_id, "realized": bool(body.get("realized")), "workflow": body.get("definition"), + # Which variable rerun_job(new_seed=True) would draw into, null when + # the workflow has none - see rerun_job on why that matters + "seed_variable": body.get("seed_variable"), "next": "Pass `workflow` to save_workflow to keep it in the catalog " "under a name, or edit it and pass it to run_workflow as " "inline_workflow.", @@ -139,14 +142,22 @@ def cancel_job(client, job_id): return client.post_json(api_path("api", "jobs", job_id, "cancel")) -def rerun_job(client, job_id, acknowledged_cost=False): +def rerun_job(client, job_id, acknowledged_cost=False, new_seed=False): """Queue a fresh job from a previous job's stored spec. This costs the same GPU time as `run_workflow` and passes through the same gate - a rerun is a run, and the gate would be worth nothing if a job id bought - a way around it.""" + a way around it. + + `new_seed` draws a fresh seed into the workflow's seed variable. Without + it the arguments repeat exactly, and a seeded workflow's rerun is served + whole from the step cache - the earlier run's files, republished in a + fraction of a second, with `reused: true`. Ask for a new seed when the + point is a different image rather than the same one again.""" if not acknowledged_cost: raise DwApiError(COST_REFUSAL) - return client.post_json(api_path("api", "jobs", job_id, "rerun")) + return client.post_json( + api_path("api", "jobs", job_id, "rerun"), {"new_seed": new_seed} + ) def move_job(client, job_id, direction): diff --git a/dw_mcp/server.py b/dw_mcp/server.py index 409979d3..533b6c19 100644 --- a/dw_mcp/server.py +++ b/dw_mcp/server.py @@ -646,12 +646,24 @@ def cancel_job(job_id: str) -> dict: Deliberately not gated - it ends a cost rather than starting one.""" return diagnose.cancel_job(client, job_id) - def rerun_job(job_id: str, acknowledged_cost: bool = False) -> dict: + def rerun_job( + job_id: str, acknowledged_cost: bool = False, new_seed: bool = False + ) -> dict: """Queue a fresh job from a previous job's stored specification. THIS COSTS GPU TIME: a rerun is a run - it occupies the machine for minutes and the engine runs one job at a time. Tell the user what - will run and get their go-ahead, then pass acknowledged_cost=true.""" - return diagnose.rerun_job(client, job_id, acknowledged_cost=acknowledged_cost) + will run and get their go-ahead, then pass acknowledged_cost=true. + + Pass new_seed=true for a different image: a workflow that pins its + seed reruns to the same pixels, and the step cache serves that whole + run from the earlier one's files (marked `reused`) in a fraction of a + second rather than generating anything.""" + return diagnose.rerun_job( + client, + job_id, + acknowledged_cost=acknowledged_cost, + new_seed=new_seed, + ) def move_job( job_id: str, direction: Literal["up", "down", "front", "back"] diff --git a/tests/test_rerun_new_seed.py b/tests/test_rerun_new_seed.py new file mode 100644 index 00000000..de27b09f --- /dev/null +++ b/tests/test_rerun_new_seed.py @@ -0,0 +1,153 @@ +"""Rerunning a seeded workflow. + +A workflow that pins its seed is a workflow the step cache can serve whole: +the same definition, the same arguments and the same output root make every +step a hit, so "Run again" finishes in a third of a second and republishes +the earlier run's files with `reused: true`. That is the cache working, but +it leaves the button promising something it did not do. A rerun can instead +draw a new seed, which is what "run it again" means for a generation - and +the workflow says which variable to draw it into. +""" + +import json + +import pytest +from fastapi.testclient import TestClient + +from dw.server.app import create_app +from dw.server.jobs import JobManager + +from .test_server import ScriptedWorkerManager, success_script, valid_workflow +from .test_server import wait_for_status + + +def seeded_workflow(job_id="seeded"): + """A workflow whose seed points at a declared variable - the shape every + catalog template uses, `workflows/models/z-image.json` included.""" + workflow = valid_workflow(job_id) + workflow["seed"] = "variable:seed" + workflow["variables"] = {**workflow["variables"], "seed": 42} + return workflow + + +def literal_seed_workflow(job_id="pinned"): + workflow = valid_workflow(job_id) + workflow["seed"] = 7 + return workflow + + +@pytest.fixture +def server(tmp_path): + (tmp_path / "workflows").mkdir(exist_ok=True) + + def make(script=success_script): + manager = JobManager( + str(tmp_path / "outputs"), + worker_manager=ScriptedWorkerManager(script), + history_path=str(tmp_path / "jobs.sqlite"), + ) + app = create_app( + workflow_dir=str(tmp_path / "workflows"), + output_dir=str(tmp_path / "outputs"), + job_manager=manager, + ) + return TestClient(app, base_url="http://localhost") + + return make + + +def submit(client, workflow, arguments=None): + job = client.post( + "/api/jobs", json={"workflow": workflow, "arguments": arguments or {}} + ).json() + wait_for_status(client, job["id"], ["succeeded"]) + return job["id"] + + +def test_the_workflow_route_names_the_seed_variable(server): + """Which variable the seed reads is the client's cue that a new-seed + rerun is available at all - it is a property of the workflow, so the + server answers it rather than the page parsing definitions.""" + with server() as client: + seeded = submit(client, seeded_workflow()) + literal = submit(client, literal_seed_workflow()) + seedless = submit(client, valid_workflow()) + + assert client.get(f"/api/jobs/{seeded}/workflow").json()["seed_variable"] == ( + "seed" + ) + # A literal seed cannot be overridden by an argument, and a seedless + # workflow already draws a fresh seed every run - neither offers one + assert client.get(f"/api/jobs/{literal}/workflow").json()["seed_variable"] is ( + None + ) + assert client.get(f"/api/jobs/{seedless}/workflow").json()["seed_variable"] is ( + None + ) + + +def test_a_rerun_with_a_new_seed_draws_one_into_that_variable(server): + with server() as client: + original = submit(client, seeded_workflow(), {"prompt": "a cat"}) + + rerun = client.post(f"/api/jobs/{original}/rerun", json={"new_seed": True}) + assert rerun.status_code == 201 + arguments = rerun.json()["arguments"] + + assert arguments["prompt"] == "a cat" # everything else rides along + assert isinstance(arguments["seed"], int) + assert arguments["seed"] != 42 + # JSON numbers are IEEE doubles in every browser that reads this back + assert arguments["seed"] < 2**53 + + +def test_two_new_seed_reruns_do_not_draw_the_same_seed(server): + with server() as client: + original = submit(client, seeded_workflow()) + seeds = { + client.post(f"/api/jobs/{original}/rerun", json={"new_seed": True}).json()[ + "arguments" + ]["seed"] + for _ in range(5) + } + assert len(seeds) == 5 + + +def test_a_plain_rerun_still_repeats_the_original_arguments(server): + """The default is unchanged: same spec, same arguments, and the step + cache is free to serve it.""" + with server() as client: + original = submit(client, seeded_workflow(), {"prompt": "a cat"}) + rerun = client.post(f"/api/jobs/{original}/rerun") + assert rerun.status_code == 201 + assert rerun.json()["arguments"] == {"prompt": "a cat"} + + +def test_a_new_seed_rerun_of_a_workflow_that_has_no_seed_variable_is_refused(server): + """Refused rather than quietly run as an ordinary rerun: the caller + asked for a different image and would otherwise get the cached one.""" + with server() as client: + literal = submit(client, literal_seed_workflow()) + response = client.post(f"/api/jobs/{literal}/rerun", json={"new_seed": True}) + assert response.status_code == 400 + assert "seed" in response.json()["detail"].lower() + + +def test_a_job_launched_from_a_file_finds_its_seed_variable(tmp_path, server): + """The seed reference lives in the workflow as written, not in the + realized copy the run wrote - which pins the seed to the integer it + used, and would look like a literal.""" + workflow_dir = tmp_path / "workflows" + workflow_dir.mkdir(exist_ok=True) + (workflow_dir / "Seeded.json").write_text(json.dumps(seeded_workflow())) + + with server() as client: + job = client.post("/api/jobs", json={"workflow_path": "Seeded"}).json() + wait_for_status(client, job["id"], ["succeeded"]) + + assert client.get(f"/api/jobs/{job['id']}/workflow").json()[ + "seed_variable" + ] == ("seed") + rerun = client.post(f"/api/jobs/{job['id']}/rerun", json={"new_seed": True}) + assert rerun.status_code == 201 + assert isinstance(rerun.json()["arguments"]["seed"], int) diff --git a/tests/test_server_workspaces.py b/tests/test_server_workspaces.py index 1e3110b9..07892bde 100644 --- a/tests/test_server_workspaces.py +++ b/tests/test_server_workspaces.py @@ -738,3 +738,33 @@ def test_history_migrates_rows_that_predate_workspaces(tmp_path): ).fetchone() assert "workspace" in columns assert stored[0] == "default" + + +def test_the_listing_reports_each_workspace_s_disk_usage(server, workspace_root): + """Every workspace in the listing carries roughly how much disk it holds. + + The shared prompt library counts once - against the default workspace, + whose folder it is - rather than once per workspace, and a named + workspace counts only what is under its own root. + """ + from dw.workspace import forget_workspace_usage + + with server() as client: + assert client.post("/api/workspaces", json={"name": "shots"}).status_code == 201 + with open(os.path.join(workspace_root.prompts, "scenic.json"), "w") as f: + f.write("x" * 500) + shots_assets = os.path.join(workspace_root.root, "shots", "assets") + with open(os.path.join(shots_assets, "big.bin"), "wb") as f: + f.write(b"0" * 4096) + + # the listing caches its walk for a minute, and the writes above + # landed after the create already primed it + forget_workspace_usage() + spaces = { + w["name"]: w["usage"] + for w in client.get("/api/workspaces").json()["workspaces"] + } + + assert spaces["shots"] == {"files": 1, "bytes": 4096} + assert spaces["default"]["files"] == 1 + assert spaces["default"]["bytes"] == 500 diff --git a/ui/src/lib/api.ts b/ui/src/lib/api.ts index 6539ec61..237c151f 100644 --- a/ui/src/lib/api.ts +++ b/ui/src/lib/api.ts @@ -270,11 +270,23 @@ export const api = { * submitted. 404s when the job named a workflow file that is no longer * readable. */ getJobWorkflow: (id: string) => - request<{ id: string; definition: Record; realized: boolean }>( - `/api/jobs/${id}/workflow`, - ), - rerunJob: (id: string) => - request(`/api/jobs/${id}/rerun`, { method: 'POST' }), + request<{ + id: string + definition: Record + realized: boolean + /** The variable a new-seed rerun would draw into, null when the + * workflow has none - the cue for whether to offer that at all. */ + seed_variable: string | null + }>(`/api/jobs/${id}/workflow`), + /** Queue the job again. `newSeed` draws a fresh seed into the workflow's + * seed variable; without it the arguments repeat exactly, which the step + * cache serves from the earlier run rather than generating anything. */ + rerunJob: (id: string, newSeed = false) => + request(`/api/jobs/${id}/rerun`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ new_seed: newSeed }), + }), listTasks: () => request<{ commands: string[] @@ -376,6 +388,7 @@ export const api = { assets: string | null outputs: string prompts: string | null + usage?: { files: number; bytes: number } }[] }>('/api/workspaces'), createWorkspace: (name: string) => diff --git a/ui/src/lib/pages/JobPage.svelte b/ui/src/lib/pages/JobPage.svelte index 73ae0515..fd5dd254 100644 --- a/ui/src/lib/pages/JobPage.svelte +++ b/ui/src/lib/pages/JobPage.svelte @@ -1,11 +1,12 @@