studio: crash-durable SQLite render queue + engine selector
SqlitePromptQueue (middleware/tasks_queue.py): drop-in for PromptQueue backed
by one SQLite file (stdlib sqlite3, WAL, zero external processes). Opt in with
STUDIO_QUEUE_DB; precedence STUDIO_QUEUE_DB > STUDIO_PERSIST_QUEUE > memory so
local default is unchanged. The put/get/task_done seam maps to a durable job
state machine: idempotent submit (INSERT OR IGNORE on prompt_id), exactly-once
claim (BEGIN IMMEDIATE), retry-on-error, and crash-recovery via a lease + reap
that re-runs an abandoned claim (idempotent — SaveImage suffixes increment).
Multiple Studio processes on the same DB compete safely.
Engine selector (middleware/engine_selector.py): GET /v1/engines lists execution
targets for the org (local + registered compute_config workers), PUT
/v1/engines/default sets a per-org default stored on ComputeProfile; route_prompt
honors it (local = unchanged). Leased cloud machines are a future engine class
(interface stubbed). Frontend picker is a TODO.
docs/federation.md rewritten to the durable-queue reality with one future
paragraph (remote Tasks backend + Go/unified-binary HIP-0106 migration + leased
machines). Tests: middleware_test/{tasks_queue,engine_selector}_test.py, incl.
crash-recovery across queue instances.
This commit is contained in:
hanzo-dev
@@ -60,9 +60,29 @@ shared Hanzo IAM (standard OIDC code flow). No embedding.
|
||||
the prompt worker from the queue item's `org_id`.
|
||||
- **GPU federation** (`middleware/{worker_client,prompt_router,compute_config,
|
||||
visor_client}.py`): `/v1/workers/register` (heartbeat), `/v1/workers` (list),
|
||||
`/v1/worker/execute` (push). Worker↔coordinator trust via
|
||||
`X-Worker-Token` (`STUDIO_WORKER_TOKEN`, KMS-sourced). NAT'd local boxes (GB10)
|
||||
join outbound — pull channel is specced in `docs/federation.md`.
|
||||
`/v1/worker/execute` (in-cluster push fast path). Worker↔coordinator trust via
|
||||
`X-Worker-Token` (`STUDIO_WORKER_TOKEN`, KMS-sourced).
|
||||
- **Durable render queue** (`middleware/tasks_queue.py`, `docs/federation.md`):
|
||||
`SqlitePromptQueue` — crash-durable render queue in a single SQLite file
|
||||
(stdlib `sqlite3`, WAL, **zero external processes**). `STUDIO_QUEUE_DB=<path>`
|
||||
→ `server.py` swaps `PromptQueue` for it (precedence **STUDIO_QUEUE_DB >
|
||||
STUDIO_PERSIST_QUEUE > memory**; default local behavior unchanged). The
|
||||
`PromptQueue` seam maps to a durable job state machine: `put→INSERT (pending,
|
||||
idempotent on prompt_id)`, `get→claim next pending under BEGIN IMMEDIATE
|
||||
(exactly-once, leased)`, `task_done→done | retry/failed`. Crash-recovery: a
|
||||
lease + reap returns an abandoned claim to `pending` (on every get(), a
|
||||
heartbeat thread, and on boot) so another process re-runs it — idempotent
|
||||
(SaveImage suffixes increment). Multiple Studio processes on the same DB file
|
||||
compete safely. Tested: `middleware_test/tasks_queue_test.py` (incl.
|
||||
crash-recovery across instances). Future (federation.md §6): a remote queue
|
||||
backend served by Hanzo Tasks + a Go/unified-binary (HIP-0106) migration.
|
||||
- **Engine selector** (`middleware/engine_selector.py`): `/v1/engines` lists
|
||||
execution targets for the org (`local` + registered compute_config workers) and
|
||||
`PUT /v1/engines/default` sets a per-org default (stored on
|
||||
`ComputeProfile.default_engine`). `prompt_router.route_prompt` routes to the
|
||||
chosen worker; `local` = unchanged. Leased cloud machines (Visor
|
||||
`GET /v1/machines`) are a future engine class (interface stubbed, not built).
|
||||
Frontend picker is a TODO. Tested: `middleware_test/engine_selector_test.py`.
|
||||
|
||||
## Tests
|
||||
`tests-unit/{folder_paths_test,middleware_test,app_test,prompt_server_test}` —
|
||||
|
||||
+117
-130
@@ -1,168 +1,155 @@
|
||||
# Hanzo Studio — GPU Federation
|
||||
|
||||
How `studio.hanzo.ai` schedules generation jobs across a heterogeneous GPU pool:
|
||||
in-cluster cloud pods **and** local boxes (e.g. a GB10) that join over an
|
||||
**outbound-only** connection. No inbound tunnel to the local box is ever
|
||||
required.
|
||||
How `studio.hanzo.ai` runs generation jobs durably and spreads them across a
|
||||
heterogeneous GPU pool: in-cluster cloud pods **and** local boxes (e.g. a GB10)
|
||||
that join over an **outbound-only** connection. No inbound tunnel to a local box
|
||||
is ever required.
|
||||
|
||||
This document is the contract. Sections are marked **[implemented]** (code in
|
||||
this repo today) or **[specced]** (agreed design, not yet wired). Nothing here is
|
||||
a stub in code — specced pieces live in this doc until they are built.
|
||||
Sections are marked **[implemented]** (code in this repo today), **[fast-path]**
|
||||
(retained in-cluster optimization), or **[future]**.
|
||||
|
||||
## 1. Roles
|
||||
|
||||
- **Coordinator** — the `studio.hanzo.ai` deployment. Full UI, IAM auth,
|
||||
multi-tenant storage. Owns the per-org worker registry and the job schedule.
|
||||
Runs a small CPU footprint (it routes; it does not need a GPU).
|
||||
- **Worker** — a headless Studio started with `--worker-mode`. Reports its
|
||||
device (GPU model, VRAM), registers with the coordinator, and executes prompts
|
||||
it is given. A worker belongs to exactly one org (`STUDIO_ORG_ID`).
|
||||
multi-tenant storage. Accepts `/prompt`, owns the durable render queue and the
|
||||
per-org worker registry + engine selection.
|
||||
- **Worker** — a Studio process that executes prompts. Either the same box as the
|
||||
coordinator (local render), an in-cluster pod reached by push (§2a), or a box
|
||||
sharing the coordinator's render queue.
|
||||
|
||||
A worker is the *same binary* as the coordinator; `--worker-mode` only drops the
|
||||
UI/websocket surface and turns on registration + the execute endpoint. This is
|
||||
the ComfyUI-Distributed lineage (workers are full ComfyUI instances the control
|
||||
plane drives) narrowed to **job-level** dispatch and hardened for multi-tenant +
|
||||
NAT traversal.
|
||||
A worker is the *same binary* as the coordinator.
|
||||
|
||||
## 2. Two dispatch models
|
||||
## 2. Durable render queue **[implemented]**
|
||||
|
||||
Reachability, not preference, decides which model a worker uses.
|
||||
The render queue is crash-durable: pending and in-flight prompts survive a
|
||||
process crash, and a killed render is re-run instead of being lost. Storage is a
|
||||
single **SQLite** file (stdlib `sqlite3`, WAL) — **zero external processes**, per
|
||||
the house rule (SQLite embedded by default; Postgres only for prod multi-instance).
|
||||
|
||||
### 2a. Push — for in-cluster workers **[implemented]**
|
||||
Enable with `STUDIO_QUEUE_DB=<path>`. `server.py` then swaps `PromptQueue` for
|
||||
`SqlitePromptQueue`; the existing prompt-executor thread is the claim loop
|
||||
unchanged. The `PromptQueue` seam maps onto a durable job state machine:
|
||||
|
||||
The coordinator can open a connection to the worker (both are pods in
|
||||
`hanzo-k8s`, or the worker has a routable `WORKER_EXTERNAL_URL`).
|
||||
```
|
||||
put(item) → INSERT job (pending) # idempotent on prompt_id
|
||||
get(timeout) → claim next pending (claimed, leased) # BEGIN IMMEDIATE
|
||||
task_done(success) → job → done
|
||||
task_done(error) → job → pending (retry) | failed # per STUDIO_TASKS_MAX_ATTEMPTS
|
||||
```
|
||||
|
||||
Guarantees:
|
||||
|
||||
- **Exactly-once submit** — the row is keyed by `prompt_id` (`INSERT OR IGNORE`);
|
||||
a resubmit of the same prompt is a no-op.
|
||||
- **Exactly-once claim** — `get()` claims under `BEGIN IMMEDIATE`, so with several
|
||||
Studio processes on the same `STUDIO_QUEUE_DB` exactly one claims each job.
|
||||
- **Crash-recovery** — a claim carries a lease. A background heartbeat renews it
|
||||
while the render runs; if the process dies, the lease expires and the reap step
|
||||
(run on every `get()`, in the heartbeat thread, and on boot) returns the job to
|
||||
`pending` for another claim. A render retry is idempotent — SaveImage suffixes
|
||||
increment.
|
||||
|
||||
Files: `middleware/tasks_queue.py` (`SqlitePromptQueue`), `server.py` (queue
|
||||
selection). Precedence:
|
||||
|
||||
| Precedence | Trigger | Backend | Durability |
|
||||
|---|---|---|---|
|
||||
| 1 | `STUDIO_QUEUE_DB` set | `SqlitePromptQueue` (one file) | Crash-durable, exactly-once, retry, multi-process |
|
||||
| 2 | `STUDIO_PERSIST_QUEUE=1` | `PromptQueue` + JSON snapshot | Single-box crash-survival (re-queue on boot) |
|
||||
| 3 | (default) | `PromptQueue` (in-memory) | None — same as upstream ComfyUI |
|
||||
|
||||
Env: `STUDIO_QUEUE_DB` (path), `STUDIO_TASKS_QUEUE` (default `studio-render`),
|
||||
`STUDIO_TASKS_MAX_ATTEMPTS` (default 3), `STUDIO_TASKS_LEASE_MS` (default 60000,
|
||||
heartbeat renews at lease/3), `STUDIO_WORKER_ID` (default `<host>-<pid>`).
|
||||
|
||||
### 2a. Push — in-cluster fast path **[fast-path]**
|
||||
|
||||
When the coordinator can already reach a worker (both pods in `hanzo-k8s`, or a
|
||||
routable `WORKER_EXTERNAL_URL`), it may forward a prompt directly:
|
||||
|
||||
```
|
||||
client → POST /prompt (coordinator)
|
||||
prompt_router.route_prompt(org, body)
|
||||
→ get_available_gpu_worker(org) # compute_config registry
|
||||
→ POST {worker.url}/v1/worker/execute # forward_to_worker
|
||||
worker queues + executes locally, returns {prompt_id, ...}
|
||||
```
|
||||
|
||||
Files: `middleware/prompt_router.py`, `middleware/worker_client.py`
|
||||
(`add_worker_routes` → `/v1/worker/execute`), `server.py` (`/prompt` calls
|
||||
`route_prompt`).
|
||||
Files: `middleware/prompt_router.py`, `middleware/worker_client.py`. This path is
|
||||
not durable on its own (a forwarded prompt lost to a worker crash is not retried);
|
||||
keep it for latency-sensitive in-cluster dispatch. It requires the coordinator to
|
||||
reach the worker, so it cannot serve a NAT'd box.
|
||||
|
||||
Limitation: requires the coordinator to reach the worker. A GB10 behind NAT is
|
||||
**not** reachable, so push cannot serve it.
|
||||
## 3. Engine selector **[implemented]**
|
||||
|
||||
### 2b. Pull — for NAT'd local boxes (the GB10 pattern) **[specced]**
|
||||
An org chooses which execution target ("engine") its prompts run on. Engines are
|
||||
derived from what already exists in the per-org compute registry:
|
||||
|
||||
The worker dials **out** and keeps the connection; the coordinator pushes jobs
|
||||
down that already-open channel. No inbound port on the box.
|
||||
- `local` — this Studio instance (always present), and
|
||||
- each registered org GPU worker (from `compute.json`'s worker registry).
|
||||
|
||||
Primary transport is a WebSocket the worker opens to the coordinator:
|
||||
Endpoints (coordinator):
|
||||
|
||||
```
|
||||
worker → WSS {coordinator}/v1/workers/connect (Authorization: worker token)
|
||||
← {type:"job", job_id, prompt, extra_data, org_id}
|
||||
→ {type:"progress", job_id, value, max}
|
||||
→ {type:"result", job_id, status, outputs:[{name, s3_url}...]}
|
||||
```
|
||||
| Method | Path | Purpose |
|
||||
|---|---|---|
|
||||
| GET | `/v1/engines` | List engines for the org + the current default. |
|
||||
| PUT | `/v1/engines/default` | Set the org's default engine. Body `{"engine": "<id>"}`. |
|
||||
|
||||
HTTP long-poll fallback for environments without WSS:
|
||||
The default is stored on the org's compute profile (`ComputeProfile.default_engine`).
|
||||
`prompt_router.route_prompt` honors it: if the default names a live worker, the
|
||||
prompt is forwarded there; `local` (the default) leaves current behavior
|
||||
unchanged. A stale default (worker deregistered) falls back to `local`, never 500.
|
||||
|
||||
```
|
||||
worker → POST /v1/jobs/claim {worker_id, org_id} # long-poll, ≤30s
|
||||
← 204 (no work) | 200 {job_id, prompt, extra_data}
|
||||
worker → POST /v1/jobs/{job_id}/result {status, outputs}
|
||||
```
|
||||
Files: `middleware/engine_selector.py`, `middleware/compute_config.py`,
|
||||
`middleware/prompt_router.py`. **Frontend TODO:** a queue-panel/settings dropdown
|
||||
to pick the engine — the API + per-org default ship now; the UI control is a
|
||||
follow-up.
|
||||
|
||||
The coordinator keeps a per-org, per-worker job queue (a natural extension of
|
||||
`compute_config`), enqueues on `/prompt` when the target worker is pull-mode, and
|
||||
hands the job to whichever channel the worker holds open. `route_prompt` gains a
|
||||
third branch (`{"action":"queued"}` → `202`) alongside the existing
|
||||
forward/provisioning/unavailable branches.
|
||||
## 4. Security **[implemented]**
|
||||
|
||||
Not implemented because it does not "drop out" of ComfyUI's `--listen`
|
||||
architecture — it needs a real coordinator-side queue and a WS control channel.
|
||||
Building it does not change the endpoints below.
|
||||
- **User auth** — all user-facing routes require an IAM JWT (JWKS-verified); org
|
||||
comes from the `owner` claim. See `middleware/iam_auth_middleware.py`.
|
||||
- **Worker trust** — the IAM-exempt `/v1/worker/execute` and `/v1/workers/register`
|
||||
require `X-Worker-Token == STUDIO_WORKER_TOKEN` (`worker_client.verify_worker_token`).
|
||||
KMS-sourced in cloud; empty in a single-trust-domain local dev, where the check
|
||||
is skipped.
|
||||
- **Org isolation** — a render carries its `org_id` in `extra_data`; the worker
|
||||
binds it as the execution org so outputs land in that org's namespace (§5). A
|
||||
worker only serves its own `STUDIO_ORG_ID`.
|
||||
|
||||
## 3. Endpoints
|
||||
## 5. Output persistence **[future]**
|
||||
|
||||
All under `/v1/` (also served under the `/api/v1/` alias ComfyUI adds for the
|
||||
frontend proxy). Worker↔coordinator endpoints are exempt from user IAM and
|
||||
carry the coordinator shared secret instead (`X-Worker-Token`, §5).
|
||||
Workers execute against their own output dir (org-scoped via
|
||||
`folder_paths.set_execution_org`). For the coordinator UI to serve results across
|
||||
boxes, workers upload outputs to `s3.lux.cloud` (hanzo bucket), key
|
||||
`studio/{org_id}/{prompt_id}/{filename}`, and the coordinator serves via `/view`
|
||||
(or a signed redirect). Until then, results remain on the executing worker —
|
||||
fine for the single-box and in-cluster cases.
|
||||
|
||||
| Method | Path | Dir | Status | Purpose |
|
||||
|---|---|---|---|---|
|
||||
| POST | `/v1/workers/register` | worker→coord | **impl** | Register / heartbeat. Body: `{worker_id, url, org_id, device, gpu_model, vram_gb, status}`. |
|
||||
| GET | `/v1/workers` | client→coord | **impl** | List the org's workers + liveness. |
|
||||
| POST | `/v1/worker/execute` | coord→worker | **impl** | Push a prompt to a reachable worker. Same body as `/prompt`. |
|
||||
| WSS | `/v1/workers/connect` | worker→coord | **spec** | Persistent pull channel for NAT'd workers. |
|
||||
| POST | `/v1/jobs/claim` | worker→coord | **spec** | HTTP long-poll fallback to claim a job. |
|
||||
| POST | `/v1/jobs/{id}/result` | worker→coord | **spec** | Report result + artifact URLs. |
|
||||
| GET | `/v1/compute/config` | client→coord | **impl** | Read the org's compute profile. |
|
||||
| PUT | `/v1/compute/config` | client→coord | **impl** | Update profile (GPU tier, auto-provision). |
|
||||
| POST | `/v1/compute/provision` | client→coord | **impl** | Autoscale a cloud GPU worker via Visor. |
|
||||
## 6. Future direction
|
||||
|
||||
(Compute-profile + Visor autoscale routes are still under `/compute/*` in
|
||||
`server.py`; the `/v1/compute/*` labels above are the canonical names to fold to
|
||||
when those handlers are next touched. Worker-registry + execute already moved.)
|
||||
|
||||
## 4. Registration & scheduling **[implemented]**
|
||||
|
||||
- Worker sends a heartbeat every 30s (`WorkerClient._heartbeat_loop`).
|
||||
- Coordinator upserts it into the org's `compute.json`
|
||||
(`compute_config.register_worker`), pruning workers silent >300s.
|
||||
Liveness for scheduling is 90s (`WorkerInfo.is_alive`).
|
||||
- `get_available_gpu_worker(org)` returns the first ready, alive CUDA worker.
|
||||
- `prompt_router.route_prompt` chooses local vs worker from the prompt's
|
||||
`device_preference` (`auto`|`cpu`|`gpu`) and the org's `gpu_enabled`.
|
||||
- If no worker and the org has `auto_provision`, Visor launches a cloud GPU VM
|
||||
that boots a worker (`middleware/visor_client.py`).
|
||||
|
||||
## 5. Security
|
||||
|
||||
- **User auth** — all user-facing routes require an IAM JWT (JWKS-verified);
|
||||
org comes from the `owner` claim. See `middleware/iam_auth_middleware.py`.
|
||||
- **Worker trust** — worker↔coordinator endpoints are IAM-exempt (they carry no
|
||||
user) and instead require `X-Worker-Token == STUDIO_WORKER_TOKEN`
|
||||
(`worker_client.verify_worker_token`). The token is KMS-sourced in cloud;
|
||||
empty in a single-trust-domain local dev, where the check is skipped.
|
||||
**[specced]** forward direction: replace the shared secret with an IAM
|
||||
service-account token (`hanzo-studio` client-credentials grant) so each worker
|
||||
is individually attributable and revocable.
|
||||
- **Org isolation** — a job carries its `org_id`; the worker binds it as the
|
||||
execution org so outputs land in that org's namespace (§6). A worker only ever
|
||||
claims jobs for its own `STUDIO_ORG_ID`.
|
||||
- **No inbound** — pull-mode boxes expose nothing; the coordinator is reached
|
||||
over TLS via `hanzoai/ingress`.
|
||||
|
||||
## 6. Output persistence **[specced]**
|
||||
|
||||
Push/pull both execute on the worker, which writes to *its own* output dir
|
||||
(org-scoped via `folder_paths.set_execution_org`). For the coordinator UI to
|
||||
show results, workers must ship artifacts to shared storage:
|
||||
|
||||
- On completion, upload outputs to `s3.lux.cloud` (hanzo bucket),
|
||||
key `studio/{org_id}/{prompt_id}/{filename}`.
|
||||
- Report the object URLs in the `result` message; the coordinator records them
|
||||
in history so `/view` (or a signed redirect) serves them.
|
||||
|
||||
Until this lands, push-mode results remain on the worker and are addressable
|
||||
only if the worker is reachable — acceptable for the in-cluster pool, blocking
|
||||
for the NAT pull pool. This is the top build item after the pull channel.
|
||||
The durable queue, engine selection, and dispatch are deliberately kept in a thin
|
||||
Python layer at the `PromptQueue`/`prompt_router` seams so the backend can be
|
||||
swapped without touching the render path. Two moves are planned: (a) a **remote
|
||||
queue backend** — the same job state machine served by **Hanzo Tasks**
|
||||
(`github.com/hanzoai/tasks`) over its HTTP surface, so multiple coordinators and
|
||||
NAT'd `--worker`-style claimers share one durable queue instead of a local file;
|
||||
and (b) folding that scheduling/queue layer into a **Go subsystem inside the
|
||||
`hanzoai/cloud` unified binary** (HIP-0106), leaving Studio's Python as a pure
|
||||
execution worker. The engine selector will likewise grow a third class,
|
||||
**leased cloud machines** provisioned via the platform's Visor compute surface
|
||||
(`GET {cloud}/v1/machines`) — the `/v1/engines` shape already accommodates the
|
||||
extra entries; the client interface is stubbed in `engine_selector.py` and wired
|
||||
when the console lands.
|
||||
|
||||
## 7. Joining a local box (GB10)
|
||||
|
||||
```bash
|
||||
python main.py \
|
||||
--worker-mode \
|
||||
--coordinator-url https://studio.hanzo.ai \
|
||||
--worker-id gb10-1
|
||||
# env: STUDIO_ORG_ID=<org> STUDIO_WORKER_TOKEN=<from KMS>
|
||||
```
|
||||
Two paths today:
|
||||
|
||||
Today (push) this works only if the coordinator can reach the box (set
|
||||
`WORKER_EXTERNAL_URL`). Once §2b lands, the box needs **no** inbound access: it
|
||||
registers, opens the WSS channel, and pulls jobs for its org.
|
||||
- **Shared queue** (durable): point the box at the same `STUDIO_QUEUE_DB` (shared
|
||||
volume) as the coordinator; its prompt-executor thread claims jobs directly.
|
||||
Exactly-once claim + crash-recovery hold across processes.
|
||||
- **Push** (in-cluster, fast): `python main.py --worker-mode --coordinator-url … --worker-id gb10-1`
|
||||
with `STUDIO_WORKER_TOKEN`; requires the coordinator to reach the box.
|
||||
|
||||
## 8. Build order
|
||||
|
||||
1. Pull WSS channel `/v1/workers/connect` + coordinator per-worker queue (§2b).
|
||||
2. Artifact upload to `s3.lux.cloud` + result recording (§6).
|
||||
3. IAM service-account worker identity replacing the shared secret (§5).
|
||||
4. Fold `/compute/*` profile routes to `/v1/compute/*` (§3).
|
||||
The remote-queue backend (§6) generalizes the shared-queue path to boxes that
|
||||
cannot share a filesystem, outbound-only.
|
||||
|
||||
@@ -63,6 +63,7 @@ class ComputeProfile:
|
||||
workers: list = field(default_factory=list)
|
||||
auto_provision: bool = False
|
||||
max_concurrent: int = 1
|
||||
default_engine: str = "local" # engine id renders route to (see engine_selector)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return asdict(self)
|
||||
@@ -139,6 +140,9 @@ def update_config(org_id: str, updates: dict) -> ComputeProfile:
|
||||
mc = int(updates["max_concurrent"])
|
||||
profile.max_concurrent = max(1, min(mc, 32))
|
||||
|
||||
if "default_engine" in updates:
|
||||
profile.default_engine = str(updates["default_engine"])
|
||||
|
||||
# Derive gpu_enabled from profile if not explicitly set
|
||||
if "active_profile" in updates and "gpu_enabled" not in updates:
|
||||
profile.gpu_enabled = profile.active_profile in ("gpu-basic", "gpu-pro", "custom")
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
"""
|
||||
Engine selector for Hanzo Studio.
|
||||
|
||||
Lists the execution targets ("engines") available to an org and records the
|
||||
org's default. Thin wrapper over the existing per-org compute registry
|
||||
(middleware/compute_config.py): an engine is either
|
||||
|
||||
- `local` — this Studio instance, always present, and
|
||||
- a registered org GPU worker (from compute.json's worker registry).
|
||||
|
||||
Selecting an engine routes prompts to it; prompt_router.py is the existing
|
||||
routing seam and already forwards to a worker by URL, so the default engine
|
||||
just names which target a prompt without an explicit device_preference lands on.
|
||||
|
||||
Leased cloud machines (via the platform's Visor compute surface,
|
||||
`GET {cloud}/v1/machines`) are a future engine class — the client interface is
|
||||
noted below but not built this iteration; the API shape here already
|
||||
accommodates the extra entries.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from middleware import compute_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Future: leased-machine engines. A VisorMachinesClient would implement
|
||||
# list_machines(org_id) -> list[dict] hitting GET {STUDIO_CLOUD_API_URL}/v1/machines
|
||||
# and its rows would be appended to list_engines() as {"type": "leased", ...}.
|
||||
# Console wiring + the client land in a later iteration; not built now.
|
||||
|
||||
|
||||
def list_engines(org_id: str) -> dict:
|
||||
"""Available engines for org_id plus the current default."""
|
||||
profile = compute_config.load_config(org_id)
|
||||
engines = [{
|
||||
"id": "local",
|
||||
"type": "local",
|
||||
"name": "This instance",
|
||||
"status": "ready",
|
||||
"device": "local",
|
||||
}]
|
||||
for w in profile.workers:
|
||||
engines.append({
|
||||
"id": w.worker_id,
|
||||
"type": "worker",
|
||||
"name": w.gpu_model or w.worker_id,
|
||||
"status": "ready" if (w.status == "ready" and w.is_alive()) else "offline",
|
||||
"device": w.device,
|
||||
"vram_gb": w.vram_gb,
|
||||
"url": w.url,
|
||||
})
|
||||
default = profile.default_engine or "local"
|
||||
if default != "local" and not any(e["id"] == default for e in engines):
|
||||
# The default engine deregistered — fall back to local, don't 500.
|
||||
default = "local"
|
||||
return {"engines": engines, "default": default}
|
||||
|
||||
|
||||
def resolve_engine_worker(org_id: str):
|
||||
"""If the org's default engine is a live worker, return it; else None (local)."""
|
||||
profile = compute_config.load_config(org_id)
|
||||
default = profile.default_engine or "local"
|
||||
if default == "local":
|
||||
return None
|
||||
for w in profile.workers:
|
||||
if w.worker_id == default and w.is_alive():
|
||||
return w
|
||||
return None
|
||||
|
||||
|
||||
def add_engine_routes(routes: web.RouteTableDef, server):
|
||||
"""Register /v1/engines on the coordinator."""
|
||||
|
||||
@routes.get("/v1/engines")
|
||||
async def get_engines(request):
|
||||
org_id = server._get_org_id(request)
|
||||
return web.json_response(list_engines(org_id))
|
||||
|
||||
@routes.put("/v1/engines/default")
|
||||
async def set_default_engine(request):
|
||||
org_id = server._get_org_id(request)
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return web.json_response({"error": "Invalid JSON"}, status=400)
|
||||
engine_id = body.get("engine")
|
||||
if not engine_id:
|
||||
return web.json_response({"error": "engine is required"}, status=400)
|
||||
valid = {e["id"] for e in list_engines(org_id)["engines"]}
|
||||
if engine_id not in valid:
|
||||
return web.json_response(
|
||||
{"error": f"unknown engine {engine_id!r}", "engines": sorted(valid)}, status=400
|
||||
)
|
||||
compute_config.update_config(org_id, {"default_engine": engine_id})
|
||||
return web.json_response(list_engines(org_id))
|
||||
@@ -59,6 +59,17 @@ async def route_prompt(
|
||||
if device_pref == "cpu":
|
||||
return None
|
||||
|
||||
# Engine selector: if the org picked a specific worker as its default
|
||||
# engine, route there. `local` (the default) leaves current behavior
|
||||
# unchanged. See middleware/engine_selector.py.
|
||||
from middleware.engine_selector import resolve_engine_worker
|
||||
engine_worker = resolve_engine_worker(org_id)
|
||||
if engine_worker is not None:
|
||||
result = await forward_to_worker(engine_worker, json_data)
|
||||
if result is not None:
|
||||
return {"action": "forward", "worker": engine_worker, "response": result}
|
||||
logger.warning("Default engine %s failed, falling through", engine_worker.worker_id)
|
||||
|
||||
config = load_config(org_id)
|
||||
|
||||
# No GPU enabled in profile — execute locally
|
||||
|
||||
@@ -0,0 +1,305 @@
|
||||
"""
|
||||
SqlitePromptQueue — Hanzo Studio's crash-durable render queue.
|
||||
|
||||
Drop-in for execution.PromptQueue, backed by a single SQLite file (stdlib
|
||||
sqlite3, WAL). When STUDIO_QUEUE_DB is set, render jobs survive a process
|
||||
crash: pending and in-flight prompts live in the DB, and a killed render is
|
||||
rescheduled for another claim instead of being lost. Zero external processes.
|
||||
|
||||
The PromptQueue seam maps straight onto a durable job state machine:
|
||||
|
||||
put(item) → INSERT job (pending) # idempotent on prompt_id
|
||||
get(timeout) → claim next pending (claimed, leased)
|
||||
task_done(...success) → job → done
|
||||
task_done(...error) → job → pending (retry) | failed (per max attempts)
|
||||
|
||||
The existing prompt_worker thread is the claim loop unchanged: it calls get()
|
||||
(claim) and task_done() (report). A background heartbeat renews the lease of
|
||||
the in-flight job while the render runs; if this process dies, the lease
|
||||
expires and the reap step (run on every get() and by the heartbeat thread)
|
||||
returns the job to `pending` so it is re-claimed. A render retry is
|
||||
idempotent — SaveImage suffixes increment.
|
||||
|
||||
Multiple Studio processes pointing at the same STUDIO_QUEUE_DB compete safely:
|
||||
claim runs under BEGIN IMMEDIATE, so exactly one process claims each job.
|
||||
|
||||
History and flags stay in the in-memory base (the Studio UI reads them); only
|
||||
the queue mechanics become durable. This supersedes STUDIO_PERSIST_QUEUE (JSON
|
||||
snapshot) as the local durable path; that remains the zero-dep fallback.
|
||||
Standard library only — no new hard dependency.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
|
||||
from execution import PromptQueue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_QUEUE = "studio-render"
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS render_jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
queue TEXT NOT NULL,
|
||||
payload TEXT NOT NULL,
|
||||
state TEXT NOT NULL, -- pending|claimed|done|failed
|
||||
attempt INTEGER NOT NULL DEFAULT 1,
|
||||
max_attempts INTEGER NOT NULL DEFAULT 3,
|
||||
worker TEXT,
|
||||
lease_expiry REAL,
|
||||
result TEXT,
|
||||
cause TEXT,
|
||||
created_at REAL NOT NULL,
|
||||
updated_at REAL NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS ix_render_jobs_ready ON render_jobs(queue, state, created_at);
|
||||
"""
|
||||
|
||||
|
||||
def serialize_item(item: tuple) -> str:
|
||||
"""Render-queue tuple → JSON job payload."""
|
||||
number, prompt_id, prompt, extra_data, outputs, sensitive = item
|
||||
return json.dumps({
|
||||
"number": number,
|
||||
"prompt_id": prompt_id,
|
||||
"prompt": prompt,
|
||||
"extra_data": extra_data,
|
||||
"outputs": outputs,
|
||||
"sensitive": sensitive,
|
||||
})
|
||||
|
||||
|
||||
def deserialize_item(payload: str) -> tuple:
|
||||
"""JSON job payload → render-queue tuple (inverse of serialize_item)."""
|
||||
d = json.loads(payload)
|
||||
return (
|
||||
d.get("number", 0),
|
||||
d["prompt_id"],
|
||||
d["prompt"],
|
||||
d.get("extra_data") or {},
|
||||
d.get("outputs") or [],
|
||||
d.get("sensitive") or {},
|
||||
)
|
||||
|
||||
|
||||
def _worker_identity() -> str:
|
||||
return os.environ.get("STUDIO_WORKER_ID") or f"{socket.gethostname()}-{os.getpid()}"
|
||||
|
||||
|
||||
class SqlitePromptQueue(PromptQueue):
|
||||
"""PromptQueue whose durability lives in a SQLite file."""
|
||||
|
||||
def __init__(self, server, db_path: str | None = None):
|
||||
super().__init__(server)
|
||||
# SQLite is the durable backend — disable the JSON snapshot path and
|
||||
# drop anything a stray STUDIO_PERSIST_QUEUE restore left behind
|
||||
# (precedence: SQLite > PERSIST > memory).
|
||||
self._persist_path = None
|
||||
self.queue = []
|
||||
|
||||
self.db_path = db_path or os.environ["STUDIO_QUEUE_DB"]
|
||||
self.queue_name = os.environ.get("STUDIO_TASKS_QUEUE", DEFAULT_QUEUE)
|
||||
self.worker_id = _worker_identity()
|
||||
self.max_attempts = int(os.environ.get("STUDIO_TASKS_MAX_ATTEMPTS", "3"))
|
||||
self.lease_s = float(os.environ.get("STUDIO_TASKS_LEASE_MS", "60000")) / 1000.0
|
||||
self.poll_interval = float(os.environ.get("STUDIO_TASKS_POLL_INTERVAL", "1.0"))
|
||||
|
||||
os.makedirs(os.path.dirname(os.path.abspath(self.db_path)), exist_ok=True)
|
||||
self._db_lock = threading.Lock()
|
||||
self._conn = sqlite3.connect(self.db_path, check_same_thread=False, isolation_level=None)
|
||||
self._conn.execute("PRAGMA journal_mode=WAL")
|
||||
self._conn.execute("PRAGMA synchronous=NORMAL")
|
||||
self._conn.execute("PRAGMA busy_timeout=5000")
|
||||
self._conn.executescript(_SCHEMA)
|
||||
|
||||
# item_id → job_id so task_done reports to the right row.
|
||||
self._job_by_item: dict[int, str] = {}
|
||||
self._pending = 0
|
||||
self._reap() # recover leases orphaned by a previous crash
|
||||
|
||||
logger.info(
|
||||
"[sqlite-queue] durable render queue: db=%s queue=%s worker=%s lease=%.0fs",
|
||||
self.db_path, self.queue_name, self.worker_id, self.lease_s,
|
||||
)
|
||||
|
||||
self._hb_interval = max(self.lease_s / 3.0, 2.0)
|
||||
self._hb_stop = threading.Event()
|
||||
self._hb_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
|
||||
self._hb_thread.start()
|
||||
|
||||
# ── producer: /prompt → durable job ──────────────────────────────
|
||||
def put(self, item):
|
||||
prompt_id = item[1]
|
||||
now = time.time()
|
||||
with self._db_lock:
|
||||
self._conn.execute(
|
||||
"INSERT OR IGNORE INTO render_jobs"
|
||||
"(id, queue, payload, state, attempt, max_attempts, created_at, updated_at)"
|
||||
" VALUES(?,?,?,'pending',1,?,?,?)",
|
||||
(prompt_id, self.queue_name, serialize_item(item), self.max_attempts, now, now),
|
||||
)
|
||||
self.server.queue_updated()
|
||||
|
||||
# ── consumer: prompt_worker → claim ──────────────────────────────
|
||||
def get(self, timeout=None):
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
while True:
|
||||
self._reap()
|
||||
claimed = self._claim()
|
||||
if claimed is not None:
|
||||
return claimed
|
||||
if deadline is not None and time.monotonic() >= deadline:
|
||||
return None
|
||||
wait = self.poll_interval
|
||||
if deadline is not None:
|
||||
wait = min(wait, max(deadline - time.monotonic(), 0.0))
|
||||
if self._hb_stop.wait(wait):
|
||||
return None
|
||||
|
||||
def _claim(self):
|
||||
now = time.time()
|
||||
with self._db_lock:
|
||||
self._conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
row = self._conn.execute(
|
||||
"SELECT id, payload FROM render_jobs"
|
||||
" WHERE queue=? AND state='pending' ORDER BY created_at LIMIT 1",
|
||||
(self.queue_name,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
self._conn.execute("COMMIT")
|
||||
return None
|
||||
job_id, payload = row
|
||||
self._conn.execute(
|
||||
"UPDATE render_jobs SET state='claimed', worker=?, lease_expiry=?, updated_at=?"
|
||||
" WHERE id=?",
|
||||
(self.worker_id, now + self.lease_s, now, job_id),
|
||||
)
|
||||
self._conn.execute("COMMIT")
|
||||
except Exception:
|
||||
self._conn.execute("ROLLBACK")
|
||||
raise
|
||||
try:
|
||||
item = deserialize_item(payload)
|
||||
except Exception as e: # corrupt payload — fail terminally, don't loop.
|
||||
logger.error("[sqlite-queue] undecodable job %s: %s", job_id, e)
|
||||
self._terminal(job_id, "failed", cause=f"undecodable payload: {e}")
|
||||
return None
|
||||
with self.mutex:
|
||||
item_id = self.task_counter
|
||||
self.currently_running[item_id] = copy.deepcopy(item)
|
||||
self._job_by_item[item_id] = job_id
|
||||
self.task_counter += 1
|
||||
self.server.queue_updated()
|
||||
return (item, item_id)
|
||||
|
||||
# ── completion → durable state ───────────────────────────────────
|
||||
def task_done(self, item_id, history_result, status, process_item=None):
|
||||
with self.mutex:
|
||||
job_id = self._job_by_item.pop(item_id, None)
|
||||
if job_id is not None:
|
||||
success = status is None or status.status_str == "success"
|
||||
if success:
|
||||
self._terminal(job_id, "done",
|
||||
result=json.dumps({"outputs": _summarize_outputs(history_result)}))
|
||||
else:
|
||||
cause = "; ".join(status.messages) if status and status.messages else "render error"
|
||||
self._fail(job_id, cause)
|
||||
super().task_done(item_id, history_result, status, process_item)
|
||||
|
||||
def _terminal(self, job_id, state, result=None, cause=None):
|
||||
with self._db_lock:
|
||||
self._conn.execute(
|
||||
"UPDATE render_jobs SET state=?, result=?, cause=?, worker=NULL,"
|
||||
" lease_expiry=NULL, updated_at=? WHERE id=?",
|
||||
(state, result, cause, time.time(), job_id),
|
||||
)
|
||||
|
||||
def _fail(self, job_id, cause):
|
||||
# Retry-aware: another attempt if the policy allows, else terminal.
|
||||
with self._db_lock:
|
||||
self._conn.execute(
|
||||
"UPDATE render_jobs SET"
|
||||
" attempt = attempt + 1,"
|
||||
" state = CASE WHEN attempt + 1 > max_attempts THEN 'failed' ELSE 'pending' END,"
|
||||
" worker=NULL, lease_expiry=NULL, cause=?, updated_at=?"
|
||||
" WHERE id=? AND state='claimed'",
|
||||
(cause, time.time(), job_id),
|
||||
)
|
||||
|
||||
# ── crash-recovery: reclaim expired leases ───────────────────────
|
||||
def _reap(self):
|
||||
now = time.time()
|
||||
with self._db_lock:
|
||||
self._conn.execute(
|
||||
"UPDATE render_jobs SET"
|
||||
" attempt = attempt + 1,"
|
||||
" state = CASE WHEN attempt + 1 > max_attempts THEN 'failed' ELSE 'pending' END,"
|
||||
" cause = CASE WHEN attempt + 1 > max_attempts THEN 'lease expired' ELSE cause END,"
|
||||
" worker=NULL, lease_expiry=NULL, updated_at=?"
|
||||
" WHERE state='claimed' AND lease_expiry IS NOT NULL AND lease_expiry < ?",
|
||||
(now, now),
|
||||
)
|
||||
|
||||
# ── heartbeat / pending refresh ──────────────────────────────────
|
||||
def _heartbeat_loop(self):
|
||||
while not self._hb_stop.wait(self._hb_interval):
|
||||
now = time.time()
|
||||
with self.mutex:
|
||||
inflight = list(self._job_by_item.values())
|
||||
if inflight:
|
||||
with self._db_lock:
|
||||
for job_id in inflight:
|
||||
self._conn.execute(
|
||||
"UPDATE render_jobs SET lease_expiry=?, updated_at=?"
|
||||
" WHERE id=? AND state='claimed'",
|
||||
(now + self.lease_s, now, job_id),
|
||||
)
|
||||
self._reap()
|
||||
with self._db_lock:
|
||||
(self._pending,) = self._conn.execute(
|
||||
"SELECT COUNT(*) FROM render_jobs WHERE queue=? AND state='pending'",
|
||||
(self.queue_name,),
|
||||
).fetchone()
|
||||
|
||||
def get_tasks_remaining(self):
|
||||
with self.mutex:
|
||||
return len(self.currently_running) + self._pending
|
||||
|
||||
def job_state(self, job_id: str):
|
||||
"""Diagnostic: current (state, attempt) of a job, or None."""
|
||||
with self._db_lock:
|
||||
row = self._conn.execute(
|
||||
"SELECT state, attempt FROM render_jobs WHERE id=?", (job_id,)
|
||||
).fetchone()
|
||||
return row
|
||||
|
||||
def stop(self):
|
||||
self._hb_stop.set()
|
||||
with self._db_lock:
|
||||
self._conn.close()
|
||||
|
||||
|
||||
def _summarize_outputs(history_result) -> list:
|
||||
"""Extract output filenames from an execution's history_result."""
|
||||
files = []
|
||||
try:
|
||||
outputs = (history_result or {}).get("outputs", {})
|
||||
for node_output in outputs.values():
|
||||
for key in ("images", "gifs", "audio", "video"):
|
||||
for entry in node_output.get(key, []) or []:
|
||||
name = entry.get("filename")
|
||||
if name:
|
||||
files.append({"filename": name, "subfolder": entry.get("subfolder", ""),
|
||||
"type": entry.get("type", "output")})
|
||||
except Exception:
|
||||
pass
|
||||
return files
|
||||
@@ -215,7 +215,15 @@ class PromptServer():
|
||||
self.node_replace_manager = NodeReplaceManager()
|
||||
self.internal_routes = InternalRoutes(self)
|
||||
self.supports = ["custom_nodes_from_web"]
|
||||
self.prompt_queue = execution.PromptQueue(self)
|
||||
# Queue backend precedence: SQLite (durable, zero external process) >
|
||||
# PERSIST (JSON snapshot) > memory. STUDIO_QUEUE_DB opts into the
|
||||
# crash-durable SQLite queue; without it Studio behaves exactly as
|
||||
# before (in-memory + optional JSON snapshot).
|
||||
if os.environ.get("STUDIO_QUEUE_DB"):
|
||||
from middleware.tasks_queue import SqlitePromptQueue
|
||||
self.prompt_queue = SqlitePromptQueue(self)
|
||||
else:
|
||||
self.prompt_queue = execution.PromptQueue(self)
|
||||
self.loop = loop
|
||||
self.messages = asyncio.Queue()
|
||||
self.client_session:Optional[aiohttp.ClientSession] = None
|
||||
@@ -1292,6 +1300,8 @@ class PromptServer():
|
||||
self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items())
|
||||
self.subgraph_manager.add_routes(self.routes, nodes.LOADED_MODULE_DIRS.items())
|
||||
self.node_replace_manager.add_routes(self.routes)
|
||||
from middleware.engine_selector import add_engine_routes
|
||||
add_engine_routes(self.routes, self)
|
||||
self.app.add_subapp('/internal', self.internal_routes.get_app())
|
||||
|
||||
# Prefix every route with /api for easier matching for delegation.
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
Standalone tests for the engine selector (/v1/engines backing logic).
|
||||
Thin wrapper over the per-org compute registry — no GPU, no torch, no network.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import folder_paths
|
||||
from middleware import compute_config, engine_selector
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def org(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(folder_paths, "get_user_directory", lambda: str(tmp_path))
|
||||
return "acme"
|
||||
|
||||
|
||||
def test_local_always_present(org):
|
||||
out = engine_selector.list_engines(org)
|
||||
ids = [e["id"] for e in out["engines"]]
|
||||
assert ids == ["local"]
|
||||
assert out["default"] == "local"
|
||||
|
||||
|
||||
def test_registered_worker_listed(org):
|
||||
compute_config.register_worker(org, {
|
||||
"worker_id": "gb10-1", "url": "http://gb10:8188",
|
||||
"device": "cuda", "status": "ready", "gpu_model": "GB10", "vram_gb": 96,
|
||||
})
|
||||
out = engine_selector.list_engines(org)
|
||||
ids = {e["id"] for e in out["engines"]}
|
||||
assert ids == {"local", "gb10-1"}
|
||||
gb10 = next(e for e in out["engines"] if e["id"] == "gb10-1")
|
||||
assert gb10["type"] == "worker" and gb10["status"] == "ready" and gb10["device"] == "cuda"
|
||||
|
||||
|
||||
def test_set_and_resolve_default(org):
|
||||
compute_config.register_worker(org, {
|
||||
"worker_id": "gb10-1", "url": "http://gb10:8188", "device": "cuda", "status": "ready",
|
||||
})
|
||||
compute_config.update_config(org, {"default_engine": "gb10-1"})
|
||||
assert engine_selector.list_engines(org)["default"] == "gb10-1"
|
||||
w = engine_selector.resolve_engine_worker(org)
|
||||
assert w is not None and w.worker_id == "gb10-1"
|
||||
|
||||
|
||||
def test_default_local_resolves_none(org):
|
||||
assert engine_selector.resolve_engine_worker(org) is None
|
||||
|
||||
|
||||
def test_stale_default_falls_back_to_local(org):
|
||||
# Default names a worker that never registered → list reports local, no 500.
|
||||
compute_config.update_config(org, {"default_engine": "ghost"})
|
||||
assert engine_selector.list_engines(org)["default"] == "local"
|
||||
assert engine_selector.resolve_engine_worker(org) is None
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
Standalone tests for SqlitePromptQueue — Hanzo Studio's crash-durable render
|
||||
queue. A fake torch-free `execution` module supplies the PromptQueue base, so
|
||||
this runs with no GPU, no torch, and no external process (stdlib sqlite3 only).
|
||||
|
||||
Covers the queue seam (put/get/task_done), exactly-once claim, retry-on-error,
|
||||
and — the headline — crash-recovery: a job claimed and then abandoned by a
|
||||
"crashed" queue instance is reaped and re-run to completion by a fresh instance
|
||||
on the same DB file.
|
||||
"""
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from collections import namedtuple
|
||||
|
||||
import pytest
|
||||
|
||||
ExecStatus = namedtuple("ExecStatus", ["status_str", "completed", "messages"])
|
||||
|
||||
|
||||
def _fake_execution_module():
|
||||
mod = types.ModuleType("execution")
|
||||
|
||||
class PromptQueue:
|
||||
def __init__(self, server):
|
||||
self.server = server
|
||||
self.mutex = threading.RLock()
|
||||
self.task_counter = 0
|
||||
self.queue = []
|
||||
self.currently_running = {}
|
||||
self.history = {}
|
||||
self.flags = {}
|
||||
self._persist_path = None
|
||||
|
||||
def task_done(self, item_id, history_result, status, process_item=None):
|
||||
self.currently_running.pop(item_id, None)
|
||||
|
||||
mod.PromptQueue = PromptQueue
|
||||
return mod
|
||||
|
||||
|
||||
class FakeServer:
|
||||
def __init__(self):
|
||||
self.updates = 0
|
||||
|
||||
def queue_updated(self):
|
||||
self.updates += 1
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tqmod(monkeypatch):
|
||||
monkeypatch.setenv("STUDIO_TASKS_POLL_INTERVAL", "0.05")
|
||||
monkeypatch.setenv("STUDIO_TASKS_LEASE_MS", "300") # short lease → fast reap
|
||||
monkeypatch.setenv("STUDIO_TASKS_MAX_ATTEMPTS", "3")
|
||||
saved = {k: sys.modules.get(k) for k in ("execution", "middleware.tasks_queue")}
|
||||
sys.modules["execution"] = _fake_execution_module()
|
||||
sys.modules.pop("middleware.tasks_queue", None)
|
||||
import middleware.tasks_queue as tq
|
||||
try:
|
||||
yield tq
|
||||
finally:
|
||||
for k, v in saved.items():
|
||||
if v is None:
|
||||
sys.modules.pop(k, None)
|
||||
else:
|
||||
sys.modules[k] = v
|
||||
|
||||
|
||||
def _item(pid="prompt-1"):
|
||||
return (0, pid, {"1": {"class_type": "EmptyImage"}}, {"org_id": "acme"}, ["1"], {})
|
||||
|
||||
|
||||
def _queue(tqmod, tmp_path, freeze_hb=True):
|
||||
import os
|
||||
os.environ["STUDIO_QUEUE_DB"] = str(tmp_path / "queue.db")
|
||||
q = tqmod.SqlitePromptQueue(FakeServer())
|
||||
if freeze_hb:
|
||||
q._hb_stop.set() # deterministic: no background reaper/heartbeat
|
||||
return q
|
||||
|
||||
|
||||
def test_serde_roundtrip(tqmod):
|
||||
assert tqmod.deserialize_item(tqmod.serialize_item(_item())) == _item()
|
||||
|
||||
|
||||
def test_put_get_task_done(tqmod, tmp_path):
|
||||
q = _queue(tqmod, tmp_path)
|
||||
q.put(_item())
|
||||
assert q.server.updates == 1
|
||||
item, item_id = q.get(timeout=1.0)
|
||||
assert item == _item()
|
||||
assert q.currently_running[item_id] == _item()
|
||||
assert q.job_state("prompt-1")[0] == "claimed"
|
||||
q.task_done(item_id, {"outputs": {"2": {"images": [{"filename": "x_00001_.png"}]}}},
|
||||
ExecStatus("success", True, []))
|
||||
assert q.job_state("prompt-1")[0] == "done"
|
||||
assert item_id not in q._job_by_item
|
||||
|
||||
|
||||
def test_put_is_idempotent(tqmod, tmp_path):
|
||||
q = _queue(tqmod, tmp_path)
|
||||
q.put(_item())
|
||||
q.put(_item()) # same prompt_id → INSERT OR IGNORE, no duplicate
|
||||
q.get(timeout=1.0)
|
||||
assert q.get(timeout=0.2) is None # only one job existed
|
||||
|
||||
|
||||
def test_exactly_once_claim(tqmod, tmp_path):
|
||||
q = _queue(tqmod, tmp_path)
|
||||
q.put(_item())
|
||||
first = q.get(timeout=1.0)
|
||||
assert first is not None
|
||||
assert q.get(timeout=0.2) is None # already claimed — no double dispatch
|
||||
|
||||
|
||||
def test_error_retries_then_fails(tqmod, tmp_path):
|
||||
q = _queue(tqmod, tmp_path)
|
||||
q.put(_item()) # max_attempts=3
|
||||
# attempt 1 → error → back to pending (attempt 2)
|
||||
_, item_id = q.get(timeout=1.0)
|
||||
q.task_done(item_id, {}, ExecStatus("error", False, ["boom-1"]))
|
||||
assert q.job_state("prompt-1") == ("pending", 2)
|
||||
# attempt 2 → error → pending (attempt 3)
|
||||
_, item_id = q.get(timeout=1.0)
|
||||
q.task_done(item_id, {}, ExecStatus("error", False, ["boom-2"]))
|
||||
assert q.job_state("prompt-1") == ("pending", 3)
|
||||
# attempt 3 → error → terminal failed
|
||||
_, item_id = q.get(timeout=1.0)
|
||||
q.task_done(item_id, {}, ExecStatus("error", False, ["boom-3"]))
|
||||
assert q.job_state("prompt-1")[0] == "failed"
|
||||
|
||||
|
||||
def test_crash_recovery_reap(tqmod, tmp_path):
|
||||
import time
|
||||
q = _queue(tqmod, tmp_path)
|
||||
q.put(_item())
|
||||
_, _item_id = q.get(timeout=1.0) # claimed, leased
|
||||
assert q.job_state("prompt-1")[0] == "claimed"
|
||||
time.sleep(0.4) # lease (0.3s) expires: worker "crashed"
|
||||
assert q.get(timeout=1.0) is not None # get() reaps + re-claims for a new attempt
|
||||
assert q.job_state("prompt-1") == ("claimed", 2)
|
||||
|
||||
|
||||
def test_crash_recovery_across_instances(tqmod, tmp_path):
|
||||
"""The headline: a job in flight when a queue instance dies is picked up
|
||||
and completed by a fresh instance on the same DB file."""
|
||||
import time
|
||||
db = str(tmp_path / "queue.db")
|
||||
import os
|
||||
os.environ["STUDIO_QUEUE_DB"] = db
|
||||
|
||||
# Instance A: submit, claim, then "crash" (drop the instance without completing).
|
||||
a = tqmod.SqlitePromptQueue(FakeServer())
|
||||
a._hb_stop.set()
|
||||
a.put(_item())
|
||||
a.get(timeout=1.0)
|
||||
assert a.job_state("prompt-1")[0] == "claimed"
|
||||
a.stop() # process dies mid-render
|
||||
|
||||
time.sleep(0.4) # lease expires
|
||||
|
||||
# Instance B: same DB. Recovers the orphaned job and finishes it.
|
||||
b = tqmod.SqlitePromptQueue(FakeServer()) # __init__ reaps on boot
|
||||
b._hb_stop.set()
|
||||
item, item_id = b.get(timeout=1.0)
|
||||
assert item == _item()
|
||||
assert b.job_state("prompt-1") == ("claimed", 2) # attempt 2 on the new worker
|
||||
b.task_done(item_id, {"outputs": {}}, ExecStatus("success", True, []))
|
||||
assert b.job_state("prompt-1")[0] == "done"
|
||||
b.stop()
|
||||
Reference in New Issue
Block a user