trending/VoiceStudio
1
0
Fork 0
mirror of https://github.com/debpalash/VoiceStudio.git synced 2026-10-02 01:26:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(workers): one gateway for GPU calls, and results too big for the wire

Browse Source 
Two phases of the remote-GPU plan, landing together because neither is
useful alone: on a 4090 any render long enough to exercise the progress
lease also outgrows the 8 MiB message cap, so a gateway that routes work
remotely without an artifact transport just moves where the failure
happens.

**The gateway** (`services/gpu_gateway.py`) is the single owner of GPU
calling, model status, downloads and model load, for both targets —
`prewarm`, `run`, `status`, `download`. prewarm and run stay separate
because collapsing them loses the two-phase load/generate budget split
(#1033/#1037) that the worker protocol already mirrors. Admission moves
in here too: the old `check_gpu_admission` call read *local* pool stats,
so under Remote it would 429 on local saturation while the remote GPU
sat idle.

**Artifacts** now move out of band above a negotiated threshold. Bytes
land in an attempt-scoped `.part` file, are verified against a declared
sha256, and are renamed into place only on an explicit last chunk — a
transfer that arrives short, reordered, or simply stops commits nothing.
A resume rehashes what is already on disk, or the digest would attest
only to the tail, which is the exact case a resume exists to protect.

Two failure modes found while verifying this, both fixed with
mutation-checked regressions:

  * an oversized payload with no session (mid-reconnect, or a control
    plane too old to serve UploadResult) has nowhere to go. It must not
    enter `_pending` — an over-cap frame is re-sent on every reconnect,
    killing the session each time and stranding every other task — but
    it must stay retryable, unlike the size gate's TERMINAL verdict:
    nothing about the render is wrong, only the route to it.

  * the upload resume loop was bounded by "did the offset change", which
    a receiver alternating between two byte counts satisfies forever.
    The worker is single-slot by default, so that is not one lost upload
    but the machine, doing nothing else, until someone restarts it.
    Bounded by a round count instead.

The control stream is split into control and bulk queues so the
heartbeat this whole liveness model rests on cannot queue behind a
payload — `result_json` has no size cliff to catch it, and the next bulk
message added to the protocol would have reintroduced the stall
silently.

Live streaming stays on the control plane and now says so once per
socket: that route exists to put audio in the user's ear before the
sentence finishes, and paying queue admission plus a round trip per
utterance would spend the one thing it is for. Silence would have been
worse than the limit — the header badge would read "gpu2" while this
machine did all the work.

Backend 5236 passed, frontend 1807 passed. End-to-end verification on
real hardware has NOT been re-run since these changes; the CHANGELOG
claim for the Synthesize button waits on that.
This commit is contained in:
velixio
2026-08-11 13:30:26 +05:30
parent c643706d07
commit b54cd28403
43 changed files with 6386 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/routers/generation.py
+467 -115
View File
@@ -851,6 +851,7 @@ def _persist_profile_ref_text(profile_id: str, ref_text: str) -> None:
async def _finalize_generation(
audio_tensor, sample_rate, *, text, history_mode, ref_audio_path,
language, instruct, resolved_profile_id, used_seed, start_time,
already_marked=False,
):
"""Shared tail of a successful generation: watermark → save WAV →
history row (self-healing) → retention prune → event emit.
@@ -860,6 +861,12 @@ async def _finalize_generation(
watermark, filename, history row, retention behavior — is identical
regardless of how the audio was delivered to the client.
``already_marked`` is for audio that arrives provenance-marked: a remote
worker marks at the tensor stage before it encodes (with ``force=True``,
so the *requesting* user's preference governs, not the GPU owner's), and
embedding a second AudioSeal payload over the first degrades detection of
both. The take users keep carries exactly one whole-take mark either way.
Returns ``(watermarked_tensor, meta)`` where ``meta`` carries
``id`` / ``filename`` / ``duration`` / ``gen_time``.
"""
@@ -874,13 +881,14 @@ async def _finalize_generation(
# Dispatched to the dedicated watermark pool, not the GPU pool (#1190):
# AudioSeal embedding is CPU work that holds no VRAM, so occupying a GPU
# worker with it only delays the next generate on 1-worker hosts.
from services.watermark import mark_synthetic
from services.model_manager import get_watermark_pool
audio_tensor = await loop.run_in_executor(
get_watermark_pool(),
functools.partial(mark_synthetic, audio_tensor, sample_rate,
context="generate.finalize"),
)
if not already_marked:
from services.watermark import mark_synthetic
from services.model_manager import get_watermark_pool
audio_tensor = await loop.run_in_executor(
get_watermark_pool(),
functools.partial(mark_synthetic, audio_tensor, sample_rate,
context="generate.finalize"),
)
gen_time = round(time.time() - start_time, 2)
audio_id = str(uuid.uuid4())[:8]
@@ -962,6 +970,144 @@ def _pcm16_b64(wav_tensor) -> str:
return base64.b64encode(pcm.cpu().numpy().tobytes()).decode("ascii")
# ── Remote GPU: this route is the producer the scheduler never had ─────────
#
# Picking a remote worker used to change a badge and nothing else — every
# render still ran on this machine, which is the whole reported bug. The
# decision is taken ONCE per request, through `services/gpu_gateway.py`, and
# BEFORE anything local is loaded — for two reasons that are not
# interchangeable: a render bound for the user's 4090 must not first pull a
# multi-GB model into this machine's RAM, and it must not be refused by a gate
# that asked whether THIS host has the accelerator the engine needs (a
# CUDA-only engine on a Mac control plane is exactly the case remote workers
# exist for).
_REMOTE_OP = "tts"
# The gateway's coarse phase → the sentence a user reads while someone else's
# GPU works. A five-minute remote render otherwise shows the same bare spinner
# as a local one, with no way to tell "queued behind another task" from
# "downloading 5 GB of weights" from "actually generating".
_REMOTE_PHASE_LABELS = {
"queued": "queued on {target}",
"loading": "loading model on {target}",
"running": "generating on {target}",
"uploading": "receiving audio from {target}",
}
class _LocalDecision:
"""Stand-in for ``worker.routing.Decision`` meaning "run here".
Used only when the gateway cannot be imported at all, so a build without
it still renders instead of 500-ing.
"""
remote = False
worker_id = None
label = "Local"
reason = ""
_LOCAL_DECISION = _LocalDecision()
def _routing_decision():
"""Local or remote for this request — resolved once, never re-asked.
Asked once because the target is user-settable at any moment: a decision
that flipped between prewarm and dispatch would either warm an engine
nothing will use or dispatch remotely after paying a local cold load.
"""
try:
from services import gpu_gateway
return gpu_gateway.decide(_REMOTE_OP)
except Exception: # noqa: BLE001 — routing is advisory; local always works
logger.debug("remote routing unavailable; running locally", exc_info=True)
return _LOCAL_DECISION
def _remote_only_local_call(target_label, reason=""):
"""The local branch of a render whose local half was deliberately skipped.
``gpu_gateway.run`` always takes a local callable — it is where rule 1
(pre-dispatch unavailability) lands. But this route skips every local
preparation step once the decision is remote, precisely so a job bound for
the 4090 does not first load gigabytes here, so there is no local render
left to fall back to.
The causes rule 1 actually covers — worker offline, disabled, not
approved, breaker open, remote workers switched off — are already answered
by ``decide()`` BEFORE that skip, and come back as a local decision with a
named reason. What is left is the narrow window where dispatch itself is
refused (a full queue, a task dropped between submit and wait). Saying so
and offering the local re-run is honest; silently returning nothing is not.
"""
from services.gpu_gateway import RemoteJobFailed
def _refuse():
raise RemoteJobFailed(
reason or f"{target_label} could not take this render",
worker_label=target_label,
code="REMOTE_NOT_DISPATCHED",
hint="Run it on this machine instead, or pick another GPU.",
)
return _refuse
def _remote_progress_frame(state, target):
"""One gateway ``on_state`` payload → the NDJSON event the UI renders."""
phase = str((state or {}).get("phase") or "running")
try:
pct = max(0, min(100, round(float((state or {}).get("progress") or 0.0) * 100)))
except (TypeError, ValueError):
pct = 0
detail = _REMOTE_PHASE_LABELS.get(phase, _REMOTE_PHASE_LABELS["running"])
detail = detail.format(target=target)
if phase == "running" and pct:
detail = f"{detail} ({pct}%)"
return {
"type": "progress", "stage": phase, "percent": pct,
"target": target, "detail": detail,
}
def _apply_routing_headers(headers, engine_notice, decision):
"""Say where this render ran, on the notice channel that already exists.
``X-OmniVoice-Routing`` / ``-Routing-Reason`` are already set for the #21
engine routing gate and already consumed as a de-duped one-time toast, so
"this ran on gpu2" and "your 4090 was asleep, this ran here" travel the
same wire rather than inventing a second one.
The engine notice wins on a local render: "the engine fell back to CPU"
explains the slowness the user is looking at, while the worker notice for
a local render is the quieter of the two. A remote render has no engine
notice at all — that gate answers for THIS host, and this host did nothing.
"""
from services.engine_routing import header_safe_reason
notice = engine_notice
if decision is not None:
try:
from services.gpu_gateway import notice_for
worker_notice = notice_for(decision)
except Exception: # noqa: BLE001 — a notice must never fail a render
worker_notice = None
if worker_notice and (getattr(decision, "remote", False) or not notice):
notice = worker_notice
if not notice:
return headers
headers["X-OmniVoice-Routing"] = notice[0]
safe = header_safe_reason(notice[1]) if notice[1] else ""
if safe:
headers["X-OmniVoice-Routing-Reason"] = safe
return headers
@router.post("/generate")
async def generate_speech(
text: str = Form(...),
@@ -1038,70 +1184,85 @@ async def generate_speech(
from core.run_sentinel import touch_activity
touch_activity("generate", engine_id)
# Single-active-engine memory discipline: hand back any OTHER resident TTS
# engine's model before loading this one, so switching engines (or a
# per-request engine= override, which bypasses /engines/select entirely)
# doesn't stack two multi-GB models in memory — the accumulation behind the
# 16 GB-Mac OOM deaths. No-op when nothing else is resident, so steady-state
# single-engine use pays nothing. Opt out: OMNIVOICE_SINGLE_ENGINE_RESIDENT=0.
from services.engine_memory import evict_other_tts_engines
await evict_other_tts_engines(engine_id)
# Non-blocking breadcrumb: if free memory is already low before this load,
# log it. A later OOM kill (the 16 GB-Mac class) then has a trail pointing
# at the load that tipped it, instead of a silent process death. Never
# blocks — the OS can reclaim cache, and a hard refuse would brick
# legitimate loads.
try:
from services.memory_budget import log_if_low
log_if_low(f"TTS load ({engine_id})")
except Exception:
pass
# VRAM eviction runs in get_model()'s warm-return path now, so every native
# TTS generate (this route, WS TTS, dub, batch, audiobook) is covered.
# ── Where does this render run? Asked once, here, because every line
# between this point and the dispatch below is preparation of THIS
# machine's GPU — model eviction, a multi-GB load, a host-capability gate.
# None of it applies to a render that belongs on the user's other box, and
# running it anyway is how "I selected gpu2" ended up meaning "the Mac did
# the work after loading the model twice".
_decision = _routing_decision()
_remote = bool(getattr(_decision, "remote", False))
_target_label = getattr(_decision, "label", "") or "the chosen worker"
_model = None
_backend = None
if backend_cls is OmniVoiceBackend:
# VoiceStudio keeps its native path: it carries the full advanced
# parameter surface (t_shift, layer/position/class controls) that the
# generic adapter protocol doesn't. Byte-identical to the old behavior.
_model = await get_model()
else:
try:
ok, msg = backend_cls.is_available()
except Exception as exc:
ok, msg = False, f"{type(exc).__name__}: {exc}"
if not ok:
raise HTTPException(
status_code=400,
detail=f"TTS engine '{engine_id}' is not available: {_mask_hf_tokens(msg)}",
)
# Reuse the per-process instance cache shared with the engine
# health-check route so weights load once, not per request.
from api.routers.engines import _get_engine_instance
_backend = _get_engine_instance(backend_cls)
# ── Routing gate (#21 — no silent CPU fallback). Computed ONCE per request
# (host caps are constant; the per-request engine= override bypasses the
# /engines/select gate, so this is the only place it's enforced for synth).
from core.device_caps import detect_host_caps
from services.engine_routing import resolve_routing, routing_notice
# The engine's declared VRAM floor (#1226) — used by the routing gate and,
# below, to let a generate TIMEOUT name the same shortfall. Resolved once:
# every other job on this GPU pool (reference transcribe, assemble) leaves
# it at 0, so only TTS generates can get the under-provisioned wording.
_engine_min_vram_gb = getattr(backend_cls, "min_vram_gb", 0.0)
_routing = resolve_routing(
getattr(backend_cls, "gpu_compat", ("cpu",)), detect_host_caps(),
_engine_min_vram_gb,
)
if _routing["routing_status"] == "unavailable":
# The engine needs an accelerator this host lacks and has no CPU path.
raise HTTPException(status_code=400, detail=_routing["routing_reason"])
_routing_notice = routing_notice(_routing) # (status, reason) or None
_routing_notice = None
if not _remote:
# Single-active-engine memory discipline: hand back any OTHER resident
# TTS engine's model before loading this one, so switching engines (or
# a per-request engine= override, which bypasses /engines/select
# entirely) doesn't stack two multi-GB models in memory — the
# accumulation behind the 16 GB-Mac OOM deaths. No-op when nothing else
# is resident, so steady-state single-engine use pays nothing. Opt out:
# OMNIVOICE_SINGLE_ENGINE_RESIDENT=0.
from services.engine_memory import evict_other_tts_engines
await evict_other_tts_engines(engine_id)
# Non-blocking breadcrumb: if free memory is already low before this
# load, log it. A later OOM kill (the 16 GB-Mac class) then has a trail
# pointing at the load that tipped it, instead of a silent process
# death. Never blocks — the OS can reclaim cache, and a hard refuse
# would brick legitimate loads.
try:
from services.memory_budget import log_if_low
log_if_low(f"TTS load ({engine_id})")
except Exception:
pass
# VRAM eviction runs in get_model()'s warm-return path now, so every
# native TTS generate (this route, WS TTS, dub, batch, audiobook) is
# covered.
if backend_cls is OmniVoiceBackend:
# VoiceStudio keeps its native path: it carries the full advanced
# parameter surface (t_shift, layer/position/class controls) that
# the generic adapter protocol doesn't. Byte-identical behavior.
_model = await get_model()
else:
try:
ok, msg = backend_cls.is_available()
except Exception as exc:
ok, msg = False, f"{type(exc).__name__}: {exc}"
if not ok:
raise HTTPException(
status_code=400,
detail=f"TTS engine '{engine_id}' is not available: {_mask_hf_tokens(msg)}",
)
# Reuse the per-process instance cache shared with the engine
# health-check route so weights load once, not per request.
from api.routers.engines import _get_engine_instance
_backend = _get_engine_instance(backend_cls)
# ── Routing gate (#21 — no silent CPU fallback). Computed ONCE per
# request (host caps are constant; the per-request engine= override
# bypasses the /engines/select gate, so this is the only place it's
# enforced for synth). Local only, and deliberately: it asks what THIS
# host can accelerate, and a remote render is precisely the case where
# that answer is none of the question — a CUDA-only engine sent to a
# 4090 from a Mac control plane would be refused by a gate describing
# a machine that is about to do nothing.
from core.device_caps import detect_host_caps
from services.engine_routing import resolve_routing, routing_notice
_routing = resolve_routing(
getattr(backend_cls, "gpu_compat", ("cpu",)), detect_host_caps(),
_engine_min_vram_gb,
)
if _routing["routing_status"] == "unavailable":
# The engine needs an accelerator this host lacks and has no CPU path.
raise HTTPException(status_code=400, detail=_routing["routing_reason"])
_routing_notice = routing_notice(_routing) # (status, reason) or None
# ── #1033/#1037: warm the engine under the LOAD budget, not the generate
# budget. A cold adapter lazily loads (and possibly downloads multi-GB
@@ -1111,18 +1272,17 @@ async def generate_speech(
# measured it: 0% GPU util for the full 300s). Model loading gets its own,
# larger budget (OMNIVOICE_MODEL_LOAD_TIMEOUT, default 1200s) — the same
# split get_model() already has for the native engine. Once warm, this is
# a no-op per request.
# a no-op per request. A remote render gets the same two-phase split from
# the worker, under the assignment's own model-load deadline.
if _backend is not None:
from services.model_manager import _model_load_timeout
from services import gpu_gateway
try:
await run_on_gpu_pool_guarded(
_backend.ensure_ready,
what=f"TTS engine '{engine_id}' model load",
timeout=_model_load_timeout(),
await gpu_gateway.prewarm(
_REMOTE_OP, backend=_backend, engine=engine_id, decision=_decision,
)
# Builtin TimeoutError base, not GpuJobTimeoutError — reload-proof
# class identity (see the twin catch in openai_compat.py).
except TimeoutError as exc:
except (TimeoutError, gpu_gateway.ModelLoadTimeout) as exc:
logger.warning("engine load exceeded the model-load budget: %s", exc)
raise HTTPException(
status_code=503,
@@ -1308,6 +1468,74 @@ async def generate_speech(
start_time = time.time()
# ── The remote assignment ───────────────────────────────────────────────
# Built even for a local render (it costs a dict) so the gateway owns the
# branch rather than this route owning two of them.
#
# The worker runs the ENTIRE render as one op — sentence split, per-chunk
# generate at ``seed + i``, crossfaded concat, effect chain, provenance
# mark — because dispatching a chunk at a time would pay a round trip, a
# progress lease and a slot per sentence against a worker whose
# concurrency defaults to 1. So every knob that shapes the local render has
# to be on the wire: a missing one is not an error, it is remote audio that
# quietly differs from local audio (no sentence splitting, no per-chunk
# seed variation, no crossfade).
from services import gpu_gateway
from services.watermark import is_enabled as _watermark_enabled
_remote_params = {
"text": text,
"language": None if (language and language.lower() == "auto") else language,
"ref_audio": ref_audio_path,
"ref_text": ref_text,
"instruct": instruct,
"duration": duration,
"speed": speed,
"num_step": num_step,
"guidance_scale": guidance_scale,
"denoise": denoise,
"postprocess_output": postprocess_output,
"t_shift": t_shift,
"layer_penalty_factor": layer_penalty_factor,
"position_temperature": position_temperature,
"class_temperature": class_temperature,
"seed": used_seed,
"max_chunk_chars": max_chunk_chars,
"crossfade_ms": crossfade_ms,
"effect_preset": effect_preset,
# The requesting user's provenance preference, not the GPU owner's.
"watermark": bool(_watermark_enabled()),
}
_remote_call = gpu_gateway.RemoteCall(
engine=engine_id, operation=_REMOTE_OP, params=_remote_params,
)
async def _render_on_worker(on_state=None):
"""One whole render on the chosen worker → ``(tensor, sample_rate)``.
The audio comes back already effect-chained and provenance-marked: the
worker mirrors the local order (split → generate → concat → effects →
mark) so a remote take and a local take of the same request differ
only in which GPU produced them.
"""
waveform, sample_rate = await gpu_gateway.run(
_REMOTE_OP,
local=gpu_gateway.LocalCall(
_remote_only_local_call(_target_label),
what="TTS generate",
timeout=_generate_timeout_s(text),
min_vram_gb=_engine_min_vram_gb,
),
remote=_remote_call,
decision=_decision,
on_state=on_state,
)
if getattr(waveform, "ndim", 2) == 1:
# `_safe_torchaudio_save` and the local paths deal in
# (channels, samples); a mono artifact reads back flat.
waveform = waveform.unsqueeze(0)
return waveform, sample_rate
# ── Streaming preview (feat: streaming-tts-preview) ─────────────────────
# Long scripts used to mean staring at a spinner until the ENTIRE render
# finished. With stream=true the existing text chunks (the Wave 1.2
@@ -1322,6 +1550,116 @@ async def generate_speech(
# seed / normalization) already ran, so per-chunk jobs spend the generate
# budget on generation only — and each chunk gets its own budget, so a
# long script can't time out merely for being long.
if stream and _remote:
# ── Remote: the streaming PREVIEW is off, the render still streams ──
# Progressive playback needs per-chunk dispatch, and per-chunk dispatch
# to a worker means a round trip, a progress lease and a slot for every
# sentence, serialised by a default concurrency of 1. So the render
# goes as ONE op and there is no first chunk to play early.
#
# The NDJSON channel stays open anyway, because the desktop UI asks for
# it whenever auto-play is on — which is the default. Answering with
# the classic WAV shape here would make the client fall back to a
# LOCAL re-render, i.e. exactly the bug this phase exists to fix: the
# user picks gpu2, clicks Synthesize, and their laptop does the work.
# What flows down it instead is coarse progress from the worker, then
# the finished take as a single chunk.
_remote_headers = _apply_routing_headers(
{"X-Seed": str(used_seed) if used_seed is not None else "",
"Cache-Control": "no-cache"},
None, _decision,
)
_progress_q: asyncio.Queue = asyncio.Queue()
def _push_progress(event):
# Called from the control plane's own loop; never let a progress
# frame break a render that is otherwise going fine.
try:
_progress_q.put_nowait(dict(event or {}))
except Exception: # noqa: BLE001
logger.debug("dropped a remote progress frame", exc_info=True)
async def _remote_stream_events():
import json
def _line(obj) -> bytes:
return (json.dumps(obj, separators=(",", ":")) + "\n").encode("utf-8")
render = asyncio.ensure_future(_render_on_worker(_push_progress))
try:
# Relay progress until the render settles, then flush whatever
# arrived in the gap so the last "generating (98%)" is not lost.
while not render.done():
getter = asyncio.ensure_future(_progress_q.get())
done, _pending = await asyncio.wait(
{render, getter}, return_when=asyncio.FIRST_COMPLETED
)
if getter in done:
yield _line(_remote_progress_frame(getter.result(), _target_label))
continue
getter.cancel()
while not _progress_q.empty():
yield _line(_remote_progress_frame(_progress_q.get_nowait(),
_target_label))
audio_tensor, sample_rate = await render
yield _line({
"type": "start", "sample_rate": sample_rate, "channels": 1,
"format": "pcm16", "total_chunks": 1, "crossfade_ms": 0,
"seed": used_seed,
})
# No second provenance mark: the worker marked at the tensor
# stage before encoding, with the requesting user's preference
# forced, and stacking a second AudioSeal payload over the
# first degrades detection of both.
yield _line({"type": "chunk", "seq": 0, "pcm": _pcm16_b64(audio_tensor)})
_, meta = await _finalize_generation(
audio_tensor, sample_rate, text=text, history_mode=history_mode,
ref_audio_path=ref_audio_path, language=language,
instruct=instruct, resolved_profile_id=resolved_profile_id,
used_seed=used_seed, start_time=start_time, already_marked=True,
)
# #1330's dropped-chunk warning has no remote carrier yet: the
# gateway hands back audio, not the worker's render metadata.
# Reported as a cross-stream gap rather than faked as zero.
yield _line({
"type": "done", "id": meta["id"], "audio_path": meta["filename"],
"duration": meta["duration"], "gen_time": meta["gen_time"],
"seed": used_seed, "sample_rate": sample_rate,
})
except (asyncio.CancelledError, GeneratorExit):
# The user hit stop, or the request was abandoned. Cancelling
# the render is what tells the worker to release its slot —
# otherwise the 4090 keeps rendering audio nobody will hear,
# holding what is often its only slot until the lease lapses.
render.cancel()
raise
except ValueError:
logger.error("Remote generation request rejected")
from core.public_errors import stream_failure
yield _line({"type": "error", **stream_failure("invalid_request")})
except Exception:
# Mid-job remote failure is NOT quietly redone here: the client
# treats a retryable error as "surface it", so the user decides
# whether to spend the same minutes again on this machine.
logger.error("Remote generation failed", exc_info=True)
from core.public_errors import stream_failure
yield _line({"type": "error", **stream_failure("generation_failed")})
finally:
if not render.done():
render.cancel()
if cleanup_ref and ref_audio_path:
with contextlib.suppress(OSError):
os.remove(ref_audio_path)
return StreamingResponse(
_remote_stream_events(),
media_type="application/x-ndjson",
headers=_remote_headers,
)
if stream:
from omnivoice.utils.text import parse_pause_markers
from services.chunked_tts import split_text_into_chunks
@@ -1425,7 +1763,7 @@ async def generate_speech(
_backend, text, language, ref_audio_path, ref_text,
instruct, duration, num_step, guidance_scale, speed,
denoise, postprocess_output, used_seed, effect_preset,
max_chunk_chars, crossfade_ms,
max_chunk_chars, crossfade_ms, dropped_sink=_dropped_sink,
),
what="TTS generate",
min_vram_gb=_engine_min_vram_gb,
@@ -1441,7 +1779,7 @@ async def generate_speech(
t_shift, denoise, postprocess_output,
layer_penalty_factor, position_temperature,
class_temperature, used_seed, effect_preset,
max_chunk_chars, crossfade_ms,
max_chunk_chars, crossfade_ms, dropped_sink=_dropped_sink,
),
what="TTS generate",
min_vram_gb=_engine_min_vram_gb,
@@ -1562,18 +1900,13 @@ async def generate_speech(
with contextlib.suppress(OSError):
os.remove(ref_audio_path)
_stream_headers = {
# Routing notice (#21): known before the stream starts, so it rides the
# same headers the classic path uses — and now also carries "your
# chosen worker was unavailable, this ran here".
_stream_headers = _apply_routing_headers({
"X-Seed": str(used_seed) if used_seed is not None else "",
"Cache-Control": "no-cache",
}
# Routing notice (#21): known before the stream starts, so it rides the
# same headers the classic path uses.
if _routing_notice:
from services.engine_routing import header_safe_reason
_stream_headers["X-OmniVoice-Routing"] = _routing_notice[0]
_hr = header_safe_reason(_routing_notice[1])
if _hr:
_stream_headers["X-OmniVoice-Routing-Reason"] = _hr
}, _routing_notice, _decision)
return StreamingResponse(
_stream_events(),
media_type="application/x-ndjson",
@@ -1586,47 +1919,54 @@ async def generate_speech(
# so. A warning in a log the user never opens is a record of the bug, not a
# fix for it.
_dropped_text: list = []
_already_marked = False
try:
if _backend is not None:
# Bounded + pool-reset on hang so a wedged generate can't starve the
# GPU pool and brick the backend ("can't reach backend", #730 class).
audio_tensor = await run_on_gpu_pool_guarded(
functools.partial(
if _remote:
# One op, one worker, the whole render — including the chunk loop.
audio_tensor, sample_rate = await _render_on_worker()
_already_marked = True
else:
# The gateway owns the dispatch on both branches. Locally it still
# lands in run_on_gpu_pool_guarded, so the #730 bound + pool reset
# that keeps a wedged generate from bricking the backend is
# unchanged.
if _backend is not None:
_local_render = functools.partial(
_run_backend_inference,
_backend, text, language, ref_audio_path, ref_text, instruct,
duration, num_step, guidance_scale, speed, denoise,
postprocess_output, used_seed, effect_preset,
max_chunk_chars, crossfade_ms, dropped_sink=_dropped_text,
),
what="TTS generate",
min_vram_gb=_engine_min_vram_gb,
timeout=_generate_timeout_s(text),
)
# Read after generation: engines with lazy model loading report
# their real rate only once weights are up.
sample_rate = _backend.sample_rate
else:
audio_tensor = await run_on_gpu_pool_guarded(
functools.partial(
)
else:
_local_render = functools.partial(
_run_inference,
_model, text, language, ref_audio_path, ref_text, instruct, duration,
num_step, guidance_scale, speed, t_shift, denoise,
postprocess_output, layer_penalty_factor, position_temperature,
class_temperature, used_seed, effect_preset,
max_chunk_chars, crossfade_ms, dropped_sink=_dropped_text,
)
audio_tensor = await gpu_gateway.run(
_REMOTE_OP,
local=gpu_gateway.LocalCall(
_local_render, what="TTS generate",
timeout=_generate_timeout_s(text),
min_vram_gb=_engine_min_vram_gb,
),
what="TTS generate",
min_vram_gb=_engine_min_vram_gb,
timeout=_generate_timeout_s(text),
decision=_decision,
)
sample_rate = _model.sampling_rate
# Read after generation: engines with lazy model loading report
# their real rate only once weights are up.
sample_rate = (_backend.sample_rate if _backend is not None
else _model.sampling_rate)
# Watermark → save → history → prune → emit, shared with the streaming
# path (see _finalize_generation) so both flows produce identical takes.
audio_tensor, _meta = await _finalize_generation(
audio_tensor, sample_rate, text=text, history_mode=history_mode,
ref_audio_path=ref_audio_path, language=language, instruct=instruct,
resolved_profile_id=resolved_profile_id, used_seed=used_seed,
start_time=start_time,
start_time=start_time, already_marked=_already_marked,
)
audio_id = _meta["id"]
audio_filename = _meta["filename"]
@@ -1660,14 +2000,10 @@ async def generate_speech(
_lost = header_safe_reason(" | ".join(t for t in _dropped_text if t))
if _lost:
_resp_headers["X-OmniVoice-Dropped-Text"] = _lost
# Routing notice (#21): cpu_fallback or accelerated-with-caveat only;
# the WAV body is binary so the header channel is the carrier.
if _routing_notice:
from services.engine_routing import header_safe_reason
_resp_headers["X-OmniVoice-Routing"] = _routing_notice[0]
_hr = header_safe_reason(_routing_notice[1])
if _hr:
_resp_headers["X-OmniVoice-Routing-Reason"] = _hr
# Routing notice (#21): cpu_fallback, accelerated-with-caveat, or the
# machine this render ran on. The WAV body is binary so the header
# channel is the carrier.
_apply_routing_headers(_resp_headers, _routing_notice, _decision)
return StreamingResponse(
_stream_wav(),
media_type="audio/wav",
@@ -1675,6 +2011,22 @@ async def generate_speech(
)
except HTTPException:
raise
except gpu_gateway.RemoteJobFailed as e:
# Rule 2 of the fallback policy: a single-shot interactive render that
# failed ON the worker is reported, not silently redone here. Minutes
# already went somewhere else, the user is watching, and quietly
# re-rendering on the slower machine turns a 20-second wait into a
# four-minute one with no explanation. The header names the target so
# the client can offer "run it on this machine instead" — a resubmit
# the user chose, with a wait they were told about.
logger.error("Remote generate failed on %s: %s", _target_label, e)
raise HTTPException(
status_code=503,
detail=f"{e} {e.hint or 'Run it on this machine instead, or pick another GPU.'}",
headers={"X-OmniVoice-Retryable": "true",
"X-OmniVoice-Routing": "remote_failed",
"Retry-After": "10"},
) from e
except GpuPoolBusyError as e:
# Saturation, not failure (#1190): the job never started, so the caller
# can retry the identical request. Retry-After + the retryable marker
backend/api/routers/tts_stream.py
+42
View File
@@ -62,6 +62,11 @@ async def ws_tts(websocket: WebSocket):
await websocket.accept()
logger.info("TTS streaming WebSocket connected")
# Said once per socket, not once per utterance: a conversational client
# sends many requests down one connection and a repeated notice would be
# noise. See `_announce_local_only`.
announced_local_only = False
try:
while True:
# Wait for a text request from the client
@@ -83,6 +88,43 @@ async def ws_tts(websocket: WebSocket):
t0 = time.perf_counter()
text = data["text"]
# Remote GPU: this socket stays on this machine, and says so.
#
# /generate's port trades progressive playback for the remote
# render — the classic path was always a single wait, so spending
# it on a faster GPU is a straight win. This route is the opposite
# shape: it exists to put audio in the user's ear before the
# sentence has finished synthesizing, and sending each utterance to
# a worker would pay queue admission, a round trip and cold-load
# risk per utterance, for the one surface where latency IS the
# feature.
#
# Silence would be worse than the limitation: the header badge
# would read "gpu2" while this machine does 100% of the work, the
# same class of lie the op-aware picker exists to stop. Said once
# per socket — a conversational client sends many requests down one
# connection — and BEFORE engine resolution, so an engine that
# cannot load still tells the user where it would have run.
if not announced_local_only:
announced_local_only = True
try:
from worker import routing as worker_routing
target = worker_routing.decide(op="tts")
except Exception: # noqa: BLE001 — advisory; never break audio
target = None
if target is not None and target.remote:
from core.scrub import scrub_text as _scrub
await websocket.send_json({
"type": "routing",
"status": "local_stream",
"reason": _scrub(
f"{target.label} is your GPU target, but live "
f"streaming runs on this machine"
),
})
try:
# Resolve engine
from services.tts_backend import (
backend/services/gpu_gateway.py
+850
View File
@@ -0,0 +1,850 @@
"""One door to a GPU — this machine's, or the one the user picked.
Every GPU call in this app used to be the same three lines inlined at ~30 call
sites: resolve a backend, ``run_on_gpu_pool_guarded`` it, translate the pool's
exceptions into an HTTP answer. That shape has exactly one destination baked
into it, so "run this on my 4090" could never be more than a badge. This module
is the seam that makes the destination a *parameter*:
decision = gpu_gateway.decide("tts") # local, or the chosen worker
await gpu_gateway.prewarm("tts", backend=b, decision=decision)
audio = await gpu_gateway.run("tts", local=..., remote=..., decision=decision)
Four calls, and each of them answers for both targets:
* :func:`prewarm` — the model-load budget
* :func:`run` — the generate budget
* :func:`status` — supported / installed / downloaded / resident
* :func:`download` — fetching weights
Design decisions that are load-bearing, and why they are not obvious:
**``prewarm`` and ``run`` are separate calls.** Collapsing them loses the
two-phase split documented at ``tts_backend.ensure_ready`` (#1033/#1037): a cold
adapter that loads lazily inside ``generate()`` spends the *generate* budget on
a multi-GB download and dies with "too heavy for the available compute". The
protocol mirrors the same split (``TaskModelLoading`` and
``Deadlines.model_load_seconds``), so keeping the two calls apart is what lets
one policy serve both targets.
**The local branch calls ``run_on_gpu_pool_guarded``; the remote branch does
not.** That function means "submit a zero-arg blocking callable to the local
thread pool". Its error taxonomy (``GpuPoolBusyError`` / ``GpuJobTimeoutError``,
plus a pool ``reset()``) describes local saturation, and its ``started.set()``
handshake is a second timeout regime that would race the attempt ``Deadlines``.
A remote job is bounded by the lease and the phase budgets instead.
**This is not a ``RemoteBackend(TTSBackend)``.** ``generate()`` is synchronous
and returns a tensor, so a remote implementation would block a pool thread on
an async round-trip *while holding a GPU-pool slot* — a hard deadlock at
``OMNIVOICE_GPU_WORKERS=1``, which is the default on the machines that most
want to offload. It would also cover none of the non-TTS GPU work.
**Admission control lives here.** ``check_gpu_admission`` reads *local* pool
stats; called unconditionally it would answer 429 "the local GPU pool is
saturated" while the remote 4090 sat idle. It runs on the local branch only.
**Fallback is three rules, not one** (see ``worker/routing.py``'s header):
1. *Pre-dispatch* unavailability — the worker is offline, disabled, paused,
the queue is full, or nothing ever accepted the task — runs locally,
quietly, with the named reason. Nothing ran remotely, so nothing is lost.
2. *Mid-job* failure on a single-shot interactive op raises
:class:`RemoteJobFailed`. Silently redoing minutes of work on the slower
machine, with no explanation, is not a kindness.
3. *Multi-unit* jobs (audiobook chapters, batches) pass a :class:`JobRun`;
after N consecutive remote failures the job latches local for the rest of
its units and reports **one** aggregated notice, instead of 160 identical
error rows because a 4090 went to sleep at chapter 40.
"""
from __future__ import annotations
import asyncio
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Callable, Optional
from worker import routing
from worker.routing import LOCAL, Decision
logger = logging.getLogger("omnivoice.gateway")
# How often the awaiting coroutine samples a remote task to report coarse
# progress. Polling rather than `scheduler.on_change`: that listener list has
# no unregister (scheduler.py), so one subscription per job would leak for the
# life of the process.
_POLL_SECONDS = 0.5
# Consecutive remote failures a multi-unit job tolerates before it stops trying
# the remote worker. One is a blip (a dropped stream, a worker restart); two in
# a row is a machine that has gone away, and the remaining 160 chapters should
# not each pay a full deadline to discover that.
_MULTI_UNIT_FAILURE_LIMIT = 2
# Coarse phases the UI can render for a remote job. `workers.py`'s task view is
# poll-only, so without these a five-minute remote render shows the same bare
# spinner as a local one and looks wedged.
PHASE_QUEUED = "queued"
PHASE_LOADING = "loading"
PHASE_RUNNING = "running"
PHASE_UPLOADING = "uploading"
# ── Errors ─────────────────────────────────────────────────────────────────
class GatewayError(RuntimeError):
"""Base for every error this module raises on its own behalf."""
class ModelLoadTimeout(GatewayError):
"""A local engine did not finish loading inside the model-load budget."""
class RemoteJobFailed(GatewayError):
"""Remote work started and then failed. Rule 2: this is not a fallback.
Carries what a caller needs to offer "Run locally instead" — the same
request with ``target=local`` — rather than a bare 500.
"""
def __init__(
self,
message: str,
*,
worker_label: str = "",
task_id: str = "",
code: str = "",
hint: str = "",
) -> None:
super().__init__(message)
self.worker_label = worker_label
self.task_id = task_id
self.code = code
self.hint = hint
# Nothing about this failure implicates the local machine, so a retry
# here is genuinely likely to work. Callers surface it as one click.
self.retry_local = True
class RemoteUnsupported(GatewayError):
"""Asked for something the remote path cannot do yet.
Deliberately not a quiet local fallback: downloading weights onto *this*
machine when the user asked for them on the 4090 is not the same operation,
and pretending it is leaves the remote box exactly as unprepared as before.
"""
class _NotDispatched(Exception):
"""Internal: the remote target never started the work. Rule 1 applies."""
def __init__(self, reason: str) -> None:
super().__init__(reason)
self.reason = reason
# ── Call descriptions ──────────────────────────────────────────────────────
@dataclass(frozen=True)
class LocalCall:
"""The local branch: a zero-arg blocking callable for the GPU pool.
``fn`` must take no arguments — wrap with ``functools.partial``, exactly as
``run_on_gpu_pool_guarded`` already requires.
"""
fn: Callable[[], Any]
what: str = "GPU job"
timeout: Optional[float] = None
queue_timeout: Optional[float] = None
# The engine's declared VRAM floor; only shapes the timeout message.
min_vram_gb: float = 0.0
@dataclass(frozen=True)
class RemoteResult:
"""A committed remote result, as this side holds it.
The artifact is a path inside the control plane's own artifact directory —
minted from the attempt record, never from anything on the wire.
"""
task_id: str
worker_id: str
worker_label: str
path: Optional[str] = None
meta: dict = field(default_factory=dict)
def read(self) -> bytes:
if not self.path:
raise RemoteJobFailed(
f"{self.worker_label or 'The worker'} reported success but sent no audio.",
worker_label=self.worker_label,
task_id=self.task_id,
code="RESULT_MISSING",
)
with open(self.path, "rb") as handle:
return handle.read()
@dataclass(frozen=True)
class RemoteCall:
"""The remote branch: one task for the scheduler, and how to read it back.
``decode`` converts the committed artifact into whatever the local branch
returns, so ``run`` has one return type regardless of where the work ran.
Left unset for ``tts``/``clone`` it defaults to :func:`decode_audio_artifact`
— the audio ops are the only ones with a remote producer today, and a
caller that has to branch on the target has gained nothing from this module.
"""
engine: str
params: dict = field(default_factory=dict)
operation: str = "tts"
# Stable, opaque, engine-scoped ("indextts:default") — never a repo id or a
# path. Empty means "any model this engine advertises".
model_id: str = ""
deadline_seconds: Optional[float] = None
idempotency_key: Optional[str] = None
decode: Optional[Callable[[RemoteResult], Any]] = None
# ── Multi-unit jobs (rule 3) ───────────────────────────────────────────────
class JobRun:
"""State for a job made of many units, so rule 3 can be applied once.
Created per audiobook / batch / dub run and passed to every ``run`` call it
makes. It counts *consecutive* remote failures — an intermittent blip
should not permanently demote a working worker — and latches local once the
limit is reached, because the alternative is paying a full remote deadline
per remaining unit to rediscover the same dead machine.
"""
def __init__(self, op: str, *, limit: int = _MULTI_UNIT_FAILURE_LIMIT) -> None:
self.op = op
self.limit = max(1, int(limit))
self.consecutive_failures = 0
self.remote_failures = 0
self.local_units = 0
self.remote_units = 0
self.latched_local = False
self.worker_label = ""
self.last_reason = ""
def record_success(self) -> None:
self.consecutive_failures = 0
self.remote_units += 1
def record_failure(self, reason: str, *, worker_label: str = "") -> bool:
"""Charge one failed unit. ``True`` if this unit may fall back locally.
Always true today: the first failed unit already falls back rather than
failing the row, and the counter decides whether *later* units still try
the worker at all.
"""
self.consecutive_failures += 1
self.remote_failures += 1
self.last_reason = reason
self.worker_label = worker_label or self.worker_label
if self.consecutive_failures >= self.limit:
self.latched_local = True
return True
def record_local(self) -> None:
self.local_units += 1
def notice(self) -> Optional[tuple[str, str]]:
"""The single aggregated notice for the whole job, or ``None``."""
if not self.remote_failures:
return None
who = self.worker_label or "the remote worker"
if self.latched_local:
return (
"local_fallback",
f"{who} failed {self.remote_failures} time(s) "
f"({self.last_reason}) — the rest of this job ran locally.",
)
return (
"local_fallback",
f"{self.remote_failures} item(s) ran locally after {who} failed "
f"({self.last_reason}).",
)
# ── Routing ────────────────────────────────────────────────────────────────
def decide(op: str, *, control_plane=None) -> Decision:
"""Where this job runs, decided **once**.
Callers pass the result to ``prewarm`` and ``run`` rather than calling this
per step: the target is user-settable at any moment, and a decision that
flipped between the two would either warm an engine nothing will use or
dispatch remotely after paying a local cold load.
"""
return routing.decide(control_plane, op=op)
def notice_for(decision: Decision) -> Optional[tuple[str, str]]:
"""``(status, reason)`` for the ``X-OmniVoice-Routing`` header channel.
Same two-tuple shape ``engine_routing.routing_notice`` produces, so the
existing header plumbing and the de-duped toast at ``routingNotice.js``
carry it without a second channel. ``None`` when there is nothing to say —
a user who chose Local does not need to be told their audio ran locally.
"""
if decision.remote:
return ("remote", f"running on {decision.label}")
if decision.reason and decision.reason != "chosen":
return ("local_fallback", decision.reason)
return None
# ── Prewarm: the model-load budget ─────────────────────────────────────────
async def prewarm(
op: str,
*,
backend=None,
engine: str = "",
decision: Optional[Decision] = None,
timeout: Optional[float] = None,
executor=None,
control_plane=None,
) -> Decision:
"""Make the model ready under the LOAD budget, on whichever GPU will run it.
Returns the decision so the caller can hand the *same* one to :func:`run`.
On the remote branch this deliberately does nothing local. Warming an
engine here before dispatching to another machine costs minutes and VRAM on
the box that is not doing the work; the worker loads under its own
``model_load_seconds`` and reports ``TaskModelLoading`` while it does, which
is the whole point of the phase existing on the wire.
"""
decision = decision or decide(op, control_plane=control_plane)
if decision.remote:
return decision
if backend is None:
# The native model path warms itself through `model_manager.get_model`;
# there is nothing engine-shaped to call.
return decision
call = LocalCall(
backend.ensure_ready,
what=f"TTS engine '{engine or getattr(backend, 'engine_id', '') or 'model'}' model load",
timeout=timeout if timeout is not None else _model_load_timeout(),
)
try:
await _run_local(call, executor=executor)
except TimeoutError as exc:
# Builtin TimeoutError, not GpuJobTimeoutError — reload-proof class
# identity, the same catch generation.py and openai_compat.py use.
# GpuPoolBusyError is a TimeoutError too and means the opposite thing
# (the load never started, nothing was spent, retry as-is); it is
# identified by its `retry_after` rather than by class, so a module
# reload cannot turn saturation into a bogus load timeout.
if hasattr(exc, "retry_after"):
raise
raise ModelLoadTimeout(
f"{call.what} did not finish within its budget. The first load of an "
f"engine can include a multi-GB download; check the connection or set "
f"a Hugging Face mirror in Settings, then retry."
) from exc
return decision
# ── Run: the generate budget ───────────────────────────────────────────────
async def run(
op: str,
*,
local: LocalCall,
remote: Optional[RemoteCall] = None,
decision: Optional[Decision] = None,
job: Optional[JobRun] = None,
admit: bool = False,
on_state: Optional[Callable[[dict], None]] = None,
executor=None,
control_plane=None,
) -> Any:
"""Run one unit of GPU work, here or on the chosen worker.
``local`` is always required — it is both the local branch and the landing
ground for rules 1 and 3, and a gateway that could not run anything locally
would turn every offline worker into a failed request.
"""
if job is not None and job.latched_local:
decision = Decision(remote=False, reason=job.last_reason or "the remote worker failed")
decision = decision or decide(op, control_plane=control_plane)
if decision.remote and remote is not None:
try:
value = await _run_remote(
remote, decision, on_state=on_state, control_plane=control_plane
)
except _NotDispatched as exc:
# Rule 1. Nothing ran remotely, so this is the quiet fallback the
# picker already promises — no compute was spent anywhere.
logger.info("Remote dispatch declined (%s); running locally", exc.reason)
if job is not None:
job.record_local()
return await _run_local(local, admit=admit, executor=executor)
except RemoteJobFailed as exc:
# Rules 2 and 3. Work started on the worker and did not finish.
if job is None:
raise
job.record_failure(str(exc), worker_label=exc.worker_label)
logger.warning(
"Remote unit failed on %s (%s); running this unit locally",
exc.worker_label or "the worker", exc,
)
job.record_local()
return await _run_local(local, admit=admit, executor=executor)
if job is not None:
job.record_success()
return value
if job is not None:
job.record_local()
return await _run_local(local, admit=admit, executor=executor)
async def _run_local(call: LocalCall, *, admit: bool = False, executor=None) -> Any:
"""The local branch: admission, then the guarded pool."""
from services.model_manager import ( # noqa: PLC0415 — torch lives down here
check_gpu_admission,
run_on_gpu_pool_guarded,
)
if admit:
# Only ever on this branch: these are local pool statistics, and a 429
# about local saturation while the remote GPU idles is a lie.
check_gpu_admission(what=call.what, executor=executor)
return await run_on_gpu_pool_guarded(
call.fn,
what=call.what,
timeout=call.timeout,
queue_timeout=call.queue_timeout,
min_vram_gb=call.min_vram_gb,
executor=executor,
)
async def _run_remote(
call: RemoteCall,
decision: Decision,
*,
on_state: Optional[Callable[[dict], None]] = None,
control_plane=None,
) -> Any:
"""The remote branch: submit, await, decode.
Raises ``_NotDispatched`` while nothing has run yet (rule 1) and
:class:`RemoteJobFailed` once a worker accepted the work (rules 2/3). That
boundary is the whole fallback policy: "no compute was spent" is the only
honest licence to silently redo the job somewhere else.
"""
from worker.scheduler import QueueFull, SchedulerStopped # noqa: PLC0415
plane = _plane(control_plane)
if plane is None or not getattr(plane, "running", False):
raise _NotDispatched("remote workers are turned off")
scheduler = plane.scheduler
if scheduler is None:
raise _NotDispatched("the control plane has no scheduler")
params = dict(call.params or {})
deadline = call.deadline_seconds
if deadline is None:
deadline = _default_deadline(call.operation, params.get("text"))
try:
task = scheduler.submit(
operation=call.operation,
engine=call.engine,
model_id=call.model_id,
params=params,
idempotency_key=call.idempotency_key,
deadline_seconds=deadline,
)
except QueueFull as exc:
raise _NotDispatched(str(exc)) from exc
_emit(on_state, {"phase": PHASE_QUEUED, "progress": 0.0, "stage": "",
"worker": decision.label, "task_id": task.task_id})
try:
settled = await _await_task(
scheduler, task.task_id, timeout=deadline,
on_state=on_state, label=decision.label,
)
except KeyError as exc:
# The scheduler no longer holds the task (purged, or restored into a
# different instance). Nothing ran here, so rule 1 applies.
raise _NotDispatched("the remote task was dropped before it ran") from exc
except SchedulerStopped as exc:
raise _classify(scheduler.get(task.task_id) or task, str(exc), decision) from exc
except TimeoutError as exc:
raise _classify(
scheduler.get(task.task_id) or task,
f"it did not finish within {float(deadline):g}s",
decision,
) from exc
return _decode(call, settled, decision)
async def _await_task(scheduler, task_id: str, *, timeout: float, on_state, label: str):
"""Await a terminal task, reporting coarse progress, cancelling if we leave.
Every exit that is not a terminal task cancels the remote task, because a
worker holds its slot — often its only one — until this side says
otherwise, and the sweeper only enforces deadlines on tasks that are still
queued. Without this, ``useTTS.js``'s AbortController abandons the request
while the 4090 keeps rendering audio nobody will ever read.
The one exception is shutdown: ``SchedulerStopped`` means this side is
quitting, the worker was never told to stop and may still be rendering, so
recording a cancellation would be a claim about someone else's GPU that we
are in no position to make.
"""
from worker.scheduler import SchedulerStopped # noqa: PLC0415
waiter = asyncio.ensure_future(scheduler.wait(task_id, timeout=timeout))
last: Optional[tuple] = None
try:
while True:
done, _pending = await asyncio.wait({waiter}, timeout=_POLL_SECONDS)
if not done:
last = _report(scheduler, task_id, on_state, label, last)
continue
# Raises here for a deadline or a shutdown; both are handled below.
return waiter.result()
except SchedulerStopped:
waiter.cancel()
raise
except asyncio.CancelledError:
waiter.cancel()
_cancel(scheduler, task_id, "the client stopped waiting")
raise
except BaseException:
waiter.cancel()
_cancel(scheduler, task_id, "the task passed its deadline")
raise
def _report(scheduler, task_id: str, on_state, label: str, last: Optional[tuple]):
"""Emit a coarse phase when it changes. Never raises."""
if on_state is None:
return last
try:
task = scheduler.get(task_id)
except Exception:
return last
if task is None:
return last
attempt = task.active_attempt
phase = _PHASES.get(getattr(task.state, "value", ""), PHASE_QUEUED)
progress = round(float(getattr(attempt, "progress", 0.0) or 0.0), 2)
stage = getattr(attempt, "stage", "") or ""
current = (phase, progress, stage)
if current == last:
return last
_emit(on_state, {"phase": phase, "progress": progress, "stage": stage,
"worker": label, "task_id": task_id})
return current
_PHASES = {
"queued": PHASE_QUEUED,
"assigned": PHASE_QUEUED,
"accepted": PHASE_QUEUED,
"model_loading": PHASE_LOADING,
"running": PHASE_RUNNING,
"result_uploading": PHASE_UPLOADING,
}
def _emit(on_state, payload: dict) -> None:
if on_state is None:
return
try:
on_state(payload)
except Exception:
logger.debug("Remote progress listener failed", exc_info=True)
def _cancel(scheduler, task_id: str, reason: str) -> None:
try:
scheduler.cancel(task_id, reason=reason)
except Exception:
logger.exception("Could not cancel abandoned remote task %s", task_id)
def _decode(call: RemoteCall, task, decision: Decision) -> Any:
"""Turn a settled task into the local branch's return value, or fail."""
state = getattr(task.state, "value", str(task.state))
if state != "completed":
raise _classify(task, _reason(task), decision)
result = RemoteResult(
task_id=task.task_id,
worker_id=decision.worker_id or "",
worker_label=decision.label,
path=task.result_ref,
meta={"engine": task.engine, "model_id": task.model_id},
)
if result.path is None or not os.path.exists(result.path):
# Completed with nothing to read. Treated as a mid-job failure, not a
# quiet fallback: the worker spent the compute, and a caller told
# "nothing ran" would be misled about where its minutes went.
raise RemoteJobFailed(
f"{decision.label} finished the job but its audio did not arrive.",
worker_label=decision.label,
task_id=task.task_id,
code="RESULT_MISSING",
hint="Run it locally instead, or check the worker's connection.",
)
decoder = call.decode or (
decode_audio_artifact if call.operation in ("tts", "clone") else None
)
if decoder is None:
return result
try:
return decoder(result)
except RemoteJobFailed:
raise
except Exception as exc: # noqa: BLE001 — any decode failure is one class
# A truncated or unreadable artifact is a mid-job failure, not a quiet
# fallback: the compute happened, and a multi-unit job must be able to
# count it against the worker like any other.
raise RemoteJobFailed(
f"{decision.label} returned audio this app could not read: {exc}",
worker_label=decision.label,
task_id=task.task_id,
code="RESULT_UNREADABLE",
) from exc
def _classify(task, reason: str, decision: Decision):
"""``_NotDispatched`` while nothing ran; ``RemoteJobFailed`` once it did."""
if not _work_started(task):
return _NotDispatched(reason)
error = getattr(task, "error", None)
return RemoteJobFailed(
f"{decision.label} did not finish this job: {reason}",
worker_label=decision.label,
task_id=task.task_id,
code=getattr(error, "code", "") or "",
hint=getattr(error, "hint", "") or "",
)
def _work_started(task) -> bool:
"""Did any worker actually accept this task?
``accepted_at`` rather than "an attempt exists": an assignment that was
rejected for capacity, or that died in a dispatch race before the worker
answered, cost nothing anywhere and is exactly the case rule 1 exists for.
"""
for attempt in getattr(task, "attempts", []) or []:
if getattr(attempt, "accepted_at", None) is not None:
return True
if getattr(attempt, "started_at", None) is not None:
return True
return False
def _reason(task) -> str:
error = getattr(task, "error", None)
message = getattr(error, "message", "") if error is not None else ""
if message:
return message
state = getattr(task.state, "value", str(task.state))
return {
"cancelled": "the task was cancelled",
"timeout": "the task passed its deadline",
}.get(state, "the task failed")
def decode_audio_artifact(result: RemoteResult):
"""``(waveform, sample_rate)`` from a remote WAV artifact.
The local branch returns a tensor plus the engine's ``sample_rate``; this
returns the same pair read from the artifact's own header rather than from
an assumed 24 kHz — VoxCPM2 renders at 48 kHz, and guessing plays it back
at half speed.
"""
import io # noqa: PLC0415
import soundfile as sf # noqa: PLC0415
data, sample_rate = sf.read(io.BytesIO(result.read()), dtype="float32", always_2d=False)
try:
import torch # noqa: PLC0415
waveform = torch.from_numpy(data)
except Exception: # noqa: BLE001 — a torch-less host still gets its audio
waveform = data
return waveform, int(sample_rate)
# ── Status: what each target can actually run ──────────────────────────────
async def status(
engine: Optional[str] = None,
*,
decision: Optional[Decision] = None,
op: str = "tts",
control_plane=None,
) -> dict:
"""The four facts per model — supported / installed / downloaded / resident
— for whichever machine would run the work.
One shape for both targets, from one producer: ``worker.capabilities``
already derives them from ``tts_backend`` for the local host, and a remote
worker reports the same records through ``Register``. Asking the local
engine layer about a remote machine is how a UI ends up offering an engine
that only exists here.
"""
decision = decision or decide(op, control_plane=control_plane)
if not decision.remote:
return {
"target": LOCAL,
"remote": False,
"label": decision.label,
"reason": decision.reason,
"models": _filtered(_local_capabilities(), engine),
}
plane = _plane(control_plane)
worker = None
pool = getattr(plane, "pool", None) if plane is not None else None
if pool is not None:
worker = pool.get(decision.worker_id)
if worker is None:
# Reachability changed between decide() and here.
return {
"target": decision.worker_id or LOCAL,
"remote": False,
"label": decision.label,
"reason": "the chosen worker is not connected",
"models": _filtered(_local_capabilities(), engine),
}
return {
"target": decision.worker_id,
"remote": True,
"label": decision.label,
"reason": decision.reason,
"models": _filtered(list(worker.record.capabilities or []), engine),
}
def _local_capabilities() -> list[dict]:
from worker import capabilities # noqa: PLC0415
return capabilities.discover(include_unavailable=True)
def _filtered(models: list[dict], engine: Optional[str]) -> list[dict]:
if not engine:
return models
return [m for m in models if m.get("engine") == engine]
# ── Download: weights, onto the machine that needs them ────────────────────
async def download(
repo_id: str,
*,
decision: Optional[Decision] = None,
op: str = "tts",
control_plane=None,
) -> dict:
"""Fetch a catalog model onto the target machine.
Remote downloads are not implemented yet, and this refuses rather than
falling back: downloading onto *this* machine when the user asked for the
weights on the 4090 leaves the remote box exactly as unprepared, having
reported success.
"""
decision = decision or decide(op, control_plane=control_plane)
if decision.remote:
raise RemoteUnsupported(
f"Model downloads do not run on {decision.label} yet — the weights "
f"would land on this machine instead. Install the model on "
f"{decision.label} directly, or switch the GPU target to Local."
)
from api.routers.setup.download import ( # noqa: PLC0415
InstallModelRequest,
install_model,
)
from api.routers.setup.models import KNOWN_MODELS # noqa: PLC0415
if repo_id not in {m.get("repo_id") for m in KNOWN_MODELS}:
# The wire and the UI both carry catalog ids only; anything else is a
# path by another name, and paths are what the protocol forbids.
raise GatewayError(f"Unknown model: {repo_id!r}.")
return await install_model(InstallModelRequest(repo_id=repo_id))
# ── Plumbing ───────────────────────────────────────────────────────────────
def _plane(control_plane=None):
if control_plane is not None:
return control_plane
try:
from worker.service import control_plane as default_plane # noqa: PLC0415
return default_plane
except Exception:
logger.debug("No control plane available", exc_info=True)
return None
def _default_deadline(operation: str, text: Optional[str]) -> float:
"""Worst-case wall time for one attempt, from the shared deadline policy.
Same budget the assignment itself carries, so the awaiting side cannot give
up on a worker that is still inside the time this side granted it.
"""
from worker import deadlines # noqa: PLC0415
return float(deadlines.for_task(operation, text=text).total_seconds)
def _model_load_timeout() -> float:
from services.model_manager import _model_load_timeout as resolve # noqa: PLC0415
return float(resolve())
__all__ = [
"GatewayError",
"JobRun",
"LOCAL",
"LocalCall",
"ModelLoadTimeout",
"RemoteCall",
"RemoteJobFailed",
"RemoteResult",
"RemoteUnsupported",
"decide",
"decode_audio_artifact",
"download",
"notice_for",
"prewarm",
"run",
"status",
]
backend/worker/executor.py
+261 -1
View File
@@ -17,8 +17,10 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import logging
import os
from typing import Any, Awaitable, Callable, Optional
from worker.errors import ErrorClass, WorkerError
@@ -29,6 +31,17 @@ logger = logging.getLogger("omnivoice.worker")
# uploaded separately so it cannot head-of-line block heartbeats.
INLINE_LIMIT_BYTES = 256 * 1024
# Where the control plane reports that it could not stage an input. Mirrors
# ``codec._INPUT_ERRORS_KEY``; the two are pinned together by a test rather
# than by an import, because this module must not depend on the transport.
INPUT_ERRORS_PARAM = "input_errors"
# Fetched inputs are cached by content hash, so the second clone of a voice
# transfers nothing. Bounded, because a cache with no ceiling is the same disk
# leak on the worker that unpurged artifacts were on the control plane.
INPUT_CACHE_LIMIT_BYTES = 2 * 1024 * 1024 * 1024
_FALLBACK_INPUT_FETCH_SECONDS = 600.0
# on_progress(fraction: float, stage: str)
# on_model_loading(fraction: float, detail: str)
#
@@ -39,6 +52,13 @@ INLINE_LIMIT_BYTES = 256 * 1024
ProgressReporter = Callable[[float, str], Awaitable[None]]
LoadReporter = Callable[[float, str], Awaitable[None]]
# fetch_input(ref: ArtifactRef, destination: str) -> Awaitable[Any]
#
# Supplied by the transport, which owns the ``DownloadArtifact`` stream and the
# session credentials it needs. The executor decides *what* to fetch and where
# it lands; it does not know there is a network.
InputFetcher = Callable[[Any, str], Awaitable[Any]]
# Used when an assignment carries no deadlines (the HTTP mirror, and tests).
# Generous on purpose: the server lease is the real bound, and a worker-side
# timeout that fires first turns a slow-but-healthy job into a hard failure.
@@ -58,9 +78,13 @@ class TaskExecutor:
*,
on_progress: Optional[ProgressReporter] = None,
on_model_loading: Optional[LoadReporter] = None,
fetch_input: Optional[InputFetcher] = None,
input_dir: Optional[str] = None,
) -> None:
self._on_progress = on_progress
self._on_model_loading = on_model_loading
self._fetch_input = fetch_input
self._input_dir = input_dir
async def execute(
self,
@@ -68,6 +92,7 @@ class TaskExecutor:
*,
on_progress: Optional[ProgressReporter] = None,
on_model_loading: Optional[LoadReporter] = None,
fetch_input: Optional[InputFetcher] = None,
) -> dict:
"""Run one assignment and return ``{"meta": {...}, "payload": bytes}``.
@@ -82,6 +107,9 @@ class TaskExecutor:
"""
operation = (assignment.operation or "tts").lower()
params = _parse_params(assignment.params_json)
params = await self._materialize_inputs(
assignment, params, fetch_input or self._fetch_input
)
handler = {
"tts": self._run_tts,
@@ -148,6 +176,93 @@ class TaskExecutor:
await report.progress(1.0, "done")
return {"meta": meta, "payload": payload}
# ── Inputs ────────────────────────────────────────────────────────────
async def _materialize_inputs(self, assignment, params: dict, fetch) -> dict:
"""Turn declared inputs into local files, then point the params at them.
The control plane sends artifact ids, never paths — its own paths mean
nothing here. So a clone arrives with ``ref_audio`` set to an id, and
the audio itself only exists once this has fetched it. Getting that
wrong does not fail loudly: the engine renders in the default voice and
the user gets audio that is simply not their clone.
"""
errors = params.get(INPUT_ERRORS_PARAM)
if errors:
detail = "; ".join(str(e) for e in errors) if isinstance(errors, list) else str(errors)
raise TaskFailure(
WorkerError(
error_class=ErrorClass.TERMINAL,
code="INPUT_UNAVAILABLE",
message=f"The task's input files could not be prepared: {detail}",
hint="Check that the reference audio still exists, then try again.",
)
)
refs = [ref for ref in (getattr(assignment, "inputs", None) or []) if ref.artifact_id]
if not refs:
return params
if fetch is None:
raise TaskFailure(
WorkerError(
error_class=ErrorClass.CAPABILITY,
code="INPUT_TRANSFER_UNSUPPORTED",
message="This worker cannot fetch task inputs.",
hint="Update the worker, or run this task on a worker that can.",
)
)
_, run_budget = _budgets(assignment)
local: dict[str, str] = {}
for ref in refs:
local[ref.artifact_id] = await self._bounded(
self._fetch_one(ref, fetch),
timeout=min(run_budget, _FALLBACK_INPUT_FETCH_SECONDS),
code="INPUT_FETCH_TIMEOUT",
what=f"Fetching '{ref.filename or ref.artifact_id}'",
)
return _rewrite_params(params, local)
async def _fetch_one(self, ref, fetch) -> str:
"""The local copy of one input, downloaded only if we lack it.
Content-addressed: the name is the hash the control plane computed, so
a second clone of the same voice — or a retry of this very task on this
worker — costs no transfer at all.
"""
directory = self._input_dir or default_input_dir()
os.makedirs(directory, exist_ok=True)
destination = os.path.join(directory, _cache_name(ref))
if _already_held(destination, ref):
_touch(destination)
return destination
partial = f"{destination}.part"
try:
await fetch(ref, partial)
except TaskFailure:
raise
except Exception as exc:
_discard(partial)
raise TaskFailure(
WorkerError(
# Transient on purpose: an id we cannot resolve now is far
# more often a dropped stream than a permanently missing
# file, and one wasted retry beats failing real work.
error_class=ErrorClass.TRANSIENT,
code="INPUT_FETCH_FAILED",
message=f"Could not fetch '{ref.filename or ref.artifact_id}': {exc}",
hint="The control plane may have restarted; the task will be retried.",
)
) from exc
# Off the loop: hashing a source video on the event loop thread would
# stall every heartbeat this worker owes the control plane.
await asyncio.to_thread(_verify, partial, ref)
os.replace(partial, destination)
await asyncio.to_thread(_prune_input_cache, directory)
return destination
# ── Engine plumbing ───────────────────────────────────────────────────
@staticmethod
@@ -328,6 +443,142 @@ def _budgets(assignment) -> tuple[float, float]:
)
def default_input_dir() -> str:
"""Where fetched inputs are cached on this worker.
Under the app's own data directory when there is one — a worker is the
ordinary backend in worker mode — and the system temp dir otherwise, so a
stripped-down install still runs instead of failing on a missing path.
"""
try:
from core.config import DATA_DIR # noqa: PLC0415
return os.path.join(str(DATA_DIR), "workers", "inputs")
except Exception: # pragma: no cover — no app data dir on this host
import tempfile # noqa: PLC0415
return os.path.join(tempfile.gettempdir(), "omnivoice-worker-inputs")
def _cache_name(ref) -> str:
"""A safe, content-addressed local name for one input.
Never the wire filename: that is remote input, and joining it onto a
directory is how a peer writes outside it. The hash the control plane sent
is the identity; the extension is kept only when it is a plain one,
because an engine that shells out to ffmpeg reads the suffix.
"""
digest = "".join(c for c in (getattr(ref, "sha256", "") or "") if c in "0123456789abcdef")
if len(digest) != 64:
digest = hashlib.sha256((ref.artifact_id or "").encode("utf-8")).hexdigest()
suffix = os.path.splitext(os.path.basename(str(getattr(ref, "filename", "") or "")))[1].lower()
if not (1 < len(suffix) <= 9 and suffix[1:].isalnum()):
suffix = ""
return f"{digest}{suffix}"
def _already_held(path: str, ref) -> bool:
"""Do we already have this exact input?
Size alone: the name is the content hash and the only writer is an atomic
rename, so a file of the right size at this name cannot be different bytes.
"""
try:
expected = int(getattr(ref, "size_bytes", 0) or 0)
return os.path.isfile(path) and (not expected or os.path.getsize(path) == expected)
except OSError: # pragma: no cover
return False
def _touch(path: str) -> None:
try:
os.utime(path, None)
except OSError: # pragma: no cover
pass
def _discard(path: str) -> None:
try:
os.remove(path)
except OSError:
pass
def _verify(path: str, ref) -> None:
"""Refuse a transfer that does not match what was announced.
A truncated reference clip does not fail — it clones three seconds of
silence — so the check has to happen before the file is committed.
"""
expected_size = int(getattr(ref, "size_bytes", 0) or 0)
expected_hash = (getattr(ref, "sha256", "") or "").lower()
try:
actual_size = os.path.getsize(path)
except OSError as exc:
_discard(path)
raise TaskFailure(
WorkerError(
error_class=ErrorClass.TRANSIENT,
code="INPUT_FETCH_FAILED",
message=f"The input '{ref.filename or ref.artifact_id}' did not arrive.",
hint="The task will be retried.",
)
) from exc
actual_hash = ""
if expected_hash:
digest = hashlib.sha256()
with open(path, "rb") as handle:
for block in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(block)
actual_hash = digest.hexdigest()
if (expected_size and actual_size != expected_size) or (
expected_hash and actual_hash != expected_hash
):
_discard(path)
raise TaskFailure(
WorkerError(
error_class=ErrorClass.TRANSIENT,
code="INPUT_CORRUPT",
message=f"The input '{ref.filename or ref.artifact_id}' arrived damaged.",
hint="The transfer will be retried.",
)
)
def _prune_input_cache(directory: str, limit_bytes: int = INPUT_CACHE_LIMIT_BYTES) -> None:
"""Keep the input cache under its ceiling, oldest first."""
try:
entries = []
total = 0
for name in os.listdir(directory):
path = os.path.join(directory, name)
if not os.path.isfile(path):
continue
stat = os.stat(path)
entries.append((stat.st_mtime, stat.st_size, path))
total += stat.st_size
for _mtime, size, path in sorted(entries):
if total <= limit_bytes:
break
os.remove(path)
total -= size
except OSError: # pragma: no cover — a full cache is not a failed task
logger.debug("Could not prune the worker input cache", exc_info=True)
def _rewrite_params(params: dict, local: dict[str, str]):
"""Replace every artifact id in the params with its local path."""
if isinstance(params, dict):
return {key: _rewrite_params(value, local) for key, value in params.items()}
if isinstance(params, list):
return [_rewrite_params(item, local) for item in params]
if isinstance(params, str):
return local.get(params, params)
return params
def _mark(audio, sample_rate: int, params: dict):
"""Provenance-mark synthetic audio before it is encoded (EU AI Act 50(2)).
@@ -370,4 +621,13 @@ def encode_inline(payload: bytes) -> str:
return base64.b64encode(payload).decode("ascii")
__all__ = ["INLINE_LIMIT_BYTES", "TaskExecutor", "TaskFailure", "UnsupportedOperation"]
__all__ = [
"INLINE_LIMIT_BYTES",
"INPUT_CACHE_LIMIT_BYTES",
"INPUT_ERRORS_PARAM",
"InputFetcher",
"TaskExecutor",
"TaskFailure",
"UnsupportedOperation",
"default_input_dir",
]
backend/worker/task_store.py
+314 -3
View File
@@ -17,12 +17,18 @@ flips the task to completed, so the ack can only follow a durable fact.
"""
from __future__ import annotations
import hashlib
import json
import logging
import mimetypes
import os
import re
import shutil
import time
from typing import Iterable, Optional
from typing import Iterable, Iterator, Optional
from core.db import db_conn
from core.path_security import UnsafePath, resolve_within, safe_filename
from worker.clock import resolve
from worker.errors import ErrorClass, WorkerError
from worker.lifecycle import Attempt, AttemptState, PriorityClass, Task, TaskState
@@ -92,6 +98,272 @@ def _row_to_task(row, attempts: list[Attempt]) -> Task:
return task
# ── Input artifacts ────────────────────────────────────────────────────────
#
# A worker is another machine. Every file-valued parameter — reference audio
# for a clone, a source video for a dub — lives in ``VOICES_DIR`` or a tempdir
# on the *control plane*, so sending its path is sending a string that names
# nothing on the far side. That is why remote cloning could not work: the
# assignment carried ``ref_audio=/Users/…/voices/x.wav`` and the worker either
# failed to open it or, worse, rendered with the default voice.
#
# Staging copies those files into the artifact directory the control plane
# already serves over ``DownloadArtifact``, which refuses anything outside it.
# The copy is named by the SHA-256 of its contents, so cloning the same voice
# a hundred times keeps exactly one copy on disk and lets the worker's own
# cache skip the transfer entirely on every clone after the first.
INPUT_PARAM_KEYS: tuple[str, ...] = (
"ref_audio",
"reference_audio",
"prompt_audio",
"prompt_wav",
"source_audio",
"audio_path",
"source_video",
"video_path",
)
# Where staged inputs live under the artifact root, and the key under which a
# task records what was staged for it. The record is what makes the purge
# exact: an input is deletable only when no surviving task still refers to it.
INPUTS_DIRNAME = "inputs"
INPUTS_PARAM_KEY = "inputs"
_HASH_CHUNK_BYTES = 1024 * 1024
_SAFE_EXTENSION = re.compile(r"^\.[A-Za-z0-9]{1,8}$")
class InputStagingError(RuntimeError):
"""A task input could not be staged for transfer to a worker.
Raised rather than swallowed: a clone whose reference audio silently went
missing does not fail, it renders someone else's voice.
"""
def artifact_root(*, create_dir: bool = True) -> str:
"""The directory the control plane serves artifacts from.
Imported lazily: ``worker.service`` owns the layout, and a module-level
import here would tie the durable store to the lifecycle module that
starts the gRPC server.
"""
from worker.service import paths # noqa: PLC0415 — layout owner, not a dependency
root = paths()["artifacts"]
if create_dir:
os.makedirs(os.path.join(root, INPUTS_DIRNAME), exist_ok=True)
return root
def _extension(source: str) -> str:
"""The source extension when it is a plain one, else nothing.
Kept for the worker's benefit — soundfile sniffs content, but an engine
that shells out to ffmpeg reads the suffix — and sanitised because the
name is about to become a filesystem path.
"""
suffix = os.path.splitext(str(source))[1]
return suffix.lower() if _SAFE_EXTENSION.match(suffix) else ""
def _digest(path: str) -> tuple[str, int]:
"""(sha256, size) read in chunks — a source video is not a bytes object."""
digest = hashlib.sha256()
size = 0
with open(path, "rb") as handle:
while True:
block = handle.read(_HASH_CHUNK_BYTES)
if not block:
break
digest.update(block)
size += len(block)
return digest.hexdigest(), size
def stage_input(source: str, *, root: Optional[str] = None, now: Optional[float] = None) -> dict:
"""Copy one input into the artifact store, keyed by its content hash.
Returns the record that ends up on the task row. ``source`` is kept in it
so a local fallback still has the original file, and stripped before the
record reaches the wire.
"""
stamp = resolve(now)
base = root or artifact_root()
try:
digest, size = _digest(source)
except OSError as exc:
raise InputStagingError(f"Could not read the task input {source!r}: {exc}") from exc
artifact_id = os.path.join(INPUTS_DIRNAME, f"{digest}{_extension(source)}")
try:
destination = resolve_within(base, artifact_id)
except UnsafePath as exc: # pragma: no cover — the id is ours, hex only
raise InputStagingError(f"Refusing to stage {source!r} outside the artifact store") from exc
try:
destination.parent.mkdir(parents=True, exist_ok=True)
# Same size at a content-addressed name means the same bytes: the only
# writer is the rename below, so a truncated file cannot exist here.
if not (destination.is_file() and destination.stat().st_size == size):
partial = destination.with_name(destination.name + ".part")
shutil.copyfile(source, partial)
os.replace(partial, destination)
# Freshness, not decoration: the purge dates an unreferenced input by
# its mtime, so re-using a staged voice has to renew it.
os.utime(destination, (stamp, stamp))
except OSError as exc:
raise InputStagingError(f"Could not stage the task input {source!r}: {exc}") from exc
filename = os.path.basename(str(source)) or f"{digest}{_extension(source)}"
return {
"artifact_id": artifact_id,
"path": str(destination),
"source": str(source),
"filename": filename,
"sha256": digest,
"size_bytes": size,
"content_type": mimetypes.guess_type(filename)[0] or "application/octet-stream",
}
def _iter_input_values(params: dict) -> Iterator[tuple[str, Optional[int], str]]:
"""``(key, index, value)`` for every parameter that could name a file."""
for key in INPUT_PARAM_KEYS:
value = params.get(key)
if isinstance(value, str):
yield key, None, value
elif isinstance(value, list):
for index, item in enumerate(value):
if isinstance(item, str):
yield key, index, item
def ensure_staged(
task: Task, *, root: Optional[str] = None, now: Optional[float] = None
) -> list[dict]:
"""Stage every file-valued parameter of *task*, once.
Idempotent by design — it runs at submission (so the durable row records
what a later purge must keep) and again at dispatch (so a task built
without the store, or a scheduler running unpersisted, still gets inputs
the worker can fetch). Already-staged keys are skipped, so the second call
does no I/O.
"""
params = task.params if isinstance(task.params, dict) else {}
recorded = params.get(INPUTS_PARAM_KEY)
entries: list[dict] = [e for e in recorded if isinstance(e, dict)] if isinstance(recorded, list) else []
if root:
# A task may have been staged when it was submitted under the default
# store, then dispatched by a servicer configured with another store.
# Recorded metadata is not proof that this servicer can serve it.
refreshed: list[dict] = []
for entry in entries:
artifact_id = str(entry.get("artifact_id") or "")
try:
available = bool(artifact_id and resolve_within(root, artifact_id).is_file())
except UnsafePath:
available = False
if available:
refreshed.append(entry)
continue
source = str(entry.get("source") or "")
if source and os.path.isfile(source):
replacement = stage_input(source, root=root, now=now)
replacement.update(key=entry.get("key"), index=entry.get("index"))
refreshed.append(replacement)
else:
raise InputStagingError(
f"The staged task input {artifact_id!r} is unavailable in this artifact store."
)
entries = refreshed
params[INPUTS_PARAM_KEY] = entries
covered = {(e.get("key"), e.get("index")) for e in entries}
for key, index, value in _iter_input_values(params):
if (key, index) in covered or not value:
continue
# Not every value of these keys is a file: an engine may take a voice
# id here. Only what exists on this disk is an input.
if not os.path.isfile(value):
continue
entry = stage_input(value, root=root, now=now)
entry["key"] = key
entry["index"] = index
entries.append(entry)
covered.add((key, index))
if entries:
params[INPUTS_PARAM_KEY] = entries
task.params = params
return entries
def _referenced_artifacts(conn) -> set[str]:
"""Every staged input still named by a surviving task row."""
referenced: set[str] = set()
for row in conn.execute("SELECT params_json FROM remote_tasks").fetchall():
try:
params = json.loads(row["params_json"] or "{}")
entries = params.get(INPUTS_PARAM_KEY) or []
except (ValueError, AttributeError):
continue
for entry in entries:
if isinstance(entry, dict) and entry.get("artifact_id"):
referenced.add(str(entry["artifact_id"]))
return referenced
def purge_artifacts(
task_ids: Iterable[str], referenced: set[str], *, cutoff: float, root: Optional[str] = None
) -> int:
"""Delete the results of purged tasks and every input nothing points at.
Both directions, deliberately: results are attempt-scoped and die with
their task, while a content-hashed input is shared, so it may only go once
no surviving task refers to it *and* it is older than the same cutoff the
rows were judged by. Nothing here raises — a purge that fails is a disk
that stays fuller than we wanted, not a failed request.
"""
removed = 0
try:
base = root or artifact_root(create_dir=False)
except Exception: # pragma: no cover — no data dir at all
logger.debug("No artifact root to purge", exc_info=True)
return 0
if not os.path.isdir(base):
return 0
for task_id in task_ids:
try:
path = resolve_within(base, safe_filename(task_id))
except UnsafePath:
continue
if os.path.isdir(path):
shutil.rmtree(path, ignore_errors=True)
removed += 1
inputs_dir = os.path.join(base, INPUTS_DIRNAME)
try:
names = os.listdir(inputs_dir)
except OSError:
return removed
for name in names:
artifact_id = os.path.join(INPUTS_DIRNAME, name)
if artifact_id in referenced:
continue
path = os.path.join(inputs_dir, name)
try:
if not os.path.isfile(path) or os.path.getmtime(path) >= cutoff:
continue
os.remove(path)
removed += 1
except OSError:
logger.debug("Could not purge the staged input %s", name, exc_info=True)
return removed
# ── Writes ─────────────────────────────────────────────────────────────────
@@ -100,12 +372,17 @@ def create(task: Task, *, project_id: Optional[str] = None, now: Optional[float]
Idempotent on ``idempotency_key``: a client that retries its HTTP request
gets the original task back rather than a second render of the same text.
Inputs are staged before the row is written, so the durable record names
the artifacts the task owns. Persisting first would leave a task whose
reference audio no purge can account for.
"""
stamp = resolve(now)
if task.idempotency_key:
existing = get_by_idempotency_key(task.idempotency_key)
if existing is not None:
return existing
ensure_staged(task, now=stamp)
with db_conn() as conn:
conn.execute(
"INSERT INTO remote_tasks "
@@ -292,10 +569,29 @@ def list_tasks(*, states: Optional[Iterable[TaskState]] = None, limit: int = 100
return [_row_to_task(r, _attempts_for(conn, r["id"])) for r in rows]
def purge_finished(*, older_than_seconds: float = 7 * 24 * 3600, now: Optional[float] = None) -> int:
def purge_finished(
*,
older_than_seconds: float = 7 * 24 * 3600,
now: Optional[float] = None,
root: Optional[str] = None,
) -> int:
"""Drop old finished tasks — rows *and* the bytes they own.
Rows only was a leak with no ceiling: every remote render leaves a result
artifact on disk, and every remote clone leaves a copy of the reference
audio. Neither was ever deleted, so the feature grew the user's disk for
as long as they used it.
"""
cutoff = resolve(now) - older_than_seconds
terminal = ", ".join(f"'{s.value}'" for s in TaskState if s.terminal)
with db_conn() as conn:
doomed = [
row["id"]
for row in conn.execute(
f"SELECT id FROM remote_tasks WHERE state IN ({terminal}) AND finished_at < ?",
(cutoff,),
).fetchall()
]
conn.execute(
f"DELETE FROM remote_task_attempts WHERE task_id IN "
f"(SELECT id FROM remote_tasks WHERE state IN ({terminal}) AND finished_at < ?)",
@@ -304,17 +600,32 @@ def purge_finished(*, older_than_seconds: float = 7 * 24 * 3600, now: Optional[f
cur = conn.execute(
f"DELETE FROM remote_tasks WHERE state IN ({terminal}) AND finished_at < ?", (cutoff,)
)
return cur.rowcount
removed = cur.rowcount
# Read the survivors inside the same transaction that deleted the
# rows: an input is only unreferenced relative to what is left.
referenced = _referenced_artifacts(conn)
# Filesystem work outside the transaction — a slow rmtree must not hold
# SQLite's write lock against the dispatch loop.
purge_artifacts(doomed, referenced, cutoff=cutoff, root=root)
return removed
__all__ = [
"INPUTS_DIRNAME",
"INPUTS_PARAM_KEY",
"INPUT_PARAM_KEYS",
"InputStagingError",
"artifact_root",
"commit_result",
"create",
"ensure_staged",
"get",
"get_by_idempotency_key",
"is_committed",
"list_tasks",
"load_unfinished",
"purge_artifacts",
"purge_finished",
"save",
"stage_input",
]
backend/worker/transport/client.py
+351 -22
View File
@@ -22,10 +22,19 @@ silent for a progress lease, and the longest silence in a task's life — the
cold model load — happens *after* the worker says it started. So every running
task carries a timer that renews the lease, marked ``keepalive`` so the server
can tell "still working" from "still ticking" and bound it by the phase budget.
**Bulk bytes never ride the control stream.** A result above the negotiated
inline threshold goes over UploadResult on a second RPC, and the control
stream carries only its ``ArtifactRef``. What is left on that stream is split
again into control and bulk queues, because the heartbeat this whole liveness
model rests on must not queue behind a payload — including the one payload,
``result_json``, that has no size cliff to catch it.
"""
from __future__ import annotations
import asyncio
import contextlib
import hashlib
import inspect
import json
import logging
@@ -59,6 +68,28 @@ _HEARTBEAT_SECONDS = 20.0
# see _oversized_result_error for why that has to be a failure and not a retry.
MAX_MESSAGE_BYTES = 8 * 1024 * 1024
# Room left for result_json, the ref, and protobuf framing when a payload does
# ride inline. The inline decision is made on the payload alone, so without a
# reserve a payload sized exactly at the frame cap would overflow it.
_INLINE_FRAME_HEADROOM_BYTES = 64 * 1024
# One upload chunk. Small enough that a chunk boundary — and therefore a lease
# renewal — comes round often on a slow uplink, large enough that a 100 MB dub
# is a hundred frames rather than a hundred thousand.
_UPLOAD_CHUNK_BYTES = 1024 * 1024
# How many times a single upload may be asked to resume before the worker
# calls the receiver broken. Generous for a genuinely flaky uplink — each
# resume restarts from a real byte count, so honest progress needs very few —
# and low enough that a receiver stuck in a resume loop costs one attempt
# rather than this worker's whole session.
_MAX_UPLOAD_RESUMES = 8
# The progress stage a result upload reports under. The control plane keys
# RESULT_UPLOADING (and its much longer delivery budget) off this, so it is a
# wire constant, not a cosmetic label.
UPLOAD_STAGE = "uploading"
# Used when an assignment carries no lease (an older control plane, or a test).
# Mirrors deadlines.py's _HEARTBEAT_GRACE_S * 4.
_DEFAULT_PROGRESS_LEASE_SECONDS = 120.0
@@ -68,7 +99,7 @@ _DEFAULT_PROGRESS_LEASE_SECONDS = 120.0
_MIN_KEEPALIVE_INTERVAL_SECONDS = 0.05
# Reporter keywords the client offers the executor, per task.
_REPORTER_KWARGS = frozenset({"on_progress", "on_model_loading"})
_EXECUTOR_KWARGS = frozenset({"on_progress", "on_model_loading", "fetch_input"})
def keepalive_interval(lease_seconds: float) -> float:
@@ -96,8 +127,8 @@ def _accepted_reporter_kwargs(execute: Callable) -> frozenset[str]:
except (TypeError, ValueError): # C-implemented or otherwise unintrospectable
return frozenset()
if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()):
return _REPORTER_KWARGS
return frozenset(name for name in _REPORTER_KWARGS if name in parameters)
return _EXECUTOR_KWARGS
return frozenset(name for name in _EXECUTOR_KWARGS if name in parameters)
def backoff_delay(attempt: int, *, jitter: Optional[Callable[[], float]] = None) -> float:
@@ -113,17 +144,72 @@ def backoff_delay(attempt: int, *, jitter: Optional[Callable[[], float]] = None)
return ceiling * roll
class _Outbox:
"""Two queues behind one interface: control frames overtake bulk ones.
The liveness model is built on the heartbeat arriving every 20 s, but a
single FIFO puts that heartbeat *behind* whatever result frame is being
written — and a result frame is the one message with no small upper bound
on its size. The worker then looks dead while it is in fact busy delivering
exactly the work it was asked for.
Splitting by class rather than shrinking the payload is the durable fix:
the upload path below already moves the big bytes off this stream, but
``result_json`` has no size cliff to catch, and the next bulk message added
to the protocol would reintroduce the stall silently.
Strict priority, not fair queuing: control frames are small, bounded in
number by the number of running tasks, and only ever *reduce* work — there
is nothing here for a bulk frame to be starved by for long.
"""
def __init__(self) -> None:
self.control: asyncio.Queue[pb.WorkerMessage] = asyncio.Queue()
self.bulk: asyncio.Queue[pb.WorkerMessage] = asyncio.Queue()
self._arrival = asyncio.Event()
def put_nowait(self, message: pb.WorkerMessage, *, bulk: bool = False) -> None:
(self.bulk if bulk else self.control).put_nowait(message)
self._arrival.set()
async def put(self, message: pb.WorkerMessage, *, bulk: bool = False) -> None:
self.put_nowait(message, bulk=bulk)
async def get(self) -> pb.WorkerMessage:
while True:
if not self.control.empty():
return self.control.get_nowait()
if not self.bulk.empty():
return self.bulk.get_nowait()
# Cleared before the wait and set by every put, with no await in
# between: on a single loop that ordering cannot lose a wakeup.
self._arrival.clear()
await self._arrival.wait()
def qsize(self) -> int:
return self.control.qsize() + self.bulk.qsize()
def empty(self) -> bool:
return self.qsize() == 0
@dataclass
class PendingResult:
"""A finished result the server has not acknowledged yet.
Held until RESULT_ACK arrives, across reconnects. Dropping it early is how
a completed forty-minute dub disappears with no error anywhere.
Anything over the inline threshold is uploaded first and represented here
by its ``ArtifactRef`` alone — the bytes are already durable on the control
plane, so a redelivery costs one small frame instead of re-sending a
payload that may not even fit in one (#B9).
"""
ref: pb.TaskRef
result_json: str = ""
inline_payload: bytes = b""
artifacts: list[pb.ArtifactRef] = field(default_factory=list)
usage: Optional[pb.UsageReport] = None
@@ -179,12 +265,18 @@ class WorkerClient:
# which would make key-based identity pointless.
self._on_registered = on_registered
self._reporter_kwargs = _accepted_reporter_kwargs(execute)
self._outbox: asyncio.Queue[pb.WorkerMessage] = asyncio.Queue()
self._outbox = _Outbox()
self._pending: dict[str, PendingResult] = {}
self._running: dict[str, asyncio.Task] = {}
self._keepalives: dict[str, asyncio.Task] = {}
self._epoch = 0
self._session_token = ""
# Negotiated by ConfigUpdate; None means "use the executor's own
# preference", so the threshold is never spelled twice.
self._inline_threshold: Optional[int] = None
# The live stub, kept so the result upload can use a second RPC on the
# same channel rather than the control stream.
self._stub = None
self._stop = asyncio.Event()
# ── Connection ────────────────────────────────────────────────────────
@@ -256,11 +348,18 @@ class WorkerClient:
heartbeat = asyncio.create_task(
self._heartbeat_loop(response.heartbeat_interval_seconds or _HEARTBEAT_SECONDS)
)
# Published only once the session is established: an upload before
# this point would carry a token the server has not issued yet.
self._stub = stub
try:
async for message in stream:
await self._on_server_message(message)
finally:
heartbeat.cancel()
# The channel closes with this block, so a stub kept past it
# would fail every upload with a confusing "channel closed"
# instead of the honest "no session".
self._stub = None
async def _register(self, stub) -> pb.RegisterResponse:
challenge = identity.new_challenge()
@@ -306,8 +405,14 @@ class WorkerClient:
message = await self._outbox.get()
yield message
async def _send(self, message: pb.WorkerMessage) -> None:
await self._outbox.put(message)
async def _send(self, message: pb.WorkerMessage, *, bulk: bool = False) -> None:
"""Enqueue a frame. ``bulk`` is for anything that can be large.
Only result frames qualify today. Everything else — heartbeat, pong,
progress, accept/reject, started/failed — is the control plane's view
of whether this worker is alive, and must not queue behind a payload.
"""
await self._outbox.put(message, bulk=bulk)
async def _heartbeat_loop(self, interval: float) -> None:
while True:
@@ -335,7 +440,7 @@ class WorkerClient:
"""Re-send anything the server never acknowledged."""
for pending in list(self._pending.values()):
logger.info("Redelivering unacknowledged result for task %s", pending.ref.task_id)
await self._send(_result_message(pending))
await self._send(_result_message(pending), bulk=True)
async def _cancel_zombies(self, authoritative: set[str]) -> None:
"""Stop work the control plane no longer believes in."""
@@ -374,6 +479,11 @@ class WorkerClient:
elif kind == "config":
if message.config.max_concurrent_tasks:
self.config.max_concurrent_tasks = message.config.max_concurrent_tasks
if message.config.inline_result_threshold_bytes:
# Negotiated, so the two sides cannot drift: the control plane
# is the one that knows how much it is willing to take on the
# control stream, and it may lower this at any time.
self._inline_threshold = int(message.config.inline_result_threshold_bytes)
elif kind == "ping":
# Answer immediately; the server times the round trip.
await self._send(pb.WorkerMessage(pong=pb.Pong(nonce=message.ping.nonce)))
@@ -421,23 +531,37 @@ class WorkerClient:
),
name=f"worker-keepalive-{assignment.ref.attempt_id}",
)
result = await self._execute(assignment, **self._reporters(assignment))
result = await self._execute(assignment, **self._executor_kwargs(assignment))
meta = result.get("meta", {}) or {}
payload = result.get("payload", b"") or b""
artifacts: list[pb.ArtifactRef] = []
if self._should_upload(payload):
# The keepalive timer is still armed here on purpose: the
# upload runs under the same attempt, and the renewals it
# sends per chunk (below) are what buy it the delivery budget.
artifacts, payload = await self._deliver_out_of_band(
assignment.ref, payload, meta
)
# Stopped before the terminal frame so no keepalive can arrive
# claiming an attempt the server has already settled.
self._stop_keepalive(key)
pending = PendingResult(
ref=assignment.ref,
result_json=json.dumps(result.get("meta", {})),
inline_payload=result.get("payload", b"") or b"",
result_json=json.dumps(meta),
inline_payload=payload,
artifacts=artifacts,
)
oversized = _oversized_result_error(pending)
if oversized is not None:
# Deliberately NOT recorded in _pending. An over-cap frame is
# rejected identically on every reconnect, so remembering it
# would redeliver a payload that can never be accepted and
# tear the session down each time (#B9) — taking every other
# task on this worker with it.
# Reachable now only through an enormous result_json: bulk
# bytes take the upload path above. Deliberately NOT recorded
# in _pending — an over-cap frame is rejected identically on
# every reconnect, so remembering it would redeliver a frame
# that can never be accepted and tear the session down each
# time (#B9), taking every other task on this worker with it.
logger.warning(
"Result for task %s is too large to deliver inline; failing it",
assignment.ref.task_id,
@@ -448,7 +572,7 @@ class WorkerClient:
# Recorded BEFORE sending: if the connection dies mid-send we must
# still know to redeliver.
self._pending[key] = pending
await self._send(_result_message(pending))
await self._send(_result_message(pending), bulk=True)
except asyncio.CancelledError:
raise
except Exception as exc:
@@ -474,6 +598,175 @@ class WorkerClient:
pb.WorkerMessage(failed=pb.TaskFailed(ref=ref, error=codec.error_to_pb(error)))
)
# ── Result delivery ───────────────────────────────────────────────────
def inline_limit(self) -> int:
"""How many payload bytes may ride the control stream.
The negotiated value when the control plane has stated one, otherwise
the executor's own preference — read from the executor rather than
copied, so there is exactly one default in the tree.
Clamped to what a frame can actually hold in either case: a control
plane that negotiates a threshold above the frame ceiling would
otherwise turn every large result into RESULT_TOO_LARGE, which is the
precise failure this phase exists to remove.
"""
if self._inline_threshold is not None:
limit = self._inline_threshold
else:
from worker.executor import INLINE_LIMIT_BYTES # noqa: PLC0415
limit = INLINE_LIMIT_BYTES
return max(0, min(int(limit), MAX_MESSAGE_BYTES - _INLINE_FRAME_HEADROOM_BYTES))
def _should_upload(self, payload: bytes) -> bool:
return bool(payload) and len(payload) > self.inline_limit()
async def _deliver_out_of_band(
self, ref: pb.TaskRef, payload: bytes, meta: dict
) -> tuple[list[pb.ArtifactRef], bytes]:
"""Upload the payload, returning ``([ref], b"")`` on success.
Falls back to inline delivery — ``([], payload)`` — only when the
payload would still fit in a frame. That fallback is what keeps an
older control plane (no UploadResult) and a one-off network stumble
from destroying a render that already succeeded; above the frame
ceiling there is no such option, and the attempt fails TRANSIENT so a
retry can find a working path rather than looping on a dead one.
"""
try:
return [await self._upload_result(ref, payload, meta)], b""
except asyncio.CancelledError:
raise
except Exception as exc:
if len(payload) <= MAX_MESSAGE_BYTES - _INLINE_FRAME_HEADROOM_BYTES:
logger.warning(
"Uploading the result for task %s failed (%s); sending it inline instead",
ref.task_id,
exc,
)
return [], payload
from worker.executor import TaskFailure # noqa: PLC0415
raise TaskFailure(
WorkerError(
error_class=ErrorClass.TRANSIENT,
code="RESULT_UPLOAD_FAILED",
message=(
f"The result ({len(payload) / (1024 * 1024):.1f} MiB) could not be "
f"uploaded to the control plane: {exc}"
),
hint="Check the connection between this worker and the control plane.",
)
) from exc
async def _upload_result(
self, ref: pb.TaskRef, payload: bytes, meta: dict
) -> pb.ArtifactRef:
"""Stream one result over UploadResult and return its committed ref.
``sha256`` and ``size_bytes`` are stated up front so the receiver can
refuse a transfer that arrives short or corrupted instead of renaming a
truncated file into place and calling the task done.
"""
stub = self._stub
if stub is None:
raise RuntimeError("no session is established")
artifact = pb.ArtifactRef(
task_id=ref.task_id,
attempt_id=ref.attempt_id,
filename=str(meta.get("filename") or f"{ref.attempt_id}.wav"),
content_type=str(meta.get("content_type") or "audio/wav"),
size_bytes=len(payload),
sha256=hashlib.sha256(payload).hexdigest(),
session_token=self._session_token,
)
# Sent before the first chunk so the control plane can move the attempt
# into RESULT_UPLOADING — and onto its delivery budget — before a slow
# uplink has had a chance to burn the ordinary progress lease.
await self._report_upload(ref, 0.0)
offset = 0
for _ in range(_MAX_UPLOAD_RESUMES):
ack = await stub.UploadResult(self._result_chunks(ref, artifact, payload, offset))
if ack.committed:
break
resumed = int(ack.bytes_received)
if ack.error.code and ack.error.code != "OFFSET_MISMATCH":
raise RuntimeError(ack.error.message or "the control plane refused the upload")
if resumed < 0 or resumed > len(payload) or resumed == offset:
raise RuntimeError(ack.error.message or "the control plane could not resume the upload")
offset = resumed
else:
# Bounded, because "did this make progress" cannot be answered by
# comparing against the previous offset alone: a receiver that
# alternates between two byte counts satisfies `resumed != offset`
# forever, and one that advances a few bytes per round would retry
# once per byte of a 100 MB dub. Either way the worker stops
# rendering anything else while it spins.
raise RuntimeError(
f"the control plane asked to resume the upload more than "
f"{_MAX_UPLOAD_RESUMES} times without committing it"
)
if ack.bytes_received and ack.bytes_received != len(payload):
raise RuntimeError(
f"the control plane received {ack.bytes_received} of {len(payload)} bytes"
)
# The final renewal cannot come from the chunk loop: the receiver stops
# pulling at ``last``, so the generator is closed before the code after
# that yield ever runs.
await self._report_upload(ref, 1.0)
committed = pb.ArtifactRef()
committed.CopyFrom(artifact)
if ack.artifact_id:
committed.artifact_id = ack.artifact_id
# The control stream is already authenticated; echoing the session
# token back on it would only widen where the token is written.
committed.ClearField("session_token")
return committed
async def _result_chunks(
self, ref: pb.TaskRef, artifact: pb.ArtifactRef, payload: bytes, offset: int = 0
):
"""Chunks in order, each ``offset`` equal to the bytes already sent.
The receiver checks that equality against the length it holds, so this
is a contract and not a hint. Exactly one chunk carries ``last``, and
only that one licenses the commit.
"""
total = len(payload)
while offset < total:
data = payload[offset : offset + _UPLOAD_CHUNK_BYTES]
offset += len(data)
yield pb.ResultChunk(
ref=artifact,
offset=offset - len(data),
data=data,
last=offset >= total,
session_token=self._session_token,
)
# Per chunk, not per timer: a lease renewed by real transfer
# progress cannot keep an attempt alive over a stalled upload.
await self._report_upload(ref, offset / total)
async def _report_upload(self, ref: pb.TaskRef, fraction: float) -> None:
await self._send(
pb.WorkerMessage(
progress=pb.TaskProgress(
ref=ref,
progress=float(fraction),
stage=UPLOAD_STAGE,
# Upload bytes are liveness, not synthesis progress. Mark
# them keepalive so the server applies the phase ceiling
# and does not replace an already-finished 100% with 0%.
keepalive=True,
)
)
)
# ── Liveness ──────────────────────────────────────────────────────────
async def _keepalive_loop(self, ref: pb.TaskRef, interval: float) -> None:
@@ -494,7 +787,37 @@ class WorkerClient:
if timer is not None:
timer.cancel()
def _reporters(self, assignment: pb.TaskAssignment) -> dict[str, Callable]:
async def _fetch_input(self, ref: pb.ArtifactRef, destination: str) -> None:
"""Download one declared input with authenticated, ordered chunks."""
if self._stub is None:
raise RuntimeError("no session is established")
request = pb.ArtifactRef()
request.CopyFrom(ref)
request.session_token = self._session_token
offset = 0
complete = False
try:
with open(destination, "wb") as handle:
async for chunk in self._stub.DownloadArtifact(request):
if int(chunk.offset) != offset:
raise RuntimeError(
f"input offset {chunk.offset} did not match {offset} bytes received"
)
handle.write(chunk.data)
offset += len(chunk.data)
if chunk.last:
complete = True
break
except Exception:
with contextlib.suppress(OSError):
os.remove(destination)
raise
if not complete:
with contextlib.suppress(OSError):
os.remove(destination)
raise RuntimeError("input download ended before its final chunk")
def _executor_kwargs(self, assignment: pb.TaskAssignment) -> dict[str, Callable]:
"""Per-task progress callbacks for the executor.
Bound to this assignment's ref rather than installed on the executor
@@ -528,7 +851,11 @@ class WorkerClient:
)
)
available = {"on_progress": on_progress, "on_model_loading": on_model_loading}
available = {
"on_progress": on_progress,
"on_model_loading": on_model_loading,
"fetch_input": self._fetch_input,
}
return {k: v for k, v in available.items() if k in self._reporter_kwargs}
@@ -544,6 +871,7 @@ def _result_message(pending: PendingResult) -> pb.WorkerMessage:
ref=pending.ref,
result_json=pending.result_json,
inline_payload=pending.inline_payload,
artifacts=pending.artifacts,
)
)
@@ -558,10 +886,10 @@ def _oversized_result_error(pending: PendingResult) -> Optional[WorkerError]:
worker in the fleet would produce the same frame and be rejected the same
way. Retrying it burns the whole fleet's slots to arrive back here.
The executor also flags ``meta["inline"]`` against its own, much smaller
preference — that flag selects the Phase 3 upload path and is deliberately
not a failure condition, because everything between the two thresholds
delivers correctly today.
A last line of defence rather than the size policy it once was: bulk bytes
now take the UploadResult path (``WorkerClient.inline_limit``), so what
still reaches this is a ``result_json`` — a transcript, a segment list —
that on its own will not fit in a frame.
"""
size = _result_message(pending).ByteSize()
if size <= MAX_MESSAGE_BYTES:
@@ -619,6 +947,7 @@ def config_from_token(
__all__ = [
"MAX_MESSAGE_BYTES",
"UPLOAD_STAGE",
"PendingResult",
"WorkerClient",
"WorkerConfig",
backend/worker/transport/codec.py
+114 -2
View File
@@ -6,6 +6,8 @@ here makes decisions; it converts.
"""
from __future__ import annotations
import logging
import os
from typing import Optional
from worker.capacity import derive_concurrency
@@ -14,6 +16,12 @@ from worker.errors import ErrorClass, WorkerError
from worker.lifecycle import Attempt, PriorityClass, Task
from worker.protocol.gen import worker_v1_pb2 as pb
logger = logging.getLogger("omnivoice.worker")
# Where a staging failure is reported to the worker. Named here because both
# sides of the wire read it: the executor turns it into a terminal error.
_INPUT_ERRORS_KEY = "input_errors"
# Domain ErrorClass ↔ protobuf enum. Explicit rather than by-name so renaming
# one side cannot silently change the meaning of a wire value.
_ERROR_TO_PB = {
@@ -69,20 +77,33 @@ def deadlines_to_pb(budget: Deadlines) -> pb.Deadlines:
)
def assignment_to_pb(task: Task, attempt: Attempt, budget: Deadlines) -> pb.TaskAssignment:
def assignment_to_pb(
task: Task, attempt: Attempt, budget: Deadlines, *, artifact_root: Optional[str] = None
) -> pb.TaskAssignment:
"""Build the wire assignment.
``params_json`` carries the operation's parameters opaquely; the transport
has no business knowing what a dub or a clone needs.
The one thing it cannot stay opaque about is a **path**. A parameter like
``ref_audio`` names a file on the control plane's disk, which is not a
thing on the worker's — so a clone assignment used to arrive naming a file
that did not exist, and the worker either failed to open it or rendered
the default voice. Every file-valued parameter is therefore staged into
the artifact store, declared in ``inputs`` for the worker to fetch over
``DownloadArtifact``, and replaced in ``params_json`` by its artifact id.
No local path ever crosses the wire.
"""
import json # noqa: PLC0415 — only needed on this path
entries, errors = _staged_inputs(task, artifact_root)
return pb.TaskAssignment(
ref=ref_for(attempt),
operation=task.operation,
engine=task.engine,
model_id=task.model_id,
params_json=json.dumps(task.params),
params_json=json.dumps(remote_params(task.params, entries, errors)),
inputs=[input_ref(entry, task, attempt) for entry in entries],
deadlines=deadlines_to_pb(budget),
priority_class=int(task.priority),
attempt_number=attempt.attempt_number,
@@ -90,6 +111,95 @@ def assignment_to_pb(task: Task, attempt: Attempt, budget: Deadlines) -> pb.Task
)
def input_ref(entry: dict, task: Task, attempt: Attempt) -> pb.ArtifactRef:
"""One staged input as the worker will ask for it back.
``sha256`` and ``size_bytes`` are populated rather than left at their
defaults because they are what lets the worker verify the transfer and,
more usefully, recognise a reference clip it already holds.
"""
return pb.ArtifactRef(
artifact_id=str(entry.get("artifact_id") or ""),
task_id=task.task_id,
attempt_id=attempt.attempt_id,
filename=str(entry.get("filename") or ""),
content_type=str(entry.get("content_type") or ""),
size_bytes=int(entry.get("size_bytes") or 0),
sha256=str(entry.get("sha256") or ""),
)
def remote_params(params: dict, entries: list[dict], errors: list[str]) -> dict:
"""The parameters as the worker should see them.
Two rules. The staging bookkeeping (which holds control-plane paths) is
stripped. And any remaining file-valued parameter is *removed* rather than
passed through: an unstaged local path is worse than an absent one,
because absent fails loudly while a dead path can silently produce audio
in the wrong voice.
"""
from worker.task_store import INPUT_PARAM_KEYS, INPUTS_PARAM_KEY # noqa: PLC0415
remote = {
key: value
for key, value in (params or {}).items()
if key not in (INPUTS_PARAM_KEY, _INPUT_ERRORS_KEY)
}
mapped: dict[str, dict[Optional[int], str]] = {}
for entry in entries:
artifact_id = str(entry.get("artifact_id") or "")
if artifact_id:
mapped.setdefault(str(entry.get("key") or ""), {})[entry.get("index")] = artifact_id
for key in INPUT_PARAM_KEYS:
if key not in remote:
continue
by_index = mapped.get(key, {})
value = remote[key]
if isinstance(value, list):
rewritten = [
by_index.get(index, item)
for index, item in enumerate(value)
if index in by_index or not _is_local_path(item)
]
remote[key] = rewritten
elif isinstance(value, str):
if None in by_index:
remote[key] = by_index[None]
elif _is_local_path(value):
remote.pop(key)
if errors:
remote[_INPUT_ERRORS_KEY] = errors
return remote
def _is_local_path(value) -> bool:
"""Does this value name a place on this machine rather than a plain id?"""
if not isinstance(value, str) or not value:
return False
return os.path.isabs(value) or os.sep in value or "/" in value or os.path.exists(value)
def _staged_inputs(task: Task, artifact_root: Optional[str]) -> tuple[list[dict], list[str]]:
"""Stage this task's inputs, or say why they could not be staged.
A staging failure must not take down the dispatch loop, and it must not
fall back to sending the path: the assignment goes out with an explicit
error the worker turns into a terminal failure the user can read.
"""
from worker import task_store # noqa: PLC0415 — control-plane only
try:
entries = task_store.ensure_staged(task, root=artifact_root)
except task_store.InputStagingError as exc:
logger.warning("Could not stage inputs for task %s: %s", task.task_id, exc)
return [], [str(exc)]
except Exception as exc: # pragma: no cover — defensive
logger.warning("Input staging failed for task %s", task.task_id, exc_info=True)
return [], [f"Task inputs could not be prepared: {exc}"]
return [e for e in entries if e.get("artifact_id")], []
def capability_to_pb(cap: dict) -> pb.ModelCapability:
"""Convert a discovered capability.
@@ -199,7 +309,9 @@ __all__ = [
"error_to_pb",
"host_from_pb",
"host_to_pb",
"input_ref",
"priority_from_pb",
"remote_params",
"ref_for",
"task_ref",
]
backend/worker/transport/server.py
+425 -56
View File
@@ -8,6 +8,8 @@ domain:
* fencing — one active session per worker, newest epoch wins, stale epochs
dropped rather than merged
* ordering — persist a result before acknowledging it
* integrity — an artifact is verified against its declared digest before it
is renamed into place, and only an explicit last chunk commits one
Everything else is delegated. If this file starts making scheduling decisions,
something has been put in the wrong place.
@@ -21,11 +23,12 @@ alive are exactly what must never queue behind anything else.
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
import os
import time
from typing import Optional
from typing import Callable, Optional
import grpc
@@ -54,10 +57,188 @@ SESSION_METADATA_KEY = "x-omnivoice-session"
# payload here head-of-line blocks the heartbeats that prove the worker alive.
INLINE_RESULT_THRESHOLD = 256 * 1024
# Ceilings on what a remote peer may stream into our filesystem. They are not
# derived from anything the worker says: ``ArtifactRef.size_bytes`` narrows the
# cap when it is declared, but can never widen it. A gibibyte is roughly six
# hours of the 24 kHz PCM16 WAV the executor writes — past any single render,
# far short of a disk.
MAX_ARTIFACT_BYTES = 1024**3
# And a budget across every artifact one task delivers, so retries and
# redeliveries cannot walk past the per-artifact cap one upload at a time.
MAX_TASK_ARTIFACT_BYTES = 2 * 1024**3
_HEARTBEAT_INTERVAL_SECONDS = 20
# How often the control plane times a round trip to each worker. Frequent
# enough that the latency shown in the UI is current, rare enough to be free.
_PING_INTERVAL_SECONDS = 5.0
# Read size when serving an input. Two orders of magnitude under the 8 MiB
# message cap, so a large input is many small frames rather than one that the
# receiver refuses outright.
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
_REHASH_BLOCK_BYTES = 1024 * 1024
def _upload_refused(
code: str,
message: str,
*,
bytes_received: int = 0,
error_class: int = pb.ERROR_CLASS_PROTOCOL,
) -> pb.ResultAck:
"""A terminal ack that commits nothing and says why.
Refusals are answered rather than aborted: the ack is the only frame this
RPC ever sends back, so aborting the call would leave the worker knowing
the upload failed and nothing about whether to retry, resume, or re-render.
"""
return pb.ResultAck(
bytes_received=bytes_received,
committed=False,
error=pb.Error(error_class=error_class, code=code, message=message),
)
class _Upload:
"""One in-progress result transfer.
Bytes land in an attempt-scoped ``.part`` file and are renamed into place
only once the declared digest matches what actually arrived, so a
truncated, reordered, or corrupted transfer can never be mistaken for a
finished result. Every rule here exists because the sender is remote: the
offset is checked against what we hold rather than trusted as a hint, the
total is capped whether or not a size was declared, and an iterator that
simply stops commits nothing.
"""
def __init__(
self,
*,
attempt: Attempt,
artifact_id: str,
final: str,
limit: int,
declared_size: int,
declared_sha256: str,
on_commit: Callable[[Attempt, int], None],
) -> None:
self.attempt = attempt
self.artifact_id = artifact_id
self.final = final
self.part = f"{final}.part"
self.limit = limit
self.declared_size = declared_size
self.declared_sha256 = declared_sha256.strip().lower()
self._on_commit = on_commit
self._digest = hashlib.sha256()
self._handle = None
self.received = 0
def held_bytes(self) -> int:
try:
return os.path.getsize(self.part)
except OSError:
return 0
async def start(self, offset: int) -> Optional[pb.ResultAck]:
"""Open the part file at ``offset``, or refuse with what we hold."""
held = self.held_bytes()
if offset == 0:
self._handle = open(self.part, "wb")
return None
if offset == held and 0 < held <= self.limit:
# The digest has to cover the bytes already on disk, or the
# verification at commit would attest only to the resumed tail —
# which is exactly the case a resume exists to protect.
await asyncio.to_thread(self._rehash_held)
self.received = held
self._handle = open(self.part, "ab")
return None
return _upload_refused(
"OFFSET_MISMATCH",
"Resume from the byte count in this ack.",
bytes_received=held,
error_class=pb.ERROR_CLASS_TRANSIENT,
)
def _rehash_held(self) -> None:
with open(self.part, "rb") as fh:
for block in iter(lambda: fh.read(_REHASH_BLOCK_BYTES), b""):
self._digest.update(block)
def write(self, chunk) -> Optional[pb.ResultAck]:
"""Append one chunk. Non-None means the transfer is over."""
if int(chunk.offset) != self.received:
# Not a resume point: a gap or an overlap inside a live stream is
# a sender that has lost track of what it sent, and appending it
# would produce a file that hashes to nothing anybody expected.
return _upload_refused(
"OFFSET_MISMATCH",
"Resume from the byte count in this ack.",
bytes_received=self.received,
error_class=pb.ERROR_CLASS_TRANSIENT,
)
data = bytes(chunk.data)
if self.received + len(data) > self.limit:
self.discard()
return _upload_refused(
"ARTIFACT_TOO_LARGE",
"This result is larger than the control plane accepts.",
)
self._handle.write(data)
self._digest.update(data)
self.received += len(data)
return None
def commit(self) -> pb.ResultAck:
"""Verify, then rename. Never the other way round."""
self.close()
if self.declared_size and self.received != self.declared_size:
self.discard()
return _upload_refused(
"SIZE_MISMATCH",
"The transfer did not deliver the number of bytes it declared.",
error_class=pb.ERROR_CLASS_TRANSIENT,
)
if self._digest.hexdigest() != self.declared_sha256:
# Keeping the part file would let the next resume append onto
# bytes already known to be wrong.
self.discard()
return _upload_refused(
"DIGEST_MISMATCH",
"The uploaded result does not match its declared sha256.",
error_class=pb.ERROR_CLASS_TRANSIENT,
)
os.replace(self.part, self.final)
self._on_commit(self.attempt, self.received)
return pb.ResultAck(
artifact_id=self.artifact_id, bytes_received=self.received, committed=True
)
def incomplete(self) -> pb.ResultAck:
"""The stream ended with no terminal chunk.
The part file survives for a resume and nothing is renamed. This used
to return ``committed=True`` over whatever bytes happened to arrive.
"""
self.close()
return _upload_refused(
"UPLOAD_INCOMPLETE",
"The upload ended before its last chunk; resume from the byte count in this ack.",
bytes_received=self.received,
error_class=pb.ERROR_CLASS_TRANSIENT,
)
def close(self) -> None:
if self._handle is not None:
self._handle.close()
self._handle = None
def discard(self) -> None:
self.close()
try:
os.remove(self.part)
except OSError:
pass
class _Session:
@@ -94,6 +275,10 @@ class WorkerServicer(pb_grpc.WorkerServiceServicer):
self.cert_fingerprint = cert_fingerprint
self._sessions: dict[str, _Session] = {}
self._by_token: dict[str, _Session] = {}
# task_id → attempt_id → committed artifact bytes. Per attempt rather
# than a running total, so a redelivered upload of the same attempt
# replaces its own entry instead of spending the task's budget twice.
self._artifact_bytes: dict[str, dict[str, int]] = {}
os.makedirs(artifact_dir, exist_ok=True)
# ── Registration ──────────────────────────────────────────────────────
@@ -261,6 +446,19 @@ class WorkerServicer(pb_grpc.WorkerServiceServicer):
return
session.stream_open = True
from worker.executor import INLINE_LIMIT_BYTES # noqa: PLC0415
await session.send(
pb.ServerMessage(
config=pb.ConfigUpdate(
heartbeat_interval_seconds=_HEARTBEAT_INTERVAL_SECONDS,
max_concurrent_tasks=max(
1, self.pool.get(session.worker_id).capacity.max_concurrent_tasks
),
inline_result_threshold_bytes=INLINE_LIMIT_BYTES,
)
)
)
writer = asyncio.create_task(self._write_loop(session, context))
reader = asyncio.create_task(self._read_loop(session, request_iterator))
pinger = asyncio.create_task(self._ping_loop(session))
@@ -601,7 +799,10 @@ class WorkerServicer(pb_grpc.WorkerServiceServicer):
await session.send(
pb.ServerMessage(
assignment=codec.assignment_to_pb(
assignment.task, assignment.attempt, assignment.deadlines
assignment.task,
assignment.attempt,
assignment.deadlines,
artifact_root=self.artifact_dir,
)
)
)
@@ -633,76 +834,245 @@ class WorkerServicer(pb_grpc.WorkerServiceServicer):
async def UploadResult(self, request_iterator, context) -> pb.ResultAck:
"""Receive a result artifact in chunks, resumably.
Written to an attempt-scoped ``.part`` file and only renamed into place
once the last chunk lands, so a partial transfer can never be mistaken
for a finished result.
Nothing the sender says is taken on trust. Every chunk must name the
exact offset this control plane already holds, the total is bounded per
artifact and per task, the digest declared in ``ArtifactRef.sha256``
must match the bytes that arrived, and only an explicit ``last`` chunk
renames the ``.part`` file into place. An iterator that simply stops
leaves the partial file for a resume and commits nothing — it used to
commit, which is how a truncated transfer became a finished render.
Resume is real, and this is where it is reported. The call is
client-streaming with a single terminal ack, so there is no mid-stream
channel for "bytes already held": a chunk whose offset disagrees with
what we hold is answered with ``committed=False`` and
``bytes_received`` set to the authoritative held count, and the worker
restarts from there. That ack is the bytes-held probe the proto
promised and no RPC provided.
"""
artifact_id = ""
received = 0
handle = None
path = ""
final = ""
upload: Optional[_Upload] = None
try:
async for chunk in request_iterator:
if handle is None:
ref = chunk.ref
session = self._session_for(context, ref)
if session is None:
await context.abort(
grpc.StatusCode.UNAUTHENTICATED, "Unknown or expired session."
)
return pb.ResultAck(committed=False)
# Same rule as an inline result: the destination is minted
# from our own attempt record, never assembled from the ids
# in the request.
attempt = self._attempt_for(session, ref)
final = (
self._artifact_path(attempt.task_id, attempt.attempt_id)
if attempt is not None
else None
)
if final is None:
await context.abort(
grpc.StatusCode.PERMISSION_DENIED,
"No such attempt is running for this worker.",
)
return pb.ResultAck(committed=False)
artifact_id = final
path = f"{final}.part"
# Resume where a previous transfer stopped.
mode = "ab" if chunk.offset and os.path.exists(path) else "wb"
handle = open(path, mode)
received = handle.tell()
handle.write(chunk.data)
received += len(chunk.data)
if upload is None:
upload, refusal = await self._open_upload(context, chunk)
if upload is None:
if refusal is None:
await context.abort(
grpc.StatusCode.UNAUTHENTICATED,
"Unknown or expired session.",
)
return _upload_refused(
"UNAUTHENTICATED", "Unknown or expired session."
)
return refusal
refusal = upload.write(chunk)
if refusal is not None:
return refusal
self._renew_upload_lease(upload.attempt)
if chunk.last:
break
return upload.commit()
finally:
if handle is not None:
handle.close()
if path and final:
os.replace(path, final)
artifact_id = final
return pb.ResultAck(artifact_id=artifact_id, bytes_received=received, committed=True)
if upload is not None:
upload.close()
if upload is None:
return _upload_refused("EMPTY_UPLOAD", "The upload carried no chunks.")
return upload.incomplete()
async def _open_upload(self, context, chunk) -> tuple[Optional[_Upload], Optional[pb.ResultAck]]:
"""Authorise the first chunk and open its destination.
``(None, None)`` means unauthenticated — the one failure answered with
a gRPC abort rather than an ack, because a caller we cannot identify
has no business being told anything about the task it named.
"""
ref = chunk.ref
session = self._session_for(context, ref) or self._session_for(context, chunk)
if session is None:
return None, None
# Same rule as an inline result: the destination is minted from our own
# attempt record, never assembled from the ids in the request.
attempt = self._attempt_for(session, ref)
final = (
self._artifact_path(attempt.task_id, attempt.attempt_id)
if attempt is not None
else None
)
if attempt is None or final is None:
return None, _upload_refused(
"UNKNOWN_ATTEMPT", "No such attempt is running for this worker."
)
if not ref.sha256:
# Refused before a single byte is accepted. An upload with no
# declared digest cannot be verified, and committing it would make
# the whole verification path decorative.
return None, _upload_refused(
"DIGEST_REQUIRED", "Declare ArtifactRef.sha256 before uploading a result."
)
declared = int(ref.size_bytes)
if declared > MAX_ARTIFACT_BYTES:
return None, _upload_refused(
"ARTIFACT_TOO_LARGE", "This result is larger than the control plane accepts."
)
# A declared size narrows the cap; an undeclared one gets the ceiling.
limit = declared or MAX_ARTIFACT_BYTES
if self._artifact_bytes_spent(attempt.task_id) + limit > MAX_TASK_ARTIFACT_BYTES:
return None, _upload_refused(
"TASK_BUDGET_EXCEEDED",
"This task has delivered as many artifact bytes as it is allowed.",
)
# Last, and only once the request is known to be one we would accept:
# a refusal must not leave a task parked in RESULT_UPLOADING with no
# transfer under way. Cancelled, timed out, or already committed by
# another attempt all fail here, because accepting these bytes would
# overwrite the artifact of whichever attempt actually won.
if not self._begin_uploading(attempt):
return None, _upload_refused(
"ATTEMPT_NOT_LIVE", "This attempt is no longer accepting a result."
)
upload = _Upload(
attempt=attempt,
artifact_id=self._artifact_id_for(final),
final=final,
limit=limit,
declared_size=declared,
declared_sha256=ref.sha256,
on_commit=self._record_artifact_bytes,
)
refusal = await upload.start(int(chunk.offset))
if refusal is not None:
upload.close()
return None, refusal
return upload, None
def _begin_uploading(self, attempt: Attempt) -> bool:
"""Put the task into RESULT_UPLOADING for the length of the transfer.
Without this transition ``Task.uploading`` has no callers, so
RESULT_UPLOADING is unreachable and the entire delivery of a large
result runs under the 120 s progress lease while the 900 s
``result_delivery_seconds`` budget sits unused because nothing ever
entered the state it applies to.
"""
task = self.scheduler.get(attempt.task_id)
if task is None or task.state.terminal or attempt.state.terminal:
return False
try:
task.uploading(attempt.attempt_id, session_epoch=attempt.session_epoch)
except Exception:
logger.debug(
"Refusing an upload for attempt %s: not in a state that can deliver",
attempt.attempt_id,
exc_info=True,
)
return False
self._renew_upload_lease(attempt)
return True
def _renew_upload_lease(self, attempt: Attempt) -> None:
"""Renew the lease from upload progress, under the delivery budget.
Routed through the scheduler rather than computed here: it owns the
phase budgets, and ``on_progress(keepalive=True)`` already caps a
renewal at the current phase's ceiling — which, now that the task is in
RESULT_UPLOADING, is ``result_delivery_seconds``. A keepalive and not a
progress frame: bytes on the wire prove the worker is alive, not that
the render advanced, and overwriting a finished 100% with a transfer's
zero is a UI that goes backwards.
"""
self.scheduler.on_progress(
attempt.task_id,
attempt.attempt_id,
progress=0.0,
keepalive=True,
epoch=attempt.session_epoch,
)
def _artifact_bytes_spent(self, task_id: str) -> int:
"""How much of this task's artifact budget is already committed."""
for known in list(self._artifact_bytes):
task = self.scheduler.get(known)
if task is None or task.state.terminal:
self._artifact_bytes.pop(known, None)
return sum(self._artifact_bytes.get(task_id, {}).values())
def _record_artifact_bytes(self, attempt: Attempt, count: int) -> None:
self._artifact_bytes.setdefault(attempt.task_id, {})[attempt.attempt_id] = count
def _artifact_id_for(self, path: str) -> str:
"""The store-relative id a worker may name this artifact by.
Relative, not the absolute path it lives at: the id travels back on the
control stream as ``TaskResult.artifacts[0].artifact_id`` and is
re-resolved against the artifact directory, and handing a remote peer
our filesystem layout buys nothing that resolution does not already do.
"""
try:
return os.path.relpath(path, self.artifact_dir)
except ValueError: # different drive on Windows; cannot happen, but
return path
async def DownloadArtifact(self, request: pb.ArtifactRef, context):
"""Stream a task input (reference audio, source video) to a worker."""
if not self._authorized(context, request):
"""Stream a task input (reference audio, source video) to a worker.
Bound to the attempt that needs the input, not merely to a live
session. Until artifacts started flowing inwards, ``inputs`` carried
nothing and "any authenticated worker may read any staged file" was a
distinction without a difference; from here those files are the user's
own reference audio, staged from their voice library.
"""
session = self._session_for(context, request)
if session is None:
await context.abort(grpc.StatusCode.UNAUTHENTICATED, "Unknown or expired session.")
return
if not self._may_read_input(session, request):
await context.abort(
grpc.StatusCode.PERMISSION_DENIED,
"This input belongs to a task that is not running on this worker.",
)
return
path = self._resolve_input(request.artifact_id)
if path is None:
await context.abort(grpc.StatusCode.NOT_FOUND, "Artifact not found.")
return
# A ref minted here rather than the caller's echoed back: the request
# carries the worker's session token, and nothing goes back out that
# did not have to go out.
served = pb.ArtifactRef(
artifact_id=request.artifact_id,
task_id=request.task_id,
attempt_id=request.attempt_id,
filename=os.path.basename(path),
content_type=request.content_type,
size_bytes=os.path.getsize(path),
)
offset = 0
with open(path, "rb") as fh:
while True:
data = fh.read(64 * 1024)
data = fh.read(_DOWNLOAD_CHUNK_BYTES)
if not data:
break
yield pb.ArtifactChunk(ref=request, offset=offset, data=data, last=False)
yield pb.ArtifactChunk(ref=served, offset=offset, data=data, last=False)
offset += len(data)
yield pb.ArtifactChunk(ref=request, offset=offset, data=b"", last=True)
yield pb.ArtifactChunk(ref=served, offset=offset, data=b"", last=True)
def _may_read_input(self, session: _Session, ref) -> bool:
"""Does this session hold a live attempt of the task naming the input?
Per-task rather than per-artifact on purpose: staged inputs are
content-hashed so repeated clones of one voice share a single copy, and
an id that is deliberately shared cannot itself carry the
authorisation. The attempt does.
"""
task = self.scheduler.get(ref.task_id) if ref.task_id else None
if task is None:
return False
if ref.attempt_id:
attempt, foreign = self._attempt_and_owner(session, ref)
return attempt is not None and not foreign and attempt.state.live
return any(
attempt.worker_id == session.worker_id and attempt.state.live
for attempt in task.attempts
)
def _session_for(self, context, ref) -> Optional[_Session]:
"""The live session a transfer belongs to, by ref token or by metadata."""
@@ -712,9 +1082,6 @@ class WorkerServicer(pb_grpc.WorkerServiceServicer):
return None if session.session.expired() else session
return self._session_from_metadata(context)
def _authorized(self, context, ref) -> bool:
return self._session_for(context, ref) is not None
def _resolve_input(self, artifact_id: str) -> Optional[str]:
"""Resolve an input reference to a path inside the artifact directory.
@@ -776,6 +1143,8 @@ def _peer_address(context) -> str:
__all__ = [
"INLINE_RESULT_THRESHOLD",
"MAX_ARTIFACT_BYTES",
"MAX_TASK_ARTIFACT_BYTES",
"MIN_SUPPORTED_VERSION",
"PROTOCOL_VERSION",
"SESSION_METADATA_KEY",
frontend/src/components/GpuTarget.jsx
+97 -16
View File
@@ -12,6 +12,13 @@
*
* `Local` has no rename control. It is not a machine — it is this machine —
* and there is nothing to name.
*
* The answer is also per operation, because a worker is not remote for
* everything: work reaches a worker only where this side has a producer for
* it, and those are ported one at a time. Without that, the badge would read
* "gpu2 ● ready" on the Dub, Audiobook and Transcripts tabs while 100% of
* that work runs here — the same lie the resolved-answer rule exists to
* prevent, in a place the user cannot even see it happen.
*/
import React, { useCallback, useRef, useState } from 'react';
import { createPortal } from 'react-dom';
@@ -20,6 +27,7 @@ import toast from 'react-hot-toast';
import { useTranslation } from 'react-i18next';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { apiFetch } from '../api/client';
import { useAppStore } from '../store';
// Two cadences: a slow tick that just keeps status honest, and a fast one
// while work is in flight so the task count reads as live rather than lagging
@@ -27,6 +35,29 @@ import { apiFetch } from '../api/client';
const IDLE_REFRESH_MS = 5000;
const BUSY_REFRESH_MS = 1000;
/**
* What the workspace in front of the user actually submits.
*
* A workspace that submits no GPU job of its own — the launchpad, Settings,
* the gallery, OmniDrive — maps to no operation, which asks about the target
* itself rather than about one job. That is exactly what the menu wants when
* you open it from anywhere to pick a machine.
*
* Ids are the control plane's (`worker/routing.py`), not the UI's: `mode` is
* a navigation id and these are units of work, so `studio` and the legacy
* `clone`/`design` modes all submit `tts`.
*/
const OP_BY_MODE = {
generate: 'tts',
studio: 'tts',
clone: 'tts',
design: 'tts',
dub: 'dub',
audiobook: 'audiobook',
stories: 'longform',
transcriptions: 'asr',
};
/** ready → green, busy → amber, offline → red. */
const DOT = {
ready: 'bg-emerald-400',
@@ -90,9 +121,14 @@ export default function GpuTarget() {
});
}, []);
// The surface being rendered, not the machine — see OP_BY_MODE. Part of the
// query key so switching tabs re-resolves instead of showing the previous
// tab's answer until the next poll.
const op = OP_BY_MODE[useAppStore((s) => s.mode)] || '';
const { data } = useQuery({
queryKey: ['workers', 'target'],
queryFn: () => request('/workers/target'),
queryKey: ['workers', 'target', op],
queryFn: () => request(op ? `/workers/target?op=${encodeURIComponent(op)}` : '/workers/target'),
refetchInterval: (query) => {
const t = (query.state?.data?.targets || []).find((x) => x.id === query.state?.data?.target);
return t && t.active_tasks > 0 ? BUSY_REFRESH_MS : IDLE_REFRESH_MS;
@@ -114,18 +150,48 @@ export default function GpuTarget() {
? targets.find((x) => x.id === active.worker_id)
: targets.find((x) => x.is_local);
const label = active?.remote ? active.label : t('gpu.local', { defaultValue: 'Local' });
const fellBack = !active?.remote && chosen !== 'local';
// What a worker can be sent at all. Absent (an older control plane, or a
// response that predates op-awareness) means "don't claim anything" — the
// coverage line and the unported reason simply do not render.
const remoteOps = data?.remote_operations || [];
const opLabel = (id) => t(`gpu.ops.${id}`, { defaultValue: id });
// Chosen a worker, on a surface nothing would ever send it. Not a failure
// and not the worker's fault, so it is deliberately NOT `fellBack`: no
// amber, because there is nothing wrong to warn about.
const opIsLocalOnly =
Boolean(op) && chosen !== 'local' && remoteOps.length > 0 && !remoteOps.includes(op);
const coverage = remoteOps.length
? t('gpu.coverage', {
ops: remoteOps.map(opLabel).join(', '),
defaultValue: '{{ops}} only',
})
: '';
const reason = opIsLocalOnly
? t('gpu.opLocal', {
op: opLabel(op),
defaultValue: 'Local — {{op}} does not run remotely yet',
})
: active?.reason || '';
const fellBack = !active?.remote && chosen !== 'local' && !opIsLocalOnly;
// When the chosen worker is unreachable the work runs here, but the DOT must
// report the worker's state — a green dot beside "Local" would hide that
// the machine you picked is down.
const dotStatus = fellBack ? chosenTarget?.status || 'offline' : activeTarget?.status || 'ready';
// the machine you picked is down. Same on an unported surface: the machine
// is fine, and the reason line is what says why it is idle.
const dotStatus =
fellBack || opIsLocalOnly ? chosenTarget?.status || 'offline' : activeTarget?.status || 'ready';
const chipLatency = latencyLabel(active?.remote ? activeTarget : null);
const choose = async (id) => {
setOpen(false);
try {
const next = await request('/workers/target', { method: 'POST', body: { target: id } });
queryClient.setQueryData(['workers', 'target'], next);
// POST answers for the target as a whole. Writing that into an
// op-scoped cache entry would paint "gpu2 ● ready" on a tab whose work
// is local until the next poll corrected it — so it seeds the cache
// only where the two questions are the same one, and the invalidate
// below refreshes the rest.
if (!op) queryClient.setQueryData(['workers', 'target', ''], next);
queryClient.invalidateQueries({ queryKey: ['workers'] });
} catch (e) {
toast.error(e?.message || String(e));
@@ -138,7 +204,7 @@ export default function GpuTarget() {
ref={buttonRef}
type="button"
onClick={toggle}
title={active?.reason || undefined}
title={reason || undefined}
aria-label={t('gpu.picker', { defaultValue: 'Where jobs run' })}
className="inline-flex items-center gap-1.5 rounded px-2 py-1 text-xs opacity-80 hover:opacity-100"
>
@@ -186,24 +252,39 @@ export default function GpuTarget() {
</span>
{/* The address disambiguates two machines a user named
similarly; the detail says why one cannot be picked; the
task count is what makes "busy" mean something. */}
task count is what makes "busy" mean something; the
coverage is what stops the entry from implying the
machine takes everything. Coverage is dropped while the
worker is unusable — "offline · TTS only" answers a
question the user is not asking yet. */}
{!target.is_local && (
<span className="block truncate opacity-60">
{target.detail
? target.detail
: target.active_tasks > 0
? `${target.endpoint} · ${target.active_tasks}/${target.max_tasks} ${t(
'gpu.tasks',
{ defaultValue: 'tasks' },
)}`
: target.endpoint}
: [
target.active_tasks > 0
? `${target.endpoint} · ${target.active_tasks}/${target.max_tasks} ${t(
'gpu.tasks',
{ defaultValue: 'tasks' },
)}`
: target.endpoint,
coverage,
]
.filter(Boolean)
.join(' · ')}
</span>
)}
</span>
</button>
))}
{fellBack && active?.reason && (
<p className="m-0 px-2 py-1 text-[11px] text-amber-400">{active.reason}</p>
{reason && (fellBack || opIsLocalOnly) && (
<p
className={`m-0 px-2 py-1 text-[11px] ${
fellBack ? 'text-amber-400' : 'opacity-60'
}`}
>
{reason}
</p>
)}
</div>
</>,
frontend/src/components/GpuTarget.test.jsx
+117 -1
View File
@@ -13,6 +13,7 @@ const respond = (body, { ok = true, status = 200 } = {}) => ({ ok, status, json:
import toast from 'react-hot-toast';
import GpuTarget from './GpuTarget';
import { useAppStore } from '../store';
const LOCAL = { id: 'local', label: 'Local', available: true, connected: true, is_local: true };
const DESKTOP = {
@@ -33,7 +34,12 @@ function renderPicker() {
);
}
beforeEach(() => vi.clearAllMocks());
beforeEach(() => {
vi.clearAllMocks();
// The picker resolves against the workspace in front of the user, so every
// test states which one it is rather than inheriting the last one's.
useAppStore.setState({ mode: 'studio' });
});
describe('GpuTarget', () => {
it('renders nothing when no worker is enrolled', async () => {
@@ -262,6 +268,116 @@ describe('GpuTarget', () => {
expect(container.querySelector('.bg-red-400')).toBeTruthy();
});
// ── Op awareness ────────────────────────────────────────────────────────
//
// A worker is not remote for everything: work reaches it only where this
// side has a producer, and those are ported one at a time. Without asking
// per operation the badge reads "gpu2 ● ready" on the Dub tab while 100% of
// dubbing runs locally — the exact lie the resolved-answer rule exists to
// prevent, in a place the user cannot see it happen.
const paths = () => apiFetch.mock.calls.filter(([, o]) => !o?.method).map(([p]) => p);
it('asks routing about the operation the current workspace submits', async () => {
useAppStore.setState({ mode: 'dub' });
apiFetch.mockResolvedValue(
respond({ target: 'w1', op: 'dub', active: { remote: false }, targets: [LOCAL, DESKTOP] }),
);
renderPicker();
await waitFor(() => expect(paths()).toContain('/workers/target?op=dub'));
});
it('asks about the target itself where the workspace submits no job', async () => {
// The launchpad renders no GPU work; the menu is being opened to pick a
// machine, not to ask about one job.
useAppStore.setState({ mode: 'launchpad' });
apiFetch.mockResolvedValue(
respond({ target: 'w1', op: '', active: { remote: false }, targets: [LOCAL, DESKTOP] }),
);
renderPicker();
await waitFor(() => expect(paths()).toContain('/workers/target'));
});
it('reads Local, in the user language, on a surface that has no remote path', async () => {
useAppStore.setState({ mode: 'dub' });
apiFetch.mockResolvedValue(
respond({
target: 'w1',
op: 'dub',
// The worker is healthy and chosen — it simply receives no dubbing.
active: { remote: false, label: 'Local', reason: 'dubbing does not run remotely yet' },
remote_operations: ['tts'],
targets: [LOCAL, READY],
}),
);
renderPicker();
expect(await screen.findByText('Local')).toBeInTheDocument();
expect(screen.queryByText('desktop-4090')).not.toBeInTheDocument();
fireEvent.click(screen.getByRole('button'));
// Localized here, not the control plane's English sentence.
expect(
await screen.findByText('Local — dubbing does not run remotely yet'),
).toBeInTheDocument();
});
it('does not warn in amber when the surface simply has no remote path', async () => {
// Nothing is wrong: the machine is fine and this work was never going to
// leave. Amber is reserved for "the worker you picked is down".
useAppStore.setState({ mode: 'dub' });
apiFetch.mockResolvedValue(
respond({
target: 'w1',
op: 'dub',
active: { remote: false, label: 'Local', reason: 'dubbing does not run remotely yet' },
remote_operations: ['tts'],
targets: [LOCAL, READY],
}),
);
const { container } = renderPicker();
await screen.findByText('Local');
expect(container.querySelector('.text-amber-400')).toBeNull();
});
it('says what a worker actually takes, in the menu', async () => {
apiFetch.mockResolvedValue(
respond({
target: 'local',
op: 'tts',
active: { remote: false },
remote_operations: ['tts'],
targets: [LOCAL, READY],
}),
);
renderPicker();
fireEvent.click(await screen.findByRole('button'));
expect(await screen.findByText(/192\.168\.0\.222:2222 · TTS only/)).toBeInTheDocument();
});
it('claims no coverage when the control plane reports none', async () => {
// An older control plane answers without `remote_operations`. Saying
// "TTS only" there would be an invention, and greying the surface out
// would break a working setup.
useAppStore.setState({ mode: 'dub' });
apiFetch.mockResolvedValue(
respond({
target: 'w1',
active: { remote: true, worker_id: 'w1', label: 'desktop-4090' },
targets: [LOCAL, READY],
}),
);
renderPicker();
expect(await screen.findByText('desktop-4090')).toBeInTheDocument();
fireEvent.click(screen.getByRole('button'));
expect(screen.queryByText(/only/i)).not.toBeInTheDocument();
});
it('shows the address and live task count for a busy worker', async () => {
apiFetch.mockResolvedValue(
respond({
frontend/src/hooks/useTTS.js
+25 -1
View File
@@ -5,6 +5,7 @@ import { pickDesignSeed } from '../utils/seed';
import { playBlobAudio, playPing } from '../utils/media';
import {
StreamingPreviewError,
resolveRemoteTtsTarget,
streamGenerateSpeech,
supportsStreamingPreview,
} from '../utils/streamingTts';
@@ -22,6 +23,11 @@ const t = i18next.t.bind(i18next);
// session (module scope, no localStorage); resets on full reload.
let _lastRoutingStatus = null;
// Same de-dup, for "progressive playback is off because your GPU is the one
// across the room". Keyed by worker so switching machines re-announces, while
// ten renders in a row on the same worker say it once.
let _lastStreamingOffWorker = null;
/**
* Encapsulates TTS generation logic, streaming response handling,
* audio ingestion (with trim gate), and preset/tag helpers.
@@ -268,8 +274,26 @@ export default function useTTS({ selectedProfile, setSelectedProfile, loadHistor
// Any MID-stream failure falls back to the classic whole-file flow with
// no user-visible difference beyond the old wait; pre-stream HTTP errors
// (ApiError) throw straight to the shared catch, exactly like before.
//
// Remote GPU is the one case where it must NOT run: the stream is
// rendered by this process, so taking it would silently ignore the
// worker the user picked — a local render dressed as a remote one. The
// classic path below is the one that goes remote, so it wins, and the
// user is told why their progressive playback stopped rather than left
// to conclude the app got slower.
let streamed = false;
if (useAppStore.getState().autoPlayPreview && supportsStreamingPreview()) {
const wantsStreaming = useAppStore.getState().autoPlayPreview && supportsStreamingPreview();
const remoteTarget = wantsStreaming
? await resolveRemoteTtsTarget({ signal: ac.signal })
: null;
if (remoteTarget) {
const who = remoteTarget.label || remoteTarget.workerId || '';
if (who !== _lastStreamingOffWorker) {
_lastStreamingOffWorker = who;
toast(t('tts.streamingOffRemote', { label: who }), { icon: '🖥️', duration: 6000 });
}
}
if (wantsStreaming && !remoteTarget) {
try {
await streamGenerateSpeech(formData, {
signal: ac.signal,
frontend/src/i18n/locales/ar.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "تنبيه: جزء واحد ({{count}}) من النص لم يُنتج أي صوت، لذا هذه النسخة ينقصها. جرّب إعادة التوليد — عادةً ما تحل بذرة مختلفة المشكلة.",
"droppedChunks_other": "تنبيه: {{count}} أجزاء من النص لم تُنتج أي صوت، لذا هذه النسخة تنقصها. جرّب إعادة التوليد — عادةً ما تحل بذرة مختلفة المشكلة.",
"droppedChunksWithText_one": "تنبيه: جزء واحد ({{count}}) من النص لم يُنتج أي صوت وينقص هذه النسخة — «{{text}}». جرّب إعادة التوليد؛ عادةً ما تحل بذرة مختلفة المشكلة.",
"droppedChunksWithText_other": "تنبيه: {{count}} أجزاء من النص لم تُنتج أي صوت وتنقص هذه النسخة — «{{text}}». جرّب إعادة التوليد؛ عادةً ما تحل بذرة مختلفة المشكلة."
"droppedChunksWithText_other": "تنبيه: {{count}} أجزاء من النص لم تُنتج أي صوت وتنقص هذه النسخة — «{{text}}». جرّب إعادة التوليد؛ عادةً ما تحل بذرة مختلفة المشكلة.",
"streamingOffRemote": "المعاينة التدريجية متوقفة أثناء معالجة {{label}} — سيتم تشغيل التسجيل النهائي فور وصوله."
},
"voiceSelector": {
"engineDefault": "المحرك الافتراضي",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "محلي",
"picker": "أين تُنفَّذ المهام",
"tasks": "مهام"
"tasks": "مهام",
"coverage": "{{ops}} فقط",
"opLocal": "محلي — {{op}} لا يعمل عن بُعد بعد",
"ops": {
"tts": "TTS",
"clone": "استنساخ الصوت",
"dub": "الدبلجة",
"audiobook": "معالجة الكتب الصوتية",
"longform": "سرد القصص",
"asr": "التفريغ النصي"
}
}
}
frontend/src/i18n/locales/de.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Hinweis: {{count}} Teil deines Textes hat kein Audio erzeugt, in dieser Aufnahme fehlt er also. Generiere es neu — ein anderer Seed behebt das meistens.",
"droppedChunks_other": "Hinweis: {{count}} Teile deines Textes haben kein Audio erzeugt, in dieser Aufnahme fehlen sie also. Generiere es neu — ein anderer Seed behebt das meistens.",
"droppedChunksWithText_one": "Hinweis: {{count}} Teil deines Textes hat kein Audio erzeugt und fehlt in dieser Aufnahme — „{{text}}“. Generiere es neu; ein anderer Seed behebt das meistens.",
"droppedChunksWithText_other": "Hinweis: {{count}} Teile deines Textes haben kein Audio erzeugt und fehlen in dieser Aufnahme — „{{text}}“. Generiere es neu; ein anderer Seed behebt das meistens."
"droppedChunksWithText_other": "Hinweis: {{count}} Teile deines Textes haben kein Audio erzeugt und fehlen in dieser Aufnahme — „{{text}}“. Generiere es neu; ein anderer Seed behebt das meistens.",
"streamingOffRemote": "Die progressive Vorschau ist deaktiviert, während {{label}} rendert — die fertige Aufnahme wird abgespielt, sobald sie eintrifft."
},
"voiceSelector": {
"engineDefault": "Motorstandard",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Lokal",
"picker": "Wo Aufträge laufen",
"tasks": "Aufgaben"
"tasks": "Aufgaben",
"coverage": "nur {{ops}}",
"opLocal": "Lokal — {{op}} läuft noch nicht auf Remote-Workern",
"ops": {
"tts": "TTS",
"clone": "Stimmklonen",
"dub": "Synchronisation",
"audiobook": "Hörbuch-Rendering",
"longform": "Story-Vertonung",
"asr": "Transkription"
}
}
}
frontend/src/i18n/locales/en.json
+13 -2
View File
@@ -2584,7 +2584,8 @@
"droppedChunks_one": "Heads up: {{count}} part of your text produced no audio, so this take is missing it. Try re-generating — a different seed usually fixes it.",
"droppedChunks_other": "Heads up: {{count}} parts of your text produced no audio, so this take is missing them. Try re-generating — a different seed usually fixes it.",
"droppedChunksWithText_one": "Heads up: {{count}} part of your text produced no audio, so this take is missing it — “{{text}}”. Try re-generating; a different seed usually fixes it.",
"droppedChunksWithText_other": "Heads up: {{count}} parts of your text produced no audio, so this take is missing them — “{{text}}”. Try re-generating; a different seed usually fixes it."
"droppedChunksWithText_other": "Heads up: {{count}} parts of your text produced no audio, so this take is missing them — “{{text}}”. Try re-generating; a different seed usually fixes it.",
"streamingOffRemote": "Progressive preview is off while {{label}} renders this — the finished take plays as soon as it lands."
},
"voiceSelector": {
"engineDefault": "Engine default",
@@ -2902,6 +2903,16 @@
"gpu": {
"local": "Local",
"picker": "Where jobs run",
"tasks": "tasks"
"tasks": "tasks",
"coverage": "{{ops}} only",
"opLocal": "Local — {{op}} does not run remotely yet",
"ops": {
"tts": "TTS",
"clone": "voice cloning",
"dub": "dubbing",
"audiobook": "audiobook rendering",
"longform": "story narration",
"asr": "transcription"
}
}
}
frontend/src/i18n/locales/es.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Aviso: {{count}} parte de tu texto no generó audio, así que falta en esta toma. Prueba a generarlo de nuevo: normalmente basta con otra semilla.",
"droppedChunks_other": "Aviso: {{count}} partes de tu texto no generaron audio, así que faltan en esta toma. Prueba a generarlo de nuevo: normalmente basta con otra semilla.",
"droppedChunksWithText_one": "Aviso: {{count}} parte de tu texto no generó audio y falta en esta toma: «{{text}}». Prueba a generarlo de nuevo; normalmente basta con otra semilla.",
"droppedChunksWithText_other": "Aviso: {{count}} partes de tu texto no generaron audio y faltan en esta toma: «{{text}}». Prueba a generarlo de nuevo; normalmente basta con otra semilla."
"droppedChunksWithText_other": "Aviso: {{count}} partes de tu texto no generaron audio y faltan en esta toma: «{{text}}». Prueba a generarlo de nuevo; normalmente basta con otra semilla.",
"streamingOffRemote": "La vista previa progresiva está desactivada mientras {{label}} genera el audio: la toma final se reproducirá en cuanto llegue."
},
"voiceSelector": {
"engineDefault": "Valor predeterminado del motor",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Local",
"picker": "Dónde se ejecutan los trabajos",
"tasks": "tareas"
"tasks": "tareas",
"coverage": "solo {{ops}}",
"opLocal": "Local: {{op}} aún no se ejecuta en remoto",
"ops": {
"tts": "TTS",
"clone": "clonación de voz",
"dub": "doblaje",
"audiobook": "renderizado de audiolibros",
"longform": "narración de historias",
"asr": "transcripción"
}
}
}
frontend/src/i18n/locales/fr.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Attention : {{count}} partie de votre texte n'a produit aucun audio, elle manque donc dans cette prise. Relancez la génération — une autre graine suffit généralement.",
"droppedChunks_other": "Attention : {{count}} parties de votre texte n'ont produit aucun audio, elles manquent donc dans cette prise. Relancez la génération — une autre graine suffit généralement.",
"droppedChunksWithText_one": "Attention : {{count}} partie de votre texte n'a produit aucun audio et manque dans cette prise — « {{text}} ». Relancez la génération ; une autre graine suffit généralement.",
"droppedChunksWithText_other": "Attention : {{count}} parties de votre texte n'ont produit aucun audio et manquent dans cette prise — « {{text}} ». Relancez la génération ; une autre graine suffit généralement."
"droppedChunksWithText_other": "Attention : {{count}} parties de votre texte n'ont produit aucun audio et manquent dans cette prise — « {{text}} ». Relancez la génération ; une autre graine suffit généralement.",
"streamingOffRemote": "L'aperçu progressif est désactivé pendant que {{label}} effectue le rendu — la prise finale sera lue dès son arrivée."
},
"voiceSelector": {
"engineDefault": "Moteur par défaut",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Local",
"picker": "Où les tâches s'exécutent",
"tasks": "tâches"
"tasks": "tâches",
"coverage": "{{ops}} uniquement",
"opLocal": "Local — {{op}} ne s'exécute pas encore à distance",
"ops": {
"tts": "TTS",
"clone": "clonage vocal",
"dub": "doublage",
"audiobook": "rendu de livre audio",
"longform": "narration d'histoires",
"asr": "transcription"
}
}
}
frontend/src/i18n/locales/hi.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "ध्यान दें: आपके टेक्स्ट का {{count}} हिस्सा कोई ऑडियो नहीं बना पाया, इसलिए यह टेक उससे वंचित है। दोबारा जेनरेट करें — आमतौर पर दूसरा सीड इसे ठीक कर देता है।",
"droppedChunks_other": "ध्यान दें: आपके टेक्स्ट के {{count}} हिस्से कोई ऑडियो नहीं बना पाए, इसलिए यह टेक उनसे वंचित है। दोबारा जेनरेट करें — आमतौर पर दूसरा सीड इसे ठीक कर देता है।",
"droppedChunksWithText_one": "ध्यान दें: आपके टेक्स्ट का {{count}} हिस्सा कोई ऑडियो नहीं बना पाया और इस टेक में नहीं है — “{{text}}”। दोबारा जेनरेट करें; आमतौर पर दूसरा सीड इसे ठीक कर देता है।",
"droppedChunksWithText_other": "ध्यान दें: आपके टेक्स्ट के {{count}} हिस्से कोई ऑडियो नहीं बना पाए और इस टेक में नहीं हैं — “{{text}}”। दोबारा जेनरेट करें; आमतौर पर दूसरा सीड इसे ठीक कर देता है।"
"droppedChunksWithText_other": "ध्यान दें: आपके टेक्स्ट के {{count}} हिस्से कोई ऑडियो नहीं बना पाए और इस टेक में नहीं हैं — “{{text}}”। दोबारा जेनरेट करें; आमतौर पर दूसरा सीड इसे ठीक कर देता है।",
"streamingOffRemote": "जब तक {{label}} रेंडर कर रहा है, प्रोग्रेसिव प्रीव्यू बंद रहेगा — तैयार ऑडियो आते ही चल जाएगा।"
},
"voiceSelector": {
"engineDefault": "इंजन डिफ़ॉल्ट",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "लोकल",
"picker": "काम कहाँ चलते हैं",
"tasks": "कार्य"
"tasks": "कार्य",
"coverage": "केवल {{ops}}",
"opLocal": "लोकल — {{op}} अभी रिमोट पर नहीं चलता",
"ops": {
"tts": "TTS",
"clone": "वॉइस क्लोनिंग",
"dub": "डबिंग",
"audiobook": "ऑडियोबुक रेंडरिंग",
"longform": "कहानी वर्णन",
"asr": "ट्रांसक्रिप्शन"
}
}
}
frontend/src/i18n/locales/id.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Perhatian: {{count}} bagian teks Anda tidak menghasilkan audio, jadi hasil ini kehilangan bagian tersebut. Coba hasilkan ulang — seed yang berbeda biasanya memperbaikinya.",
"droppedChunks_other": "Perhatian: {{count}} bagian teks Anda tidak menghasilkan audio, jadi hasil ini kehilangan bagian-bagian tersebut. Coba hasilkan ulang — seed yang berbeda biasanya memperbaikinya.",
"droppedChunksWithText_one": "Perhatian: {{count}} bagian teks Anda tidak menghasilkan audio dan hilang dari hasil ini — “{{text}}”. Coba hasilkan ulang; seed yang berbeda biasanya memperbaikinya.",
"droppedChunksWithText_other": "Perhatian: {{count}} bagian teks Anda tidak menghasilkan audio dan hilang dari hasil ini — “{{text}}”. Coba hasilkan ulang; seed yang berbeda biasanya memperbaikinya."
"droppedChunksWithText_other": "Perhatian: {{count}} bagian teks Anda tidak menghasilkan audio dan hilang dari hasil ini — “{{text}}”. Coba hasilkan ulang; seed yang berbeda biasanya memperbaikinya.",
"streamingOffRemote": "Pratinjau progresif dimatikan selama {{label}} merender — hasil akhirnya diputar begitu tiba."
},
"voiceSelector": {
"engineDefault": "Default mesin",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Lokal",
"picker": "Tempat pekerjaan dijalankan",
"tasks": "tugas"
"tasks": "tugas",
"coverage": "hanya {{ops}}",
"opLocal": "Lokal — {{op}} belum berjalan dari jarak jauh",
"ops": {
"tts": "TTS",
"clone": "kloning suara",
"dub": "sulih suara",
"audiobook": "render buku audio",
"longform": "narasi cerita",
"asr": "transkripsi"
}
}
}
frontend/src/i18n/locales/it.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Attenzione: {{count}} parte del testo non ha prodotto audio, quindi manca in questa registrazione. Riprova a generare: di solito basta un seed diverso.",
"droppedChunks_other": "Attenzione: {{count}} parti del testo non hanno prodotto audio, quindi mancano in questa registrazione. Riprova a generare: di solito basta un seed diverso.",
"droppedChunksWithText_one": "Attenzione: {{count}} parte del testo non ha prodotto audio e manca in questa registrazione — «{{text}}». Riprova a generare; di solito basta un seed diverso.",
"droppedChunksWithText_other": "Attenzione: {{count}} parti del testo non hanno prodotto audio e mancano in questa registrazione — «{{text}}». Riprova a generare; di solito basta un seed diverso."
"droppedChunksWithText_other": "Attenzione: {{count}} parti del testo non hanno prodotto audio e mancano in questa registrazione — «{{text}}». Riprova a generare; di solito basta un seed diverso.",
"streamingOffRemote": "L'anteprima progressiva è disattivata mentre {{label}} esegue il rendering: la traccia finale verrà riprodotta appena arriva."
},
"voiceSelector": {
"engineDefault": "Motore difettoso",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Locale",
"picker": "Dove vengono eseguiti i lavori",
"tasks": "attività"
"tasks": "attività",
"coverage": "solo {{ops}}",
"opLocal": "Locale — {{op}} non viene ancora eseguito da remoto",
"ops": {
"tts": "TTS",
"clone": "clonazione vocale",
"dub": "doppiaggio",
"audiobook": "rendering di audiolibri",
"longform": "narrazione di storie",
"asr": "trascrizione"
}
}
}
frontend/src/i18n/locales/ja.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "ご注意: テキストのうち {{count}} 箇所が音声になりませんでした。このテイクにはその部分がありません。生成し直してください — 別のシードで解決することがほとんどです。",
"droppedChunks_other": "ご注意: テキストのうち {{count}} 箇所が音声になりませんでした。このテイクにはそれらの部分がありません。生成し直してください — 別のシードで解決することがほとんどです。",
"droppedChunksWithText_one": "ご注意: テキストのうち {{count}} 箇所が音声にならず、このテイクから欠けています —「{{text}}」。生成し直してください。別のシードで解決することがほとんどです。",
"droppedChunksWithText_other": "ご注意: テキストのうち {{count}} 箇所が音声にならず、このテイクから欠けています —「{{text}}」。生成し直してください。別のシードで解決することがほとんどです。"
"droppedChunksWithText_other": "ご注意: テキストのうち {{count}} 箇所が音声にならず、このテイクから欠けています —「{{text}}」。生成し直してください。別のシードで解決することがほとんどです。",
"streamingOffRemote": "{{label}} がレンダリング中はプログレッシブ再生を無効にします。完成した音声は届き次第再生されます。"
},
"voiceSelector": {
"engineDefault": "エンジンのデフォルト",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "ローカル",
"picker": "ジョブの実行場所",
"tasks": "タスク"
"tasks": "タスク",
"coverage": "{{ops}} のみ",
"opLocal": "ローカル — {{op}}はまだリモートで実行できません",
"ops": {
"tts": "TTS",
"clone": "音声クローン",
"dub": "吹き替え",
"audiobook": "オーディオブック生成",
"longform": "ストーリー朗読",
"asr": "文字起こし"
}
}
}
frontend/src/i18n/locales/ko.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "알림: 텍스트 중 {{count}}개 부분이 오디오로 생성되지 않아 이 테이크에서 빠졌습니다. 다시 생성해 보세요 — 보통 다른 시드로 해결됩니다.",
"droppedChunks_other": "알림: 텍스트 중 {{count}}개 부분이 오디오로 생성되지 않아 이 테이크에서 빠졌습니다. 다시 생성해 보세요 — 보통 다른 시드로 해결됩니다.",
"droppedChunksWithText_one": "알림: 텍스트 중 {{count}}개 부분이 오디오로 생성되지 않아 이 테이크에서 빠졌습니다 — “{{text}}”. 다시 생성해 보세요. 보통 다른 시드로 해결됩니다.",
"droppedChunksWithText_other": "알림: 텍스트 중 {{count}}개 부분이 오디오로 생성되지 않아 이 테이크에서 빠졌습니다 — “{{text}}”. 다시 생성해 보세요. 보통 다른 시드로 해결됩니다."
"droppedChunksWithText_other": "알림: 텍스트 중 {{count}}개 부분이 오디오로 생성되지 않아 이 테이크에서 빠졌습니다 — “{{text}}”. 다시 생성해 보세요. 보통 다른 시드로 해결됩니다.",
"streamingOffRemote": "{{label}}이(가) 렌더링하는 동안 점진적 미리듣기가 꺼집니다. 완성된 결과는 도착하는 즉시 재생됩니다."
},
"voiceSelector": {
"engineDefault": "엔진 기본값",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "로컬",
"picker": "작업 실행 위치",
"tasks": "작업"
"tasks": "작업",
"coverage": "{{ops}}만",
"opLocal": "로컬 — {{op}}은(는) 아직 원격에서 실행되지 않습니다",
"ops": {
"tts": "TTS",
"clone": "음성 복제",
"dub": "더빙",
"audiobook": "오디오북 렌더링",
"longform": "스토리 내레이션",
"asr": "전사"
}
}
}
frontend/src/i18n/locales/nl.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Let op: {{count}} deel van je tekst leverde geen audio op en ontbreekt dus in deze opname. Genereer opnieuw — een andere seed lost dit meestal op.",
"droppedChunks_other": "Let op: {{count}} delen van je tekst leverden geen audio op en ontbreken dus in deze opname. Genereer opnieuw — een andere seed lost dit meestal op.",
"droppedChunksWithText_one": "Let op: {{count}} deel van je tekst leverde geen audio op en ontbreekt in deze opname — “{{text}}”. Genereer opnieuw; een andere seed lost dit meestal op.",
"droppedChunksWithText_other": "Let op: {{count}} delen van je tekst leverden geen audio op en ontbreken in deze opname — “{{text}}”. Genereer opnieuw; een andere seed lost dit meestal op."
"droppedChunksWithText_other": "Let op: {{count}} delen van je tekst leverden geen audio op en ontbreken in deze opname — “{{text}}”. Genereer opnieuw; een andere seed lost dit meestal op.",
"streamingOffRemote": "De progressieve preview staat uit terwijl {{label}} rendert — de voltooide opname speelt zodra die binnen is."
},
"voiceSelector": {
"engineDefault": "Standaard motor",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Lokaal",
"picker": "Waar taken draaien",
"tasks": "taken"
"tasks": "taken",
"coverage": "alleen {{ops}}",
"opLocal": "Lokaal — {{op}} draait nog niet op afstand",
"ops": {
"tts": "TTS",
"clone": "stemklonen",
"dub": "nasynchronisatie",
"audiobook": "audioboeken renderen",
"longform": "verhaalvertelling",
"asr": "transcriptie"
}
}
}
frontend/src/i18n/locales/pl.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Uwaga: {{count}} fragment tekstu nie wygenerował dźwięku, więc brakuje go w tym nagraniu. Wygeneruj ponownie — inne ziarno zwykle rozwiązuje problem.",
"droppedChunks_other": "Uwaga: fragmenty tekstu ({{count}}) nie wygenerowały dźwięku, więc brakuje ich w tym nagraniu. Wygeneruj ponownie — inne ziarno zwykle rozwiązuje problem.",
"droppedChunksWithText_one": "Uwaga: {{count}} fragment tekstu nie wygenerował dźwięku i brakuje go w tym nagraniu — „{{text}}”. Wygeneruj ponownie; inne ziarno zwykle rozwiązuje problem.",
"droppedChunksWithText_other": "Uwaga: fragmenty tekstu ({{count}}) nie wygenerowały dźwięku i brakuje ich w tym nagraniu — „{{text}}”. Wygeneruj ponownie; inne ziarno zwykle rozwiązuje problem."
"droppedChunksWithText_other": "Uwaga: fragmenty tekstu ({{count}}) nie wygenerowały dźwięku i brakuje ich w tym nagraniu — „{{text}}”. Wygeneruj ponownie; inne ziarno zwykle rozwiązuje problem.",
"streamingOffRemote": "Podgląd progresywny jest wyłączony, gdy renderuje {{label}} — gotowe nagranie odtworzy się zaraz po nadejściu."
},
"voiceSelector": {
"engineDefault": "Domyślne ustawienie silnika",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Lokalnie",
"picker": "Gdzie wykonywane są zadania",
"tasks": "zadania"
"tasks": "zadania",
"coverage": "tylko {{ops}}",
"opLocal": "Lokalnie — {{op}} nie działa jeszcze zdalnie",
"ops": {
"tts": "TTS",
"clone": "klonowanie głosu",
"dub": "dubbing",
"audiobook": "renderowanie audiobooków",
"longform": "narracja opowiadań",
"asr": "transkrypcja"
}
}
}
frontend/src/i18n/locales/pt.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Atenção: {{count}} parte do seu texto não gerou áudio, então falta nesta gravação. Tente gerar de novo — outra semente costuma resolver.",
"droppedChunks_other": "Atenção: {{count}} partes do seu texto não geraram áudio, então faltam nesta gravação. Tente gerar de novo — outra semente costuma resolver.",
"droppedChunksWithText_one": "Atenção: {{count}} parte do seu texto não gerou áudio e falta nesta gravação — “{{text}}”. Tente gerar de novo; outra semente costuma resolver.",
"droppedChunksWithText_other": "Atenção: {{count}} partes do seu texto não geraram áudio e faltam nesta gravação — “{{text}}”. Tente gerar de novo; outra semente costuma resolver."
"droppedChunksWithText_other": "Atenção: {{count}} partes do seu texto não geraram áudio e faltam nesta gravação — “{{text}}”. Tente gerar de novo; outra semente costuma resolver.",
"streamingOffRemote": "A pré-visualização progressiva está desativada enquanto {{label}} renderiza — a gravação final será reproduzida assim que chegar."
},
"voiceSelector": {
"engineDefault": "Padrão do mecanismo",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Local",
"picker": "Onde os trabalhos são executados",
"tasks": "tarefas"
"tasks": "tarefas",
"coverage": "apenas {{ops}}",
"opLocal": "Local — {{op}} ainda não é executado remotamente",
"ops": {
"tts": "TTS",
"clone": "clonagem de voz",
"dub": "dublagem",
"audiobook": "renderização de audiolivros",
"longform": "narração de histórias",
"asr": "transcrição"
}
}
}
frontend/src/i18n/locales/ru.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Внимание: {{count}} фрагмент текста не дал звука, поэтому его нет в этой записи. Попробуйте сгенерировать заново — обычно помогает другое зерно.",
"droppedChunks_other": "Внимание: фрагментов текста без звука: {{count}} — их нет в этой записи. Попробуйте сгенерировать заново, обычно помогает другое зерно.",
"droppedChunksWithText_one": "Внимание: {{count}} фрагмент текста не дал звука, и его нет в этой записи — «{{text}}». Попробуйте сгенерировать заново; обычно помогает другое зерно.",
"droppedChunksWithText_other": "Внимание: фрагментов текста без звука: {{count}} — их нет в этой записи: «{{text}}». Попробуйте сгенерировать заново; обычно помогает другое зерно."
"droppedChunksWithText_other": "Внимание: фрагментов текста без звука: {{count}} — их нет в этой записи: «{{text}}». Попробуйте сгенерировать заново; обычно помогает другое зерно.",
"streamingOffRemote": "Прогрессивное воспроизведение отключено, пока {{label}} выполняет рендеринг — готовая запись начнёт играть, как только придёт."
},
"voiceSelector": {
"engineDefault": "Двигатель по умолчанию",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Локально",
"picker": "Где выполняются задачи",
"tasks": "задачи"
"tasks": "задачи",
"coverage": "только {{ops}}",
"opLocal": "Локально — {{op}} пока не выполняется удалённо",
"ops": {
"tts": "TTS",
"clone": "клонирование голоса",
"dub": "дубляж",
"audiobook": "рендеринг аудиокниг",
"longform": "озвучивание историй",
"asr": "расшифровка"
}
}
}
frontend/src/i18n/locales/sv.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Obs: {{count}} del av din text gav inget ljud, så den saknas i den här tagningen. Generera om — ett annat frö brukar lösa det.",
"droppedChunks_other": "Obs: {{count}} delar av din text gav inget ljud, så de saknas i den här tagningen. Generera om — ett annat frö brukar lösa det.",
"droppedChunksWithText_one": "Obs: {{count}} del av din text gav inget ljud och saknas i den här tagningen — ”{{text}}”. Generera om; ett annat frö brukar lösa det.",
"droppedChunksWithText_other": "Obs: {{count}} delar av din text gav inget ljud och saknas i den här tagningen — ”{{text}}”. Generera om; ett annat frö brukar lösa det."
"droppedChunksWithText_other": "Obs: {{count}} delar av din text gav inget ljud och saknas i den här tagningen — ”{{text}}”. Generera om; ett annat frö brukar lösa det.",
"streamingOffRemote": "Progressiv förhandsvisning är av medan {{label}} renderar — den färdiga tagningen spelas upp så snart den kommer."
},
"voiceSelector": {
"engineDefault": "Standard för motor",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Lokalt",
"picker": "Var jobben körs",
"tasks": "uppgifter"
"tasks": "uppgifter",
"coverage": "endast {{ops}}",
"opLocal": "Lokalt — {{op}} körs inte på distans än",
"ops": {
"tts": "TTS",
"clone": "röstkloning",
"dub": "dubbning",
"audiobook": "ljudboksrendering",
"longform": "berättarröst",
"asr": "transkribering"
}
}
}
frontend/src/i18n/locales/th.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "โปรดทราบ: ข้อความของคุณ {{count}} ส่วนไม่ได้สร้างเสียง ไฟล์นี้จึงขาดส่วนนั้นไป ลองสร้างใหม่ — โดยปกติการเปลี่ยน seed จะแก้ได้",
"droppedChunks_other": "โปรดทราบ: ข้อความของคุณ {{count}} ส่วนไม่ได้สร้างเสียง ไฟล์นี้จึงขาดส่วนเหล่านั้นไป ลองสร้างใหม่ — โดยปกติการเปลี่ยน seed จะแก้ได้",
"droppedChunksWithText_one": "โปรดทราบ: ข้อความของคุณ {{count}} ส่วนไม่ได้สร้างเสียงและหายไปจากไฟล์นี้ — “{{text}}” ลองสร้างใหม่ โดยปกติการเปลี่ยน seed จะแก้ได้",
"droppedChunksWithText_other": "โปรดทราบ: ข้อความของคุณ {{count}} ส่วนไม่ได้สร้างเสียงและหายไปจากไฟล์นี้ — “{{text}}” ลองสร้างใหม่ โดยปกติการเปลี่ยน seed จะแก้ได้"
"droppedChunksWithText_other": "โปรดทราบ: ข้อความของคุณ {{count}} ส่วนไม่ได้สร้างเสียงและหายไปจากไฟล์นี้ — “{{text}}” ลองสร้างใหม่ โดยปกติการเปลี่ยน seed จะแก้ได้",
"streamingOffRemote": "ปิดการแสดงตัวอย่างแบบต่อเนื่องขณะที่ {{label}} กำลังเรนเดอร์ — ไฟล์ที่เสร็จแล้วจะเล่นทันทีที่มาถึง"
},
"voiceSelector": {
"engineDefault": "ค่าเริ่มต้นของเครื่องยนต์",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "ในเครื่อง",
"picker": "ตำแหน่งที่งานทำงาน",
"tasks": "งาน"
"tasks": "งาน",
"coverage": "เฉพาะ {{ops}}",
"opLocal": "ในเครื่อง — {{op}} ยังไม่ทำงานบนเครื่องระยะไกล",
"ops": {
"tts": "TTS",
"clone": "การโคลนเสียง",
"dub": "การพากย์",
"audiobook": "การเรนเดอร์หนังสือเสียง",
"longform": "การบรรยายเรื่อง",
"asr": "การถอดเสียง"
}
}
}
frontend/src/i18n/locales/tr.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Dikkat: Metninizin {{count}} bölümü ses üretmedi, bu nedenle bu kayıtta eksik. Yeniden oluşturmayı deneyin — genellikle farklı bir tohum sorunu çözer.",
"droppedChunks_other": "Dikkat: Metninizin {{count}} bölümü ses üretmedi, bu nedenle bu kayıtta eksikler. Yeniden oluşturmayı deneyin — genellikle farklı bir tohum sorunu çözer.",
"droppedChunksWithText_one": "Dikkat: Metninizin {{count}} bölümü ses üretmedi ve bu kayıtta eksik — “{{text}}”. Yeniden oluşturmayı deneyin; genellikle farklı bir tohum sorunu çözer.",
"droppedChunksWithText_other": "Dikkat: Metninizin {{count}} bölümü ses üretmedi ve bu kayıtta eksikler — “{{text}}”. Yeniden oluşturmayı deneyin; genellikle farklı bir tohum sorunu çözer."
"droppedChunksWithText_other": "Dikkat: Metninizin {{count}} bölümü ses üretmedi ve bu kayıtta eksikler — “{{text}}”. Yeniden oluşturmayı deneyin; genellikle farklı bir tohum sorunu çözer.",
"streamingOffRemote": "{{label}} işlerken aşamalı önizleme kapalıdır — tamamlanan kayıt geldiği anda çalınır."
},
"voiceSelector": {
"engineDefault": "Motor varsayılanı",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Yerel",
"picker": "İşlerin çalıştığı yer",
"tasks": "görev"
"tasks": "görev",
"coverage": "yalnızca {{ops}}",
"opLocal": "Yerel — {{op}} henüz uzaktan çalışmıyor",
"ops": {
"tts": "TTS",
"clone": "ses klonlama",
"dub": "dublaj",
"audiobook": "sesli kitap işleme",
"longform": "hikâye seslendirme",
"asr": "deşifre"
}
}
}
frontend/src/i18n/locales/uk.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Увага: {{count}} фрагмент тексту не дав звуку, тому його немає в цьому записі. Спробуйте згенерувати ще раз — зазвичай допомагає інше зерно.",
"droppedChunks_other": "Увага: фрагментів тексту без звуку: {{count}} — їх немає в цьому записі. Спробуйте згенерувати ще раз, зазвичай допомагає інше зерно.",
"droppedChunksWithText_one": "Увага: {{count}} фрагмент тексту не дав звуку, і його немає в цьому записі — «{{text}}». Спробуйте згенерувати ще раз; зазвичай допомагає інше зерно.",
"droppedChunksWithText_other": "Увага: фрагментів тексту без звуку: {{count}} — їх немає в цьому записі: «{{text}}». Спробуйте згенерувати ще раз; зазвичай допомагає інше зерно."
"droppedChunksWithText_other": "Увага: фрагментів тексту без звуку: {{count}} — їх немає в цьому записі: «{{text}}». Спробуйте згенерувати ще раз; зазвичай допомагає інше зерно.",
"streamingOffRemote": "Прогресивний попередній перегляд вимкнено, доки {{label}} виконує рендеринг — готовий запис відтвориться, щойно надійде."
},
"voiceSelector": {
"engineDefault": "Двигун за замовчуванням",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Локально",
"picker": "Де виконуються завдання",
"tasks": "завдання"
"tasks": "завдання",
"coverage": "лише {{ops}}",
"opLocal": "Локально — {{op}} поки не виконується віддалено",
"ops": {
"tts": "TTS",
"clone": "клонування голосу",
"dub": "дубляж",
"audiobook": "рендеринг аудіокниг",
"longform": "озвучення історій",
"asr": "транскрибування"
}
}
}
frontend/src/i18n/locales/vi.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "Lưu ý: {{count}} phần trong văn bản của bạn không tạo ra âm thanh, nên bản thu này bị thiếu phần đó. Hãy thử tạo lại — thường chỉ cần một seed khác.",
"droppedChunks_other": "Lưu ý: {{count}} phần trong văn bản của bạn không tạo ra âm thanh, nên bản thu này bị thiếu các phần đó. Hãy thử tạo lại — thường chỉ cần một seed khác.",
"droppedChunksWithText_one": "Lưu ý: {{count}} phần trong văn bản của bạn không tạo ra âm thanh và bị thiếu trong bản thu này — “{{text}}”. Hãy thử tạo lại; thường chỉ cần một seed khác.",
"droppedChunksWithText_other": "Lưu ý: {{count}} phần trong văn bản của bạn không tạo ra âm thanh và bị thiếu trong bản thu này — “{{text}}”. Hãy thử tạo lại; thường chỉ cần một seed khác."
"droppedChunksWithText_other": "Lưu ý: {{count}} phần trong văn bản của bạn không tạo ra âm thanh và bị thiếu trong bản thu này — “{{text}}”. Hãy thử tạo lại; thường chỉ cần một seed khác.",
"streamingOffRemote": "Xem trước lũy tiến bị tắt trong khi {{label}} kết xuất — bản hoàn chỉnh sẽ phát ngay khi có."
},
"voiceSelector": {
"engineDefault": "Mặc định công cụ",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "Cục bộ",
"picker": "Nơi tác vụ chạy",
"tasks": "tác vụ"
"tasks": "tác vụ",
"coverage": "chỉ {{ops}}",
"opLocal": "Cục bộ — {{op}} chưa chạy từ xa",
"ops": {
"tts": "TTS",
"clone": "nhân bản giọng nói",
"dub": "lồng tiếng",
"audiobook": "kết xuất sách nói",
"longform": "kể chuyện",
"asr": "phiên âm"
}
}
}
frontend/src/i18n/locales/zh-CN.json
+13 -2
View File
@@ -2307,7 +2307,8 @@
"droppedChunks_one": "请注意：文本中有 {{count}} 处没有生成音频，因此这段录音缺少该部分。请重新生成——换一个随机种子通常就能解决。",
"droppedChunks_other": "请注意：文本中有 {{count}} 处没有生成音频，因此这段录音缺少这些部分。请重新生成——换一个随机种子通常就能解决。",
"droppedChunksWithText_one": "请注意：文本中有 {{count}} 处没有生成音频，这段录音缺少该部分——“{{text}}”。请重新生成，换一个随机种子通常就能解决。",
"droppedChunksWithText_other": "请注意：文本中有 {{count}} 处没有生成音频，这段录音缺少这些部分——“{{text}}”。请重新生成，换一个随机种子通常就能解决。"
"droppedChunksWithText_other": "请注意：文本中有 {{count}} 处没有生成音频，这段录音缺少这些部分——“{{text}}”。请重新生成，换一个随机种子通常就能解决。",
"streamingOffRemote": "{{label}} 渲染期间已关闭渐进式预览，完成后会立即播放。"
},
"voiceSelector": {
"engineDefault": "引擎默认",
@@ -2382,6 +2383,16 @@
"gpu": {
"local": "本机",
"picker": "任务运行位置",
"tasks": "任务"
"tasks": "任务",
"coverage": "仅 {{ops}}",
"opLocal": "本机 — {{op}}尚不支持远程运行",
"ops": {
"tts": "TTS",
"clone": "声音克隆",
"dub": "视频配音",
"audiobook": "有声书渲染",
"longform": "故事朗读",
"asr": "转写"
}
}
}
frontend/src/i18n/locales/zh-TW.json
+13 -2
View File
@@ -2300,7 +2300,8 @@
"droppedChunks_one": "請注意：文字中有 {{count}} 處未產生音訊，因此這段錄音缺少該部分。請重新產生——換一個隨機種子通常就能解決。",
"droppedChunks_other": "請注意：文字中有 {{count}} 處未產生音訊，因此這段錄音缺少這些部分。請重新產生——換一個隨機種子通常就能解決。",
"droppedChunksWithText_one": "請注意：文字中有 {{count}} 處未產生音訊，這段錄音缺少該部分——「{{text}}」。請重新產生，換一個隨機種子通常就能解決。",
"droppedChunksWithText_other": "請注意：文字中有 {{count}} 處未產生音訊，這段錄音缺少這些部分——「{{text}}」。請重新產生，換一個隨機種子通常就能解決。"
"droppedChunksWithText_other": "請注意：文字中有 {{count}} 處未產生音訊，這段錄音缺少這些部分——「{{text}}」。請重新產生，換一個隨機種子通常就能解決。",
"streamingOffRemote": "{{label}} 算繪期間已關閉漸進式預覽，完成後會立即播放。"
},
"voiceSelector": {
"engineDefault": "引擎預設",
@@ -2375,6 +2376,16 @@
"gpu": {
"local": "本機",
"picker": "工作執行位置",
"tasks": "工作"
"tasks": "工作",
"coverage": "僅 {{ops}}",
"opLocal": "本機 — {{op}}尚未支援遠端執行",
"ops": {
"tts": "TTS",
"clone": "聲音複製",
"dub": "影片配音",
"audiobook": "有聲書算繪",
"longform": "故事朗讀",
"asr": "轉錄"
}
}
}
frontend/src/test/streamingTts.test.js
+43
View File
@@ -15,6 +15,7 @@ const {
streamGenerateSpeech,
createStreamingChunkPlayer,
supportsStreamingPreview,
resolveRemoteTtsTarget,
decodePcm16Base64,
peaksFromChunkList,
StreamingPreviewError,
@@ -185,6 +186,48 @@ describe('supportsStreamingPreview', () => {
});
});
// ── resolveRemoteTtsTarget ──────────────────────────────────────────────────
//
// The stream is rendered by THIS process, so a progressive preview on a
// remote target would quietly run the job on this machine after the user
// picked their 4090. This is the check that stops that, and it must fail
// open: a picker that cannot be reached must never block a local render.
describe('resolveRemoteTtsTarget', () => {
const json = (body, ok = true) => ({ ok, json: async () => body });
it('asks routing for the tts operation specifically', async () => {
apiFetch.mockResolvedValue(json({ active: { remote: false } }));
await resolveRemoteTtsTarget();
expect(apiFetch.mock.calls[0][0]).toBe('/workers/target?op=tts');
// A dead backend must fail this probe in one round trip, not stall the
// click behind the transport retry ladder.
expect(apiFetch.mock.calls[0][1]).toMatchObject({ retryTransport: false });
});
it('reports the worker when the resolved target is remote', async () => {
apiFetch.mockResolvedValue(
json({ active: { remote: true, worker_id: 'w1', label: 'desktop-4090' } }),
);
expect(await resolveRemoteTtsTarget()).toEqual({ workerId: 'w1', label: 'desktop-4090' });
});
it('answers local for a fallback decision, so streaming stays available', async () => {
apiFetch.mockResolvedValue(
json({ active: { remote: false, reason: 'desktop-4090 is offline — running locally' } }),
);
expect(await resolveRemoteTtsTarget()).toBeNull();
});
it('answers local when the endpoint errors or the backend is unreachable', async () => {
apiFetch.mockResolvedValue(json({ detail: 'nope' }, false));
expect(await resolveRemoteTtsTarget()).toBeNull();
apiFetch.mockRejectedValue(new Error('connection refused'));
expect(await resolveRemoteTtsTarget()).toBeNull();
});
});
// ── streamGenerateSpeech ────────────────────────────────────────────────────
describe('streamGenerateSpeech', () => {
frontend/src/test/useTTSAutoplayPref.test.jsx
+77
View File
@@ -3,6 +3,13 @@ import { renderHook, act } from '@testing-library/react';
import useTTS from '../hooks/useTTS';
import { useAppStore } from '../store';
import { playBlobAudio } from '../utils/media';
import {
resolveRemoteTtsTarget,
streamGenerateSpeech,
supportsStreamingPreview,
} from '../utils/streamingTts';
import { generateSpeech } from '../api/generate';
import toast from 'react-hot-toast';
// #1032: Settings → Appearance "Auto-play preview" ("play the output as soon
// as a render finishes", #666/#667) only gated the WaveformPlayer preview
@@ -42,6 +49,32 @@ vi.mock('../api/generate', async (importOriginal) => {
};
});
// Streaming renders in THIS process, so with a worker selected the classic
// path is the only one that reaches it. The delivery path is chosen here, so
// this is where "did it actually go remote?" is decided.
vi.mock('../utils/streamingTts', async (importOriginal) => {
const actual = await importOriginal();
return {
...actual,
// Defaults to the real answer (jsdom has no Web Audio → classic path), so
// the auto-play tests above keep exercising what they always did; the
// delivery-path tests below turn it on explicitly.
supportsStreamingPreview: vi.fn(actual.supportsStreamingPreview),
streamGenerateSpeech: vi.fn().mockResolvedValue({ id: 'x', audio_path: 'x.wav' }),
resolveRemoteTtsTarget: vi.fn().mockResolvedValue(null),
};
});
vi.mock('react-hot-toast', () => {
const fn = vi.fn();
fn.error = vi.fn();
fn.success = vi.fn();
fn.dismiss = vi.fn();
fn.loading = vi.fn();
fn.custom = vi.fn();
return { default: fn, toast: fn, Toaster: () => null };
});
const hookProps = () => ({
selectedProfile: null,
setSelectedProfile: vi.fn(),
@@ -79,3 +112,47 @@ describe('useTTS auto-play pref (#1032)', () => {
expect(playBlobAudio).not.toHaveBeenCalled();
});
});
describe('useTTS delivery path vs the chosen GPU', () => {
beforeEach(() => {
useAppStore.setState({ autoPlayPreview: true });
vi.mocked(supportsStreamingPreview).mockReturnValue(true);
vi.mocked(streamGenerateSpeech).mockClear();
vi.mocked(generateSpeech).mockClear();
vi.mocked(toast).mockClear();
vi.mocked(resolveRemoteTtsTarget).mockResolvedValue(null);
});
it('streams progressively when the work runs on this machine', async () => {
await runGenerate();
expect(streamGenerateSpeech).toHaveBeenCalledTimes(1);
expect(generateSpeech).not.toHaveBeenCalled();
});
it('takes the classic path when the resolved target is a worker', async () => {
// Streaming would have rendered here — a local job wearing the badge of
// the 4090 the user picked. The classic path is the one that goes remote.
vi.mocked(resolveRemoteTtsTarget).mockResolvedValue({
workerId: 'w1',
label: 'desktop-4090',
});
await runGenerate();
expect(streamGenerateSpeech).not.toHaveBeenCalled();
expect(generateSpeech).toHaveBeenCalledTimes(1);
});
it('says why progressive playback stopped, once per worker', async () => {
// Silently dropping the feature reads as "the app got slower"; the user
// has to be able to connect it to the choice they made.
vi.mocked(resolveRemoteTtsTarget).mockResolvedValue({ workerId: 'w9', label: 'gpu2' });
await runGenerate();
await runGenerate();
const said = vi.mocked(toast).mock.calls.map(([msg]) => msg);
const notices = said.filter((m) => typeof m === 'string' && m.includes('gpu2'));
expect(notices).toHaveLength(1);
expect(notices[0]).toMatch(/gpu2/);
expect(notices[0]).not.toMatch(/streamingOffRemote/); // a real string, not the key
});
});
frontend/src/utils/streamingTts.js
+33
View File
@@ -26,6 +26,7 @@
* falling back would only duplicate the failure).
*/
import { generateSpeech, withTtsInflight } from '../api/generate';
import { apiFetch } from '../api/client';
import { claimTrackedPlayback } from './playback';
/** True when the webview can progressively play PCM chunks (Web Audio). */
@@ -33,6 +34,38 @@ export const supportsStreamingPreview = () =>
typeof window !== 'undefined' &&
typeof (window.AudioContext || window.webkitAudioContext) === 'function';
/**
* The resolved GPU target for synthesis, when it is a remote worker.
*
* `/generate?stream=true` renders in THIS process — there is no streaming
* path to a worker — so a progressive preview would quietly run the job on
* this machine after the user picked their 4090. That is the one thing the
* whole picker exists to prevent, so the caller takes the classic (remote-
* capable) path instead and says why.
*
* Resolved per generate rather than once at mount, deliberately: a worker can
* go to sleep between two clicks, and this asks routing the same question the
* generation path asks, so the two cannot disagree.
*
* Any failure answers "local". Remote workers are opt-in and the endpoint is
* loopback-only; a picker that cannot be reached must never stand between a
* user and an ordinary local render, and transport retries are off so a dead
* backend fails here in one round trip instead of stalling the click.
*
* @returns {Promise<{workerId: string|null, label: string}|null>}
*/
export async function resolveRemoteTtsTarget({ signal } = {}) {
try {
const res = await apiFetch('/workers/target?op=tts', { signal, retryTransport: false });
if (!res?.ok) return null;
const active = (await res.json())?.active;
if (!active?.remote) return null;
return { workerId: active.worker_id || null, label: active.label || '' };
} catch {
return null;
}
}
/** A failure AFTER the stream started — the signal to fall back to the
* classic whole-file /generate flow. */
export class StreamingPreviewError extends Error {
tests/test_dropped_audio_reaches_the_user_1330.py
+3 -2
View File
@@ -155,8 +155,9 @@ def test_the_streaming_path_emits_a_warning_frame_before_done():
# Order matters: a consumer that stops reading at `done` would never see a
# warning emitted after it.
src = _generation_src()
warn = src.index('"type": "warning", "code": "dropped_chunks"')
done = src.index('"type": "done", "id": meta["id"]')
local_stream = src.index("if _dropped_sink:")
warn = src.index('"type": "warning", "code": "dropped_chunks"', local_stream)
done = src.index('"type": "done", "id": meta["id"]', warn)
assert warn < done
tests/test_generation_remote_routing.py
+467
View File
@@ -0,0 +1,467 @@
"""Clicking Synthesize with a remote GPU selected must reach that GPU.
Before this, `/generate` had no idea remote workers existed: `routing.decide()`
only ever painted a header badge and every render ran on the control plane's
own device. These tests pin the four things that made "I picked gpu2" true:
1. The render is dispatched through the GPU gateway with the REMOTE
decision, and nothing local is loaded on the way — no `get_model()`, no
engine instance, no host-capability gate that would refuse a CUDA-only
engine on a Mac.
2. The whole chunked render travels as ONE op. The assignment carries the
seed, chunk size and crossfade so the worker can reproduce
split → generate(seed + i) → crossfaded concat → effect chain; the
control plane does not pre-split and does not dispatch per chunk.
3. `stream=true` — what the desktop UI sends whenever auto-play is on, i.e.
by default — keeps its NDJSON channel but stops previewing per chunk:
coarse worker progress, then the finished take as one chunk. Answering
with the classic WAV shape here would send the client back to a LOCAL
re-render, which is the entire reported bug.
4. The take is provenance-marked exactly once (the worker marks before it
encodes) and the response says where the work ran.
The gateway itself (`services/gpu_gateway.py`) is stubbed: these are tests of
the CALL SITE, and the fake records exactly what the real gateway would be
handed.
"""
import base64
import importlib
import io
import json
import os
import sqlite3
os.environ.setdefault("OMNIVOICE_MODEL", "test")
os.environ.setdefault("OMNIVOICE_DISABLE_FILE_LOG", "1")
import pytest
import soundfile as sf
import torch
LONG_TEXT = (
"The first sentence sets the scene tonight. "
"A second sentence carries the middle part. "
"The third sentence wraps everything up now."
)
@pytest.fixture(autouse=True)
def _hermetic_store(tmp_path, monkeypatch):
"""Pin OUTPUTS_DIR (both module bindings) and the history DB per test.
Same reasoning as tests/test_generate_streaming.py: the router saves
through `api.routers.generation.OUTPUTS_DIR` while a read-back goes via
`core.config.OUTPUTS_DIR`, and a full-suite run can split the two.
"""
import api.routers.generation as gen
import core.config as cfg
outdir = tmp_path / "outputs"
outdir.mkdir()
monkeypatch.setattr(cfg, "OUTPUTS_DIR", str(outdir))
monkeypatch.setattr(gen, "OUTPUTS_DIR", str(outdir))
dbf = tmp_path / "history.db"
def _get_db():
conn = sqlite3.connect(str(dbf))
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys=ON")
return conn
monkeypatch.setitem(gen.ensure_schema.__globals__, "get_db", _get_db)
gen.ensure_schema()
@pytest.fixture()
def client():
from fastapi.testclient import TestClient
from main import app
return TestClient(app, client=("127.0.0.1", 50000))
@pytest.fixture(autouse=True)
def _no_local_gpu(monkeypatch):
"""Any local render attempt in a REMOTE test is the bug under test."""
import api.routers.generation as gen
async def _boom():
raise AssertionError("get_model() ran for a remote render")
monkeypatch.setattr(gen, "get_model", _boom)
def _remote_decision(label="gpu2"):
from worker.routing import Decision
return Decision(remote=True, worker_id="0123456789ab", label=label,
reason="chosen")
def _local_decision(reason="chosen"):
from worker.routing import Decision
return Decision(remote=False, label="Local", reason=reason)
def _worker_wav(seconds=0.4, sample_rate=24000, amplitude=0.31):
"""WAV bytes shaped like what a worker's `_encode` puts on the wire."""
n = int(seconds * sample_rate)
wave = (torch.linspace(-1.0, 1.0, n) * amplitude).numpy()
buf = io.BytesIO()
sf.write(buf, wave, sample_rate, format="WAV", subtype="PCM_16")
return buf.getvalue()
class FakeGateway:
"""Stands in for the scheduler half of `services/gpu_gateway.py`.
Same signature and same return contract as the real `run` — a
`(waveform, sample_rate)` pair for `tts`, from either branch — so a call
site that works against this works against the real one. What it does NOT
do is talk to a control plane; that is the gateway's own test surface.
"""
def __init__(self, *, payload=None, states=(), raises=None):
self.calls = []
self._payload = payload if payload is not None else _worker_wav()
self._states = list(states)
self._raises = raises
async def run(self, op, *, local, remote=None, decision=None, job=None,
admit=False, on_state=None, executor=None, control_plane=None):
self.calls.append({
"op": op, "local": local, "remote": remote, "decision": decision,
"job": job, "admit": admit,
})
if self._raises is not None:
raise self._raises
for state in self._states:
if on_state is not None:
on_state(state)
if getattr(decision, "remote", False) and remote is not None:
import services.gpu_gateway as real
return real.decode_audio_artifact(
real.RemoteResult(task_id="t1", worker_id="w1",
worker_label=getattr(decision, "label", "gpu2"),
path=self._artifact())
)
# The real gateway hands `local.fn` to run_on_gpu_pool_guarded; the
# pool is not what these tests are about, so call it directly.
return local.fn()
def _artifact(self):
import tempfile
handle = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
handle.write(self._payload)
handle.close()
return handle.name
def _install(monkeypatch, gateway, decision):
"""Point the route's gateway calls at the fake, with a fixed decision."""
import services.gpu_gateway as real
monkeypatch.setattr(real, "run", gateway.run)
monkeypatch.setattr(real, "decide", lambda op, **k: decision)
return real
def _post(client, **overrides):
data = {"text": LONG_TEXT, "engine": "omnivoice", "seed": "4242",
"max_chunk_chars": "40", "crossfade_ms": "70"}
data.update(overrides)
return client.post("/generate", data=data)
def _stream_events(client, **overrides):
data = {"text": LONG_TEXT, "engine": "omnivoice", "seed": "4242",
"stream": "true"}
data.update(overrides)
events = []
with client.stream("POST", "/generate", data=data) as r:
assert r.status_code == 200, r.read()
assert r.headers["content-type"].startswith("application/x-ndjson")
headers = dict(r.headers)
for line in r.iter_lines():
if line.strip():
events.append(json.loads(line))
return headers, events
# ── 1. The classic path actually leaves this machine ────────────────────────
def test_remote_target_dispatches_through_the_gateway(client, monkeypatch):
"""The render goes to the worker, and the WAV that comes back is served."""
payload = _worker_wav(seconds=0.5, amplitude=0.27)
gateway = FakeGateway(payload=payload)
_install(monkeypatch, gateway, _remote_decision("gpu2"))
r = _post(client)
assert r.status_code == 200, r.text
assert r.headers["content-type"] == "audio/wav"
assert len(gateway.calls) == 1
call = gateway.calls[0]
assert call["op"] == "tts"
assert call["decision"].remote is True
# Audio identity: the served take is the worker's audio, not a local one.
served, sr = sf.read(io.BytesIO(r.content), dtype="float32", always_2d=True)
expected, expected_sr = sf.read(io.BytesIO(payload), dtype="float32",
always_2d=True)
assert sr == expected_sr
assert served.shape == expected.shape
assert abs(float(served.max()) - float(expected.max())) < 1e-3
def test_remote_render_says_where_it_ran(client, monkeypatch):
"""The existing #21 notice channel carries "this ran on gpu2"."""
_install(monkeypatch, FakeGateway(), _remote_decision("gpu2"))
r = _post(client)
assert r.headers.get("X-OmniVoice-Routing") == "remote"
assert "gpu2" in r.headers.get("X-OmniVoice-Routing-Reason", "")
def test_unavailable_worker_falls_back_quietly_and_names_the_machine(
client, monkeypatch
):
"""Rule 1: pre-dispatch unavailability renders here, with the reason.
The gateway still owns the call — the local branch runs through it — so
the route keeps exactly one dispatch site.
"""
import api.routers.generation as gen
class _Model:
sampling_rate = 24000
async def _get_model():
return _Model()
decision = _local_decision("gpu2 is offline — running locally")
gateway = FakeGateway()
_install(monkeypatch, gateway, decision)
monkeypatch.setattr(gen, "get_model", _get_model)
monkeypatch.setattr(gen, "_run_inference",
lambda *a, **k: torch.zeros(1, 2400))
r = _post(client)
assert r.status_code == 200, r.text
assert r.headers.get("X-OmniVoice-Routing") == "local_fallback"
assert "gpu2 is offline" in r.headers.get("X-OmniVoice-Routing-Reason", "")
assert gateway.calls[0]["local"] is not None
assert gateway.calls[0]["decision"].remote is False
# ── 2. The whole chunked render travels as one op ───────────────────────────
def test_remote_assignment_carries_the_whole_chunked_render(client, monkeypatch):
"""Per-chunk dispatch is rejected, so every knob the chunk loop reads has
to be on the wire — otherwise remote audio silently differs from local."""
gateway = FakeGateway()
_install(monkeypatch, gateway, _remote_decision())
r = _post(client, max_chunk_chars="40", crossfade_ms="70", seed="4242",
effect_preset="broadcast")
assert r.status_code == 200, r.text
assert len(gateway.calls) == 1, "one op per render, never one per chunk"
call = gateway.calls[0]["remote"]
params = call.params
assert call.engine == "omnivoice"
assert call.operation == "tts"
assert params["seed"] == 4242
assert params["max_chunk_chars"] == 40
assert params["crossfade_ms"] == 70
assert params["effect_preset"] == "broadcast"
# The FULL text, unsplit: the split belongs to the worker.
assert params["text"].startswith("The first sentence")
assert params["text"].endswith("everything up now.")
def test_reference_audio_travels_with_the_assignment(client, monkeypatch, tmp_path):
"""A clone's reference lives only on this machine, so the assignment has
to carry it — the transport is what stages it onto the worker."""
gateway = FakeGateway()
_install(monkeypatch, gateway, _remote_decision())
ref = tmp_path / "ref.wav"
sf.write(str(ref), torch.zeros(2400).numpy(), 24000, format="WAV")
asr = importlib.import_module("services.asr_backend")
monkeypatch.setattr(asr, "transcribe_reference", lambda *a, **k: None)
with open(ref, "rb") as fh:
r = client.post(
"/generate",
data={"text": "Hello there.", "engine": "omnivoice",
"ref_text": "hello"},
files={"ref_audio": ("ref.wav", fh.read(), "audio/wav")},
)
assert r.status_code == 200, r.text
call = gateway.calls[0]["remote"]
assert call.params["ref_audio"]
assert call.params["ref_text"] == "hello"
# ── 3. Streaming preview off, progress on ───────────────────────────────────
def test_remote_stream_reports_progress_then_one_chunk(client, monkeypatch):
"""stream=true stays NDJSON (the client asked for it) but stops previewing
per chunk — coarse worker progress, then the finished take."""
states = [
{"phase": "queued", "progress": 0.0, "worker": "gpu2"},
{"phase": "loading", "progress": 0.5, "worker": "gpu2"},
{"phase": "running", "progress": 0.42, "worker": "gpu2"},
]
gateway = FakeGateway(states=states)
_install(monkeypatch, gateway, _remote_decision("gpu2"))
headers, events = _stream_events(client)
kinds = [e["type"] for e in events]
assert kinds.count("chunk") == 1, "one op means one delivered chunk"
assert kinds[-1] == "done"
stages = [(e["stage"], e["detail"]) for e in events if e["type"] == "progress"]
assert [s for s, _ in stages] == ["queued", "loading", "running"]
assert stages[0][1] == "queued on gpu2"
assert stages[1][1] == "loading model on gpu2"
assert stages[2][1] == "generating on gpu2 (42%)"
assert all(e["target"] == "gpu2"
for e in events if e["type"] == "progress")
# Progress lands BEFORE the audio, or it is not progress.
assert kinds.index("progress") < kinds.index("start")
assert headers.get("x-omnivoice-routing") == "remote"
def test_remote_stream_delivers_the_workers_audio(client, monkeypatch):
"""The single chunk is the worker's render, and the take is saved."""
payload = _worker_wav(seconds=0.25, amplitude=0.4)
_install(monkeypatch, FakeGateway(payload=payload), _remote_decision())
_headers, events = _stream_events(client)
chunk = next(e for e in events if e["type"] == "chunk")
done = next(e for e in events if e["type"] == "done")
pcm = base64.b64decode(chunk["pcm"])
assert len(pcm) == 2 * int(0.25 * 24000)
from core.config import OUTPUTS_DIR
assert os.path.exists(os.path.join(OUTPUTS_DIR, done["audio_path"]))
def test_remote_stream_does_not_run_the_local_chunk_loop(client, monkeypatch):
"""The local streaming path must not run at all under a remote target —
it is what made the user's click render on their laptop."""
def _never(*a, **k):
raise AssertionError("the local chunk splitter ran for a remote render")
_install(monkeypatch, FakeGateway(), _remote_decision())
chunked = importlib.import_module("services.chunked_tts")
monkeypatch.setattr(chunked, "split_text_into_chunks", _never)
_headers, events = _stream_events(client)
assert events[-1]["type"] == "done"
# ── 4. Provenance and failure ───────────────────────────────────────────────
def test_remote_take_is_not_marked_twice(client, monkeypatch):
"""The worker marks before it encodes; a second AudioSeal payload over the
first degrades detection of both."""
import api.routers.generation as gen
marked = []
def _mark(*a, **k):
marked.append(k.get("context"))
return a[0]
watermark = importlib.import_module("services.watermark")
monkeypatch.setattr(watermark, "mark_synthetic", _mark)
_install(monkeypatch, FakeGateway(), _remote_decision())
assert _post(client).status_code == 200
assert "generate.finalize" not in marked
def test_watermark_preference_travels_with_the_assignment(client, monkeypatch):
"""The requesting user's preference governs, not the GPU owner's."""
watermark = importlib.import_module("services.watermark")
monkeypatch.setattr(watermark, "is_enabled", lambda: False)
gateway = FakeGateway()
_install(monkeypatch, gateway, _remote_decision())
assert _post(client).status_code == 200
assert gateway.calls[0]["remote"].params["watermark"] is False
def test_midjob_remote_failure_is_reported_not_silently_redone(client, monkeypatch):
"""Rule 2: minutes already spent elsewhere are not silently respent here."""
from services.gpu_gateway import RemoteJobFailed
gateway = FakeGateway(raises=RemoteJobFailed(
"gpu2 did not finish this job: the worker went away",
worker_label="gpu2", hint="Run it on this machine instead.",
))
_install(monkeypatch, gateway, _remote_decision("gpu2"))
r = _post(client)
assert r.status_code == 503, r.text
assert "gpu2" in r.json()["detail"]
assert "Run it on this machine instead." in r.json()["detail"]
assert r.headers.get("X-OmniVoice-Retryable") == "true"
assert r.headers.get("X-OmniVoice-Routing") == "remote_failed"
def test_remote_render_is_not_refused_by_this_hosts_capabilities(
client, monkeypatch
):
"""A CUDA-only engine sent to a 4090 must not be 400-ed because the
control plane is a Mac — the gate describes a machine doing nothing."""
def _unavailable(*a, **k):
raise AssertionError("the local host-capability gate ran for a remote render")
gateway = FakeGateway()
_install(monkeypatch, gateway, _remote_decision())
engine_routing = importlib.import_module("services.engine_routing")
monkeypatch.setattr(engine_routing, "resolve_routing", _unavailable)
r = _post(client)
assert r.status_code == 200, r.text
assert gateway.calls[0]["decision"].remote is True
# ── 5. The streaming socket says it stays here ──────────────────────────────
def test_ws_tts_says_it_runs_on_this_machine(client, monkeypatch):
"""/ws/tts has no remote form: latency is the whole point of the route.
Staying silent would let the header badge imply the 4090 is doing it."""
from worker import routing as worker_routing
monkeypatch.setattr(
worker_routing, "decide", lambda **k: _remote_decision("gpu2")
)
frames = []
with client.websocket_connect("/ws/tts") as ws:
ws.send_json({"text": "Hello.", "engine": "definitely-not-an-engine"})
for _ in range(3):
frame = ws.receive_json()
frames.append(frame)
if frame.get("type") in ("done", "error"):
break
routing_frames = [f for f in frames if f.get("type") == "routing"]
assert any(f.get("status") == "local_stream" for f in routing_frames), frames
local_only = next(f for f in routing_frames if f["status"] == "local_stream")
assert "gpu2" in local_only["reason"]
assert "this machine" in local_only["reason"]
tests/test_gpu_gateway.py
+640
View File
@@ -0,0 +1,640 @@
"""The single door to a GPU: same call, either machine, honest fallback.
Three properties carry this module, and each of them was a real defect before
the gateway existed:
* **The target is a parameter.** One decision governs prewarm and run, so a
job cannot pay a local model load and then dispatch remotely.
* **Fallback is three rules.** Nothing-ran falls back quietly; work-that-ran
raises; a multi-unit job falls back per unit and latches after N.
* **Local-only policy stays on the local branch.** `check_gpu_admission`
reads local pool statistics, so under a remote target it would answer 429
about local saturation while the chosen GPU sat idle.
"""
from __future__ import annotations
import asyncio
import io
from concurrent.futures import ThreadPoolExecutor
import pytest
from services import gpu_gateway
from worker.errors import ErrorClass, WorkerError
from worker.lifecycle import Attempt, AttemptState, Task, TaskState
from worker.routing import Decision
REMOTE = Decision(remote=True, worker_id="abc123456789", label="gpu2", reason="chosen")
LOCAL_CHOSEN = Decision(remote=False, reason="chosen")
LOCAL_FALLBACK = Decision(remote=False, reason="gpu2 is offline — running locally")
# ── Fakes ──────────────────────────────────────────────────────────────────
class FakeScheduler:
"""Enough scheduler to exercise submit / wait / cancel / get."""
def __init__(self, *, outcome="completed", result_ref=None, started=True,
error=None, raises=None, delay=0.0):
self.outcome = outcome
self.result_ref = result_ref
self.started = started
self.error = error
self.raises = raises
self.delay = delay
self.tasks: dict[str, Task] = {}
self.submitted: list[dict] = []
self.cancelled: list[tuple[str, str]] = []
def submit(self, **kwargs):
self.submitted.append(kwargs)
task = Task(
task_id=f"t{len(self.submitted)}",
operation=kwargs["operation"],
engine=kwargs["engine"],
model_id=kwargs.get("model_id") or "",
params=kwargs.get("params") or {},
)
self.tasks[task.task_id] = task
return task
def get(self, task_id):
return self.tasks.get(task_id)
def cancel(self, task_id, reason=""):
self.cancelled.append((task_id, reason))
return True
async def wait(self, task_id, timeout=None):
if self.delay:
await asyncio.sleep(self.delay)
if self.raises is not None:
raise self.raises
task = self.tasks[task_id]
attempt = Attempt(
attempt_id="a1", task_id=task_id, worker_id="abc123456789",
session_epoch=1, attempt_number=1,
)
if self.started:
# Accepting is what makes a failure "mid-job" — the boundary the
# whole fallback policy turns on.
attempt.accepted_at = 100.0
attempt.state = AttemptState.RUNNING
task.attempts.append(attempt)
task.state = {
"completed": TaskState.COMPLETED,
"failed": TaskState.FAILED,
"timeout": TaskState.TIMEOUT,
"cancelled": TaskState.CANCELLED,
}[self.outcome]
task.result_ref = self.result_ref
task.error = self.error
return task
class FakePlane:
def __init__(self, scheduler=None, running=True, pool=None):
self.running = running
self.scheduler = scheduler
self.pool = pool
class FakePool:
def __init__(self, worker=None):
self._worker = worker
def get(self, worker_id):
return self._worker
def local_call(value="local", *, boom=None):
def _fn():
if boom is not None:
raise boom
return value
return gpu_gateway.LocalCall(_fn, what="TTS generate")
def remote_call(**kw):
kw.setdefault("engine", "indextts")
kw.setdefault("params", {"text": "hello"})
kw.setdefault("decode", lambda result: "remote")
return gpu_gateway.RemoteCall(**kw)
@pytest.fixture
def pool_executor():
ex = ThreadPoolExecutor(max_workers=1)
yield ex
ex.shutdown(wait=False)
def wav_bytes(seconds=0.25, sample_rate=48_000):
import numpy as np
import soundfile as sf
samples = np.zeros(int(seconds * sample_rate), dtype="float32")
buffer = io.BytesIO()
sf.write(buffer, samples, sample_rate, format="WAV")
return buffer.getvalue()
# ── The local branch ───────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_local_target_runs_on_the_pool(pool_executor):
value = await gpu_gateway.run(
"tts", local=local_call("audio"), remote=remote_call(),
decision=LOCAL_CHOSEN, executor=pool_executor,
)
assert value == "audio"
@pytest.mark.asyncio
async def test_admission_is_local_only(tmp_path, pool_executor, monkeypatch):
"""A remote target must never be refused for LOCAL pool saturation.
`check_gpu_admission` reads local queue depth; calling it on the remote
branch answers "the local GPU worker pool is saturated" while the chosen
4090 is idle — the exact 429 the gateway exists to stop.
"""
calls = []
monkeypatch.setattr(
"services.model_manager.check_gpu_admission",
lambda **kw: calls.append(kw),
)
artifact = tmp_path / "a1.bin"
artifact.write_bytes(b"wav")
plane = FakePlane(FakeScheduler(result_ref=str(artifact)))
await gpu_gateway.run(
"tts", local=local_call(), remote=remote_call(), decision=REMOTE,
admit=True, control_plane=plane, executor=pool_executor,
)
assert calls == []
await gpu_gateway.run(
"tts", local=local_call(), decision=LOCAL_CHOSEN,
admit=True, executor=pool_executor,
)
assert len(calls) == 1
@pytest.mark.asyncio
async def test_prewarm_skips_the_local_load_when_the_work_goes_remote():
"""Warming here before dispatching elsewhere costs minutes and VRAM on the
machine that is not doing the work."""
loaded = []
class Backend:
def ensure_ready(self):
loaded.append(True)
decision = await gpu_gateway.prewarm("tts", backend=Backend(), decision=REMOTE)
assert decision is REMOTE
assert loaded == []
@pytest.mark.asyncio
async def test_prewarm_loads_under_the_model_load_budget(pool_executor, monkeypatch):
loaded = []
class Backend:
def ensure_ready(self):
loaded.append(True)
monkeypatch.setattr("services.model_manager._model_load_timeout", lambda: 30.0)
await gpu_gateway.prewarm(
"tts", backend=Backend(), engine="indextts",
decision=LOCAL_CHOSEN, executor=pool_executor,
)
assert loaded == [True]
@pytest.mark.asyncio
async def test_prewarm_names_a_load_timeout(pool_executor, monkeypatch):
class Backend:
def ensure_ready(self):
import time
time.sleep(0.4)
monkeypatch.setattr("services.model_manager._model_load_timeout", lambda: 0.05)
with pytest.raises(gpu_gateway.ModelLoadTimeout):
await gpu_gateway.prewarm(
"tts", backend=Backend(), engine="indextts",
decision=LOCAL_CHOSEN, executor=pool_executor,
)
@pytest.mark.asyncio
async def test_prewarm_does_not_disguise_pool_saturation(pool_executor, monkeypatch):
"""GpuPoolBusyError is a TimeoutError, and means the opposite thing."""
from services.model_manager import GpuPoolBusyError
async def _busy(*a, **kw):
raise GpuPoolBusyError("saturated", retry_after=12)
monkeypatch.setattr("services.model_manager.run_on_gpu_pool_guarded", _busy)
monkeypatch.setattr("services.model_manager._model_load_timeout", lambda: 30.0)
class Backend:
def ensure_ready(self):
pass
with pytest.raises(GpuPoolBusyError):
await gpu_gateway.prewarm(
"tts", backend=Backend(), decision=LOCAL_CHOSEN, executor=pool_executor
)
# ── The remote branch ──────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_remote_target_submits_and_decodes(tmp_path, pool_executor):
artifact = tmp_path / "a1.bin"
artifact.write_bytes(b"wav")
scheduler = FakeScheduler(result_ref=str(artifact))
plane = FakePlane(scheduler)
value = await gpu_gateway.run(
"tts",
local=local_call("LOCAL RAN"),
remote=remote_call(decode=lambda r: r.read()),
decision=REMOTE, control_plane=plane, executor=pool_executor,
)
assert value == b"wav"
assert scheduler.submitted[0]["engine"] == "indextts"
assert scheduler.submitted[0]["deadline_seconds"] > 0
@pytest.mark.asyncio
async def test_audio_decoder_reads_the_artifacts_own_sample_rate(tmp_path):
"""Assuming 24 kHz plays a 48 kHz engine back at half speed."""
artifact = tmp_path / "a1.wav"
artifact.write_bytes(wav_bytes(sample_rate=48_000))
result = gpu_gateway.RemoteResult(
task_id="t1", worker_id="w", worker_label="gpu2", path=str(artifact)
)
_waveform, sample_rate = gpu_gateway.decode_audio_artifact(result)
assert sample_rate == 48_000
# ── Fallback rule 1: nothing ran, so run it here quietly ───────────────────
@pytest.mark.asyncio
async def test_control_plane_off_falls_back_quietly(pool_executor):
value = await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(None, running=False),
executor=pool_executor,
)
assert value == "LOCAL"
@pytest.mark.asyncio
async def test_queue_full_falls_back_quietly(pool_executor):
from worker.scheduler import QueueFull
class Full(FakeScheduler):
def submit(self, **kwargs):
raise QueueFull("the queue is full")
value = await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(Full()), executor=pool_executor,
)
assert value == "LOCAL"
@pytest.mark.asyncio
async def test_failure_before_any_worker_accepted_falls_back_quietly(pool_executor):
"""A rejected or never-dispatched assignment cost nothing anywhere."""
scheduler = FakeScheduler(outcome="failed", started=False)
value = await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(scheduler), executor=pool_executor,
)
assert value == "LOCAL"
# ── Fallback rule 2: work ran, so say so ───────────────────────────────────
@pytest.mark.asyncio
async def test_mid_job_failure_raises_for_a_single_shot_op(pool_executor):
scheduler = FakeScheduler(
outcome="failed",
error=WorkerError(
error_class=ErrorClass.TIMEOUT, code="EXECUTION_TIMEOUT",
message="Synthesis exceeded its budget", hint="try a shorter input",
),
)
with pytest.raises(gpu_gateway.RemoteJobFailed) as excinfo:
await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(scheduler),
executor=pool_executor,
)
assert excinfo.value.worker_label == "gpu2"
assert excinfo.value.retry_local is True
assert "Synthesis exceeded its budget" in str(excinfo.value)
@pytest.mark.asyncio
async def test_completed_without_an_artifact_is_a_failure_not_a_fallback(pool_executor):
scheduler = FakeScheduler(outcome="completed", result_ref=None)
with pytest.raises(gpu_gateway.RemoteJobFailed):
await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(scheduler),
executor=pool_executor,
)
@pytest.mark.asyncio
async def test_unreadable_artifact_is_a_mid_job_failure(tmp_path, pool_executor):
artifact = tmp_path / "a1.bin"
artifact.write_bytes(b"not audio")
scheduler = FakeScheduler(result_ref=str(artifact))
with pytest.raises(gpu_gateway.RemoteJobFailed) as excinfo:
await gpu_gateway.run(
"tts", local=local_call(),
remote=remote_call(decode=gpu_gateway.decode_audio_artifact),
decision=REMOTE, control_plane=FakePlane(scheduler),
executor=pool_executor,
)
assert excinfo.value.code == "RESULT_UNREADABLE"
# ── Fallback rule 3: per unit, then latch, with one notice ─────────────────
@pytest.mark.asyncio
async def test_multi_unit_job_falls_back_per_unit_and_latches(tmp_path, pool_executor):
scheduler = FakeScheduler(outcome="failed")
plane = FakePlane(scheduler)
job = gpu_gateway.JobRun("tts")
values = []
for _ in range(4):
values.append(
await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, job=job, control_plane=plane,
executor=pool_executor,
)
)
assert values == ["LOCAL"] * 4
# Latched after two consecutive failures: the remaining units never paid
# another remote deadline to rediscover the same dead machine.
assert job.latched_local is True
assert len(scheduler.submitted) == 2
status, reason = job.notice()
assert status == "local_fallback"
assert "gpu2" in reason and "rest of this job ran locally" in reason
@pytest.mark.asyncio
async def test_one_bad_unit_does_not_demote_a_working_worker(tmp_path, pool_executor):
artifact = tmp_path / "a1.bin"
artifact.write_bytes(b"wav")
scheduler = FakeScheduler(outcome="failed")
plane = FakePlane(scheduler)
job = gpu_gateway.JobRun("tts")
await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, job=job, control_plane=plane, executor=pool_executor,
)
scheduler.outcome = "completed"
scheduler.result_ref = str(artifact)
value = await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, job=job, control_plane=plane, executor=pool_executor,
)
assert value == "remote"
assert job.latched_local is False
assert job.consecutive_failures == 0
def test_a_clean_job_has_nothing_to_say():
job = gpu_gateway.JobRun("tts")
job.record_success()
assert job.notice() is None
# ── Abandonment ────────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_caller_cancellation_cancels_the_remote_task(pool_executor):
"""A worker holds its only slot until this side says otherwise."""
scheduler = FakeScheduler(delay=5.0)
plane = FakePlane(scheduler)
task = asyncio.ensure_future(
gpu_gateway.run(
"tts", local=local_call(), remote=remote_call(),
decision=REMOTE, control_plane=plane, executor=pool_executor,
)
)
await asyncio.sleep(0.05)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert [t for t, _ in scheduler.cancelled] == ["t1"]
@pytest.mark.asyncio
async def test_deadline_cancels_the_remote_task_and_reports_it(pool_executor):
scheduler = FakeScheduler(raises=TimeoutError("deadline"), started=True)
scheduler.tasks = {}
plane = FakePlane(scheduler)
with pytest.raises(gpu_gateway._NotDispatched):
# No attempt was ever accepted, so rule 1 applies — but the task is
# still cancelled, which is what stops an orphaned render.
await gpu_gateway._run_remote(remote_call(), REMOTE, control_plane=plane)
assert scheduler.cancelled and scheduler.cancelled[0][0] == "t1"
@pytest.mark.asyncio
async def test_shutdown_does_not_claim_someone_elses_gpu_stopped(pool_executor):
from worker.scheduler import SchedulerStopped
scheduler = FakeScheduler(raises=SchedulerStopped("the control plane stopped"))
value = await gpu_gateway.run(
"tts", local=local_call("LOCAL"), remote=remote_call(),
decision=REMOTE, control_plane=FakePlane(scheduler), executor=pool_executor,
)
# Nothing had been accepted, so this side runs the work; and it never told
# the worker anything, because it no longer can.
assert value == "LOCAL"
assert scheduler.cancelled == []
# ── Coarse progress ────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_remote_progress_is_reported_by_phase(tmp_path, pool_executor, monkeypatch):
"""`workers.py` is poll-only, so without this a five-minute remote render
shows the same bare spinner as a local one."""
monkeypatch.setattr(gpu_gateway, "_POLL_SECONDS", 0.01)
artifact = tmp_path / "a1.bin"
artifact.write_bytes(b"wav")
scheduler = FakeScheduler(result_ref=str(artifact), delay=0.08)
plane = FakePlane(scheduler)
seen: list[dict] = []
async def _drive():
return await gpu_gateway.run(
"tts", local=local_call(), remote=remote_call(),
decision=REMOTE, control_plane=plane, on_state=seen.append,
executor=pool_executor,
)
runner = asyncio.ensure_future(_drive())
await asyncio.sleep(0.03)
task = scheduler.tasks["t1"]
attempt = Attempt(
attempt_id="a0", task_id="t1", worker_id="abc123456789",
session_epoch=1, attempt_number=1, state=AttemptState.MODEL_LOADING,
)
task.attempts.append(attempt)
task.state = TaskState.MODEL_LOADING
await asyncio.sleep(0.05)
await runner
phases = [event["phase"] for event in seen]
assert phases[0] == gpu_gateway.PHASE_QUEUED
assert gpu_gateway.PHASE_LOADING in phases
assert all(event["worker"] == "gpu2" for event in seen)
# ── Notices ────────────────────────────────────────────────────────────────
def test_notice_says_nothing_when_the_user_chose_local():
assert gpu_gateway.notice_for(LOCAL_CHOSEN) is None
def test_notice_names_the_machine_that_was_skipped():
status, reason = gpu_gateway.notice_for(LOCAL_FALLBACK)
assert status == "local_fallback"
assert "gpu2 is offline" in reason
def test_notice_is_header_safe():
"""It rides the X-OmniVoice-Routing channel, which is latin-1."""
from services.engine_routing import header_safe_reason
_status, reason = gpu_gateway.notice_for(REMOTE)
assert header_safe_reason(reason) == reason
# ── Status and downloads ───────────────────────────────────────────────────
class _Record:
capabilities = [
{"engine": "indextts", "model_id": "indextts:default", "supported": True,
"installed": True, "downloaded": True, "resident": False},
{"engine": "mlx-audio", "model_id": "mlx-audio:kokoro", "supported": True,
"installed": True, "downloaded": True, "resident": True},
]
class _Worker:
record = _Record()
@pytest.mark.asyncio
async def test_status_answers_for_the_remote_host_not_this_one(monkeypatch):
"""Asking the local engine layer about another machine is how a UI offers
an engine that only exists here."""
monkeypatch.setattr(
gpu_gateway, "_local_capabilities",
lambda: [{"engine": "local-only", "model_id": "local-only:default"}],
)
plane = FakePlane(FakeScheduler(), pool=FakePool(_Worker()))
answer = await gpu_gateway.status(decision=REMOTE, control_plane=plane)
assert answer["remote"] is True
assert [m["engine"] for m in answer["models"]] == ["indextts", "mlx-audio"]
filtered = await gpu_gateway.status(
"indextts", decision=REMOTE, control_plane=plane
)
assert [m["engine"] for m in filtered["models"]] == ["indextts"]
@pytest.mark.asyncio
async def test_status_falls_back_when_the_worker_dropped(monkeypatch):
monkeypatch.setattr(
gpu_gateway, "_local_capabilities",
lambda: [{"engine": "indextts", "model_id": "indextts:default"}],
)
plane = FakePlane(FakeScheduler(), pool=FakePool(None))
answer = await gpu_gateway.status(decision=REMOTE, control_plane=plane)
assert answer["remote"] is False
assert answer["reason"] == "the chosen worker is not connected"
@pytest.mark.asyncio
async def test_remote_download_refuses_instead_of_downloading_here():
"""Weights fetched onto the wrong machine leave the 4090 as unprepared as
before, having reported success."""
with pytest.raises(gpu_gateway.RemoteUnsupported) as excinfo:
await gpu_gateway.download("k2-fsa/OmniVoice", decision=REMOTE)
assert "gpu2" in str(excinfo.value)
@pytest.mark.asyncio
async def test_local_download_rejects_anything_outside_the_catalog():
with pytest.raises(gpu_gateway.GatewayError):
await gpu_gateway.download("../../etc/passwd", decision=LOCAL_CHOSEN)
@pytest.mark.asyncio
async def test_local_download_delegates_to_the_installer(monkeypatch):
from api.routers.setup import download as installer
from api.routers.setup import models as catalog
monkeypatch.setattr(catalog, "KNOWN_MODELS", [{"repo_id": "acme/tts"}])
seen = {}
async def _install(req):
seen["repo_id"] = req.repo_id
return {"status": "install_started", "repo_id": req.repo_id}
monkeypatch.setattr(installer, "install_model", _install)
answer = await gpu_gateway.download("acme/tts", decision=LOCAL_CHOSEN)
assert seen["repo_id"] == "acme/tts"
assert answer["status"] == "install_started"
# ── One decision per job ───────────────────────────────────────────────────
def test_decide_is_the_single_answer(monkeypatch):
"""prewarm and run take the SAME decision; re-deciding between them could
warm an engine nothing uses, or dispatch after paying a local cold load."""
calls = []
def _decide(plane, *, op=None):
calls.append(op)
return LOCAL_CHOSEN
monkeypatch.setattr("worker.routing.decide", _decide)
assert gpu_gateway.decide("tts") is LOCAL_CHOSEN
assert calls == ["tts"]
tests/test_worker_client_liveness.py
+107 -12
View File
@@ -272,32 +272,40 @@ async def test_an_executor_that_takes_no_reporters_still_runs():
@pytest.mark.asyncio
async def test_an_oversized_result_never_enters_the_redelivery_set():
"""The B9 regression: recorded-then-rejected means it is re-sent on every
reconnect, killing the session each time and stranding every other task."""
async def test_an_oversized_result_uploads_then_enters_the_redelivery_set():
"""Large payload bytes use UploadResult; only the small artifact reference
is retained for control-stream redelivery."""
oversized = b"\0" * (MAX_MESSAGE_BYTES + 1024)
async def execute(assignment, **_):
return {"meta": {"bytes": len(oversized), "inline": False}, "payload": oversized}
class Stub:
async def UploadResult(self, chunks): # noqa: N802
received = 0
async for chunk in chunks:
assert chunk.offset == received
received += len(chunk.data)
return pb.ResultAck(
artifact_id="t-1/a-1.bin", bytes_received=received, committed=True
)
client = _client(execute)
client._stub = Stub()
client._session_token = "session"
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
failed = await wire.until("failed")
result = await wire.until("result")
await _settle(client)
assert client._pending == {}
assert "result" not in wire.kinds()
assert failed.failed.error.code == "RESULT_TOO_LARGE"
# TERMINAL: the size is a property of the output, so every worker in
# the fleet would produce the same frame and be rejected the same way.
assert failed.failed.error.error_class == pb.ERROR_CLASS_TERMINAL
assert not result.result.inline_payload
assert result.result.artifacts[0].artifact_id == "t-1/a-1.bin"
assert list(client._pending) == ["t-1/a-1"]
before = len(wire.frames)
await client._redeliver_pending()
await asyncio.sleep(0.05)
assert len(wire.frames) == before
assert wire.kinds().count("result") == 2
finally:
await wire.close()
@@ -347,6 +355,93 @@ async def test_the_size_gate_measures_the_frame_not_just_the_payload():
await wire.close()
@pytest.mark.asyncio
async def test_an_oversized_payload_with_no_session_fails_without_being_remembered():
"""The upload path's own failure mode, which the size gate never sees.
``_stub`` is None whenever no session is established — between a disconnect
and the next Register, and against a control plane too old to serve
UploadResult at all. The payload is then over the frame cap with nowhere to
go, which is the shape B9 started as.
Two properties matter and they pull in opposite directions. It must not be
remembered: an over-cap frame in ``_pending`` is re-sent on every reconnect,
killing the session each time and stranding every other task. But it must
stay RETRYABLE, unlike the size gate's TERMINAL verdict — nothing about the
output is wrong here, only the route to the control plane, and the very next
attempt has a live session to upload through.
"""
oversized = b"\0" * (MAX_MESSAGE_BYTES + 1024)
async def execute(assignment, **_):
return {"meta": {"bytes": len(oversized)}, "payload": oversized}
client = _client(execute)
assert client._stub is None, "no session — there is nothing to upload through"
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
failed = await wire.until("failed")
await _settle(client)
assert client._pending == {}
assert "result" not in wire.kinds()
assert failed.failed.error.code == "RESULT_UPLOAD_FAILED"
assert failed.failed.error.error_class == pb.ERROR_CLASS_TRANSIENT
# The reconnect that used to re-send the over-cap frame sends nothing.
before = len(wire.frames)
await client._redeliver_pending()
await asyncio.sleep(0.05)
assert len(wire.frames) == before
finally:
await wire.close()
@pytest.mark.asyncio
async def test_a_receiver_that_never_commits_cannot_spin_the_upload_forever():
"""A resume loop is bounded by a count, not by "did the offset change".
The tempting guard — refuse when the receiver repeats the offset it just
gave us — passes a receiver that alternates between two byte counts, and
passes one that advances a handful of bytes per round. Both spin: the first
forever, the second once per few bytes of a payload measured in megabytes.
Neither is distinguishable from a slow-but-honest resume by looking at a
single pair of offsets, so the bound has to be on the number of rounds.
The worker is single-slot by default, so a spin here is not one lost
upload — it is the machine, doing nothing else, until someone restarts it.
"""
payload = b"\0" * (MAX_MESSAGE_BYTES + 1024)
rounds = []
async def execute(assignment, **_):
return {"meta": {"bytes": len(payload)}, "payload": payload}
class Oscillating:
async def UploadResult(self, chunks): # noqa: N802
async for _ in chunks:
pass
rounds.append(len(rounds))
# Alternates, so `resumed != offset` holds on every single round.
return pb.ResultAck(bytes_received=8 if len(rounds) % 2 else 16, committed=False)
client = _client(execute)
client._stub = Oscillating()
client._session_token = "session"
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
failed = await wire.until("failed")
await _settle(client)
assert len(rounds) <= 16, "the upload kept asking to resume"
assert client._pending == {}
assert failed.failed.error.code == "RESULT_UPLOAD_FAILED"
finally:
await wire.close()
def test_result_too_large_is_terminal_in_the_taxonomy():
"""Guards the codec mapping the wire assertion above depends on."""
from worker.transport import codec
tests/test_worker_executor_residency.py
+13 -3
View File
@@ -485,12 +485,22 @@ async def test_the_executor_reports_through_the_callbacks_the_client_injects(
def test_the_client_can_see_that_this_executor_takes_reporters():
"""The transport probes the injected executor's signature rather than
assuming it. A rename here would silently stop every progress frame — and
the only symptom is tasks dying of an expired lease."""
assuming it, so every name here is a wire contract in disguise.
Spelled out literally rather than compared against the transport's own
constant: both sides deriving the set from one symbol would agree with each
other while agreeing with nothing the executor actually accepts.
Each omission fails silently and differently. Drop ``on_progress`` or
``on_model_loading`` and no frame is ever sent — the only symptom is tasks
dying of an expired lease. Drop ``fetch_input`` and the executor cannot pull
the reference audio a clone needs, so it renders *something* and returns it
as success: a plausible wrong result, which is strictly worse.
"""
from worker.transport.client import _accepted_reporter_kwargs
assert _accepted_reporter_kwargs(TaskExecutor().execute) == frozenset(
{"on_progress", "on_model_loading"}
{"on_progress", "on_model_loading", "fetch_input"}
)
tests/test_worker_inputs.py
+493
View File
@@ -0,0 +1,493 @@
"""Getting the user's own files onto the machine that renders them (B5).
Remote cloning could not work, and did not fail either. ``assignment_to_pb``
never populated ``inputs``, so the assignment carried ``ref_audio`` as a path
on the *control plane* — ``~/…/omnivoice_data/voices/x.wav``, which names
nothing on the worker. ``DownloadArtifact`` had no caller, and would have
404'd if it had one: it serves only the artifact directory, while reference
audio lives in ``VOICES_DIR`` or a tempfile. Meanwhile ``clone`` was
advertised as supported.
The failure mode is the quiet one. An engine handed a dead reference path does
not raise — it renders in its default voice, and the user gets audio that is
simply not their clone.
So the path has three halves, and this file covers all three:
* the control plane **stages** every file-valued parameter into the artifact
store under its content hash (one copy per voice, however many clones),
* the assignment **declares** them and carries ids instead of paths,
* the worker **fetches** them and points the parameters at its own copies.
Plus the disk: ``purge_finished`` deleted rows and left every rendered result
and every staged reference clip behind, forever.
"""
from __future__ import annotations
import json
import os
import shutil
import sqlite3
import pytest
from core.path_security import resolve_within
from services import tts_backend
from worker import deadlines, executor as executor_module, task_store
from worker.errors import ErrorClass
from worker.executor import TaskExecutor, TaskFailure
from worker.lifecycle import Attempt, Task, TaskState
from worker.protocol.gen import worker_v1_pb2 as pb
from worker.transport import codec
# ── Fixtures ───────────────────────────────────────────────────────────────
@pytest.fixture
def db(tmp_path, monkeypatch):
"""See test_worker_registry.py: patch the globals the store actually reads,
because tests/backend/conftest.py purges core.* between tests."""
from worker import task_store as ts
db_globals = ts.db_conn.__wrapped__.__globals__
path = str(tmp_path / "userdata.db")
with sqlite3.connect(path) as conn:
conn.executescript(db_globals["_BASE_SCHEMA"])
monkeypatch.setitem(db_globals, "DB_PATH", path)
return path
@pytest.fixture
def artifacts(tmp_path, monkeypatch):
"""The control plane's artifact directory, isolated per test."""
root = tmp_path / "artifacts"
(root / task_store.INPUTS_DIRNAME).mkdir(parents=True)
monkeypatch.setattr(task_store, "artifact_root", lambda **_kw: str(root))
return str(root)
@pytest.fixture
def voice(tmp_path):
"""A reference clip that exists ONLY on the control plane."""
path = tmp_path / "voices" / "my-voice.wav"
path.parent.mkdir(parents=True)
path.write_bytes(b"RIFF" + b"reference audio" * 64)
return str(path)
def _task(task_id="t1", **params) -> Task:
return Task(
task_id=task_id,
operation="clone",
engine="fake-engine",
model_id="fake-model",
params={"text": "hello", **params},
)
def _attempt(task_id="t1", attempt_id="a1") -> Attempt:
return Attempt(
attempt_id=attempt_id, task_id=task_id, worker_id="w1", session_epoch=1, attempt_number=1
)
def _budget() -> deadlines.Deadlines:
return deadlines.Deadlines(
accept_seconds=30,
model_load_seconds=600,
execution_seconds=300,
progress_lease_seconds=120,
result_delivery_seconds=900,
grace_seconds=60,
)
def _assignment(task: Task, *, artifact_root: str) -> pb.TaskAssignment:
return codec.assignment_to_pb(task, _attempt(task.task_id), _budget(), artifact_root=artifact_root)
def _download_from(artifact_root: str, *, corrupt: bool = False):
"""A stand-in for ``DownloadArtifact``, resolved exactly as the server does.
``server._resolve_input`` is ``resolve_within(artifact_dir, artifact_id)``
plus an ``isfile`` check, so a ref this cannot resolve is one the real RPC
would answer with NOT_FOUND.
"""
calls: list[str] = []
async def fetch(ref, destination):
calls.append(ref.artifact_id)
source = resolve_within(artifact_root, ref.artifact_id)
if not os.path.isfile(source):
raise FileNotFoundError(ref.artifact_id)
if corrupt:
with open(destination, "wb") as handle:
handle.write(b"truncated")
return
shutil.copyfile(source, destination)
fetch.calls = calls
return fetch
# ── The worker's engine stack ──────────────────────────────────────────────
@pytest.fixture(autouse=True)
def _live_tts_backend():
"""Re-resolve the module alias post-purge — see test_worker_executor_residency."""
global tts_backend
import services.tts_backend # noqa: PLC0415
tts_backend = services.tts_backend
class _FakeBackend:
"""Records the kwargs the engine was actually called with."""
id = "fake-engine"
display_name = "Fake Engine (test)"
sample_rate = 24_000
last_kwargs: dict = {}
def ensure_ready(self) -> None:
pass
def generate(self, text, **kwargs):
import torch
type(self).last_kwargs = dict(kwargs)
return torch.zeros(240)
@pytest.fixture
def engine(monkeypatch, _live_tts_backend):
_FakeBackend.last_kwargs = {}
monkeypatch.setitem(tts_backend._REGISTRY, "fake-engine", _FakeBackend)
monkeypatch.setattr(tts_backend, "_ENGINE_INSTANCES", {})
monkeypatch.setattr(tts_backend, "_ENGINE_LAST_USED", {})
monkeypatch.setattr(tts_backend, "_ENGINE_IN_USE", {})
from services import watermark
monkeypatch.setattr(watermark, "mark_synthetic", lambda audio, sr, **kw: audio)
return _FakeBackend
# ── Staging, on the control plane ──────────────────────────────────────────
def test_reference_audio_is_copied_into_the_artifact_store(artifacts, voice):
task = _task(ref_audio=voice)
entries = task_store.ensure_staged(task, now=1000.0)
assert len(entries) == 1
entry = entries[0]
assert entry["key"] == "ref_audio"
staged = resolve_within(artifacts, entry["artifact_id"])
assert staged.read_bytes() == open(voice, "rb").read()
assert entry["size_bytes"] == os.path.getsize(voice)
assert entry["sha256"] in entry["artifact_id"]
# The original stays where it is: a local fallback still needs it.
assert task.params["ref_audio"] == voice
assert os.path.isfile(voice)
def test_repeated_clones_of_one_voice_share_a_single_copy(artifacts, voice, tmp_path):
copy = tmp_path / "a-different-name.wav"
shutil.copyfile(voice, copy)
first = task_store.ensure_staged(_task("t1", ref_audio=voice), now=1000.0)
second = task_store.ensure_staged(_task("t2", ref_audio=str(copy)), now=1001.0)
assert first[0]["artifact_id"] == second[0]["artifact_id"], "content hash, not filename"
inputs = os.listdir(os.path.join(artifacts, task_store.INPUTS_DIRNAME))
assert len(inputs) == 1
def test_staging_twice_stages_once(artifacts, voice, monkeypatch):
task = _task(ref_audio=voice)
task_store.ensure_staged(task, now=1000.0)
def _explode(*_a, **_kw):
raise AssertionError("re-staged an input that was already staged")
monkeypatch.setattr(task_store, "stage_input", _explode)
assert len(task_store.ensure_staged(task, now=1001.0)) == 1
def test_a_parameter_that_is_not_a_file_is_left_alone(artifacts):
task = _task(ref_audio="voice-profile-id")
assert task_store.ensure_staged(task, now=1000.0) == []
assert task_store.INPUTS_PARAM_KEY not in task.params
def test_an_unreadable_reference_is_reported_not_swallowed(artifacts, voice, monkeypatch):
def _denied(path, *_a, **_kw):
raise PermissionError("nope")
monkeypatch.setattr(task_store, "_digest", _denied)
with pytest.raises(task_store.InputStagingError):
task_store.ensure_staged(_task(ref_audio=voice), now=1000.0)
def test_submitting_a_task_stages_and_records_its_inputs(db, artifacts, voice):
task_store.create(_task(ref_audio=voice), now=1000.0)
stored = task_store.get("t1")
entries = stored.params[task_store.INPUTS_PARAM_KEY]
assert len(entries) == 1, "the durable row must name what the task owns"
assert resolve_within(artifacts, entries[0]["artifact_id"]).is_file()
# ── The assignment ─────────────────────────────────────────────────────────
def test_the_assignment_declares_its_inputs(artifacts, voice):
"""The regression: ``inputs`` was never populated, by anyone, ever."""
task = _task(ref_audio=voice)
assignment = _assignment(task, artifact_root=artifacts)
assert len(assignment.inputs) == 1
ref = assignment.inputs[0]
assert ref.artifact_id.startswith(task_store.INPUTS_DIRNAME)
assert ref.size_bytes == os.path.getsize(voice)
assert len(ref.sha256) == 64
assert ref.task_id == "t1" and ref.attempt_id == "a1"
def test_no_control_plane_path_reaches_the_worker(artifacts, voice):
task = _task(ref_audio=voice)
assignment = _assignment(task, artifact_root=artifacts)
assert voice not in assignment.params_json, "sent a path that means nothing remotely"
params = json.loads(assignment.params_json)
assert params["ref_audio"] == assignment.inputs[0].artifact_id
assert params["text"] == "hello"
# Staging bookkeeping holds control-plane paths; it stays home.
assert task_store.INPUTS_PARAM_KEY not in params
def test_an_unstageable_input_fails_the_task_instead_of_shipping_a_path(
artifacts, voice, monkeypatch
):
def _boom(*_a, **_kw):
raise task_store.InputStagingError("the reference clip vanished")
monkeypatch.setattr(task_store, "stage_input", _boom)
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
params = json.loads(assignment.params_json)
assert list(assignment.inputs) == []
assert "ref_audio" not in params, "a dead path renders the wrong voice, silently"
assert params[executor_module.INPUT_ERRORS_PARAM]
def test_the_error_key_is_the_one_the_worker_reads():
"""Two modules, one contract, no import between them."""
assert codec._INPUT_ERRORS_KEY == executor_module.INPUT_ERRORS_PARAM
# ── The worker side ────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_a_clone_whose_reference_only_exists_on_the_control_plane_succeeds(
db, artifacts, voice, engine, tmp_path
):
"""The whole point of the phase, end to end.
The reference clip exists on the control plane and nowhere else. Submit,
build the assignment, run it on a worker whose only route to the file is
``DownloadArtifact`` — and the engine must be called with a readable local
copy of the user's actual voice.
"""
task = _task(ref_audio=voice)
task_store.create(task, now=1000.0)
assignment = _assignment(task, artifact_root=artifacts)
fetch = _download_from(artifacts)
worker = TaskExecutor(fetch_input=fetch, input_dir=str(tmp_path / "worker-inputs"))
await worker.execute(assignment, fetch_input=fetch)
used = engine.last_kwargs["ref_audio"]
assert used != voice, "the worker cannot open a control-plane path"
assert os.path.isfile(used)
assert open(used, "rb").read() == open(voice, "rb").read()
assert fetch.calls == [assignment.inputs[0].artifact_id]
@pytest.mark.asyncio
async def test_the_second_clone_of_a_voice_transfers_nothing(
artifacts, voice, engine, tmp_path
):
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
fetch = _download_from(artifacts)
worker = TaskExecutor(fetch_input=fetch, input_dir=str(tmp_path / "worker-inputs"))
await worker.execute(assignment)
await worker.execute(assignment)
assert len(fetch.calls) == 1, "content-addressed cache re-downloaded a clip it held"
@pytest.mark.asyncio
async def test_a_damaged_transfer_is_refused(artifacts, voice, engine, tmp_path):
"""A truncated clip does not fail — it clones silence."""
cache = tmp_path / "worker-inputs"
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
worker = TaskExecutor(fetch_input=_download_from(artifacts, corrupt=True), input_dir=str(cache))
with pytest.raises(TaskFailure) as raised:
await worker.execute(assignment)
assert raised.value.error.code == "INPUT_CORRUPT"
assert list(cache.iterdir()) == [], "a damaged transfer must not be committed"
@pytest.mark.asyncio
async def test_an_unreachable_input_is_retryable_not_terminal(artifacts, voice, engine, tmp_path):
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
shutil.rmtree(os.path.join(artifacts, task_store.INPUTS_DIRNAME))
worker = TaskExecutor(fetch_input=_download_from(artifacts), input_dir=str(tmp_path / "in"))
with pytest.raises(TaskFailure) as raised:
await worker.execute(assignment)
assert raised.value.error.error_class is ErrorClass.TRANSIENT
@pytest.mark.asyncio
async def test_a_worker_that_cannot_fetch_says_so(artifacts, voice, engine):
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
with pytest.raises(TaskFailure) as raised:
await TaskExecutor().execute(assignment)
assert raised.value.error.error_class is ErrorClass.CAPABILITY
assert raised.value.error.code == "INPUT_TRANSFER_UNSUPPORTED"
@pytest.mark.asyncio
async def test_a_staging_error_is_terminal_on_the_worker(engine, tmp_path):
assignment = pb.TaskAssignment(
operation="clone",
engine="fake-engine",
params_json='{"text": "hi", "input_errors": ["the reference clip vanished"]}',
)
with pytest.raises(TaskFailure) as raised:
await TaskExecutor(input_dir=str(tmp_path)).execute(assignment)
assert raised.value.error.error_class is ErrorClass.TERMINAL
assert raised.value.error.code == "INPUT_UNAVAILABLE"
@pytest.mark.asyncio
async def test_a_hostile_filename_cannot_escape_the_cache(artifacts, voice, engine, tmp_path):
"""``filename`` is remote input; only the hash names the local copy."""
cache = tmp_path / "worker-inputs"
assignment = _assignment(_task(ref_audio=voice), artifact_root=artifacts)
assignment.inputs[0].filename = "../../../../pwned.wav"
worker = TaskExecutor(fetch_input=_download_from(artifacts), input_dir=str(cache))
await worker.execute(assignment)
used = engine.last_kwargs["ref_audio"]
assert os.path.dirname(os.path.realpath(used)) == os.path.realpath(str(cache))
assert not (tmp_path.parent / "pwned.wav").exists()
def test_the_worker_input_cache_has_a_ceiling(tmp_path):
directory = tmp_path / "cache"
directory.mkdir()
for index in range(5):
path = directory / f"{index}.bin"
path.write_bytes(b"x" * 100)
os.utime(path, (1000 + index, 1000 + index))
executor_module._prune_input_cache(str(directory), limit_bytes=250)
survivors = sorted(p.name for p in directory.iterdir())
assert survivors == ["3.bin", "4.bin"], "the cache must evict oldest-first"
# ── The disk ───────────────────────────────────────────────────────────────
def _finish(task: Task, *, at: float) -> None:
task.state = TaskState.COMPLETED
task.finished_at = at
def _result_artifact(artifacts: str, task_id: str) -> str:
path = os.path.join(artifacts, task_id, "a1.bin")
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as handle:
handle.write(b"rendered audio")
return path
WEEK = 7 * 24 * 3600
def test_purging_a_task_takes_its_artifacts_with_it(db, artifacts, voice):
"""The regression: rows were purged, bytes were kept — forever."""
task = _task(ref_audio=voice)
task_store.create(task, now=1000.0)
entry = task.params[task_store.INPUTS_PARAM_KEY][0]
staged = str(resolve_within(artifacts, entry["artifact_id"]))
result = _result_artifact(artifacts, "t1")
_finish(task, at=1000.0)
task_store.save(task, now=1000.0)
removed = task_store.purge_finished(now=1000.0 + WEEK + 1)
assert removed == 1
assert task_store.get("t1") is None
assert not os.path.exists(result), "every remote render leaked its output"
assert not os.path.exists(staged), "every remote clone leaked a copy of the voice"
def test_a_voice_another_task_still_uses_survives_the_purge(db, artifacts, voice):
old = _task("t1", ref_audio=voice)
task_store.create(old, now=1000.0)
_finish(old, at=1000.0)
task_store.save(old, now=1000.0)
live = _task("t2", ref_audio=voice)
task_store.create(live, now=1000.0)
staged = str(resolve_within(artifacts, live.params[task_store.INPUTS_PARAM_KEY][0]["artifact_id"]))
task_store.purge_finished(now=1000.0 + WEEK + 1)
assert task_store.get("t1") is None
assert task_store.get("t2") is not None
assert os.path.isfile(staged), "one copy is shared by every clone of that voice"
def test_a_recently_staged_input_is_never_swept(db, artifacts, voice):
task = _task(ref_audio=voice)
task_store.create(task, now=1000.0 + WEEK)
staged = str(resolve_within(artifacts, task.params[task_store.INPUTS_PARAM_KEY][0]["artifact_id"]))
_finish(task, at=1000.0)
task_store.save(task, now=1000.0)
task_store.purge_finished(now=1000.0 + WEEK + 1)
assert os.path.isfile(staged), "swept an input younger than the cutoff"
def test_purge_survives_a_missing_artifact_directory(db, tmp_path, monkeypatch):
monkeypatch.setattr(task_store, "artifact_root", lambda **_kw: str(tmp_path / "gone"))
task = _task()
task_store.create(task, now=1000.0)
_finish(task, at=1000.0)
task_store.save(task, now=1000.0)
assert task_store.purge_finished(now=1000.0 + WEEK + 1) == 1
tests/test_worker_upload_client.py
+564
View File
@@ -0,0 +1,564 @@
"""Worker-side result delivery: upload, integrity, and outbox priority.
Three failures live here, all of them in the gap between "the worker finished
the work" and "the control plane has it":
* **B4** — every result rode the control stream inline against an 8 MiB frame
ceiling, so anything past roughly three minutes of 24 kHz audio could not be
delivered at all. The executor has always computed an inline/upload decision
and the client has always discarded it; ``UploadResult`` was implemented on
the server and called by nobody.
* **The lease during delivery** — a multi-minute upload sent nothing on the
control stream, so a 120 s progress lease expired mid-transfer and killed an
attempt whose audio was already rendered.
* **Head-of-line blocking** — one FIFO outbox put the heartbeat behind
whatever bulk frame was being written, which is how a busy worker gets
declared dead.
These drive a real ``WorkerClient`` against a fake upload stub and read its
outbox: the invariant is what this side puts on which wire, and in what order,
which a real server round trip would only obscure.
"""
from __future__ import annotations
import asyncio
import hashlib
import pytest
from worker.executor import INLINE_LIMIT_BYTES
from worker.identity import WorkerKeypair
from worker.protocol.gen import worker_v1_pb2 as pb
from worker.transport.client import (
MAX_MESSAGE_BYTES,
UPLOAD_STAGE,
WorkerClient,
WorkerConfig,
_Outbox,
)
ENGINE, MODEL, OP = "indextts", "indextts:v2", "tts"
LEASE_SECONDS = 1
ARTIFACT_ID = "t-1/a-1.bin"
# Comfortably past the frame ceiling: before this change, exactly the result
# that came back as a terminal RESULT_TOO_LARGE.
OVERSIZED = b"\x7f" * (MAX_MESSAGE_BYTES + 1024)
class _FakeStub:
"""Records what UploadResult received, and can be told to refuse."""
def __init__(self, *, error: Exception | None = None, chunk_delay: float = 0.0) -> None:
self.chunks: list[pb.ResultChunk] = []
self.calls = 0
self._error = error
self._chunk_delay = chunk_delay
async def UploadResult(self, request_iterator) -> pb.ResultAck: # noqa: N802
self.calls += 1
if self._error is not None:
raise self._error
received = 0
async for chunk in request_iterator:
self.chunks.append(chunk)
received += len(chunk.data)
if self._chunk_delay:
await asyncio.sleep(self._chunk_delay)
if chunk.last:
break
return pb.ResultAck(artifact_id=ARTIFACT_ID, bytes_received=received, committed=True)
@property
def uploaded(self) -> bytes:
return b"".join(c.data for c in self.chunks)
class _ResumingStub(_FakeStub):
async def UploadResult(self, request_iterator) -> pb.ResultAck: # noqa: N802
self.calls += 1
first = None
current = []
async for chunk in request_iterator:
first = first or chunk
self.chunks.append(chunk)
current.append(chunk)
if self.calls == 1:
return pb.ResultAck(bytes_received=2 * 1024 * 1024, committed=False)
received = first.offset + sum(len(c.data) for c in current)
return pb.ResultAck(artifact_id=ARTIFACT_ID, bytes_received=received, committed=True)
def _client(execute, *, stub: _FakeStub | None = None) -> WorkerClient:
config = WorkerConfig(
endpoint="127.0.0.1:1",
cert_fingerprint="",
certificate_pem=b"",
keypair=WorkerKeypair.generate(),
worker_id="w-1",
)
client = WorkerClient(config, execute=execute)
client._session_token = "sess-1"
client._stub = stub
return client
def _returning(payload: bytes, meta: dict | None = None):
async def execute(assignment, **_):
return {"meta": dict(meta or {"ok": True}), "payload": payload}
return execute
@pytest.mark.asyncio
async def test_declared_input_is_downloaded_with_the_session_and_written_locally(tmp_path):
class Stub:
def DownloadArtifact(self, request): # noqa: N802
assert request.session_token == "sess-1"
async def chunks():
yield pb.ArtifactChunk(offset=0, data=b"reference ", last=False)
yield pb.ArtifactChunk(offset=10, data=b"audio", last=True)
return chunks()
client = _client(_returning(b""), stub=Stub())
destination = tmp_path / "voice.part"
await client._fetch_input(pb.ArtifactRef(artifact_id="inputs/voice.wav"), str(destination))
assert destination.read_bytes() == b"reference audio"
def _assignment(*, lease_seconds: int = LEASE_SECONDS) -> pb.TaskAssignment:
return pb.TaskAssignment(
ref=pb.TaskRef(task_id="t-1", attempt_id="a-1", session_epoch=1),
operation=OP,
engine=ENGINE,
model_id=MODEL,
params_json="{}",
deadlines=pb.Deadlines(
accept_seconds=20,
model_load_seconds=600,
execution_seconds=300,
progress_lease_seconds=lease_seconds,
),
)
class _Wire:
"""Drains the client's outbox, recording each frame's arrival time."""
def __init__(self, client: WorkerClient) -> None:
self.frames: list[tuple[float, pb.WorkerMessage]] = []
self._task = asyncio.create_task(self._drain(client))
async def _drain(self, client: WorkerClient) -> None:
loop = asyncio.get_running_loop()
while True:
message = await client._outbox.get()
self.frames.append((loop.time(), message))
def kinds(self) -> list[str]:
return [m.WhichOneof("payload") for _, m in self.frames]
def of(self, kind: str) -> list[pb.WorkerMessage]:
return [m for _, m in self.frames if m.WhichOneof("payload") == kind]
async def until(self, *kinds: str, timeout: float = 20.0) -> pb.WorkerMessage:
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
for _, message in self.frames:
if message.WhichOneof("payload") in kinds:
return message
await asyncio.sleep(0.02)
raise AssertionError(f"no {kinds} frame; saw {self.kinds()}")
async def close(self) -> None:
self._task.cancel()
await asyncio.gather(self._task, return_exceptions=True)
async def _drained(wire: "_Wire", kind: str, count: int, timeout: float = 5.0) -> None:
"""Wait until *count* frames of *kind* have come off the outbox."""
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
if wire.kinds().count(kind) >= count:
assert wire.kinds().count(kind) == count
return
await asyncio.sleep(0.02)
raise AssertionError(f"only {wire.kinds().count(kind)} {kind} frames, wanted {count}")
async def _settle(client: WorkerClient, timeout: float = 20.0) -> None:
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
if not client._running:
return
await asyncio.sleep(0.02)
raise AssertionError("the task never finished")
# ── B4: over the threshold, the result is uploaded rather than refused ─────
@pytest.mark.asyncio
async def test_an_oversized_result_is_uploaded_not_failed():
"""The B4 regression, and the direct reversal of the Phase 0 stopgap.
A result over the frame ceiling used to come back as a terminal
RESULT_TOO_LARGE with the audio thrown away. It must now arrive.
"""
stub = _FakeStub()
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
message = await wire.until("result", "failed")
await _settle(client)
assert message.WhichOneof("payload") == "result", (
f"delivery failed instead of uploading: {message.failed.error.code}"
)
assert stub.uploaded == OVERSIZED
assert not message.result.inline_payload
assert [a.artifact_id for a in message.result.artifacts] == [ARTIFACT_ID]
finally:
await wire.close()
@pytest.mark.asyncio
async def test_pending_holds_only_the_reference_never_the_payload():
"""#B9's other half: an over-cap frame re-sent on every reconnect tears the
session down each time. The bytes are durable on the control plane once the
upload commits, so redelivery must cost one small frame."""
stub = _FakeStub()
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
await wire.until("result")
await _settle(client)
pending = client._pending["t-1/a-1"]
assert pending.inline_payload == b""
assert [a.artifact_id for a in pending.artifacts] == [ARTIFACT_ID]
await client._redeliver_pending()
await _drained(wire, "result", 2)
# The redelivered frame is the one the server would have rejected.
assert wire.of("result")[-1].ByteSize() < MAX_MESSAGE_BYTES
assert stub.calls == 1, "redelivery must not re-upload the bytes"
finally:
await wire.close()
@pytest.mark.asyncio
async def test_a_result_under_the_threshold_still_rides_inline():
"""The upload path must not cost small results their single-frame delivery."""
stub = _FakeStub()
payload = b"\0" * 1024
client = _client(_returning(payload), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
message = await wire.until("result")
await _settle(client)
assert stub.calls == 0
assert message.result.inline_payload == payload
assert not message.result.artifacts
finally:
await wire.close()
# ── Integrity: the ref states what was sent ────────────────────────────────
@pytest.mark.asyncio
async def test_the_artifact_ref_carries_sha256_and_size():
"""Both fields exist in the proto and were populated by nobody, so the
receiver had no way to tell a truncated transfer from a finished one."""
stub = _FakeStub()
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
message = await wire.until("result")
await _settle(client)
ref = message.result.artifacts[0]
assert ref.sha256 == hashlib.sha256(OVERSIZED).hexdigest()
assert ref.size_bytes == len(OVERSIZED)
assert (ref.task_id, ref.attempt_id) == ("t-1", "a-1")
# Every chunk announced the same digest and length up front.
assert {c.ref.sha256 for c in stub.chunks} == {ref.sha256}
# The control stream is already authenticated; don't widen where the
# session token is written by echoing it back on it.
assert ref.session_token == ""
finally:
await wire.close()
@pytest.mark.asyncio
async def test_chunk_offsets_are_bytes_already_sent_and_only_the_last_commits():
"""The receiver checks ``offset`` against the length it holds and commits
only on ``last`` — both are contracts this side has to keep."""
stub = _FakeStub()
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
await wire.until("result")
await _settle(client)
assert len(stub.chunks) > 1, "an 8 MiB result must not be one chunk"
sent = 0
for chunk in stub.chunks:
assert chunk.offset == sent, f"offset {chunk.offset} != {sent} bytes already sent"
assert chunk.data, "an empty chunk carries no progress"
sent += len(chunk.data)
assert chunk.session_token == "sess-1"
assert sent == len(OVERSIZED)
assert [c.last for c in stub.chunks].count(True) == 1
assert stub.chunks[-1].last
# Every chunk has to fit in a frame with room for its own ref.
assert max(len(c.data) for c in stub.chunks) < MAX_MESSAGE_BYTES
finally:
await wire.close()
# ── The negotiated threshold ───────────────────────────────────────────────
def test_the_default_threshold_is_the_executors_and_is_not_spelled_twice():
client = _client(_returning(b""))
assert client.inline_limit() == INLINE_LIMIT_BYTES
@pytest.mark.asyncio
async def test_config_update_lowers_the_threshold_and_is_honoured():
"""``inline_result_threshold_bytes`` is proto field 4, sent by nobody and
read by nobody. A payload that inlined a moment ago must now upload."""
stub = _FakeStub()
payload = b"\0" * 4096
client = _client(_returning(payload), stub=stub)
wire = _Wire(client)
try:
await client._on_server_message(
pb.ServerMessage(config=pb.ConfigUpdate(inline_result_threshold_bytes=1024))
)
assert client.inline_limit() == 1024
await client._on_assignment(_assignment())
message = await wire.until("result")
await _settle(client)
assert stub.uploaded == payload
assert not message.result.inline_payload
finally:
await wire.close()
@pytest.mark.asyncio
async def test_a_config_update_leaves_the_other_fields_alone():
"""One negotiated key must not silently reset another."""
client = _client(_returning(b""))
client.config.max_concurrent_tasks = 3
await client._on_server_message(
pb.ServerMessage(config=pb.ConfigUpdate(inline_result_threshold_bytes=2048))
)
assert client.config.max_concurrent_tasks == 3
await client._on_server_message(
pb.ServerMessage(config=pb.ConfigUpdate(max_concurrent_tasks=2))
)
assert client.inline_limit() == 2048
assert client.config.max_concurrent_tasks == 2
def test_a_negotiated_threshold_cannot_exceed_what_a_frame_holds():
"""Otherwise a generous control plane turns every large result back into
the RESULT_TOO_LARGE this phase exists to remove."""
client = _client(_returning(b""))
client._inline_threshold = 64 * 1024 * 1024
assert client.inline_limit() < MAX_MESSAGE_BYTES
assert client._should_upload(OVERSIZED)
@pytest.mark.asyncio
async def test_a_non_committed_ack_resumes_from_the_server_offset():
stub = _ResumingStub()
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
await wire.until("result")
await _settle(client)
assert stub.calls == 2
starts = [chunk.offset for chunk in stub.chunks]
assert 2 * 1024 * 1024 in starts
finally:
await wire.close()
# ── The lease survives a slow upload ───────────────────────────────────────
@pytest.mark.asyncio
async def test_upload_progress_renews_the_lease_across_a_slow_transfer():
"""A transfer longer than the progress lease used to die mid-delivery.
Asserted as a gap invariant, because the lease expires on the interval
*between* frames — and as a stage, because the control plane keys the much
longer delivery budget off it.
"""
# Eight chunks at a third of a lease each: three leases' worth of transfer.
stub = _FakeStub(chunk_delay=LEASE_SECONDS / 3.0)
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
await wire.until("result", timeout=LEASE_SECONDS * 20)
await _settle(client)
times = [t for t, _ in wire.frames]
gaps = [b - a for a, b in zip(times, times[1:])]
assert max(gaps) < LEASE_SECONDS, (
f"went silent for {max(gaps):.2f}s under a {LEASE_SECONDS}s lease"
)
uploading = [
m.progress for m in wire.of("progress") if m.progress.stage == UPLOAD_STAGE
]
assert len(uploading) >= len(stub.chunks)
# The first one lands before any bytes, so the attempt is on its
# delivery budget before a slow uplink can burn the ordinary lease.
assert uploading[0].progress == 0.0
assert uploading[-1].progress == pytest.approx(1.0)
# Upload progress is a bounded delivery keepalive, not synthesis
# progress; otherwise these frames erase the delivery deadline and
# overwrite the completed 100% synthesis value.
assert all(p.keepalive for p in uploading)
assert [p.progress for p in uploading] == sorted(p.progress for p in uploading)
finally:
await wire.close()
# ── When the upload cannot happen ──────────────────────────────────────────
@pytest.mark.asyncio
async def test_an_upload_failure_falls_back_to_inline_when_it_still_fits():
"""An older control plane without UploadResult, or one stumble, must not
destroy a render that already succeeded."""
payload = b"\0" * (INLINE_LIMIT_BYTES * 2)
stub = _FakeStub(error=RuntimeError("UNIMPLEMENTED"))
client = _client(_returning(payload), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
message = await wire.until("result", "failed")
await _settle(client)
assert message.WhichOneof("payload") == "result"
assert message.result.inline_payload == payload
assert client._pending["t-1/a-1"].inline_payload == payload
finally:
await wire.close()
@pytest.mark.asyncio
async def test_an_upload_failure_above_the_frame_ceiling_is_transient():
"""There is no inline fallback here. TRANSIENT, not TERMINAL: the failure
is the path, not the output, so another worker can succeed — and nothing
undeliverable may enter the redelivery set."""
stub = _FakeStub(error=RuntimeError("connection reset"))
client = _client(_returning(OVERSIZED), stub=stub)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
failed = await wire.until("failed")
await _settle(client)
assert failed.failed.error.code == "RESULT_UPLOAD_FAILED"
assert failed.failed.error.error_class == pb.ERROR_CLASS_TRANSIENT
assert client._pending == {}
assert "result" not in wire.kinds()
finally:
await wire.close()
@pytest.mark.asyncio
async def test_no_session_is_a_failure_not_a_crash():
"""``_stub`` is None between connections; an oversized result finishing in
that window must be reported, not raised into the task's generic handler
as an unclassified error."""
client = _client(_returning(OVERSIZED), stub=None)
wire = _Wire(client)
try:
await client._on_assignment(_assignment())
failed = await wire.until("failed")
await _settle(client)
assert failed.failed.error.code == "RESULT_UPLOAD_FAILED"
assert client._pending == {}
finally:
await wire.close()
# ── The outbox split ───────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_control_frames_overtake_a_queued_bulk_frame():
"""The heartbeat is the whole liveness model; it must not queue behind a
payload that has no bounded size."""
outbox = _Outbox()
await outbox.put(pb.WorkerMessage(result=pb.TaskResult()), bulk=True)
await outbox.put(pb.WorkerMessage(heartbeat=pb.Heartbeat(active_tasks=1)))
await outbox.put(pb.WorkerMessage(pong=pb.Pong(nonce=7)))
order = [(await outbox.get()).WhichOneof("payload") for _ in range(3)]
assert order == ["heartbeat", "pong", "result"]
@pytest.mark.asyncio
async def test_the_outbox_blocks_rather_than_spinning_when_empty():
outbox = _Outbox()
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(outbox.get(), timeout=0.1)
waiter = asyncio.create_task(outbox.get())
await asyncio.sleep(0.05)
await outbox.put(pb.WorkerMessage(pong=pb.Pong(nonce=7)))
assert (await asyncio.wait_for(waiter, timeout=1)).WhichOneof("payload") == "pong"
@pytest.mark.asyncio
async def test_a_result_is_the_only_frame_queued_as_bulk():
"""Everything the control plane uses to decide this worker is alive has to
stay on the fast queue — including the upload's own progress."""
stub = _FakeStub()
client = _client(_returning(OVERSIZED), stub=stub)
try:
await client._on_assignment(_assignment())
deadline = asyncio.get_running_loop().time() + 20
while not client._outbox.bulk.qsize():
assert asyncio.get_running_loop().time() < deadline, "no result was ever queued"
await asyncio.sleep(0.02)
bulk = [client._outbox.bulk.get_nowait() for _ in range(client._outbox.bulk.qsize())]
control = [
client._outbox.control.get_nowait()
for _ in range(client._outbox.control.qsize())
]
assert {m.WhichOneof("payload") for m in bulk} == {"result"}
assert "result" not in {m.WhichOneof("payload") for m in control}
assert {"accepted", "started", "progress"} <= {
m.WhichOneof("payload") for m in control
}
finally:
await _settle(client)
tests/test_worker_upload_server.py
+610
View File
@@ -0,0 +1,610 @@
"""Artifact transfer at the control-plane boundary.
The upload receiver is the one place where a remote peer writes bytes into the
user's filesystem and the app afterwards calls those bytes a finished render.
Every case here is a way that could go wrong without anybody noticing: a
transfer that stops early and is committed anyway, bytes that arrive out of
order and are appended regardless, a digest nobody checks, an artifact with no
ceiling, and — in the other direction — one worker reading the reference audio
staged for another's task.
The RPCs are driven directly rather than over a real stream: what is under test
is the integrity rule, not gRPC.
"""
from __future__ import annotations
import hashlib
import os
import sqlite3
import pytest
import pytest_asyncio
from worker import deadlines as deadline_policy
from worker import identity, registry
from worker.clock import resolve
from worker.errors import ErrorClass, WorkerError
from worker.identity import WorkerKeypair
from worker.lifecycle import AttemptState, TaskState
from worker.pool import WorkerPool
from worker.protocol.gen import worker_v1_pb2 as pb
from worker.scheduler import Scheduler
from worker.transport import codec, server as server_module
from worker.transport.server import SESSION_METADATA_KEY, WorkerServicer
ENGINE, MODEL, OP = "indextts", "IndexTTS-2", "tts"
@pytest.fixture
def db(tmp_path, monkeypatch):
"""Throwaway DB, patched where the stores actually read it."""
from worker import registry as reg
db_globals = reg.db_conn.__wrapped__.__globals__
path = str(tmp_path / "userdata.db")
with sqlite3.connect(path) as conn:
conn.executescript(db_globals["_BASE_SCHEMA"])
monkeypatch.setitem(db_globals, "DB_PATH", path)
return path
class _Aborted(Exception):
"""What a real gRPC ``context.abort`` does: it raises."""
def __init__(self, code, detail: str) -> None:
super().__init__(detail)
self.code = code
self.detail = detail
class _Context:
def __init__(self, token: str = "") -> None:
self.token = token
def peer(self) -> str:
return "ipv4:127.0.0.1:5555"
def invocation_metadata(self):
return ((SESSION_METADATA_KEY, self.token),) if self.token else ()
async def abort(self, code, detail):
raise _Aborted(code, detail)
def _capabilities() -> list[dict]:
return [
{
"engine": ENGINE,
"model_id": MODEL,
"operations": [OP],
"supported": True,
"installed": True,
"downloaded": True,
"resident": False,
"backend": "cuda",
"free_memory_bytes": 24 * 1024**3,
}
]
class _Plane:
"""A servicer with one enrolled worker, driven RPC by RPC."""
def __init__(self, tmp_path) -> None:
self.artifact_dir = str(tmp_path / "artifacts")
self.pool = WorkerPool()
self.scheduler = Scheduler(self.pool)
self.servicer = WorkerServicer(
self.scheduler, self.pool, artifact_dir=self.artifact_dir
)
self.keypair = WorkerKeypair.generate()
self.worker_id = ""
self.epoch = 0
async def register(self) -> None:
token = registry.create_enrollment(endpoint="localhost:1", cert_fingerprint="fp").encode()
challenge, nonce = identity.new_challenge(), identity.new_challenge()
signature = self.keypair.sign(
identity.challenge_message(
challenge=challenge,
worker_id=self.worker_id,
session_epoch=self.epoch,
nonce=nonce,
)
)
response = await self.servicer.Register(
pb.RegisterRequest(
envelope=pb.Envelope(sequence=self.epoch),
protocol_version_min=1,
protocol_version_max=1,
enrollment_token=token,
public_key=self.keypair.public_bytes(),
challenge=challenge,
challenge_signature=signature,
nonce=nonce,
key_id=self.keypair.key_id,
host=codec.host_to_pb({"hostname": "gpu2", "os": "linux", "arch": "x86_64"}),
capabilities=[codec.capability_to_pb(c) for c in _capabilities()],
max_concurrent_tasks=2,
),
_Context(),
)
assert not response.error.code, response.error.code
self.worker_id = response.worker_id
self.epoch = response.session_epoch
@property
def token(self) -> str:
return self.servicer._sessions[self.worker_id].session.token
def running(self):
"""One task assigned to this worker and rendering."""
task = self.scheduler.submit(operation=OP, engine=ENGINE, model_id=MODEL)
assignment = self.scheduler.next_assignment()
assert assignment is not None
attempt = assignment.attempt
self.scheduler.on_accepted(task.task_id, attempt.attempt_id, epoch=attempt.session_epoch)
self.scheduler.on_started(task.task_id, attempt.attempt_id, epoch=attempt.session_epoch)
return task, attempt
def final_path(self, task, attempt) -> str:
return os.path.join(self.artifact_dir, task.task_id, f"{attempt.attempt_id}.bin")
async def upload(self, chunks) -> pb.ResultAck:
return await self.servicer.UploadResult(_aiter(chunks), _Context(self.token))
async def download(self, ref, *, context=None):
collected = []
async for chunk in self.servicer.DownloadArtifact(ref, context or _Context(self.token)):
collected.append(chunk)
return collected
@pytest_asyncio.fixture
async def plane(tmp_path, db):
p = _Plane(tmp_path)
await p.register()
return p
async def _aiter(items):
for item in items:
yield item
def _ref(plane, task, attempt, *, payload=b"", sha256=None, size=None) -> pb.ArtifactRef:
return pb.ArtifactRef(
artifact_id="",
task_id=task.task_id,
attempt_id=attempt.attempt_id,
filename="result.wav",
size_bytes=len(payload) if size is None else size,
sha256=hashlib.sha256(payload).hexdigest() if sha256 is None else sha256,
session_token=plane.token,
)
def _chunks(ref, payload: bytes, *, size: int = 4, last: bool = True):
"""Split ``payload`` into offset-correct chunks."""
out = []
for start in range(0, len(payload), size):
out.append(pb.ResultChunk(ref=ref, offset=start, data=payload[start : start + size]))
if out and last:
out[-1].last = True
return out
# ── Commit only against a verified, complete transfer ──────────────────────
@pytest.mark.asyncio
async def test_a_verified_upload_commits_under_its_own_attempt(plane):
task, attempt = plane.running()
payload = b"rendered audio bytes"
ref = _ref(plane, task, attempt, payload=payload)
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is True
assert ack.bytes_received == len(payload)
final = plane.final_path(task, attempt)
assert open(final, "rb").read() == payload
assert not os.path.exists(f"{final}.part")
# The id handed back is store-relative, and re-resolves to what was
# written — the worker never learns our filesystem layout.
assert not os.path.isabs(ack.artifact_id)
assert plane.servicer._contained_artifact(ack.artifact_id) == final
@pytest.mark.asyncio
async def test_a_stream_that_ends_without_a_last_chunk_commits_nothing(plane):
"""The iterator simply stopping is a truncated transfer, not a result.
This committed whatever had arrived, renamed it into place, and returned
committed=True — so a dropped connection two thirds of the way through a
render delivered two thirds of a render as the finished article.
"""
task, attempt = plane.running()
payload = b"half a render, and then the link died"
ref = _ref(plane, task, attempt, payload=payload)
ack = await plane.upload(_chunks(ref, payload[:12], last=False))
assert ack.committed is False
assert ack.error.code == "UPLOAD_INCOMPLETE"
assert ack.bytes_received == 12
final = plane.final_path(task, attempt)
assert not os.path.exists(final)
# Kept, so the resume has something to resume onto.
assert os.path.getsize(f"{final}.part") == 12
@pytest.mark.asyncio
async def test_a_digest_mismatch_is_never_renamed_into_place(plane):
task, attempt = plane.running()
payload = b"corrupted on the wire"
# Same length, different bytes: only the digest can tell these apart.
ref = _ref(plane, task, attempt, payload=b"what the worker sent!")
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is False
assert ack.error.code == "DIGEST_MISMATCH"
final = plane.final_path(task, attempt)
assert not os.path.exists(final)
# And the bad bytes are gone: a resume must not append onto them.
assert not os.path.exists(f"{final}.part")
@pytest.mark.asyncio
async def test_an_upload_with_no_declared_digest_is_refused_before_any_bytes(plane):
task, attempt = plane.running()
payload = b"unverifiable"
ref = _ref(plane, task, attempt, payload=payload, sha256="")
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is False
assert ack.error.code == "DIGEST_REQUIRED"
assert not os.path.exists(f"{plane.final_path(task, attempt)}.part")
@pytest.mark.asyncio
async def test_a_size_that_disagrees_with_the_bytes_delivered_is_refused(plane):
task, attempt = plane.running()
payload = b"eight..."
ref = _ref(plane, task, attempt, payload=payload, size=len(payload) + 4)
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is False
assert ack.error.code == "SIZE_MISMATCH"
assert not os.path.exists(plane.final_path(task, attempt))
# ── Offsets are checked, and the ack is the bytes-held probe ───────────────
@pytest.mark.asyncio
async def test_a_chunk_at_the_wrong_offset_is_refused_with_the_bytes_held(plane):
"""``chunk.offset`` was read as a truthiness flag and then ignored, so a
gap or an overlap was appended as though it were the next byte."""
task, attempt = plane.running()
payload = b"0123456789abcdef"
ref = _ref(plane, task, attempt, payload=payload)
stream = [
pb.ResultChunk(ref=ref, offset=0, data=payload[:4]),
pb.ResultChunk(ref=ref, offset=999, data=payload[4:], last=True),
]
ack = await plane.upload(stream)
assert ack.committed is False
assert ack.error.code == "OFFSET_MISMATCH"
# The only report of "bytes already held" this RPC can make: one terminal
# ack, carrying the offset to resume from.
assert ack.bytes_received == 4
assert not os.path.exists(plane.final_path(task, attempt))
@pytest.mark.asyncio
async def test_a_resume_hashes_the_bytes_already_on_disk(plane):
"""Otherwise the digest would attest only to the resumed tail — verifying
the half of the file that was never in doubt."""
task, attempt = plane.running()
payload = b"the first half of it | and the second half of it"
ref = _ref(plane, task, attempt, payload=payload)
dropped = await plane.upload(_chunks(ref, payload[:20], last=False))
assert dropped.bytes_received == 20
resumed = await plane.upload(
[pb.ResultChunk(ref=ref, offset=20, data=payload[20:], last=True)]
)
assert resumed.committed is True
assert resumed.bytes_received == len(payload)
assert open(plane.final_path(task, attempt), "rb").read() == payload
@pytest.mark.asyncio
async def test_a_resume_onto_corrupted_held_bytes_still_fails_verification(plane):
task, attempt = plane.running()
payload = b"the first half of it | and the second half of it"
ref = _ref(plane, task, attempt, payload=payload)
await plane.upload(_chunks(ref, b"tampered with here!!", last=False))
resumed = await plane.upload(
[pb.ResultChunk(ref=ref, offset=20, data=payload[20:], last=True)]
)
assert resumed.committed is False
assert resumed.error.code == "DIGEST_MISMATCH"
assert not os.path.exists(plane.final_path(task, attempt))
# ── Ceilings ───────────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_an_upload_past_its_declared_size_is_cut_off(plane):
"""A declared size narrows the cap; it cannot be exceeded by streaming."""
task, attempt = plane.running()
ref = _ref(plane, task, attempt, payload=b"tiny", size=4)
ack = await plane.upload(_chunks(ref, b"very much larger than four bytes"))
assert ack.committed is False
assert ack.error.code == "ARTIFACT_TOO_LARGE"
final = plane.final_path(task, attempt)
assert not os.path.exists(final)
assert not os.path.exists(f"{final}.part")
@pytest.mark.asyncio
async def test_an_undeclared_upload_is_bounded_by_the_artifact_ceiling(plane, monkeypatch):
monkeypatch.setattr(server_module, "MAX_ARTIFACT_BYTES", 8)
task, attempt = plane.running()
ref = _ref(plane, task, attempt, payload=b"sixteen bytes!!!", size=0)
ack = await plane.upload(_chunks(ref, b"sixteen bytes!!!"))
assert ack.committed is False
assert ack.error.code == "ARTIFACT_TOO_LARGE"
@pytest.mark.asyncio
async def test_the_per_task_artifact_budget_is_enforced_across_attempts(plane, monkeypatch):
"""One artifact under the cap, twice, must not add up to more than a task
is allowed to deliver."""
monkeypatch.setattr(server_module, "MAX_TASK_ARTIFACT_BYTES", 24)
task, attempt = plane.running()
payload = b"sixteen bytes!!!"
first = await plane.upload(_chunks(ref := _ref(plane, task, attempt, payload=payload), payload))
assert first.committed is True
assert ref.size_bytes == 16
# A second attempt of the same task, delivering another 16 bytes. CAPACITY
# so the retry can land on the same worker — anything else excludes it.
plane.scheduler.on_failed(
task.task_id,
attempt.attempt_id,
WorkerError(error_class=ErrorClass.CAPACITY, code="RETRY", message="again"),
epoch=attempt.session_epoch,
)
retry = plane.scheduler.next_assignment()
assert retry is not None
plane.scheduler.on_accepted(task.task_id, retry.attempt.attempt_id, epoch=retry.attempt.session_epoch)
plane.scheduler.on_started(task.task_id, retry.attempt.attempt_id, epoch=retry.attempt.session_epoch)
ack = await plane.upload(
_chunks(_ref(plane, task, retry.attempt, payload=payload), payload)
)
assert ack.committed is False
assert ack.error.code == "TASK_BUDGET_EXCEEDED"
# ── The delivery phase actually exists ─────────────────────────────────────
@pytest.mark.asyncio
async def test_an_upload_moves_the_task_into_result_uploading(plane):
"""``Task.uploading`` had zero callers, so RESULT_UPLOADING was
unreachable and every byte of delivery ran under the execution phase."""
task, attempt = plane.running()
payload = b"0123456789abcdef"
ref = _ref(plane, task, attempt, payload=payload)
seen: list[TaskState] = []
async def observed():
for chunk in _chunks(ref, payload):
yield chunk
seen.append(task.state)
ack = await plane.servicer.UploadResult(observed(), _Context(plane.token))
assert ack.committed is True
assert seen[0] is TaskState.RESULT_UPLOADING
assert attempt.state is AttemptState.UPLOADING
@pytest.mark.asyncio
async def test_the_upload_lease_runs_on_the_result_delivery_budget(plane):
"""A slow delivery is bounded by ``result_delivery_seconds`` (900s), not by
the execution budget it used to inherit — which is what made a large
upload die mid-transfer under the 120s progress lease."""
task, attempt = plane.running()
budget = deadline_policy.for_task(OP)
payload = b"0123456789abcdef"
ref = _ref(plane, task, attempt, payload=payload)
async def slow():
chunks = _chunks(ref, payload)
yield chunks[0]
# Age the delivery phase past the execution budget but well inside the
# delivery one. Under the old code there was no delivery phase, so the
# keepalive ceiling clamped the lease into the past and the next sweep
# would have failed a task that was uploading fine.
attempt.phase_started_at = resolve(None) - (budget.execution_seconds + 60)
for chunk in chunks[1:]:
yield chunk
ack = await plane.servicer.UploadResult(slow(), _Context(plane.token))
assert ack.committed is True
assert budget.result_delivery_seconds > budget.execution_seconds + 60
assert not attempt.lease_expired()
def test_upload_keepalive_does_not_move_completed_progress_backwards(plane):
task, attempt = plane.running()
plane.scheduler.on_progress(
task.task_id, attempt.attempt_id, progress=1.0, epoch=attempt.session_epoch
)
plane.servicer._renew_upload_lease(attempt)
assert attempt.progress == 1.0
@pytest.mark.asyncio
async def test_an_upload_onto_a_cancelled_task_is_refused(plane):
task, attempt = plane.running()
payload = b"too late"
ref = _ref(plane, task, attempt, payload=payload)
plane.scheduler.cancel(task.task_id)
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is False
assert ack.error.code == "ATTEMPT_NOT_LIVE"
assert not os.path.exists(plane.final_path(task, attempt))
@pytest.mark.asyncio
async def test_an_upload_for_another_workers_attempt_is_refused(plane):
task, attempt = plane.running()
payload = b"not yours"
ref = _ref(plane, task, attempt, payload=payload)
plane.servicer._sessions[plane.worker_id].worker_id = "someone-else"
ack = await plane.upload(_chunks(ref, payload))
assert ack.committed is False
assert ack.error.code == "UNKNOWN_ATTEMPT"
assert not os.path.exists(plane.final_path(task, attempt))
# ── Serving staged inputs ──────────────────────────────────────────────────
def _stage(plane, name: str, data: bytes) -> str:
"""Stand in for the input-staging step: a file inside the artifact store."""
path = os.path.join(plane.artifact_dir, "inputs", name)
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as fh:
fh.write(data)
return os.path.join("inputs", name)
@pytest.mark.asyncio
async def test_a_worker_can_read_the_input_staged_for_its_own_task(plane):
task, attempt = plane.running()
artifact_id = _stage(plane, "voice.wav", b"reference audio" * 10)
chunks = await plane.download(
pb.ArtifactRef(
artifact_id=artifact_id,
task_id=task.task_id,
attempt_id=attempt.attempt_id,
session_token=plane.token,
)
)
assert b"".join(c.data for c in chunks) == b"reference audio" * 10
assert chunks[-1].last is True
assert chunks[0].ref.size_bytes == len(b"reference audio" * 10)
# Nothing hands a session token back out that did not have to go out.
assert all(not c.ref.session_token for c in chunks)
@pytest.mark.asyncio
async def test_dispatch_stages_and_serves_from_the_servicers_artifact_root(plane, tmp_path):
voice = tmp_path / "voice.wav"
voice.write_bytes(b"reference audio")
task = plane.scheduler.submit(
operation=OP, engine=ENGINE, model_id=MODEL, params={"ref_audio": str(voice)}
)
assignment = plane.scheduler.next_assignment()
assert assignment is not None
assert await plane.servicer.dispatch(assignment)
message = await plane.servicer._sessions[plane.worker_id].outbox.get()
wire = message.assignment
assert wire.inputs
staged = os.path.join(plane.artifact_dir, wire.inputs[0].artifact_id)
assert os.path.isfile(staged)
chunks = await plane.download(
pb.ArtifactRef(
artifact_id=wire.inputs[0].artifact_id,
task_id=task.task_id,
attempt_id=assignment.attempt.attempt_id,
session_token=plane.token,
)
)
assert b"".join(chunk.data for chunk in chunks) == voice.read_bytes()
@pytest.mark.asyncio
async def test_a_worker_cannot_read_an_input_for_a_task_it_is_not_running(plane):
"""Authentication is not authorisation: from this phase on, staged inputs
are the user's own reference audio."""
task, attempt = plane.running()
artifact_id = _stage(plane, "voice.wav", b"reference audio")
plane.servicer._sessions[plane.worker_id].worker_id = "someone-else"
with pytest.raises(_Aborted) as caught:
await plane.download(
pb.ArtifactRef(
artifact_id=artifact_id,
task_id=task.task_id,
attempt_id=attempt.attempt_id,
session_token=plane.token,
)
)
assert "PERMISSION_DENIED" in str(caught.value.code)
@pytest.mark.asyncio
async def test_an_input_request_naming_no_task_is_refused(plane):
plane.running()
artifact_id = _stage(plane, "voice.wav", b"reference audio")
with pytest.raises(_Aborted) as caught:
await plane.download(
pb.ArtifactRef(artifact_id=artifact_id, session_token=plane.token)
)
assert "PERMISSION_DENIED" in str(caught.value.code)
@pytest.mark.asyncio
async def test_an_input_outside_the_artifact_store_is_not_served(plane, tmp_path):
task, attempt = plane.running()
secret = tmp_path / "secret.txt"
secret.write_text("private")
for artifact_id in ("../secret.txt", str(secret), "/etc/passwd"):
with pytest.raises(_Aborted) as caught:
await plane.download(
pb.ArtifactRef(
artifact_id=artifact_id,
task_id=task.task_id,
attempt_id=attempt.attempt_id,
session_token=plane.token,
)
)
assert "NOT_FOUND" in str(caught.value.code)