Skip to content

Commit d890661

Browse files
dkackmanclaude
andcommitted
Merge feat/462-B-worker-pool: #676 stage B - worker pool and VRAM-aware dispatcher
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 parents 0402df9 + 0a2508d commit d890661

25 files changed

Lines changed: 1161 additions & 157 deletions

‎CLAUDE.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ python -m dw.test
4141
the seam, and the test or check that enforces it. Open it before grepping - it names the
4242
module, and that module's docstring holds the detail.
4343

44-
- `dw.serve` runs every job in one persistent spawned worker (`dw/worker.py`, managed by
44+
- `dw.serve` runs every job in a persistent spawned worker (one per `--devices` card) (`dw/worker.py`, managed by
4545
`dw/worker_manager.py`) that keeps models cached between runs, so a change to engine
4646
code needs a server restart before a job sees it.
4747
- The packaged `dw/workflows/` is what a `builtin:` step names (resolved in

‎docs/ARCHITECTURE.md‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,8 @@ one to open.
107107
| No seed, no cache | `dw/workflow_run.py`: `prepare_run` | A workflow that sets no `seed` skips the cache. | `tests/test_workflow_step_cache.py::TestCacheHits::test_an_unseeded_workflow_has_no_hits` |
108108
| A different result | `dw/server/routes/jobs.py`, `dw_mcp/tools_jobs.py`: `rerun_job` | `rerun` with `new_seed` draws a fresh seed into the workflow's seed variable. | `tests/test_rerun_new_seed.py::test_a_rerun_with_a_new_seed_draws_one_into_that_variable` |
109109
| Worker protocol | `dw/worker_protocol.py`: `parse_reply` | Every command and reply is a frozen dataclass that travels as a wire dict. | `tests/test_worker_messages.py::test_from_wire_inverts_to_wire`, `tests/test_worker_messages.py::test_an_unknown_reply_type_is_kept_whole_rather_than_raised` |
110-
| Persistent worker | `dw/worker.py`, `dw/worker_manager.py`, `dw/serve.py` | Jobs run in one spawned worker that keeps models loaded, so a change to engine code needs a server restart. | — |
111-
| Worker card pinning | `dw/devices.py`: `worker_environment`, `pinned_environment`, `resolve_serve_devices`, `check_device_present`, `device_label` | A worker named with a CUDA index is spawned with `CUDA_VISIBLE_DEVICES` set to that card and `DW_DEVICE=cuda` (a bare `cuda` pins nothing), so index-0 reads inside it measure its own card; `--devices`/`devices` takes one entry, a card the box lacks is refused at startup naming the cards present, and a job's `device` is `device_label`'s `"cuda:N <card name>"`. | `tests/test_worker_manager.py::TestWorkerPinning::test_a_cuda_index_pins_the_worker_through_the_environment`, `tests/test_worker_manager.py::TestWorkerPinning::test_a_bare_cuda_sets_no_pin`, `tests/test_serve_main.py::TestConfigureDevices::test_a_refusal_exits_with_code_2`, `tests/test_serve_main.py::TestResolveServeDevices::test_a_card_the_machine_lacks_is_refused_naming_the_cards_present` |
110+
| Persistent worker | `dw/worker.py`, `dw/worker_manager.py`, `dw/serve.py` | Jobs run in spawned workers (one per card) that keep models loaded, so a change to engine code needs a server restart. | — |
111+
| Worker card pinning | `dw/devices.py`: `worker_environment`, `pinned_environment`, `resolve_serve_devices`, `check_device_present`, `device_label` | A worker named with a CUDA index is spawned with `CUDA_VISIBLE_DEVICES` set to that card and `DW_DEVICE=cuda` (a bare `cuda` pins nothing), so index-0 reads inside it measure its own card; `--devices`/`devices` takes one entry per worker (a card named twice, or a bare `cuda` in a list of several, is refused), a card the box lacks is refused at startup naming the cards present, and a job's `device` is `device_label`'s `"cuda:N <card name>"`. | `tests/test_worker_manager.py::TestWorkerPinning::test_a_cuda_index_pins_the_worker_through_the_environment`, `tests/test_worker_manager.py::TestWorkerPinning::test_a_bare_cuda_sets_no_pin`, `tests/test_serve_main.py::TestConfigureDevices::test_a_refusal_exits_with_code_2`, `tests/test_serve_main.py::TestResolveServeDevices::test_a_card_the_machine_lacks_is_refused_naming_the_cards_present`, `tests/test_serve_main.py::TestResolveServeDevices::test_two_entries_are_both_kept`, `tests/test_serve_main.py::TestResolveServeDevices::test_a_card_named_twice_is_refused` |
112112
| Failed-run reporting | `dw/worker.py`, `dw/worker_protocol.py`: `Failed`, `Cancelled` | A failed or cancelled run's reply still carries the manifest of the steps that ran. | `tests/test_worker_execute.py::test_failure_carries_the_manifest_of_the_steps_that_ran`, `tests/test_worker_execute.py::test_cancellation_carries_the_manifest_too` |
113113
| Failure path redaction | `dw/path_redaction.py`: `redact_paths`, called by `dw/worker.py` | A failed run's message and traceback name a file under the job's asset roots or output directory by its `asset:`/`output:` reference, never its absolute server path. | `tests/test_worker_execute.py::test_a_failure_names_an_asset_by_reference_not_by_server_path`, `tests/test_path_redaction.py` |
114114
| Run-time warnings | `dw/events.py`: `emit_warning` | A warning found at run time is emitted as an event, so it reaches the caller and not just the log. | `tests/test_concat_videos.py::TestWarningsReachTheCaller::test_the_level_spread_warning_is_emitted_as_an_event` |
@@ -141,7 +141,8 @@ one to open.
141141
| --- | --- | --- | --- |
142142
| App factory | `dw/server/app.py`: `create_app` | `create_app` builds the JobManager and `app.state`, installs the middleware and registers the routes, and all state lives in the JobManager. | — |
143143
| Routers | `dw/server/routes/*.py`, `dw/server/routes/__init__.py`: `ROUTERS`, `include_routers`, `include_file_routes` | Routers register in `ROUTERS` order, because a greedy `{name:path}` route must come after its specific siblings. The files router comes later, after `/mcp`. | `tests/test_server_downloads.py::test_download_workflow_sets_content_disposition_attachment` |
144-
| Job queue | `dw/server/jobs.py`: `JobManager` | One runner thread runs jobs FIFO on the single worker. A job carries its own `output_dir`, `asset_dir` and `workflow_dir`, so it stays in its workspace. | `tests/test_server_workspaces.py::TestRunning::test_rerun_stays_in_the_workspace_it_ran_in` (FIFO: —) |
144+
| Job queue | `dw/server/jobs.py`: `JobManager` | A dispatcher thread hands queued jobs, in queue order, to the pool's workers (`Worker pool dispatch`); each worker runs one job at a time. A job carries its own `output_dir`, `asset_dir` and `workflow_dir`, so it stays in its workspace. | `tests/test_server_workspaces.py::TestRunning::test_rerun_stays_in_the_workspace_it_ran_in` (FIFO: —) |
145+
| Worker pool dispatch | `dw/server/jobs.py`: `JobManager._next_dispatch`, `check_fits`, `dw/vram_estimate.py`: `required_vram_gb` | `--devices` runs one worker per card; a queued job goes to the first free card (in `--devices` order) that fits what it needs, a job that fits no free card keeps its place while a smaller one behind it takes the card, a job bigger than every card is refused at submit, and cancel or a crash touches only the worker running that job. | `tests/test_worker_pool.py::TestConcurrency::test_two_workers_run_two_jobs_at_once`, `tests/test_worker_pool.py::TestBackfill::test_a_job_too_big_for_the_free_card_waits_while_a_smaller_one_takes_it`, `tests/test_worker_pool.py::TestFitAtSubmit::test_a_job_too_big_for_every_card_is_refused_naming_the_largest`, `tests/test_worker_pool.py::TestCancel::test_cancel_reaches_the_worker_running_the_job`, `tests/test_worker_pool.py::TestCrashIsolation::test_a_crashed_worker_fails_only_its_own_job`, `tests/test_worker_pool.py::TestOomRanking::test_the_later_started_worker_is_ranked_first` |
145146
| Job-to-run link | `dw/server/job_record.py`, `dw/server/jobs.py`: `JobManager.realized` | A job records `run_id`, `run_dir` and `run_version`, and `JobManager.realized` reads that run's `workflow.json`, confined to the output root. | `tests/test_server_jobs.py::test_realized_reads_the_file_the_run_wrote`, `tests/test_server_jobs.py::test_realized_refuses_a_run_dir_that_escapes_the_output_root` |
146147
| Job history | `dw/server/job_history.py`: `JobHistory` | Finished jobs persist in `jobs.sqlite`, in WAL mode, so the Jobs view survives a restart. | `tests/test_server.py::test_job_history_survives_restart_and_reruns`, `tests/test_job_history_wal.py::test_the_jobs_database_uses_wal_mode` |
147148
| HTTP security | `dw/server/http_security.py`: `install_middleware`, `query_token_ok` | Every API route needs the token, and only a route marked `query_token_ok` takes it as `?token=`. | `tests/test_security_auth.py::TestTheTokenGate::test_every_api_spelling_needs_the_token`, `tests/test_security_auth.py::TestTheTokenGate::test_the_query_token_is_refused_where_it_is_not_allowed` |

‎docs/MCP.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ when no single workflow covers it.
223223
| `get_task(command)` | `command` | Get a task command's argument schema |
224224
| `list_models()` | — | List what the Hugging Face model cache holds, largest first |
225225
| `get_memory()` | — | Get the worker's VRAM and RAM statistics. `gpu_*` is the card; `host_memory_rss_mb` / `host_memory_peak_rss_mb` are the worker process's resident and high-water host memory, beside the machine's `host_memory_total_mb` / `host_memory_available_mb` - read both, since an offloading workflow keeps its weights in host RAM and the card says little about what it holds (a host field is absent, not null, where the platform cannot measure it). `host_pinned_reserved_mb` / `host_pinned_allocated_mb`, when present, are torch's pinned-host cache - the staging buffers group offloading moves weights through, part of `host_memory_rss_mb` and invisible in every `gpu_*` figure, which is what a worker holding GB after releasing every model is usually holding (#98). `live: true` means `info` was measured now and is the worker's own memory; only live readings are comparable with each other. `live: false` with `info: null` (and `stale: false`) means nothing has been measured because nothing is resident; `live: false` with a populated `info` is a cached earlier reading, with `reason` (`job_running`, `worker_stopped`, `worker_busy`, `worker_unreachable`) and `age_seconds` - one cached while a job loads a model understates what is resident, so ask again when the server is idle rather than comparing it with a live figure |
226-
| `get_health()` | — | Check that the server is alive, and which machine answered: `version`, `device`, whether a model process is currently resident (`worker_alive`), the job running now and the queue depth. `worker_alive: false` is the normal idle state on a server that has not run a job since startup or the last memory clear - not a fault - the on-demand worker starts with the next job (#206) |
226+
| `get_health()` | — | Check that the server is alive, and which machine answered: `version`, `device`, whether a model process is currently resident (`worker_alive`), the job running now, a `workers` entry per card with its own `current_job`, and the queue depth. `worker_alive: false` is the normal idle state on a server that has not run a job since startup or the last memory clear - not a fault - the on-demand worker starts with the next job (#206) |
227227
| `get_server_info()` | — | What this installation can do and where it keeps things: `device` (the accelerator a run will use), `version`, the `workspace` this session is working in and the workflow/asset/output/prompt `directories` of *that* workspace, the bind address and port, whether a token is required, and whether MCP is mounted. Check the device before authoring - a CUDA-only choice (bitsandbytes, `torch.compile`, flash attention) is not available on an `mps` or `cpu` server. `runtime` (#222) reports the Python version, torch version and the CUDA version torch was built against, the NVIDIA driver version (when `nvidia-smi` is reachable), and the installed versions of diffusers, transformers, accelerate, bitsandbytes, peft, safetensors and sentencepiece (`null` for one not installed) - for diagnosing an environment mismatch between boxes without shelling in |
228228
| `list_jobs(limit=20, status=None, workspace=None)` | optional `limit` (newest N), `status` (one state or a comma-separated set of `queued`, `running`, `succeeded`, `failed`, `cancelled`), `workspace` | List queued, running and recent jobs, **newest first**. Bounded by default: the unbounded listing was over a client's tool-result limit on a server with a few months of history, which made it a tool that could not be called at all. `total` says how many matched and `truncated`/`next` say so when the answer was cut - raise `limit` or narrow with `status`. Without `workspace`, a named workspace lists its own jobs and the default one lists every job the server holds |
229229
| `list_gallery(limit=50, subfolder=None, only_orphans=False, workspace=None, folder=None, version=None)` | `limit`, `subfolder`, `only_orphans`, `workspace`, `folder`, `version` | List generated output files, newest first. A name is `<workflow>/<run id>/<file>`, where `<file>` may sit in the subfolder the step chose (`final/episode.mp4`); each entry carries `folder` (the workflow) and `subfolder` (by convention `final` or `intermediate`, `''` when the step chose none, any path the workflow wrote otherwise), and `subfolder="final"` lists only deliverables. Each entry carries `run_id` and `version` - that run's ordinal among the workflow's runs, which is how one of several runs that wrote the same basename is named to a person: the web UI labels the same file `v5`. The number is assigned when the run opens and never renumbered, so deleting a run leaves a gap rather than sliding the rest down (a failed run, or a rerun that reused every step, leaves one too - it took a number and may have nothing to list), and it is `null` under the flat output layout, which has no runs. `folder` with `version` lists that one run's files, and `output:<folder>/v5/<file>` names one in a workflow; every other tool takes `name`. Each entry also carries a ready-made `url`, already scoped to the workspace that made it - a hand-built `/outputs/<name>` URL 404s for anything but the default workspace. `only_orphans=True` inverts the call: instead of files, it returns run directories with no media anywhere under them (`runs`, each `{name, mtime}`) - a run whose output was deleted before `delete_output` could remove it by name, or one that failed before writing anything; `subfolder` does not apply in this mode, and `name` is exactly what `delete_output` accepts (#170). `workspace` names the workspace for this one call without switching the session to it - the same pin `run_workflow` takes, so a job run into another workspace stays reachable from the session that queued it |
@@ -343,7 +343,7 @@ through is not a gate.
343343

344344
| Tool | What it commits |
345345
| --- | --- |
346-
| `run_workflow` | Minutes of GPU time; the engine runs one job at a time |
346+
| `run_workflow` | Minutes of GPU time; the engine runs one job per card at a time |
347347
| `rerun_job` | The same run, from a stored spec |
348348
| `download_model` | Tens of gigabytes of network and disk |
349349
| `delete_model` | Cached weights, unrecoverably — getting them back means downloading again |

‎docs/SERVER.md‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ python -m dw.serve --workspace ~/studio --examples-dir ~/src/diffusers-workflow/
1717
python -m dw.serve --host 0.0.0.0 --token "some-long-random-string" # reachable off this machine
1818
python -m dw.serve --host 0.0.0.0 --token "..." --mcp # ...and drivable by an agent on another machine
1919
python -m dw.serve --devices cuda:1 # run the worker on the second card
20+
python -m dw.serve --devices cuda:0,cuda:1 # one worker per card, two jobs at once
2021
python -m dw.serve --trust-workflows # only if nothing untrusted can reach POST /api/jobs - see Security model
2122
```
2223

@@ -33,8 +34,9 @@ load entirely.
3334
with `CUDA_VISIBLE_DEVICES` (and `DW_DEVICE=cuda`), so inside the worker the
3435
card is the only one it can see; a bare `cuda`, or no flag at all, pins
3536
nothing and behaves as before. The `devices` setting is the standing form of
36-
the same choice, and the flag wins over it. There is one entry only for now:
37-
a second (`--devices cuda:0,cuda:1`) is refused at startup, as is a card the
37+
the same choice, and the flag wins over it. A list (`--devices cuda:0,cuda:1`)
38+
runs one worker per card, each running one job at a time. A card named twice,
39+
or a bare `cuda` in a list of several, is refused at startup, as is a card the
3840
machine lacks, with the message naming the cards that are present
3941
(`cuda:0 (NVIDIA GeForce RTX 4090)`, ...).
4042

@@ -197,8 +199,26 @@ agent from another machine, plus the queue across every workspace:
197199
| `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 |
198200
| `POST /api/jobs/{id}/move` | Reorder a queued job: `{"direction": "up"\|"down"\|"front"\|"back"}`. Job listings carry each waiting job's `queue_position`. |
199201

200-
One job runs at a time (it is one GPU); submissions queue in order, and
201-
the waiting portion of the queue can be reordered.
202+
Each worker runs one job at a time, so with one device (the default) one job
203+
runs at a time; submissions queue in order, and the waiting portion of the
204+
queue can be reordered. With several `--devices` the dispatcher walks the
205+
queue in order and gives each job the first free card, in `--devices` order,
206+
that fits it. A job that fits no free card keeps its place while a smaller
207+
one behind it takes the card (backfill).
208+
209+
What a job needs: the workflow's declared `vram_estimate` projected over the
210+
job's arguments (hard: a job bigger than every card is refused at submit with
211+
a 400 naming the largest card); else the smallest `cost` entry's `vram_gb` for
212+
this device type (soft: if no card is that big, the job runs on any card, as
213+
on one card); else any card. A card's size is its GiB rounded up (a 3090
214+
counts as 24). Cancel reaches only the worker running that job, and a worker
215+
crash fails only its own job. The worker started later gets a higher Linux
216+
`oom_score_adj` (+100 over the highest other live worker), so a host-RAM
217+
squeeze kills the later job; nothing is written with one worker. There is no
218+
host-RAM gate at dispatch: two big jobs can still exhaust host RAM together.
219+
`GET /api/memory`, the probe cache and memory reporting still read the first
220+
worker (or report `job_running` when any job runs); `POST /api/memory/clear`
221+
clears every idle worker.
202222

203223
### Progress events
204224

@@ -610,7 +630,10 @@ The editor's forms come from these; they are just as usable from scripts:
610630
mistaken for the worker's memory now - `info: null` means nothing has been
611631
measured because nothing is resident. health also reports `hostname`,
612632
`device` and whether `mcp` is mounted, so a remote client can tell which
613-
machine answered
633+
machine answered. `current_job` is the longest-running job and `worker_alive`
634+
is true if any worker is alive; `workers` has one entry per card (one with a
635+
single device): `{device, name, vram_gb, current_job, alive,
636+
host_memory_rss_mb}`, the last absent when the worker process isn't running
614637
- `POST /api/memory/clear` (#221) — drops every loaded pipeline and the step
615638
cache (MCP `clear_memory`), and returns the
616639
memory reading taken right after. Refused with 409 while a job is running

0 commit comments

Comments
 (0)