FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Merge feat/693-B-device-fields (#777) · dkackman/diffusers-workflow@d1efb6b · GitHub

Repository navigation

Commit d1efb6b

Browse files
committed
Merge feat/693-B-device-fields (#777)
2 parents bc0b19b + 49895e9 commit d1efb6b

12 files changed

Lines changed: 283 additions & 79 deletions

‎docs/ARCHITECTURE.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ one to open.
110110
| 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` |
111111
| 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` |
112112
| 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. | — |
113-
| Worker card pinning | `dw/devices.py`: `worker_environment`, `pinned_environment`, `resolve_serve_devices`, `check_device_present`, `device_label`, `card_of`, `ordinal_of` | 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>"`, read back only through `card_of` and `ordinal_of`. | `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` |
113+
| Worker card pinning | `dw/devices.py`: `worker_environment`, `pinned_environment`, `resolve_serve_devices`, `check_device_present`, `label_of` | 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 stores the card it ran on as `device_ordinal` and `device_card` (#693), job and history alike: rerun affinity reads the ordinal, observed cost buckets by the card, and the `device` label `"cuda:N <card name>"` is only `label_of` joining the two - nothing splits a label back apart. | `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` |
114114
| 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` |
115115
| 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` |
116116
| 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` |

‎dw/devices.py‎

Lines changed: 13 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,13 @@
1616
the *parent* around `Process.start()` (`pinned_environment`): `dw/__init__.py`
1717
imports torch, so nothing in the child runs early enough to set them there.
1818
19-
**The server process sees every card**, so it names one by its ordinal -
20-
`device_label()` gives `"cuda:1 NVIDIA GeForce RTX 3090"`, the string a job
21-
record's `device` carries and the card observed cost buckets by
22-
(dw/server/observed_cost.py).
19+
**The server process sees every card**, so it names one by its ordinal
20+
(`device_ordinal()`, `"cuda:1"`) and its card (`card_name()`, `"NVIDIA
21+
GeForce RTX 3090"`). A job record stores the two as separate fields
22+
(`device_ordinal`, `device_card`, #693): rerun affinity reads the ordinal and
23+
observed cost buckets by the card (dw/server/observed_cost.py), so neither
24+
parses a string back apart. `label_of()` joins them into the
25+
`"cuda:1 NVIDIA GeForce RTX 3090"` a job's `device` shows.
2326
2427
A device with no explicit index (`cuda`, the default on every box that
2528
names none) is not pinned at all, so an unconfigured server spawns its
@@ -229,28 +232,13 @@ def card_name(device=None):
229232
return None
230233

231234

232-
def device_label(device=None):
233-
"""`"cuda:1 NVIDIA GeForce RTX 3090"` - the ordinal, then the card's
234-
name when there is one. What a job record's `device` holds."""
235-
ordinal = device_ordinal(device)
236-
name = card_name(ordinal)
237-
return f"{ordinal} {name}" if name else ordinal
238-
239-
240-
def card_of(label):
241-
"""The card name in a `device_label` string, or None for a label that
242-
carries none (`cpu`) or no label at all."""
243-
if not label or " " not in label:
244-
return None
245-
return label.split(" ", 1)[1]
246-
247-
248-
def ordinal_of(label):
249-
"""The ordinal in a `device_label` string ('cuda:1'), or None for no
250-
label at all - `card_of`'s other half."""
251-
if not label:
235+
def label_of(ordinal, card):
236+
"""The `device` label for a stored ordinal and card: `"cuda:1 NVIDIA
237+
GeForce RTX 3090"`, or the bare ordinal (`cpu`) for a card with no name,
238+
or None for no ordinal at all."""
239+
if not ordinal:
252240
return None
253-
return label.split(" ", 1)[0]
241+
return f"{ordinal} {card}" if card else ordinal
254242

255243

256244
def is_default_device(device=None):

‎dw/server/job_history.py‎

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,12 @@
44
the column added and every existing row backfilled to `default`, since
55
history that cannot say which workspace a job ran in stops making sense once
66
there are two.
7+
8+
The card a job ran on is three columns: `device`, the label every reader
9+
returns (`"cuda:1 NVIDIA GeForce RTX 3090"`), and `device_ordinal`/
10+
`device_card`, the same card as the two fields rerun affinity and observed
11+
cost read (#693). A database that predates the two gets them backfilled from
12+
`device` in SQL when it is first opened, so no reader parses the label.
713
"""
814

915
import json
@@ -114,6 +120,27 @@ def __init__(self, db_path):
114120
# card has been swapped
115121
if "device" not in columns:
116122
connection.execute("ALTER TABLE jobs ADD COLUMN device TEXT")
123+
# The same card as two fields - its ordinal ("cuda:1") and its
124+
# name ("NVIDIA GeForce RTX 3090") - so rerun affinity and
125+
# observed cost read them rather than parse `device` (#693).
126+
# `device` stays the label every reader returns. Rows before the
127+
# columns are backfilled from it once, in SQL: an ordinal never
128+
# holds a space, so the label splits at its first one, and a
129+
# label with none (`cpu`, `mps`) is all ordinal with no card
130+
if "device_ordinal" not in columns:
131+
connection.execute("ALTER TABLE jobs ADD COLUMN device_ordinal TEXT")
132+
if "device_card" not in columns:
133+
connection.execute("ALTER TABLE jobs ADD COLUMN device_card TEXT")
134+
if "device_ordinal" not in columns or "device_card" not in columns:
135+
connection.execute(
136+
"UPDATE jobs SET"
137+
" device_ordinal = CASE WHEN INSTR(device, ' ') > 0"
138+
" THEN SUBSTR(device, 1, INSTR(device, ' ') - 1)"
139+
" ELSE device END,"
140+
" device_card = CASE WHEN INSTR(device, ' ') > 0"
141+
" THEN SUBSTR(device, INSTR(device, ' ') + 1) END"
142+
" WHERE device IS NOT NULL"
143+
)
117144

118145
def _connect(self):
119146
# WAL mode lets a reader (the web UI polling job status, an MCP
@@ -136,8 +163,9 @@ def record(self, job):
136163
" started_at, finished_at, arguments, spec, manifest, warnings,"
137164
" error, events, workspace, workflow_name, run_id, run_dir,"
138165
" acknowledged, host_memory_peak_rss_mb,"
139-
" host_memory_job_peak_rss_mb, run_version, device) VALUES"
140-
" (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
166+
" host_memory_job_peak_rss_mb, run_version, device,"
167+
" device_ordinal, device_card) VALUES"
168+
" (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
141169
(
142170
job.id,
143171
job.workflow_name,
@@ -164,6 +192,8 @@ def record(self, job):
164192
getattr(job, "host_memory_job_peak_rss_mb", None),
165193
getattr(job, "run_version", None),
166194
getattr(job, "device", None),
195+
getattr(job, "device_ordinal", None),
196+
getattr(job, "device_card", None),
167197
),
168198
)
169199

@@ -228,6 +258,15 @@ def get(self, job_id):
228258
).fetchone()
229259
return self._to_detail(row) if row else None
230260

261+
def device_ordinal(self, job_id):
262+
"""The ordinal of the card a finished job ran on ('cuda:1'), or None
263+
for a job that never ran or that history has never seen."""
264+
with self._lock, self._connect() as connection:
265+
row = connection.execute(
266+
"SELECT device_ordinal FROM jobs WHERE id = ?", (job_id,)
267+
).fetchone()
268+
return row[0] if row else None
269+
231270
def watermark(self):
232271
"""How far the table has got - what a derived figure caches against.
233272
@@ -272,8 +311,8 @@ def finished_runs(self):
272311
"SELECT workflow_name, workspace, started_at, finished_at,"
273312
" arguments, manifest, INSTR(COALESCE(events, ''), ?) > 0,"
274313
" COALESCE(json_array_length(COALESCE(events, '[]')), 0) >= ?,"
275-
" host_memory_peak_rss_mb, host_memory_job_peak_rss_mb, device"
276-
" FROM jobs WHERE status = ? AND workflow_name IS NOT NULL"
314+
" host_memory_peak_rss_mb, host_memory_job_peak_rss_mb, device,"
315+
" device_ordinal, device_card FROM jobs WHERE status = ? AND workflow_name IS NOT NULL"
277316
" AND started_at IS NOT NULL AND finished_at IS NOT NULL",
278317
(LOADING_MARKER, EVENT_CAP, SUCCEEDED),
279318
).fetchall()
@@ -290,6 +329,8 @@ def finished_runs(self):
290329
peak_rss_mb,
291330
job_peak_rss_mb,
292331
device,
332+
ordinal,
333+
card,
293334
) in rows:
294335
key = (workspace or DEFAULT_WORKSPACE_NAME, name)
295336
grouped.setdefault(key, []).append(
@@ -304,6 +345,8 @@ def finished_runs(self):
304345
"host_memory_peak_rss_mb": peak_rss_mb,
305346
"host_memory_job_peak_rss_mb": job_peak_rss_mb,
306347
"device": device,
348+
"device_ordinal": ordinal,
349+
"device_card": card,
307350
}
308351
)
309352
return grouped

‎dw/server/job_record.py‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import time
1111
import uuid
1212

13+
from ..devices import label_of
1314
from ..download_watch import format_progress
1415
from ..workspace import DEFAULT_WORKSPACE_NAME
1516

@@ -78,10 +79,13 @@ def __init__(self, spec):
7879
self.run_id = None
7980
self.run_dir = None
8081
self.run_version = None
81-
# The card the job ran on, as `dw.devices.device_label` names it -
82-
# set when it starts running, so None while queued and forever for
83-
# a job cancelled before it ran (#462)
84-
self.device = None
82+
# The card the job ran on - its ordinal ("cuda:1") and its name
83+
# ("NVIDIA GeForce RTX 3090", None for a card with none) - set when
84+
# it starts running, so None while queued and forever for a job
85+
# cancelled before it ran (#462, #693). `device` is the two as one
86+
# label
87+
self.device_ordinal = None
88+
self.device_card = None
8589
# Which form of cost acknowledgement queued this job (#85)
8690
self.acknowledged = spec.get("acknowledged") or ACK_NONE
8791
# The worker's own high-water mark for this run, from its final
@@ -110,6 +114,11 @@ def __init__(self, spec):
110114
self.denoise_total_steps = None
111115
self.condition = threading.Condition()
112116

117+
@property
118+
def device(self):
119+
"""`"cuda:1 NVIDIA GeForce RTX 3090"`, or None before it ran."""
120+
return label_of(self.device_ordinal, self.device_card)
121+
113122
def add_event(self, event):
114123
with self.condition:
115124
# `at` is seconds since the job started (since it was created,

‎dw/server/jobs.py‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
import threading
4141

4242
from ..worker_protocol import Execute, workflow_identity
43-
from ..devices import card_of, device_ordinal, ordinal_of
43+
from ..devices import device_ordinal
4444
from ..host_memory import process_rss_mb
4545
from ..worker_manager import WorkerManager
4646
from ..workflow_run import SEED_BITS
@@ -549,10 +549,8 @@ def _original_device(self, job_id):
549549
"""The ordinal of the card `job_id` ran on ('cuda:1'), or None."""
550550
job = self.jobs.get(job_id)
551551
if job is not None:
552-
label = job.device
553-
else:
554-
label = (self.history.get(job_id) or {}).get("device")
555-
return ordinal_of(label)
552+
return job.device_ordinal
553+
return self.history.device_ordinal(job_id)
556554

557555
def queue_position(self, job_id):
558556
"""Index in the waiting queue, or None when the job is not queued."""
@@ -685,7 +683,7 @@ def _run_loop(self):
685683
slot, job = picked
686684
job.status = RUNNING
687685
job.started_at = time.time()
688-
job.device = self._job_device(slot)
686+
job.device_ordinal, job.device_card = slot.device_fields()
689687
slot.current_job_id = job.id
690688
slot.started_at = job.started_at
691689
# The worker switches to this identity for the run
@@ -839,12 +837,6 @@ def _rank_for_oom(self, slot):
839837
highest = max((other.oom_score_adj or 0) for other in others)
840838
manager.set_oom_score_adj(min(1000, highest + 100))
841839

842-
def _job_device(self, slot=None):
843-
"""The card the worker runs on, for the job record. A label that
844-
cannot be read leaves the job's `device` null rather than failing
845-
the job - the run itself does not depend on it."""
846-
return (slot or self.slots[0]).label()
847-
848840
def is_busy(self):
849841
"""True while a job is running on any card or queued - the window in
850842
which a worker may be reading model files a cache delete would rip
@@ -889,7 +881,7 @@ def workers(self):
889881
pid = slot.manager.pid() if hasattr(slot.manager, "pid") else None
890882
entry = {
891883
"device": device_ordinal(slot.device),
892-
"name": card_of(slot.label()),
884+
"name": slot.device_fields()[1],
893885
"vram_gb": slot.capacity_gb(),
894886
"current_job": job_id,
895887
"alive": slot.alive(),

‎dw/server/observed_cost.py‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -337,12 +337,12 @@ def _on_this_card(self, rows, card=None):
337337
"""The rows that ran on the card this server runs on now, or on
338338
`card` when one is named (#462).
339339
340-
A row names its card (`device`, "cuda:1 NVIDIA GeForce RTX 3090")
341-
and counts when the card name matches, whichever index it had - a
342-
3090 run tells a 3090 estimate, not a 4090 one. A row from before
343-
jobs carried a `device` ran on the box's default card, so it counts
344-
only for that card."""
345-
from ..devices import card_of, is_default_device
340+
A row names its card (`device_card`, "NVIDIA GeForce RTX 3090") and
341+
counts when the card name matches, whichever index it had - a 3090
342+
run tells a 3090 estimate, not a 4090 one. A row from before jobs
343+
carried a device (no `device_ordinal`) ran on the box's default
344+
card, so it counts only for that card."""
345+
from ..devices import is_default_device
346346

347347
if card is not None:
348348
_, name, default = self.card(card)
@@ -355,11 +355,10 @@ def _on_this_card(self, rows, card=None):
355355
default = True
356356
kept = []
357357
for row in rows:
358-
label = row.get("device")
359-
if label is None:
358+
if row.get("device_ordinal") is None:
360359
if default:
361360
kept.append(row)
362-
elif card_of(label) == name:
361+
elif row.get("device_card") == name:
363362
kept.append(row)
364363
return kept
365364

‎dw/server/pool.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,15 @@ def label(self):
6666
logger.debug("Could not read the worker's device", exc_info=True)
6767
return None
6868

69+
def device_fields(self):
70+
"""The card as a job record stores it - `("cuda:1", "NVIDIA GeForce
71+
RTX 3090")` - or `(None, None)` where it cannot be read."""
72+
try:
73+
return self.manager.device_fields()
74+
except Exception:
75+
logger.debug("Could not read the worker's device", exc_info=True)
76+
return None, None
77+
6978
def ordinal(self):
7079
"""The card as `cuda:1` - what a `device` argument names it by."""
7180
return device_ordinal(self.device)

‎dw/worker_manager.py‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,13 @@
1414
import time
1515
from typing import Optional
1616
from . import device_capacity_gb, get_device
17-
from .devices import device_label, pinned_environment, worker_environment
17+
from .devices import (
18+
card_name,
19+
device_ordinal,
20+
label_of,
21+
pinned_environment,
22+
worker_environment,
23+
)
1824
from .worker import worker_main
1925
from .worker_protocol import Cancel, Shutdown, WorkerCrashed, parse_reply
2026

@@ -42,7 +48,7 @@ def __init__(self, device: Optional[str] = None):
4248
(dw/devices.py)
4349
"""
4450
self.device = device
45-
self._device_label = None
51+
self._device_fields = None
4652
self._capacity_read = False
4753
self._capacity_gb = None
4854
self._ceiling_gb = None
@@ -77,13 +83,19 @@ def ensure_worker(self, log_level: str = "INFO"):
7783
self.oom_score_adj = None
7884
logger.info("Worker process started")
7985

86+
def device_fields(self):
87+
"""This worker's card as a job record stores it - `("cuda:1",
88+
"NVIDIA GeForce RTX 3090")`, the card None where it has no name -
89+
read once, in the server process, which sees every card by its
90+
ordinal."""
91+
if self._device_fields is None:
92+
ordinal = device_ordinal(self.device)
93+
self._device_fields = (ordinal, card_name(ordinal))
94+
return self._device_fields
95+
8096
def device_label(self):
81-
"""This worker's card as a job record names it -
82-
`"cuda:1 NVIDIA GeForce RTX 3090"` - read once, in the server
83-
process, which sees every card by its ordinal."""
84-
if self._device_label is None:
85-
self._device_label = device_label(self.device)
86-
return self._device_label
97+
"""`"cuda:1 NVIDIA GeForce RTX 3090"` - the two fields as one label."""
98+
return label_of(*self.device_fields())
8799

88100
def shutdown_worker(self):
89101
"""Gracefully shutdown worker process."""

0 commit comments

Comments
 (0)

Back | FazBrowse Home | New Git URL