Compare commits
58
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d6a2277e6a | ||
|
|
a59fdee094 | ||
|
|
3fe2018589 | ||
|
|
00bf55c96f | ||
|
|
1e4fdd12e0 | ||
|
|
3b6e15dad5 | ||
|
|
de51120d6a | ||
|
|
3e5aa90e3f | ||
|
|
8cc7c88692 | ||
|
|
e7e97e4b4e | ||
|
|
97bfbfecd1 | ||
|
|
0c8a40774f | ||
|
|
f45c885171 | ||
|
|
da8358916c | ||
|
|
0268f43e7f | ||
|
|
6d42d2255e | ||
|
|
1103898d7b | ||
|
|
28220ff2dd | ||
|
|
fa9bfd41be | ||
|
|
0e47053823 | ||
|
|
bea72cadbe | ||
|
|
6624b3fb19 | ||
|
|
fc603bcc92 | ||
|
|
2300d4d5f7 | ||
|
|
c7890d2a5a | ||
|
|
633e991edc | ||
|
|
64b3785b0b | ||
|
|
4b3c62d783 | ||
|
|
103a5dbbe4 | ||
|
|
86f6c9ec8e | ||
|
|
e1a76f9aca | ||
|
|
57263edad7 | ||
|
|
4c0a0a26b7 | ||
|
|
634936eb10 | ||
|
|
422dbd1313 | ||
|
|
d23f22d56b | ||
|
|
eaa87ff041 | ||
|
|
0ee9bc35d0 | ||
|
|
a1cd15964c | ||
|
|
895e62df7c | ||
|
|
49c175301a | ||
|
|
b92d35ac5d | ||
|
|
b6bd125f23 | ||
|
|
53331a6f5a | ||
|
|
1d445855d2 | ||
|
|
9dc01d4b8b | ||
|
|
e93b6366e6 | ||
|
|
89cee3f824 | ||
|
|
078bad8e0f | ||
|
|
032c5aba86 | ||
|
|
a8371baaa8 | ||
|
|
5a615d2c66 | ||
|
|
12480b81b1 | ||
|
|
ef1cb57944 | ||
|
|
6447788fdf | ||
|
|
afa361913c | ||
|
|
0d3c81b596 | ||
|
|
98c9e68aae |
@@ -40,6 +40,7 @@ sudo apt-get install -y \
|
||||
libwebkit2gtk-4.1-dev libgtk-3-dev libpango1.0-dev libcairo2-dev \
|
||||
libsoup-3.0-dev libgdk-pixbuf-2.0-dev \
|
||||
libayatana-appindicator3-dev librsvg2-dev libssl-dev libxdo-dev \
|
||||
gstreamer1.0-plugins-good \
|
||||
libasound2-dev build-essential curl wget file
|
||||
```
|
||||
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@
|
||||
|
||||
| Version | Supported |
|
||||
|---------|-----------|
|
||||
| 0.3.x (latest release + `main` previews) | ✅ Current — all fixes land here |
|
||||
| 0.5.x (latest release + `main` previews) | ✅ Current — all fixes land here |
|
||||
| 0.2.7 | ⚠️ Legacy stable — security fixes only, upgrade recommended |
|
||||
| < 0.2.7 | ❌ No longer supported |
|
||||
|
||||
|
||||
@@ -51,6 +51,13 @@ jobs:
|
||||
- os: ubuntu-latest
|
||||
platform: linux-x86_64
|
||||
experimental: false
|
||||
- os: ubuntu-24.04-arm
|
||||
platform: linux-aarch64
|
||||
# Apple Silicon under Asahi Linux. Experimental: the Vulkan
|
||||
# (Honeykrisp GPU) build path is new and the hosted arm64
|
||||
# runner has no GPU — it validates that the binary builds;
|
||||
# on-host Vulkan acceleration is exercised by users.
|
||||
experimental: true
|
||||
- os: windows-latest
|
||||
platform: windows-x86_64
|
||||
experimental: false
|
||||
@@ -80,11 +87,18 @@ jobs:
|
||||
# Linux-only: upstream `buildcpu.sh` enables `-DGGML_BLAS=ON` which
|
||||
# requires a system BLAS implementation at cmake configure time.
|
||||
- name: Linux system deps (BLAS for ggml-blas backend)
|
||||
if: matrix.platform == 'linux-x86_64'
|
||||
if: startsWith(matrix.platform, 'linux')
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y libopenblas-dev pkg-config
|
||||
|
||||
# linux-aarch64: let the build script's Vulkan path (Honeykrisp GPU
|
||||
# on Asahi) engage instead of silently falling back to CPU.
|
||||
- name: Vulkan dev deps (linux-aarch64 GPU backend)
|
||||
if: matrix.platform == 'linux-aarch64'
|
||||
run: |
|
||||
sudo apt-get install -y glslc libvulkan-dev spirv-headers
|
||||
|
||||
- name: Build omnivoice-tts
|
||||
shell: bash
|
||||
# Pass values through env (quoted) rather than ${{ }} interpolation
|
||||
|
||||
@@ -159,6 +159,13 @@ tests/probe/reports/
|
||||
# reports). Working notes for whoever is driving a change, not a repo artifact.
|
||||
/remote/
|
||||
|
||||
# OmniVoice GGUF runtime build artifacts (scripts/build-omnivoice-tts.sh).
|
||||
# Only 0-byte placeholders of omnivoice-tts-* are tracked; real binaries,
|
||||
# the checksums manifest and the copied libggml shared libs ship via CI.
|
||||
bin/libggml*
|
||||
bin/checksums.sha256
|
||||
bin/omnivoice-tts-linux-aarch64
|
||||
|
||||
# Dubbing-demo intermediates. The .mp4/.srt/manifest.json in this directory ARE
|
||||
# committed (they ship with the app); the per-language source WAVs are just the
|
||||
# inputs scripts/render_dub_demo_audio.py hands to scripts/build_dub_demo.sh.
|
||||
|
||||
@@ -10,6 +10,34 @@ the frozen-backend fallback mirror it for their toolchains.
|
||||
|
||||
**Highlights**
|
||||
|
||||
### Changed
|
||||
|
||||
### Added
|
||||
|
||||
### Docs
|
||||
|
||||
### Fixed
|
||||
|
||||
## [0.5.1] — 2026-08-28
|
||||
|
||||
**Highlights**
|
||||
|
||||
- OmniVoice generation on Apple Silicon now runs in a crash-isolated child, so fatal MPS memory exits no longer take down the local backend (#1697, #1698) — thanks @ndntran14!
|
||||
- Model-load GPU exhaustion now returns a sanitized, actionable dubbing error, and readiness correctly attributes the shared model status to TTS (#1695)
|
||||
- Source-mode development now restarts an isolated backend crash without tearing down the UI, while repeated crash loops still stop loudly with diagnostics (#1690)
|
||||
- Dubbing playback now keeps an audible companion source when a WebView can render the preview picture but cannot decode its audio (#1692)
|
||||
- Model Catalogue engine rows now use the available desktop width and keep identity, runtime state, and actions from crowding one another (#1689)
|
||||
- VoiceStudio now acts as a local speech platform: other apps can trigger its native dictation or connect through versioned HTTP, WebSocket, JSON-RPC, CLI, and MCP transports (#1646)
|
||||
- A timed-out in-process dub transcription no longer starts a second WhisperX/CTranslate2 call over the abandoned native worker, preventing the overlapping access that preceded Windows `0xC0000005` exits (#1669)
|
||||
- Windows debugger termination code `0x40010004` is no longer misreported as a backend crash or charged against automatic restart recovery (#1663)
|
||||
- Studio now keeps one generation reservation across page changes, preventing a remount from stacking native jobs until the backend reports capacity busy or is killed under memory pressure (#1670)
|
||||
- Uploaded dubbing videos are normalized to browser-safe H.264/AAC before preview, preventing valid VP9, AV1, or Opus media from failing with “no supported sources” (#1644)
|
||||
- Dubbing now separates spoken and target languages, preserves translations through segment cleanup, and lets failed translations be retried or skipped without restarting the batch (#1654) — thanks @Number16BusShelter!
|
||||
- Importing replacement SRT subtitles now keeps each cue bound to the best-overlapping source speaker and clone instead of resetting every line to a random default voice (#1660) — thanks @invio-a11y!
|
||||
- Uploading a Dub preview no longer blocks every backend request while ffmpeg extracts its audio (#1667) — thanks @tfreyd!
|
||||
- Docker quick starts now require the administrator key needed through container NAT instead of starting a UI whose protected actions return 403 (#1651) — thanks @wd357dui!
|
||||
- WSL2 AMD containers now use the `/dev/dxg` ROCDXG bridge with actionable GPU diagnostics instead of silently falling back to CPU (#1655) — thanks @wd357dui!
|
||||
- Ad-hoc voice-clone references now stay alive until cancelled or timed-out GPU work actually stops reading them, so prompt caching can finish instead of failing on a deleted temp file (#1668) — thanks @tfreyd!
|
||||
- Dictation now stays bound to the app where it started and recovers locally from silent recognizer output (#1175)
|
||||
- The backend now answers within a second of launch and narrates its startup step by step (#1550)
|
||||
- Reporting a bug from an outdated build now offers the latest release first (#1547)
|
||||
@@ -18,10 +46,15 @@ the frozen-backend fallback mirror it for their toolchains.
|
||||
- Dub subtitles can be retimed, inserted, and merged in either direction from the segment table (#1612) — thanks @invio-a11y!
|
||||
|
||||
### Changed
|
||||
- Model Catalogue now uses one breathable workspace canvas with simpler pane and engine-family navigation instead of nested cards and scroll regions (#1685)
|
||||
- Linux source launchers now catch missing libxdo and GStreamer audio plugins before they can cause a linker error or an aborted, blank WebKit renderer (#1680, #1682)
|
||||
- Dictation now carries one native output session from shortcut-down through final delivery, restores text, HTML, image, or file-list clipboards only when untouched, keeps Wayland copy-safe unless current-focus insertion is explicitly enabled, and retries silent Sherpa speech only through an already-installed local ASR model (#1175)
|
||||
- The backend binds its port immediately and reports startup progress live — `/health` answers 503-with-step and a new `/startup/progress` endpoint lists every step while PyTorch, API routes, and database migrations load in the background, so "starting at step X" is never mistakable for "dead"; the desktop splash narrates each step (#1550)
|
||||
|
||||
### Added
|
||||
- A bundled Rust loopback sidecar exposes dictation start/stop/toggle, focused-output sessions, discovery, and JSON-RPC; the backend adds versioned streaming events and a dependency-free CLI bridge for Herdr, coding agents, editors, desktop apps, and TUIs (#1646)
|
||||
- Headless NVIDIA and ROCm machines can now join as worker-only Docker Compose services with no published UI and durable protocol-v2 enrollment; update both machines together before reconnecting (#1638) — thanks @jkrogers9862!
|
||||
- Linux ARM64 (Asahi Apple Silicon) support for the OmniVoice GGUF engine — a `linux-aarch64` binary built with GGML Vulkan where the toolchain allows it, so Apple GPUs accelerate generation through the open-source Honeykrisp driver instead of falling back to CPU-only (#1641)
|
||||
- One-command install on every desktop OS: `curl -fsSL https://voicestudio.sh/install | sh` (macOS/Linux/WSL) or `irm https://voicestudio.sh/install | iex` (Windows) — the URL serves the right script per platform, and Windows gains a source installer (`scripts/install.ps1`) with a 3-OS CI smoke (#1626)
|
||||
- Per-line subtitle management in the dub table: a line's end time is editable alongside its start (typing a time and dragging its timeline edge now take the same path), lines merge with the previous row as well as the next (`Ctrl/Cmd+Shift+M`), and a new line can be inserted into the gap after any row (#1612) — thanks @invio-a11y!
|
||||
- CI now enforces performance regression budgets on the hot paths — operation-count tests pin streaming TTS to one synthesis per sentence and cached dub re-mixes to zero re-synthesis; fast-path guards cover zero re-decoding and ⌈N/W⌉ native batch calls when enabled (#1594)
|
||||
@@ -35,6 +68,7 @@ the frozen-backend fallback mirror it for their toolchains.
|
||||
- Opt-in 24-layer PocketTTS checkpoints via `OMNIVOICE_POCKETTTS_24L` — better prosody for it/de/es/pt at roughly 2x render time (still faster than real-time); the fast 6-layer model stays the default (#1613) — thanks @paoloantinori!
|
||||
|
||||
### Docs
|
||||
- Supported-version and install guidance now identifies 0.5.1 as the stable desktop and container release (#1687)
|
||||
- The Docker Hub overview now shows the current engine-switching demo, Model Catalogue, and gallery voice workflow (#1593)
|
||||
- The Docker Hub overview and install guide now show the v0.5 tags and the built-in API-key/share-PIN security model instead of obsolete v0.4 and no-authentication guidance (#1592)
|
||||
- The READMEs now lead with download buttons and a three-step first-clone walkthrough, and a new benchmarks page anchors measured per-engine/per-device numbers on the in-repo harness (#1555)
|
||||
@@ -42,6 +76,18 @@ the frozen-backend fallback mirror it for their toolchains.
|
||||
- The OmniVoice guide now covers combining style attributes with a reference clip (consistent instruct stabilizes cloning; the reference wins conflicts), inline pronunciation control (pinyin / CMU phonemes), and corrects the claim that the default engine can't do voice design — it can, from attributes (#1565)
|
||||
|
||||
### Fixed
|
||||
- Workspaces now measure their responsive width when the post-bootstrap shell actually mounts, so native UI scaling reflows Projects and History instead of crushing the Dubbing demo into unreadable columns (#1683)
|
||||
- Dubbing keeps the source-language selector visible after a local file is chosen, so ASR can be pinned before transcription starts (#1678) — thanks @Lonki-lomki-cloud!
|
||||
- First-run media-engine downloads become available to TTS immediately without a restart, and missing media-process failures now point to repair controls (#1677) — thanks @farhataligpt-dev!
|
||||
- Source installs on AMD GPUs honour `OMNIVOICE_TORCH_VARIANT=rocm`: `bun run desktop` now swaps in the ROCm torch wheel after `uv sync` and launches the backend without re-syncing, instead of silently reverting to the CPU-only CUDA build on every start (#1665) — thanks @uberclokr!
|
||||
- `bun run desktop` on a fresh clone no longer fails with "resource path `../../frontend/dist` doesn't exist" — the dev launcher creates the placeholder Tauri resource directory before compiling (#1664) — thanks @uberclokr!
|
||||
- macOS no longer loses TTS after the first request when Python lacks `os.waitid`; subprocess ownership now uses a safe `waitpid` fallback without risking reused process groups (#1656) — thanks @paoloantinori!
|
||||
- Desktop startup, Retry, reset, uninstall, shutdown, and crash recovery now share one backend lifecycle owner; quitting interrupts first-run installers and gracefully drains then force-cleans the full backend process tree, so overlaps cannot duplicate or orphan it (#1635) — thanks @Xohaibxobi!
|
||||
- Large Stories and Audiobook projects now persist in IndexedDB instead of overflowing the `omnivoice.app` localStorage envelope, with quota-safe migration and orderly exit/reload flushing (#1636) — thanks @leodzai!
|
||||
- OmniVoice and its crash-isolated subprocess now route to AMD ROCm GPUs instead of warning and falling back to CPU (#1629) — thanks @j4r3kb!
|
||||
- Dictation now cancels pending startup work, capture resources, sockets, and timers when the capture widget closes, preventing late work against a destroyed webview (#1645)
|
||||
- Streaming generation failures now show recognized recovery guidance and appear in Diagnostics instead of only returning a generic error (#1607)
|
||||
- The worker-capacity transport test no longer races its own setup: the 1-slot limit now goes through the enrollment handshake instead of mutating client config after connect, where the server's stream-open ConfigUpdate (carrying the registered capacity of 2) could overwrite it and fake an over-accept; failed CI twice on 2026-08-21 (#1630)
|
||||
- Moving words across a speaker boundary in a dub — merging two lines and splitting them again — no longer dubs the second half in the first speaker's voice; each half now keeps the speaker, voice, direction, gain, and language of whoever actually says it (#1612) — thanks @invio-a11y!
|
||||
- Dictation on a WebView that refuses a 16 kHz audio context (WKWebView) now low-passes before downsampling, so frequencies above 8 kHz stop folding into the speech the recognizer is fed (#1610)
|
||||
- A microphone context that cannot be resumed now reports a mic error instead of leaving the dictation pill on "Listening" while capturing nothing (#1610)
|
||||
|
||||
@@ -60,7 +60,7 @@
|
||||
| macOS 13.3+ | DMG, Apple Silicon | [Install on macOS](docs/install/macos.md) |
|
||||
| Windows 10/11 | MSI, x64 | [Install on Windows](docs/install/windows.md) |
|
||||
| Linux | AppImage, x86_64 with glibc 2.39+ | [Install on Linux](docs/install/linux.md) |
|
||||
| Docker | CUDA, ROCm, or CPU | [Run with Docker](docs/install/docker.md) |
|
||||
| Docker | CUDA, ROCm, or CPU; worker-only GPU profiles | [Run with Docker](docs/install/docker.md) |
|
||||
|
||||
Download packages from the [latest release](https://github.com/debpalash/VoiceStudio/releases/latest). First launch creates a managed Python environment and downloads the default model. Later launches reuse both.
|
||||
|
||||
@@ -254,7 +254,7 @@ FastAPI backend
|
||||
|
||||
<a id="api"></a>
|
||||
|
||||
## OpenAI-compatible API
|
||||
## Local speech platform and OpenAI-compatible API
|
||||
|
||||
Point an OpenAI-compatible audio client at the local backend:
|
||||
|
||||
@@ -267,6 +267,8 @@ Point an OpenAI-compatible audio client at the local backend:
|
||||
|---|---|
|
||||
| `POST /v1/audio/speech` | TTS to `mp3`, `opus`, `aac`, `flac`, `wav`, or `pcm`; select a profile with `voice` and an engine with `model` |
|
||||
| `POST /v1/audio/transcriptions` | STT to `json`, `text`, `verbose_json`, `srt`, or `vtt` |
|
||||
| `WS /v1/audio/transcriptions/stream` | Live PCM/WebM transcription with partial, utterance, and session-final events |
|
||||
| `GET /.well-known/voicestudio-speech` | Discover HTTP, WebSocket, MCP, and native dictation-control transports |
|
||||
| `GET /v1/audio/voices` | List local voice profiles and engines |
|
||||
|
||||
```python
|
||||
@@ -283,7 +285,12 @@ with client.audio.speech.with_streaming_response.create(
|
||||
response.stream_to_file("speech.wav")
|
||||
```
|
||||
|
||||
The full API reference is in **Settings → OpenAPI Reference**. For LAN, Tailscale, or proxy access, read [API authentication](docs/api-auth.md) before exposing the backend.
|
||||
The bundled Rust control sidecar also lets Herdr, coding agents, VS Code,
|
||||
desktop apps, and TUIs trigger the existing system-wide dictation flow or reuse
|
||||
its safe native insertion. See the [speech platform guide](docs/speech-platform.md).
|
||||
The full API reference is in **Settings → OpenAPI Reference**. For LAN,
|
||||
Tailscale, or proxy access, read [API authentication](docs/api-auth.md) before
|
||||
exposing the backend.
|
||||
|
||||
### Agent skills
|
||||
|
||||
@@ -312,7 +319,7 @@ The [notebook](notebooks/OmniVoice_Studio_Colab.ipynb) runs the app and web UI o
|
||||
| Fix setup | [Troubleshooting](docs/install/troubleshooting.md) · [model downloads](docs/downloading-models.md) · [Hugging Face token](docs/setup/huggingface-token.md) |
|
||||
| Choose an engine | [Engine guides](docs/engines/README.md) · [benchmarks](docs/benchmarks.md) · [expressive speech](docs/expressive-speech.md) |
|
||||
| Tune hardware | [Performance](docs/performance.md) · [remote workers](docs/remote-workers.md) |
|
||||
| Build integrations | [API auth](docs/api-auth.md) · [MCP](docs/mcp.md) · [examples](examples/README.md) |
|
||||
| Build integrations | [Speech platform](docs/speech-platform.md) · [API auth](docs/api-auth.md) · [MCP](docs/mcp.md) · [examples](examples/README.md) |
|
||||
| Build VoiceStudio | [Contributing](.github/CONTRIBUTING.md) · [engine acceptance](docs/engine-acceptance.md) |
|
||||
| Track changes | [Changelog](CHANGELOG.md) · [roadmap](docs/ROADMAP.md) · [latest release](https://github.com/debpalash/VoiceStudio/releases/latest) |
|
||||
| Remove everything | [Uninstall guide](docs/install/uninstall.md) |
|
||||
|
||||
@@ -54,6 +54,15 @@ def _server_mode() -> bool:
|
||||
return os.environ.get("OMNIVOICE_SERVER_MODE", "").strip().lower() in _TRUTHY
|
||||
|
||||
|
||||
def validate_server_admin_key() -> None:
|
||||
"""Reject an explicitly blank key before a server-mode app starts."""
|
||||
raw_key = os.environ.get("OMNIVOICE_API_KEY")
|
||||
if _server_mode() and raw_key is not None and not raw_key.strip():
|
||||
raise RuntimeError(
|
||||
"OMNIVOICE_API_KEY is blank; configure a non-whitespace administrator key"
|
||||
)
|
||||
|
||||
|
||||
def _configured_pin(request) -> str | None:
|
||||
"""The active share PIN (``app.state.network_share.pin``) or None. Read via
|
||||
getattr so a bare Request stub (or a request that hit before lifespan set
|
||||
|
||||
@@ -361,7 +361,7 @@ LONGFORM_NUM_STEP = 32
|
||||
LONGFORM_GUIDANCE_SCALE = 2.0
|
||||
|
||||
|
||||
def _seed_segment_rng(base_seed, text: str, nonce: int = 0) -> None:
|
||||
def _seed_segment_rng(base_seed, text: str, nonce: int = 0) -> int | None:
|
||||
"""Apply a profile's pinned seed to this synth call (#1139).
|
||||
|
||||
``_resolve_voice`` has always fetched the profile ``seed`` — but only the
|
||||
@@ -380,11 +380,13 @@ def _seed_segment_rng(base_seed, text: str, nonce: int = 0) -> None:
|
||||
must cover /generate and here together, not one path.
|
||||
"""
|
||||
if base_seed is None:
|
||||
return
|
||||
return None
|
||||
import torch
|
||||
|
||||
from services.audiobook import segment_seed
|
||||
torch.manual_seed(segment_seed(base_seed, text, nonce))
|
||||
seed = segment_seed(base_seed, text, nonce)
|
||||
torch.manual_seed(seed)
|
||||
return seed
|
||||
|
||||
|
||||
def _base_seed(opts: ExpressiveOptions, voice: dict):
|
||||
@@ -508,16 +510,21 @@ def _build_synth(
|
||||
"get_model": get_model, "language": language, "opts": opts}
|
||||
|
||||
backend = cls()
|
||||
extra = _generic_extra_kwargs(opts)
|
||||
native_proxy = bool(getattr(cls, "supports_native_omnivoice_controls", False))
|
||||
extra = (_omnivoice_sampling_kwargs(opts) if native_proxy
|
||||
else _generic_extra_kwargs(opts))
|
||||
next_nonce = _make_occ_counter(opts)
|
||||
|
||||
def synth(text, voice_id, speed=None):
|
||||
v = resolve(voice_id)
|
||||
_seed_segment_rng(_base_seed(opts, v), text, next_nonce())
|
||||
seed = _seed_segment_rng(_base_seed(opts, v), text, next_nonce())
|
||||
call_extra = dict(extra)
|
||||
if native_proxy and seed is not None:
|
||||
call_extra["seed"] = seed
|
||||
return backend.generate(
|
||||
text, language=language, ref_audio=v["ref_audio"],
|
||||
ref_text=v["ref_text"], instruct=v["instruct"], duration=None,
|
||||
speed=float(speed) if speed else 1.0, **extra,
|
||||
speed=float(speed) if speed else 1.0, **call_extra,
|
||||
)
|
||||
return {"mode": "generic", "resolve": resolve, "engine_id": engine_id,
|
||||
"synth": synth, "sample_rate": backend.sample_rate}
|
||||
|
||||
@@ -38,11 +38,14 @@ Protocol:
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||
|
||||
@@ -52,6 +55,9 @@ from services.text_polish import polish_text
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger("omnivoice.capture_ws")
|
||||
|
||||
SPEECH_PROTOCOL = "voicestudio.speech.v1"
|
||||
PLATFORM_STREAM_PATH = "/v1/audio/transcriptions/stream"
|
||||
|
||||
# How often (seconds) to run transcription on the accumulated buffer.
|
||||
# Shorter = more responsive but more GPU load.
|
||||
PARTIAL_INTERVAL_S = float(os.environ.get("OMNIVOICE_STREAM_INTERVAL", "2.0"))
|
||||
@@ -83,6 +89,39 @@ _AEC_FAR = 0x01 # playback reference frame (feed the echo model only)
|
||||
SR_MIN, SR_MAX = 8000, 96000
|
||||
|
||||
|
||||
def _is_end_control(text: str | None) -> bool:
|
||||
"""Accept the versioned JSON control frame and the legacy ``EOF`` frame."""
|
||||
if text == "EOF":
|
||||
return True
|
||||
if not text:
|
||||
return False
|
||||
try:
|
||||
message = json.loads(text)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
return False
|
||||
return isinstance(message, dict) and message.get("type") == "input_audio.end"
|
||||
|
||||
|
||||
class _PlatformWebSocket:
|
||||
"""Add v1 session metadata without changing the legacy WebSocket contract."""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
self._websocket = websocket
|
||||
self.session_id = uuid.uuid4().hex
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._websocket, name)
|
||||
|
||||
async def send_json(self, data: Any, mode: str = "text") -> None:
|
||||
if isinstance(data, dict):
|
||||
data = dict(data)
|
||||
data.setdefault("protocol", SPEECH_PROTOCOL)
|
||||
data.setdefault("session_id", self.session_id)
|
||||
if data.get("type") == "final":
|
||||
data.setdefault("final_kind", "summary")
|
||||
await self._websocket.send_json(data, mode=mode)
|
||||
|
||||
|
||||
def _bounded_sample_rate(query_params) -> int:
|
||||
try:
|
||||
sample_rate = int(query_params.get("sr", "16000"))
|
||||
@@ -194,9 +233,24 @@ def _select_sherpa_spec(websocket: WebSocket):
|
||||
return _usable_spec(mid) if mid else None
|
||||
|
||||
|
||||
@router.websocket(PLATFORM_STREAM_PATH)
|
||||
@router.websocket("/ws/transcribe")
|
||||
async def ws_transcribe(websocket: WebSocket):
|
||||
"""Stream audio in, get partial + final transcription out."""
|
||||
is_platform_stream = websocket.url.path == PLATFORM_STREAM_PATH
|
||||
if is_platform_stream:
|
||||
websocket = _PlatformWebSocket(websocket)
|
||||
# A browser can reach localhost regardless of the page's own origin.
|
||||
# Reject ambient cross-site WebSocket handshakes before the loopback-host
|
||||
# shortcut or accept(), while keeping native clients (no Origin header)
|
||||
# and configured/same-origin browser UIs working (#1646 review).
|
||||
origin = websocket.headers.get("origin")
|
||||
if origin:
|
||||
from core.csrf import origin_allowed
|
||||
|
||||
if not origin_allowed(websocket):
|
||||
await websocket.close(code=1008, reason="browser origin not allowed")
|
||||
return
|
||||
# Loopback origin guard — refuse anything not from 127.0.0.1, ::1, or
|
||||
# localhost. Privileged HTTP routers use Depends(require_admin) at router
|
||||
# level; WebSocket dependency injection differs across FastAPI versions, so we
|
||||
@@ -211,6 +265,16 @@ async def ws_transcribe(websocket: WebSocket):
|
||||
return
|
||||
|
||||
await websocket.accept()
|
||||
if is_platform_stream:
|
||||
await websocket.send_json({
|
||||
"type": "session.started",
|
||||
"input_format": (
|
||||
"audio/pcm;encoding=s16le;channels=1"
|
||||
if _requested_pcm_sample_rate(websocket.query_params) is not None
|
||||
else "audio/webm;codecs=opus"
|
||||
),
|
||||
"sample_rate": _bounded_sample_rate(websocket.query_params),
|
||||
})
|
||||
|
||||
# Live-dictation engine selection. When a sherpa-onnx model is selected
|
||||
# (via ?model= or the dictation.model_id pref) AND sherpa is installed,
|
||||
@@ -333,7 +397,7 @@ async def ws_transcribe(websocket: WebSocket):
|
||||
total_bytes += len(data)
|
||||
last_audio_time = time.monotonic()
|
||||
continue
|
||||
if msg.get("text") == "EOF":
|
||||
if _is_end_control(msg.get("text")):
|
||||
# Client signals end-of-audio but stays connected for `final`.
|
||||
running = False
|
||||
break
|
||||
@@ -656,7 +720,7 @@ async def _recv_pcm_frame(websocket: WebSocket, aec):
|
||||
return "skip", b""
|
||||
return "near", aec.process_near_end(payload)
|
||||
return "near", data
|
||||
if msg.get("text") == "EOF":
|
||||
if _is_end_control(msg.get("text")):
|
||||
return "eof", b""
|
||||
return "skip", b""
|
||||
|
||||
|
||||
+178
-33
@@ -134,6 +134,82 @@ _save_job = dub_pipeline.save_job
|
||||
# paste (or a mis-aimed binary) burn CPU in the parser.
|
||||
_MAX_SUBTITLE_PASTE_CHARS = 2_000_000
|
||||
|
||||
_SRT_REPLACED_FIELDS = {
|
||||
"id",
|
||||
"start",
|
||||
"end",
|
||||
"text",
|
||||
"text_original",
|
||||
"translations",
|
||||
"translate_error",
|
||||
"translate_degraded",
|
||||
}
|
||||
|
||||
|
||||
def _best_overlapping_segment(cue: dict, existing: list[dict]) -> dict | None:
|
||||
"""Return the prior segment with the strongest temporal overlap."""
|
||||
cue_start = float(cue.get("start") or 0.0)
|
||||
cue_end = float(cue.get("end") or cue_start)
|
||||
cue_mid = (cue_start + cue_end) / 2.0
|
||||
best = None
|
||||
best_key = None
|
||||
for index, segment in enumerate(existing):
|
||||
start = float(segment.get("start") or 0.0)
|
||||
end = float(segment.get("end") or start)
|
||||
overlap = min(cue_end, end) - max(cue_start, start)
|
||||
if overlap <= 0:
|
||||
continue
|
||||
midpoint_distance = abs(cue_mid - ((start + end) / 2.0))
|
||||
key = (overlap, -midpoint_distance, -index)
|
||||
if best_key is None or key > best_key:
|
||||
best = segment
|
||||
best_key = key
|
||||
return best
|
||||
|
||||
|
||||
def _carry_srt_voice_metadata(
|
||||
cues: list[dict],
|
||||
existing: list[dict],
|
||||
segment_clones: dict | None,
|
||||
speaker_clones: dict | None = None,
|
||||
) -> tuple[list[dict], dict]:
|
||||
"""Replace subtitle content while retaining the source cast assignment."""
|
||||
source_clones = dict(segment_clones or {})
|
||||
source_speaker_clones = dict(speaker_clones or {})
|
||||
# Replacement cues get new positional ids. Starting from the old map would
|
||||
# let an unmatched cue whose new id happens to equal an old id inherit an
|
||||
# unrelated reference. Only explicitly overlap-matched references survive.
|
||||
clones = {}
|
||||
merged_segments = []
|
||||
for new_id, cue in enumerate(cues):
|
||||
prior = _best_overlapping_segment(cue, existing)
|
||||
metadata = {
|
||||
key: value
|
||||
for key, value in (prior or {}).items()
|
||||
if key not in _SRT_REPLACED_FIELDS
|
||||
}
|
||||
merged = {
|
||||
**metadata,
|
||||
"id": new_id,
|
||||
"start": cue.get("start", 0.0),
|
||||
"end": cue.get("end", 0.0),
|
||||
"text": cue.get("text", ""),
|
||||
"text_original": cue.get("text", ""),
|
||||
}
|
||||
if not merged.get("speaker_id"):
|
||||
merged["speaker_id"] = cue.get("speaker_id") or "Speaker 1"
|
||||
if prior is not None:
|
||||
prior_id = str(prior.get("id", ""))
|
||||
clone = source_clones.get(prior_id)
|
||||
if clone is None:
|
||||
clone = source_speaker_clones.get(prior.get("speaker_id"))
|
||||
if clone is not None:
|
||||
clones[str(new_id)] = clone
|
||||
if merged.get("profile_id") == f"auto-seg:{prior_id}":
|
||||
merged["profile_id"] = f"auto-seg:{new_id}"
|
||||
merged_segments.append(merged)
|
||||
return merged_segments, clones
|
||||
|
||||
|
||||
@router.post("/dub/parse-subtitle-text")
|
||||
def dub_parse_subtitle_text(req: ParseSubtitleTextRequest):
|
||||
@@ -234,7 +310,32 @@ async def dub_import_srt(job_id: str, file: UploadFile = File(...)):
|
||||
else:
|
||||
segments = result.segments
|
||||
|
||||
prior_segments = [
|
||||
segment for segment in (job.get("segments") or []) if isinstance(segment, dict)
|
||||
]
|
||||
segments, segment_clones = _carry_srt_voice_metadata(
|
||||
segments,
|
||||
prior_segments,
|
||||
job.get("segment_clones"),
|
||||
job.get("speaker_clones"),
|
||||
)
|
||||
job["segments"] = segments
|
||||
job["segment_clones"] = segment_clones
|
||||
# A pooled speaker clone is keyed only by a display label. Replacement
|
||||
# cues can reuse that label without overlapping the original speaker, so
|
||||
# retain matched pooled references as segment-specific clones above and
|
||||
# drop the global map before rebuilding the cast.
|
||||
job["speaker_clones"] = {}
|
||||
if segment_clones:
|
||||
from services.speaker_clone import build_cast_sources
|
||||
|
||||
job["cast_sources"] = build_cast_sources(
|
||||
segments,
|
||||
None,
|
||||
segment_clones,
|
||||
)
|
||||
else:
|
||||
job.pop("cast_sources", None)
|
||||
# `source_lang` stays whatever the user (or the upload step) set; we
|
||||
# don't try to language-detect off the cue text — that's noisy and the
|
||||
# user usually knows what their .srt is.
|
||||
@@ -351,12 +452,13 @@ async def preview_upload(video: UploadFile = File(...)):
|
||||
safe_name = f"{uuid.uuid4().hex[:12]}"
|
||||
vid_path = os.path.join(PREVIEW_DIR, f"{safe_name}{ext}")
|
||||
wav_path = os.path.join(PREVIEW_DIR, f"{safe_name}.wav")
|
||||
|
||||
with open(vid_path, "wb") as f:
|
||||
f.write(await video.read())
|
||||
|
||||
has_audio = False
|
||||
if ext not in [".wav", ".mp3", ".m4a", ".aac"]:
|
||||
payload = await video.read()
|
||||
|
||||
def _write_and_extract() -> bool:
|
||||
with open(vid_path, "wb") as f:
|
||||
f.write(payload)
|
||||
if ext in {".wav", ".mp3", ".m4a", ".aac"}:
|
||||
return False
|
||||
try:
|
||||
ffmpeg_cmd = [
|
||||
find_ffmpeg(), "-y", "-i", vid_path,
|
||||
@@ -368,10 +470,16 @@ async def preview_upload(video: UploadFile = File(...)):
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
timeout=300,
|
||||
)
|
||||
has_audio = True
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.warning("FFmpeg extraction failed: %s", log_safe(e))
|
||||
pass
|
||||
return False
|
||||
|
||||
# File writes and ffmpeg are blocking operations. Keep them on the bounded
|
||||
# CPU pool so a large preview cannot stall unrelated API requests (#1667).
|
||||
has_audio = await asyncio.get_running_loop().run_in_executor(
|
||||
_cpu_pool, _write_and_extract
|
||||
)
|
||||
|
||||
return {
|
||||
"url": f"/preview/{safe_name}{ext}",
|
||||
@@ -410,12 +518,38 @@ _ingest_gen = dub_pipeline.ingest_pipeline
|
||||
#: container so a mislabelled video can't slip past the video-skipping branch.
|
||||
_AUDIO_EXTS = {".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg", ".opus", ".wma"}
|
||||
|
||||
# Source-language choices exposed by the first-party dub UI. Keeping this an
|
||||
# allow-list rejects language names and private-use BCP-47 tags before they are
|
||||
# persisted as ASR overrides. Values are normalized to lowercase below.
|
||||
_DUB_SOURCE_LANG_CODES = frozenset({
|
||||
"af", "sq", "am", "ar", "hy", "az", "eu", "be", "bn", "bs", "bg",
|
||||
"my", "ca", "cmn-hans", "cmn-hant", "hr", "cs", "da", "nl", "en",
|
||||
"et", "fi", "fr", "gl", "ka", "de", "el", "gu", "ht", "ha", "haw",
|
||||
"he", "hi", "hu", "is", "id", "it", "ja", "jw", "kn", "kk", "km",
|
||||
"ko", "ku", "ky", "lo", "la", "lv", "lt", "mk", "ms", "ml", "mt",
|
||||
"mi", "mr", "mn", "ne", "no", "ps", "fa", "pl", "pt", "pa", "ro",
|
||||
"ru", "sm", "gd", "sr", "sn", "sd", "si", "sk", "sl", "so", "es",
|
||||
"su", "sw", "sv", "tg", "ta", "te", "th", "tr", "uk", "ur", "uz",
|
||||
"vi", "cy", "xh", "yi", "yo", "zu",
|
||||
})
|
||||
|
||||
|
||||
def _source_lang_override(value: str | None) -> str | None:
|
||||
"""Normalize a user-selected source language; auto/und means detect."""
|
||||
code = (value or "").strip().lower()
|
||||
if code in {"", "auto", "und"}:
|
||||
return None
|
||||
if code not in _DUB_SOURCE_LANG_CODES:
|
||||
raise HTTPException(status_code=400, detail="Invalid source language code")
|
||||
return code
|
||||
|
||||
|
||||
@router.post("/dub/upload")
|
||||
async def dub_upload(
|
||||
video: UploadFile = File(...),
|
||||
job_id: Optional[str] = Form(None),
|
||||
input_type: str = Form("video"),
|
||||
source_lang: Optional[str] = Form(None),
|
||||
):
|
||||
"""Accept a media upload, write to disk, queue background prep task.
|
||||
|
||||
@@ -445,6 +579,7 @@ async def dub_upload(
|
||||
detail=f"Audio-only dubbing needs an audio file ({', '.join(sorted(_AUDIO_EXTS))}); got '{ext or 'no extension'}'.",
|
||||
)
|
||||
|
||||
source_lang_override = _source_lang_override(source_lang)
|
||||
os.makedirs(job_dir, exist_ok=True)
|
||||
|
||||
video_path = os.path.join(job_dir, f"original{ext}")
|
||||
@@ -456,7 +591,13 @@ async def dub_upload(
|
||||
await task_manager.add_task(
|
||||
task_id, "prep",
|
||||
_ingest_gen, job_id, job_dir,
|
||||
{"kind": "file", "path": video_path, "input_type": input_type}, filename,
|
||||
{
|
||||
"kind": "file",
|
||||
"path": video_path,
|
||||
"input_type": input_type,
|
||||
"source_lang": source_lang_override,
|
||||
},
|
||||
filename,
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=202,
|
||||
@@ -478,6 +619,7 @@ async def dub_ingest_url(req: DubIngestUrlRequest, request: Request):
|
||||
status_code=400,
|
||||
detail="URL must start with http:// or https://. Paste a full video link (e.g. https://youtube.com/watch?v=…) or drop a local file instead.",
|
||||
)
|
||||
source_lang_override = _source_lang_override(req.source_lang)
|
||||
|
||||
try:
|
||||
import yt_dlp # noqa: F401
|
||||
@@ -513,6 +655,7 @@ async def dub_ingest_url(req: DubIngestUrlRequest, request: Request):
|
||||
"fetch_subs": bool(req.fetch_subs),
|
||||
"sub_langs": req.sub_langs or None,
|
||||
"cookie_file": cookie_path,
|
||||
"source_lang": source_lang_override,
|
||||
}
|
||||
try:
|
||||
await task_manager.add_task(
|
||||
@@ -1099,21 +1242,13 @@ async def dub_transcribe_stream(
|
||||
"error_code": failure["code"],
|
||||
}
|
||||
|
||||
# Retry a failed/timed-out chunk once on a fresh pool before giving
|
||||
# up. Otherwise a transient wedge on the FIRST chunk (whisperx often
|
||||
# cold-loads its model there, the #730 hang) drops that whole window
|
||||
# and the transcript is "missing the beginning, only middle+end".
|
||||
# The retry reuses the same audio window, so a recovered chunk fills
|
||||
# the hole instead of leaving silent gaps.
|
||||
# Retry an ordinary completed failure once. A timed-out native call
|
||||
# is different: its thread is still executing and must not overlap
|
||||
# a retry against the same backend (#1669).
|
||||
part = None
|
||||
timed_out = False
|
||||
for _attempt in range(1, _CHUNK_TRANSCRIBE_ATTEMPTS + 1):
|
||||
# A wedged chunk gets the SAME guarded-timeout + pool-reset
|
||||
# semantics as the whole-file paths (#730/#851):
|
||||
# run_transcribe_guarded bounds the call, abandons the poisoned
|
||||
# pool so the retry (and any concurrent TTS work) gets a fresh
|
||||
# worker, and raises the actionable ASRTimeoutError. Run it as
|
||||
# a task and poll so we can keep yielding pings — the
|
||||
# EventSource connection drops without them.
|
||||
# Run as a task and poll so pings keep the EventSource alive.
|
||||
task = asyncio.ensure_future(run_transcribe_guarded(
|
||||
_gpu_pool, _transcribe_chunk,
|
||||
what=f"Dub chunk {i + 1}/{chunks_n}",
|
||||
@@ -1128,9 +1263,12 @@ async def dub_transcribe_stream(
|
||||
try:
|
||||
part = task.result()
|
||||
except ASRTimeoutError:
|
||||
# The guard already reset the pool; keep the actionable
|
||||
# message (it names the durable fixes, and — after repeated
|
||||
# timeouts — the crash-isolated engine escape hatch).
|
||||
# Python cannot kill an in-process native transcribe. Do
|
||||
# not swap pools and retry over the still-running call:
|
||||
# concurrent whisperx/CTranslate2 access caused the native
|
||||
# Windows access violation in #1669. Stop this transcript;
|
||||
# the worker remains honestly occupied until it exits.
|
||||
timed_out = True
|
||||
logger.error(
|
||||
"Transcribe chunk %d/%d timed out after %.0fs (attempt %d/%d, job=%s)",
|
||||
i + 1, chunks_n, transcribe_timeout_s, _attempt,
|
||||
@@ -1149,20 +1287,22 @@ async def dub_transcribe_stream(
|
||||
# error-part; the timeout path already reset the pool).
|
||||
if part is not None and not part.get("error"):
|
||||
break
|
||||
if _attempt < _CHUNK_TRANSCRIBE_ATTEMPTS:
|
||||
if timed_out:
|
||||
break
|
||||
if not timed_out and _attempt < _CHUNK_TRANSCRIBE_ATTEMPTS:
|
||||
logger.warning(
|
||||
"Retrying transcribe chunk %d/%d after failure/timeout (next attempt %d/%d, job=%s)",
|
||||
i + 1, chunks_n, _attempt + 1, _CHUNK_TRANSCRIBE_ATTEMPTS, log_safe(job_id),
|
||||
)
|
||||
# A completed exception did not wedge the worker. Resetting
|
||||
# the pool here leaked a healthy executor on every ordinary
|
||||
# decode failure; run_transcribe_guarded already resets the
|
||||
# pool on the only case that needs it: a real timeout.
|
||||
# A completed exception did not leave native work behind,
|
||||
# so retrying this same audio window is safe.
|
||||
if part.get("error"):
|
||||
chunk_errors.append(part["error"])
|
||||
if part.get("error_code"):
|
||||
chunk_error_codes.append(part["error_code"])
|
||||
logger.warning("Chunk %d/%d error: %s", i + 1, chunks_n, log_safe(part["error"]))
|
||||
if timed_out:
|
||||
break
|
||||
if detected_lang is None and part.get("language"):
|
||||
detected_lang = part["language"]
|
||||
asr_speaker_turns.extend(part.get("speaker_turns") or [])
|
||||
@@ -1669,7 +1809,9 @@ async def dub_transcribe_stream(
|
||||
except Exception as e:
|
||||
logger.warning("speaker_clone extraction skipped: %s", e)
|
||||
|
||||
job["source_lang"] = ((detected_lang or "en").split("_")[0][:2] or "en").lower()
|
||||
job["source_lang"] = job.get("source_lang_override") or (
|
||||
(detected_lang or "en").split("_")[0][:2] or "en"
|
||||
).lower()
|
||||
job["full_transcript"] = " ".join(s.get("text", "") for s in final_segs)
|
||||
_save_job(job_id, job)
|
||||
|
||||
@@ -1866,7 +2008,9 @@ async def dub_transcribe(job_id: str, num_speakers: Optional[int] = None):
|
||||
except Exception as e:
|
||||
logger.warning("Failed to unload ASR backend: %s", e)
|
||||
|
||||
job["source_lang"] = (detected_lang or "en").split("_")[0][:2].lower()
|
||||
job["source_lang"] = job.get("source_lang_override") or (
|
||||
(detected_lang or "en").split("_")[0][:2] or "en"
|
||||
).lower()
|
||||
|
||||
scene_cuts = job.get("scene_cuts") or []
|
||||
segments = segment_transcript(result, duration=job.get("duration", 0.0), scene_cuts=scene_cuts)
|
||||
@@ -1922,7 +2066,8 @@ async def dub_transcribe(job_id: str, num_speakers: Optional[int] = None):
|
||||
# Bound the whole-file transcribe (#730): a wedged whisperx/CTranslate2
|
||||
# call would otherwise hold its GPU-pool worker forever and starve
|
||||
# every other request into a "can't reach backend". run_transcribe_guarded
|
||||
# also resets the pool on timeout so capacity is restored.
|
||||
# leaves an unkillable native worker accounted for on timeout so a
|
||||
# retry cannot overlap it (#1669).
|
||||
segments_result = await run_transcribe_guarded(_gpu_pool, _transcribe, what="Dub")
|
||||
except asyncio.CancelledError:
|
||||
job["aborted"] = True
|
||||
|
||||
@@ -1599,7 +1599,10 @@ async def dub_download_audio(
|
||||
return _native_save(wav_path, save_path, dl_name, media_type="audio/wav")
|
||||
return FileResponse(
|
||||
wav_path, media_type="audio/wav",
|
||||
headers={"Content-Disposition": content_disposition(dl_name)},
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"Content-Disposition": content_disposition(dl_name),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -503,6 +503,18 @@ async def dub_generate(job_id: str, req: DubRequest):
|
||||
backend = await resolve_generation_backend(require_cloning=True)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
from core.failure import is_gpu_oom
|
||||
|
||||
if not is_gpu_oom(e):
|
||||
raise
|
||||
from core.public_errors import public_exception_response
|
||||
|
||||
payload = public_exception_response(
|
||||
e,
|
||||
fallback="The TTS model could not be loaded.",
|
||||
)
|
||||
raise HTTPException(status_code=503, detail=payload["detail"]) from e
|
||||
|
||||
async def _stream(task_id):
|
||||
total = len(req.segments)
|
||||
@@ -1291,7 +1303,15 @@ async def dub_generate(job_id: str, req: DubRequest):
|
||||
pass
|
||||
_release_audio_tensors()
|
||||
except Exception as e:
|
||||
yield f"data: {json.dumps({'type': 'error', 'segment': i, 'error': str(e)})}\n\n"
|
||||
# A task-stream error bypasses the global exception handler.
|
||||
# Never publish engine exception text here: allocator errors
|
||||
# carry process tables and arbitrary failures can carry paths,
|
||||
# tokens, or source text. The shared helper enriches recognized
|
||||
# classes using VoiceStudio-owned constants only.
|
||||
from core.public_errors import stream_generation_failure
|
||||
|
||||
error_detail = stream_generation_failure(e)["detail"]
|
||||
yield f"data: {json.dumps({'type': 'error', 'segment': i, 'error': error_detail})}\n\n"
|
||||
sr = backend.sample_rate
|
||||
all_segment_wavs.append(_store_mix_wav(seg.start, seg.end, torch.zeros(1, max(0, int(seg_duration * sr))), sr, f"mix_{seg_id}"))
|
||||
sync_scores.append(1.0)
|
||||
|
||||
@@ -8,6 +8,7 @@ import asyncio
|
||||
import tempfile
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import traceback
|
||||
from typing import Optional
|
||||
from fastapi import APIRouter, File, Form, UploadFile, HTTPException
|
||||
@@ -32,6 +33,83 @@ router = APIRouter()
|
||||
logger = logging.getLogger("omnivoice.generate")
|
||||
|
||||
|
||||
class _TempReferenceLease:
|
||||
"""Delete a request-owned reference once every abandoned reader drains."""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = path
|
||||
self._lock = threading.Lock()
|
||||
self._active = 0
|
||||
self._request_done = False
|
||||
self._deleted = False
|
||||
|
||||
def acquire(self):
|
||||
with self._lock:
|
||||
if self._request_done:
|
||||
raise RuntimeError("reference lease acquired after request cleanup")
|
||||
self._active += 1
|
||||
once_lock = threading.Lock()
|
||||
released = False
|
||||
|
||||
def release() -> None:
|
||||
nonlocal released
|
||||
with once_lock:
|
||||
if released:
|
||||
return
|
||||
released = True
|
||||
self._release()
|
||||
|
||||
return release
|
||||
|
||||
def _release(self) -> None:
|
||||
delete = False
|
||||
with self._lock:
|
||||
self._active -= 1
|
||||
if self._active < 0:
|
||||
raise RuntimeError("reference lease released too many times")
|
||||
if self._request_done and self._active == 0 and not self._deleted:
|
||||
self._deleted = True
|
||||
delete = True
|
||||
if delete:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(self.path)
|
||||
|
||||
def finish_request(self) -> None:
|
||||
delete = False
|
||||
with self._lock:
|
||||
self._request_done = True
|
||||
if self._active == 0 and not self._deleted:
|
||||
self._deleted = True
|
||||
delete = True
|
||||
if delete:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(self.path)
|
||||
|
||||
|
||||
async def _run_with_reference_lease(lease, factory):
|
||||
"""Hold an ad-hoc reference through one local GPU-pool dispatch."""
|
||||
if lease is None:
|
||||
return await factory(None)
|
||||
release = lease.acquire()
|
||||
abandoned = False
|
||||
try:
|
||||
return await factory(release)
|
||||
except GpuPoolBusyError:
|
||||
# Busy means no job started; release now. The callback may already have
|
||||
# done so, and the lease token is deliberately idempotent.
|
||||
release()
|
||||
abandoned = True
|
||||
raise
|
||||
except (asyncio.CancelledError, GpuJobTimeoutError):
|
||||
# The guard owns release now: immediately for a queued cancellation,
|
||||
# or from the worker finalizer after an in-flight job drains.
|
||||
abandoned = True
|
||||
raise
|
||||
finally:
|
||||
if not abandoned:
|
||||
release()
|
||||
|
||||
|
||||
def _profile_instruct(row):
|
||||
"""Validator-safe instruct for a stored profile row.
|
||||
|
||||
@@ -381,6 +459,31 @@ def _is_timeout_failure(e) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _is_media_process_launch_failure(exc: BaseException) -> bool:
|
||||
"""Identify an ffmpeg/ffprobe launch ENOENT without guessing from a file name."""
|
||||
if not isinstance(exc, FileNotFoundError):
|
||||
return False
|
||||
|
||||
# A regular missing reference/model file may itself be named "ffmpeg".
|
||||
# Require the innermost raise site to be Python's process launcher so that
|
||||
# basename collisions keep the normal missing-file diagnosis (#1677).
|
||||
traceback_cursor = exc.__traceback__
|
||||
if traceback_cursor is None:
|
||||
return False
|
||||
while traceback_cursor.tb_next is not None:
|
||||
traceback_cursor = traceback_cursor.tb_next
|
||||
origin_module = traceback_cursor.tb_frame.f_globals.get("__name__", "")
|
||||
if origin_module != "subprocess" and not origin_module.startswith("asyncio."):
|
||||
return False
|
||||
|
||||
filename = getattr(exc, "filename", None)
|
||||
if not filename:
|
||||
return "[winerror 2]" in str(exc).lower()
|
||||
return os.path.basename(str(filename)).lower() in {
|
||||
"ffmpeg", "ffmpeg.exe", "ffprobe", "ffprobe.exe",
|
||||
}
|
||||
|
||||
|
||||
def _oom_friendly_reraise(e):
|
||||
"""Best-effort cache flush + the user-facing OOM hint shared by both
|
||||
inference paths."""
|
||||
@@ -405,6 +508,21 @@ def _oom_friendly_reraise(e):
|
||||
# that lost its +x bit) is NOT an OOM — don't send the user to the Flush
|
||||
# button; tell them what's actually wrong.
|
||||
es = str(e)
|
||||
# #1677: Windows CreateProcess reports a missing executable as a bare
|
||||
# ``FileNotFoundError: [WinError 2] ...`` with no filename, while POSIX
|
||||
# includes the missing ffmpeg/ffprobe name. The bundled-media downloader
|
||||
# now republishes PATH as soon as it finishes, but a failed/blocked
|
||||
# download still needs an actionable recovery rather than the unknown-
|
||||
# error dead end. Keep missing reference/model files on their own path.
|
||||
for _exc in _exception_chain(e):
|
||||
if _is_media_process_launch_failure(_exc):
|
||||
raise RuntimeError(
|
||||
"A required media program couldn't be launched. Open "
|
||||
"Settings → Audio tools and use "
|
||||
"Download/Repair for the media engine, then retry. If Audio "
|
||||
"tools is already ready, repair the selected TTS engine and "
|
||||
f"restart VoiceStudio. Underlying error: {_safe_exc_text(_exc)}"
|
||||
) from e
|
||||
if isinstance(e, PermissionError) or "Permission denied" in es or "Errno 13" in es:
|
||||
raise RuntimeError(
|
||||
f"A required engine binary couldn't be executed (permission denied). "
|
||||
@@ -695,15 +813,17 @@ def _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="broadcast",
|
||||
max_chunk_chars=None, crossfade_ms=None, *, dropped_sink=None,
|
||||
max_chunk_chars=None, crossfade_ms=None, *, t_shift=None,
|
||||
layer_penalty_factor=None, position_temperature=None,
|
||||
class_temperature=None, dropped_sink=None,
|
||||
):
|
||||
"""Engine-aware twin of :func:`_run_inference` (issue #312).
|
||||
|
||||
Runs the request through a pluggable ``TTSBackend`` adapter instead of the
|
||||
VoiceStudio model directly. The adapter protocol is narrower than the
|
||||
VoiceStudio-native surface — engine-specific extras (``t_shift``,
|
||||
``layer_penalty_factor``, …) only exist on the native path, which is why
|
||||
VoiceStudio itself still goes through ``_run_inference``.
|
||||
VoiceStudio model directly. A crash-isolated OmniVoice proxy advertises
|
||||
``supports_native_omnivoice_controls`` and receives the same advanced
|
||||
controls and per-call seed as the native path; other adapters keep the
|
||||
narrower protocol unchanged.
|
||||
"""
|
||||
import torch
|
||||
try:
|
||||
@@ -718,6 +838,18 @@ def _run_backend_inference(
|
||||
instruct=instruct, num_step=num_step, guidance_scale=guidance_scale,
|
||||
speed=speed, denoise=denoise, postprocess_output=postprocess_output,
|
||||
)
|
||||
native_proxy = bool(
|
||||
getattr(backend, "supports_native_omnivoice_controls", False)
|
||||
)
|
||||
if native_proxy:
|
||||
gen_kwargs.update({
|
||||
key: value for key, value in {
|
||||
"t_shift": t_shift,
|
||||
"layer_penalty_factor": layer_penalty_factor,
|
||||
"position_temperature": position_temperature,
|
||||
"class_temperature": class_temperature,
|
||||
}.items() if value is not None
|
||||
})
|
||||
sr = backend.sample_rate
|
||||
|
||||
# Inline [pause Nms] markers (issue #276) work for every engine — the
|
||||
@@ -727,10 +859,17 @@ def _run_backend_inference(
|
||||
has_pause = len(segments) > 1 or (segments and segments[0][1] > 0)
|
||||
|
||||
if has_pause:
|
||||
first_span = True
|
||||
|
||||
def _gen_span(span_text):
|
||||
nonlocal first_span
|
||||
# Per-span duration is left to the engine; an explicit overall
|
||||
# `duration` can't be meaningfully split across spans.
|
||||
return backend.generate(span_text, duration=None, **gen_kwargs)
|
||||
span_kwargs = dict(gen_kwargs)
|
||||
if native_proxy and first_span and used_seed is not None:
|
||||
span_kwargs["seed"] = used_seed
|
||||
first_span = False
|
||||
return backend.generate(span_text, duration=None, **span_kwargs)
|
||||
audio_out = _render_with_pauses(_gen_span, segments, sr)
|
||||
else:
|
||||
# Wave 1.2: sentence-boundary chunking for long text (see
|
||||
@@ -747,12 +886,19 @@ def _run_backend_inference(
|
||||
for i, chunk_text in enumerate(text_chunks):
|
||||
if used_seed is not None:
|
||||
torch.manual_seed(used_seed + i)
|
||||
parts.append(backend.generate(chunk_text, duration=None, **gen_kwargs))
|
||||
chunk_kwargs = dict(gen_kwargs)
|
||||
if native_proxy and used_seed is not None:
|
||||
chunk_kwargs["seed"] = used_seed + i
|
||||
parts.append(backend.generate(
|
||||
chunk_text, duration=None, **chunk_kwargs
|
||||
))
|
||||
_note_generate_progress()
|
||||
audio_out = concatenate_audio_chunks(parts, sr, _xfade_ms,
|
||||
texts=text_chunks,
|
||||
sink=dropped_sink)
|
||||
else:
|
||||
if native_proxy and used_seed is not None:
|
||||
gen_kwargs["seed"] = used_seed
|
||||
audio_out = backend.generate(text, duration=duration, **gen_kwargs)
|
||||
|
||||
return _apply_effect_chain(
|
||||
@@ -1297,6 +1443,7 @@ async def generate_speech(
|
||||
|
||||
ref_audio_path = None
|
||||
cleanup_ref = False
|
||||
ref_lease = None
|
||||
used_seed = seed
|
||||
resolved_profile_id = None
|
||||
history_mode = None # profile.kind when a profile drives; else inferred at insert
|
||||
@@ -1384,6 +1531,7 @@ async def generate_speech(
|
||||
f.write(await ref_audio.read())
|
||||
ref_audio_path = f.name
|
||||
cleanup_ref = True
|
||||
ref_lease = _TempReferenceLease(ref_audio_path)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@@ -1400,13 +1548,19 @@ async def generate_speech(
|
||||
# built-in ASR fallback), so a timeout degrades to None rather than
|
||||
# failing the whole generate.
|
||||
try:
|
||||
ref_text = await run_on_gpu_pool_guarded(
|
||||
functools.partial(transcribe_reference, ref_audio_path),
|
||||
what="Reference transcribe",
|
||||
# Floor budget (#1190): a reference clip is seconds of audio,
|
||||
# so the length-scaled bonus never applies — but the timeout is
|
||||
# explicit here too, so no dispatch relies on a hidden default.
|
||||
timeout=_generate_timeout_s("", execution_device=_routing["effective_device"]),
|
||||
ref_text = await _run_with_reference_lease(
|
||||
ref_lease,
|
||||
lambda release: run_on_gpu_pool_guarded(
|
||||
functools.partial(transcribe_reference, ref_audio_path),
|
||||
what="Reference transcribe",
|
||||
# Floor budget (#1190): a reference clip is seconds of audio,
|
||||
# so the length-scaled bonus never applies — but the timeout is
|
||||
# explicit here too, so no dispatch relies on a hidden default.
|
||||
timeout=_generate_timeout_s(
|
||||
"", execution_device=_routing["effective_device"]
|
||||
),
|
||||
on_abandon=release,
|
||||
)
|
||||
)
|
||||
# TimeoutError covers both the execution bound and pool saturation:
|
||||
# this path is best-effort either way.
|
||||
@@ -1662,19 +1816,30 @@ async def generate_speech(
|
||||
"target_label": e.worker_label or _target_label,
|
||||
"hint": e.hint,
|
||||
})
|
||||
except Exception:
|
||||
except Exception as exc:
|
||||
# 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")})
|
||||
# whether to spend the same minutes again on this machine. Like
|
||||
# the local streaming path, this in-band frame stands in for the
|
||||
# global 500 handler, so it journals the scrubbed failure and
|
||||
# names a recognized cause instead of the bare generic string
|
||||
# (#1607).
|
||||
logger.error(
|
||||
"Remote generation failed (class=%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
from core.public_errors import stream_generation_failure
|
||||
from core import error_journal
|
||||
|
||||
error_journal.record(
|
||||
exc, route="/generate", trace=traceback.format_exc()
|
||||
)
|
||||
yield _line({"type": "error", **stream_generation_failure(exc)})
|
||||
finally:
|
||||
if not render.done():
|
||||
render.cancel()
|
||||
if cleanup_ref and ref_audio_path:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(ref_audio_path)
|
||||
if cleanup_ref and ref_lease is not None:
|
||||
ref_lease.finish_request()
|
||||
|
||||
return StreamingResponse(
|
||||
_remote_stream_events(),
|
||||
@@ -1716,6 +1881,17 @@ async def generate_speech(
|
||||
instruct=instruct, num_step=num_step,
|
||||
guidance_scale=guidance_scale, speed=speed,
|
||||
denoise=denoise, postprocess_output=postprocess_output,
|
||||
**({
|
||||
key: value for key, value in {
|
||||
"t_shift": t_shift,
|
||||
"layer_penalty_factor": layer_penalty_factor,
|
||||
"position_temperature": position_temperature,
|
||||
"class_temperature": class_temperature,
|
||||
"seed": used_seed + i if used_seed is not None else None,
|
||||
}.items() if value is not None
|
||||
} if getattr(
|
||||
_backend, "supports_native_omnivoice_controls", False
|
||||
) else {}),
|
||||
)
|
||||
sr = _backend.sample_rate
|
||||
skip = getattr(_backend, "applies_own_mastering", False)
|
||||
@@ -1779,33 +1955,45 @@ async def generate_speech(
|
||||
if _has_pause or len(_text_chunks) <= 1:
|
||||
# Single-shot pipeline, unchanged — streamed as one chunk.
|
||||
if _backend is not None:
|
||||
audio_tensor = await run_on_gpu_pool_guarded(
|
||||
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_sink,
|
||||
),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
timeout=_generate_timeout_s(text, execution_device=_routing["effective_device"]),
|
||||
audio_tensor = await _run_with_reference_lease(
|
||||
ref_lease,
|
||||
lambda release: run_on_gpu_pool_guarded(
|
||||
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, t_shift=t_shift,
|
||||
layer_penalty_factor=layer_penalty_factor,
|
||||
position_temperature=position_temperature,
|
||||
class_temperature=class_temperature,
|
||||
dropped_sink=_dropped_sink,
|
||||
),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
timeout=_generate_timeout_s(text, execution_device=_routing["effective_device"]),
|
||||
on_abandon=release,
|
||||
)
|
||||
)
|
||||
sample_rate = _backend.sample_rate
|
||||
else:
|
||||
audio_tensor = await run_on_gpu_pool_guarded(
|
||||
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_sink,
|
||||
),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
timeout=_generate_timeout_s(text, execution_device=_routing["effective_device"]),
|
||||
audio_tensor = await _run_with_reference_lease(
|
||||
ref_lease,
|
||||
lambda release: run_on_gpu_pool_guarded(
|
||||
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_sink,
|
||||
),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
timeout=_generate_timeout_s(text, execution_device=_routing["effective_device"]),
|
||||
on_abandon=release,
|
||||
)
|
||||
)
|
||||
sample_rate = _model.sampling_rate
|
||||
yield _line({
|
||||
@@ -1834,14 +2022,18 @@ async def generate_speech(
|
||||
for i, chunk_text in enumerate(_text_chunks):
|
||||
# Bounded per chunk + pool-reset on hang (#730 class);
|
||||
# a timeout surfaces as an "error" event below.
|
||||
raw, preview, sample_rate = await run_on_gpu_pool_guarded(
|
||||
functools.partial(_render_stream_chunk, i, chunk_text),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
# Budget scaled to THIS chunk (#1190) — the flat
|
||||
# 300s here is what made long streamed renders fail
|
||||
# even after the v0.3.22 scaled budget shipped.
|
||||
timeout=_generate_timeout_s(chunk_text, execution_device=_routing["effective_device"]),
|
||||
raw, preview, sample_rate = await _run_with_reference_lease(
|
||||
ref_lease,
|
||||
lambda release: run_on_gpu_pool_guarded(
|
||||
functools.partial(_render_stream_chunk, i, chunk_text),
|
||||
what="TTS generate",
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
# Budget scaled to THIS chunk (#1190) — the flat
|
||||
# 300s here is what made long streamed renders fail
|
||||
# even after the v0.3.22 scaled budget shipped.
|
||||
timeout=_generate_timeout_s(chunk_text, execution_device=_routing["effective_device"]),
|
||||
on_abandon=release,
|
||||
)
|
||||
)
|
||||
parts.append(raw)
|
||||
# Provenance-mark the streamed copy off the GPU pool
|
||||
@@ -1915,16 +2107,32 @@ async def generate_speech(
|
||||
logger.error("Streaming generation request rejected")
|
||||
from core.public_errors import stream_failure
|
||||
yield _line({"type": "error", **stream_failure("invalid_request")})
|
||||
except Exception:
|
||||
logger.error("Streaming generation failed unexpectedly")
|
||||
from core.public_errors import stream_failure
|
||||
yield _line({"type": "error", **stream_failure("generation_failed")})
|
||||
except Exception as exc:
|
||||
# A streaming request answers 200 and carries its failure as an
|
||||
# in-band error frame, so it never reaches the global 500
|
||||
# handler — which is where a classic /generate failure gets its
|
||||
# scrubbed journal entry (Diagnostics / recent errors) AND its
|
||||
# classified, actionable message. Both have to be reproduced
|
||||
# here or a streaming generation failure is invisible in the
|
||||
# diagnostic bundle and opaque to the user (#1607). The raw
|
||||
# exception is NOT logged: it can carry a reference-clip path or
|
||||
# a provider secret, and only the journal scrubs before storing.
|
||||
logger.error(
|
||||
"Streaming generation failed unexpectedly (class=%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
from core.public_errors import stream_generation_failure
|
||||
from core import error_journal
|
||||
|
||||
error_journal.record(
|
||||
exc, route="/generate", trace=traceback.format_exc()
|
||||
)
|
||||
yield _line({"type": "error", **stream_generation_failure(exc)})
|
||||
finally:
|
||||
# Ownership of the temp reference clip moves to this generator
|
||||
# in stream mode (the route returns before rendering starts).
|
||||
if cleanup_ref and ref_audio_path:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(ref_audio_path)
|
||||
if cleanup_ref and ref_lease is not None:
|
||||
ref_lease.finish_request()
|
||||
|
||||
# Routing notice (#21): known before the stream starts, so it rides the
|
||||
# same headers the classic path uses — and now also carries "your
|
||||
@@ -1962,7 +2170,11 @@ 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, dropped_sink=_dropped_text,
|
||||
max_chunk_chars, crossfade_ms, t_shift=t_shift,
|
||||
layer_penalty_factor=layer_penalty_factor,
|
||||
position_temperature=position_temperature,
|
||||
class_temperature=class_temperature,
|
||||
dropped_sink=_dropped_text,
|
||||
)
|
||||
else:
|
||||
_local_render = functools.partial(
|
||||
@@ -1973,14 +2185,18 @@ async def generate_speech(
|
||||
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, execution_device=_routing["effective_device"]),
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
),
|
||||
decision=_decision,
|
||||
audio_tensor = await _run_with_reference_lease(
|
||||
ref_lease,
|
||||
lambda release: gpu_gateway.run(
|
||||
_REMOTE_OP,
|
||||
local=gpu_gateway.LocalCall(
|
||||
_local_render, what="TTS generate",
|
||||
timeout=_generate_timeout_s(text, execution_device=_routing["effective_device"]),
|
||||
min_vram_gb=_engine_min_vram_gb,
|
||||
on_abandon=release,
|
||||
),
|
||||
decision=_decision,
|
||||
)
|
||||
)
|
||||
# Read after generation: engines with lazy model loading report
|
||||
# their real rate only once weights are up.
|
||||
@@ -2112,9 +2328,8 @@ async def generate_speech(
|
||||
),
|
||||
)
|
||||
finally:
|
||||
if cleanup_ref and ref_audio_path:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(ref_audio_path)
|
||||
if cleanup_ref and ref_lease is not None:
|
||||
ref_lease.finish_request()
|
||||
|
||||
def _safe_output_path(name):
|
||||
if not name:
|
||||
|
||||
@@ -76,6 +76,7 @@ _cancelled: set[str] = set()
|
||||
_active_installs: set[str] = set()
|
||||
_active_installs_lock = threading.Lock()
|
||||
_install_tasks: set[asyncio.Task] = set()
|
||||
_install_tasks_by_repo: dict[str, asyncio.Task] = {}
|
||||
|
||||
|
||||
def _download_max_workers() -> int:
|
||||
@@ -420,16 +421,11 @@ async def install_model(req: InstallModelRequest):
|
||||
f"Retry in {remaining}s or check your network."
|
||||
),
|
||||
)
|
||||
with _active_installs_lock:
|
||||
if req.repo_id in _active_installs:
|
||||
return {"status": "already_running", "repo_id": req.repo_id}
|
||||
_active_installs.add(req.repo_id)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _do():
|
||||
token = hf_progress.current_repo_id.set(req.repo_id)
|
||||
target_token = hf_progress.current_target.set("local")
|
||||
_cancelled.discard(req.repo_id) # clear any stale cancel from a prior run
|
||||
hf_progress.emit({
|
||||
"repo_id": req.repo_id,
|
||||
"filename": req.repo_id,
|
||||
@@ -688,17 +684,53 @@ async def install_model(req: InstallModelRequest):
|
||||
with _active_installs_lock:
|
||||
_active_installs.discard(req.repo_id)
|
||||
|
||||
try:
|
||||
task = loop.create_task(asyncio.to_thread(_do))
|
||||
_install_tasks.add(task)
|
||||
task.add_done_callback(_install_tasks.discard)
|
||||
except Exception:
|
||||
with _active_installs_lock:
|
||||
with _active_installs_lock:
|
||||
if req.repo_id in _active_installs:
|
||||
return {"status": "already_running", "repo_id": req.repo_id}
|
||||
_active_installs.add(req.repo_id)
|
||||
# Admission and task publication are one atomic generation boundary:
|
||||
# cancellation can never observe an admitted install without its task.
|
||||
_cancelled.discard(req.repo_id)
|
||||
try:
|
||||
task = loop.create_task(asyncio.to_thread(_do))
|
||||
_install_tasks.add(task)
|
||||
_install_tasks_by_repo[req.repo_id] = task
|
||||
except Exception:
|
||||
_active_installs.discard(req.repo_id)
|
||||
raise
|
||||
raise
|
||||
|
||||
def install_finished(completed: asyncio.Task) -> None:
|
||||
with _active_installs_lock:
|
||||
_install_tasks.discard(completed)
|
||||
if _install_tasks_by_repo.get(req.repo_id) is completed:
|
||||
_install_tasks_by_repo.pop(req.repo_id, None)
|
||||
|
||||
task.add_done_callback(install_finished)
|
||||
return {"status": "install_started", "repo_id": req.repo_id}
|
||||
|
||||
|
||||
async def cancel_install_and_wait(repo_id: str) -> None:
|
||||
"""Request cancellation and retain authority until its thread exits."""
|
||||
from worker.async_utils import drain_task # noqa: PLC0415
|
||||
|
||||
with _active_installs_lock:
|
||||
_cancelled.add(repo_id)
|
||||
_install_cooldowns.pop(repo_id, None)
|
||||
task = _install_tasks_by_repo.get(repo_id)
|
||||
if task is None:
|
||||
return
|
||||
try:
|
||||
# asyncio.to_thread cannot stop snapshot_download mid-file. Cancelling
|
||||
# its wrapper would only detach the thread, so wait until the blocking
|
||||
# call observes the flag or naturally returns.
|
||||
await drain_task(task)
|
||||
finally:
|
||||
with _active_installs_lock:
|
||||
current = _install_tasks_by_repo.get(repo_id)
|
||||
if current is None or current is task:
|
||||
_cancelled.discard(repo_id)
|
||||
|
||||
|
||||
@router.post("/models/install/cancel")
|
||||
async def cancel_install(req: InstallModelRequest):
|
||||
"""Request cancellation of an in-flight install (FDL-11).
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Discovery contract for VoiceStudio's local speech platform.
|
||||
|
||||
Interfaces should discover this document instead of hard-coding whichever
|
||||
dictation route the desktop happens to use. Endpoint URLs are relative so the
|
||||
same response works on loopback, a tailnet GPU host, and a reverse proxy.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from core.version import APP_VERSION
|
||||
|
||||
router = APIRouter(tags=["Speech Platform"])
|
||||
|
||||
SPEECH_PROTOCOL = "voicestudio.speech.v1"
|
||||
STREAM_PATH = "/v1/audio/transcriptions/stream"
|
||||
|
||||
|
||||
class EndpointCapability(BaseModel):
|
||||
path: str
|
||||
transport: Literal["http", "websocket", "mcp-streamable-http", "mcp-stdio"]
|
||||
method: str | None = None
|
||||
protocol: str | None = None
|
||||
|
||||
|
||||
class StreamInputCapability(BaseModel):
|
||||
framing: Literal["binary"] = "binary"
|
||||
formats: list[str]
|
||||
default_format: str
|
||||
sample_rate_query: str = "sr"
|
||||
end_control: dict[str, str]
|
||||
|
||||
|
||||
class StreamOutputCapability(BaseModel):
|
||||
framing: Literal["json"] = "json"
|
||||
events: list[str]
|
||||
final_kinds: list[str]
|
||||
|
||||
|
||||
class SpeechFeatureCapabilities(BaseModel):
|
||||
batch_transcription: bool = True
|
||||
streaming_transcription: bool = True
|
||||
partial_transcripts: bool = True
|
||||
utterance_finals: bool = True
|
||||
session_summary: bool = True
|
||||
word_timestamps: bool = True
|
||||
local_refinement: bool = True
|
||||
acoustic_echo_cancellation: bool = True
|
||||
native_dictation_control: bool = False
|
||||
|
||||
|
||||
class SpeechAuthCapabilities(BaseModel):
|
||||
loopback: Literal["none"] = "none"
|
||||
remote: Literal["bearer"] = "bearer"
|
||||
header: str = "Authorization: Bearer <OMNIVOICE_API_KEY>"
|
||||
browser_session_endpoint: str = "/api/auth/session"
|
||||
websocket_ticket_endpoint: str = "/api/auth/ws-ticket"
|
||||
websocket_ticket_query_parameter: Literal["ws_ticket"] = "ws_ticket"
|
||||
|
||||
|
||||
class SpeechCapabilities(BaseModel):
|
||||
schema_: Literal["voicestudio.speech-capabilities"] = Field(
|
||||
default="voicestudio.speech-capabilities",
|
||||
serialization_alias="schema",
|
||||
)
|
||||
protocol: Literal["voicestudio.speech.v1"] = SPEECH_PROTOCOL
|
||||
protocol_version: Literal["1.0"] = "1.0"
|
||||
service: str = "VoiceStudio"
|
||||
service_version: str = APP_VERSION
|
||||
local_first: bool = True
|
||||
endpoints: dict[str, EndpointCapability]
|
||||
stream_input: StreamInputCapability
|
||||
stream_output: StreamOutputCapability
|
||||
features: SpeechFeatureCapabilities
|
||||
authentication: SpeechAuthCapabilities
|
||||
|
||||
|
||||
def speech_capabilities() -> SpeechCapabilities:
|
||||
"""Return the stable, side-effect-free integration contract."""
|
||||
endpoints = {
|
||||
"capabilities": EndpointCapability(
|
||||
path="/.well-known/voicestudio-speech",
|
||||
transport="http",
|
||||
method="GET",
|
||||
),
|
||||
"batch_transcription": EndpointCapability(
|
||||
path="/v1/audio/transcriptions",
|
||||
transport="http",
|
||||
method="POST",
|
||||
protocol="openai.audio.transcriptions",
|
||||
),
|
||||
"streaming_transcription": EndpointCapability(
|
||||
path=STREAM_PATH,
|
||||
transport="websocket",
|
||||
protocol=SPEECH_PROTOCOL,
|
||||
),
|
||||
"mcp": EndpointCapability(
|
||||
path="/mcp",
|
||||
transport="mcp-streamable-http",
|
||||
method="POST",
|
||||
protocol="mcp",
|
||||
),
|
||||
"mcp_stdio": EndpointCapability(
|
||||
path="python -m backend.mcp_shim",
|
||||
transport="mcp-stdio",
|
||||
protocol="mcp",
|
||||
),
|
||||
}
|
||||
native_control = False
|
||||
try:
|
||||
control_port = int(os.environ.get("VOICESTUDIO_SPEECH_CONTROL_PORT", ""))
|
||||
except (TypeError, ValueError):
|
||||
control_port = 0
|
||||
if 0 < control_port <= 65535:
|
||||
native_control = True
|
||||
endpoints["native_dictation_control"] = EndpointCapability(
|
||||
path=f"http://127.0.0.1:{control_port}/v1/capabilities",
|
||||
transport="http",
|
||||
method="GET",
|
||||
protocol=SPEECH_PROTOCOL,
|
||||
)
|
||||
|
||||
return SpeechCapabilities(
|
||||
endpoints=endpoints,
|
||||
stream_input=StreamInputCapability(
|
||||
formats=[
|
||||
"audio/pcm;encoding=s16le;channels=1",
|
||||
"audio/webm;codecs=opus",
|
||||
],
|
||||
default_format="audio/webm;codecs=opus",
|
||||
end_control={"type": "input_audio.end"},
|
||||
),
|
||||
stream_output=StreamOutputCapability(
|
||||
events=["session.started", "status", "partial", "final", "error"],
|
||||
final_kinds=["utterance", "summary"],
|
||||
),
|
||||
features=SpeechFeatureCapabilities(
|
||||
native_dictation_control=native_control,
|
||||
),
|
||||
authentication=SpeechAuthCapabilities(),
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/.well-known/voicestudio-speech",
|
||||
response_model=SpeechCapabilities,
|
||||
response_model_by_alias=True,
|
||||
)
|
||||
@router.get(
|
||||
"/v1/audio/capabilities",
|
||||
response_model=SpeechCapabilities,
|
||||
response_model_by_alias=True,
|
||||
)
|
||||
async def get_speech_capabilities() -> SpeechCapabilities:
|
||||
"""Advertise batch, streaming, and agent-facing speech transports."""
|
||||
return speech_capabilities()
|
||||
@@ -39,6 +39,26 @@ CHUNK_SAMPLES = int(os.environ.get("OMNIVOICE_STREAM_CHUNK", "4800"))
|
||||
_perf_counter = time.perf_counter
|
||||
|
||||
|
||||
async def _resolve_stream_backend(engine_id: str | None):
|
||||
"""Resolve the live-stream engine without bypassing host isolation."""
|
||||
from services.tts_backend import (
|
||||
OmniVoiceBackend,
|
||||
active_backend_id,
|
||||
get_active_tts_backend,
|
||||
get_backend_class,
|
||||
)
|
||||
|
||||
if engine_id:
|
||||
return get_backend_class(engine_id)()
|
||||
|
||||
cls = get_backend_class(active_backend_id())
|
||||
if cls is OmniVoiceBackend:
|
||||
from services.model_manager import get_model
|
||||
|
||||
return get_active_tts_backend(model=await get_model())
|
||||
return get_active_tts_backend()
|
||||
|
||||
|
||||
class StreamTTSRequest(BaseModel):
|
||||
"""Client request for streaming TTS."""
|
||||
text: str
|
||||
@@ -132,10 +152,6 @@ async def ws_tts(websocket: WebSocket):
|
||||
|
||||
try:
|
||||
# Resolve engine
|
||||
from services.tts_backend import (
|
||||
get_active_tts_backend,
|
||||
get_backend_class,
|
||||
)
|
||||
engine_id = data.get("engine")
|
||||
# #1224: leave a breadcrumb when memory is already tight before
|
||||
# a heavy load. /generate has done this since the 16 GB-Mac
|
||||
@@ -151,13 +167,7 @@ async def ws_tts(websocket: WebSocket):
|
||||
log_if_low(f"TTS stream load ({engine_id or 'active engine'})")
|
||||
except Exception:
|
||||
pass
|
||||
if engine_id:
|
||||
cls = get_backend_class(engine_id)
|
||||
backend = cls()
|
||||
else:
|
||||
from services.model_manager import get_model
|
||||
model = await get_model()
|
||||
backend = get_active_tts_backend(model=model)
|
||||
backend = await _resolve_stream_backend(engine_id)
|
||||
|
||||
# ── Routing gate (#21 — no silent CPU fallback). WebSockets have
|
||||
# no response headers, so this uses frames: an error frame +
|
||||
|
||||
+266
-55
@@ -23,13 +23,16 @@ appears and is replaced by the GPU gateway.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from api.dependencies import require_admin
|
||||
from worker import registry, routing, service
|
||||
from worker.async_utils import drain_task, to_thread_and_defer_cancellation
|
||||
|
||||
logger = logging.getLogger("omnivoice.worker")
|
||||
|
||||
@@ -158,6 +161,19 @@ def agent_status() -> dict:
|
||||
return worker_agent.agent.status()
|
||||
|
||||
|
||||
@router.get("/agent/readiness", include_in_schema=False, response_model=None)
|
||||
def agent_readiness() -> JSONResponse:
|
||||
"""Container readiness: 200 only after this process registered as a worker."""
|
||||
from worker import agent as worker_agent # noqa: PLC0415
|
||||
|
||||
readiness = worker_agent.agent.readiness()
|
||||
return JSONResponse(
|
||||
status_code=200 if readiness["ready"] else 503,
|
||||
content=readiness,
|
||||
headers={} if readiness["ready"] else {"Retry-After": "2"},
|
||||
)
|
||||
|
||||
|
||||
def _refuse_when_env_pinned(worker_agent) -> None:
|
||||
"""OMNIVOICE_WORKER_MODE wins over the setting everywhere else.
|
||||
|
||||
@@ -176,6 +192,63 @@ def _refuse_when_env_pinned(worker_agent) -> None:
|
||||
)
|
||||
|
||||
|
||||
async def _finish_cleanup(awaitable):
|
||||
"""Run rollback to completion even if its HTTP task was cancelled."""
|
||||
task = asyncio.create_task(awaitable)
|
||||
await drain_task(task)
|
||||
return task.result()
|
||||
|
||||
|
||||
async def _set_worker_mode(worker_agent, enabled: bool) -> None:
|
||||
_result, cancelled = await to_thread_and_defer_cancellation(
|
||||
worker_agent.set_worker_mode_enabled, enabled
|
||||
)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
|
||||
async def _restore_agent_transaction(
|
||||
worker_agent, previous: dict, *, was_running: bool
|
||||
) -> None:
|
||||
"""Restore durable enrollment/settings and the exact prior live state."""
|
||||
try:
|
||||
await _finish_cleanup(worker_agent.agent.stop())
|
||||
await _finish_cleanup(worker_agent.restore_enrollment(previous))
|
||||
if was_running and not worker_agent.agent.running:
|
||||
await _finish_cleanup(worker_agent.agent.start())
|
||||
elif not was_running and worker_agent.agent.running:
|
||||
await _finish_cleanup(worker_agent.agent.stop())
|
||||
except worker_agent.EnrollmentRollbackError:
|
||||
raise
|
||||
except BaseException as exc:
|
||||
message = (
|
||||
"The previous worker state could not be restored safely. "
|
||||
"Worker mode remains stopped; fix its enrollment/settings storage, then retry."
|
||||
)
|
||||
with contextlib.suppress(BaseException):
|
||||
await _finish_cleanup(worker_agent.agent.stop())
|
||||
worker_agent.agent.last_error = message
|
||||
raise worker_agent.EnrollmentRollbackError(message) from exc
|
||||
|
||||
|
||||
def _raise_agent_transaction_failure(
|
||||
worker_agent, operation: BaseException, rollback: BaseException | None
|
||||
) -> None:
|
||||
if isinstance(operation, asyncio.CancelledError):
|
||||
if rollback is not None:
|
||||
logger.error(
|
||||
"Worker rollback failed during request cancellation",
|
||||
exc_info=(type(rollback), rollback, rollback.__traceback__),
|
||||
)
|
||||
raise operation
|
||||
if rollback is not None:
|
||||
raise HTTPException(status_code=409, detail=str(rollback)) from rollback
|
||||
if isinstance(operation, Exception):
|
||||
worker_agent.agent.last_error = str(operation)
|
||||
raise HTTPException(status_code=409, detail=str(operation)) from operation
|
||||
raise operation
|
||||
|
||||
|
||||
@router.post("/agent/join")
|
||||
async def join_control_plane(request: JoinRequest) -> dict:
|
||||
"""Redeem a join code and start working for that control plane.
|
||||
@@ -200,26 +273,37 @@ async def join_control_plane(request: JoinRequest) -> dict:
|
||||
# says it joined and never lends anything (CodeRabbit).
|
||||
_refuse_when_env_pinned(worker_agent)
|
||||
async with worker_agent.agent.lifecycle:
|
||||
# A rejoin replaces a working enrollment. Keep enough to put it back:
|
||||
# pinning the new certificate overwrites the old one on disk, so a
|
||||
# failed rejoin would otherwise leave the machine unable to reconnect
|
||||
# to the control plane it was already serving.
|
||||
previous = worker_agent.snapshot_enrollment()
|
||||
await worker_agent.agent.stop()
|
||||
try:
|
||||
previous, cancelled = await to_thread_and_defer_cancellation(
|
||||
worker_agent.snapshot_enrollment
|
||||
)
|
||||
except worker_agent.EnrollmentStateError as exc:
|
||||
worker_agent.agent.last_error = str(exc)
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
was_running = worker_agent.agent.running
|
||||
|
||||
# A rejoin stops a working agent before the replacement is accepted.
|
||||
# Stop, acceptance and the durable setting are one transaction: every
|
||||
# failure, including cancellation, restores both trust and live state.
|
||||
try:
|
||||
await worker_agent.agent.stop()
|
||||
await worker_agent.agent.start(token_text=token)
|
||||
# Success is the control plane ACCEPTING this worker, not the
|
||||
# connection being scheduled — see wait_until_registered.
|
||||
await worker_agent.agent.wait_until_registered()
|
||||
except Exception as exc:
|
||||
worker_agent.agent.last_error = str(exc)
|
||||
await worker_agent.agent.stop()
|
||||
await worker_agent.restore_enrollment(previous)
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
await _set_worker_mode(worker_agent, True)
|
||||
except BaseException as exc:
|
||||
rollback_exc = None
|
||||
try:
|
||||
await _restore_agent_transaction(
|
||||
worker_agent, previous, was_running=was_running
|
||||
)
|
||||
except BaseException as rollback_error:
|
||||
rollback_exc = rollback_error
|
||||
_raise_agent_transaction_failure(worker_agent, exc, rollback_exc)
|
||||
worker_agent.agent.last_error = ""
|
||||
# Persisted only after the join actually worked: a machine that failed
|
||||
# to enrol must not come back up trying again forever.
|
||||
worker_agent.set_worker_mode_enabled(True)
|
||||
return worker_agent.agent.status()
|
||||
|
||||
|
||||
@@ -235,19 +319,35 @@ async def set_agent_enabled(request: EnableRequest) -> dict:
|
||||
|
||||
_refuse_when_env_pinned(worker_agent)
|
||||
async with worker_agent.agent.lifecycle:
|
||||
if request.enabled:
|
||||
try:
|
||||
try:
|
||||
previous, cancelled = await to_thread_and_defer_cancellation(
|
||||
worker_agent.snapshot_enrollment
|
||||
)
|
||||
except worker_agent.EnrollmentStateError as exc:
|
||||
worker_agent.agent.last_error = str(exc)
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
was_running = worker_agent.agent.running
|
||||
|
||||
try:
|
||||
if request.enabled:
|
||||
await worker_agent.agent.start()
|
||||
await worker_agent.agent.wait_until_registered()
|
||||
except Exception as exc:
|
||||
worker_agent.agent.last_error = str(exc)
|
||||
await _set_worker_mode(worker_agent, True)
|
||||
else:
|
||||
await worker_agent.agent.stop()
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
worker_agent.agent.last_error = ""
|
||||
worker_agent.set_worker_mode_enabled(True)
|
||||
else:
|
||||
await worker_agent.agent.stop()
|
||||
worker_agent.set_worker_mode_enabled(False)
|
||||
await _set_worker_mode(worker_agent, False)
|
||||
except BaseException as exc:
|
||||
rollback_exc = None
|
||||
try:
|
||||
await _restore_agent_transaction(
|
||||
worker_agent, previous, was_running=was_running
|
||||
)
|
||||
except BaseException as rollback_error:
|
||||
rollback_exc = rollback_error
|
||||
_raise_agent_transaction_failure(worker_agent, exc, rollback_exc)
|
||||
worker_agent.agent.last_error = ""
|
||||
return worker_agent.agent.status()
|
||||
|
||||
|
||||
@@ -263,9 +363,14 @@ def create_enrollment(request: EnrollRequest) -> dict:
|
||||
status_code=409,
|
||||
detail="Remote workers are turned off. Enable them in Settings → System → Remote workers first.",
|
||||
)
|
||||
token = service.control_plane.create_enrollment(
|
||||
endpoint=request.endpoint, label=request.label, ttl_seconds=request.ttl_seconds
|
||||
)
|
||||
try:
|
||||
token = service.control_plane.create_enrollment(
|
||||
endpoint=request.endpoint,
|
||||
label=request.label,
|
||||
ttl_seconds=request.ttl_seconds,
|
||||
)
|
||||
except service.EndpointCertificateError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
return {
|
||||
"token": token.encode(),
|
||||
"endpoint": token.endpoint,
|
||||
@@ -275,23 +380,55 @@ def create_enrollment(request: EnrollRequest) -> dict:
|
||||
}
|
||||
|
||||
|
||||
def _persist_worker_update(
|
||||
worker_id: str, request: WorkerUpdate
|
||||
):
|
||||
"""Write policy on a worker thread; live publication stays loop-owned."""
|
||||
return registry.update_policy(
|
||||
worker_id,
|
||||
name=request.name,
|
||||
enabled=request.enabled,
|
||||
priority=request.priority,
|
||||
)
|
||||
|
||||
|
||||
@router.patch("/{worker_id}")
|
||||
def update_worker(worker_id: str, request: WorkerUpdate) -> dict:
|
||||
worker = registry.get(worker_id)
|
||||
if worker is None:
|
||||
async def update_worker(worker_id: str, request: WorkerUpdate) -> dict:
|
||||
pool = service.control_plane.pool if service.control_plane.running else None
|
||||
live = None
|
||||
was_pending = False
|
||||
if pool is not None:
|
||||
# Quiesce dispatch before releasing authority for the SQLite write.
|
||||
# The publication after the await restores the exact prior state, so a
|
||||
# concurrent registration handoff remains quiesced for its own reason.
|
||||
with registry.authority_guard():
|
||||
live = pool.get(worker_id)
|
||||
if live is not None:
|
||||
was_pending = live.registration_pending
|
||||
live.registration_pending = True
|
||||
updated = None
|
||||
cancelled = False
|
||||
try:
|
||||
updated, cancelled = await to_thread_and_defer_cancellation(
|
||||
_persist_worker_update, worker_id, request
|
||||
)
|
||||
finally:
|
||||
if pool is not None:
|
||||
with registry.authority_guard():
|
||||
if updated is not None:
|
||||
# Pool state, including the cached record the scheduler
|
||||
# reads, belongs to the app's event loop.
|
||||
pool.refresh_record(updated)
|
||||
current = pool.get(worker_id)
|
||||
if current is live:
|
||||
current.registration_pending = was_pending
|
||||
if updated is None:
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
raise HTTPException(status_code=404, detail="No such worker.")
|
||||
if request.name is not None:
|
||||
registry.rename(worker_id, request.name)
|
||||
if request.enabled is not None:
|
||||
registry.set_enabled(worker_id, request.enabled)
|
||||
if request.priority is not None:
|
||||
registry.set_priority(worker_id, request.priority)
|
||||
updated = registry.get(worker_id)
|
||||
# Keep the live copy in step, so the scheduler and its logs do not go on
|
||||
# using the name or priority this worker had when it connected.
|
||||
if updated is not None and service.control_plane.running:
|
||||
service.control_plane.pool.refresh_record(updated)
|
||||
return updated.to_dict() if updated else {}
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return updated.to_dict()
|
||||
|
||||
|
||||
@router.post("/{worker_id}/consent")
|
||||
@@ -305,7 +442,7 @@ def grant_consent(worker_id: str) -> dict:
|
||||
|
||||
|
||||
@router.post("/{worker_id}/resume")
|
||||
def clear_breaker(worker_id: str) -> dict:
|
||||
async def clear_breaker(worker_id: str) -> dict:
|
||||
"""Clear a paused worker's circuit breakers.
|
||||
|
||||
The user fixed the machine and knows it — a breaker with no manual clear is
|
||||
@@ -320,18 +457,53 @@ def clear_breaker(worker_id: str) -> dict:
|
||||
|
||||
|
||||
@router.delete("/{worker_id}")
|
||||
def revoke_worker(worker_id: str) -> dict:
|
||||
async def revoke_worker(worker_id: str) -> dict:
|
||||
"""Remove a worker — which means revoke its key, not hide the row.
|
||||
|
||||
Its in-flight work is released so it can be retried elsewhere rather than
|
||||
waiting out a lease on a machine that will never answer again.
|
||||
"""
|
||||
if registry.get(worker_id) is None:
|
||||
pool = service.control_plane.pool if service.control_plane.running else None
|
||||
live = None
|
||||
was_pending = False
|
||||
if pool is not None:
|
||||
with registry.authority_guard():
|
||||
live = pool.get(worker_id)
|
||||
if live is not None:
|
||||
was_pending = live.registration_pending
|
||||
live.registration_pending = True
|
||||
try:
|
||||
revoked, cancelled = await to_thread_and_defer_cancellation(
|
||||
registry.revoke, worker_id
|
||||
)
|
||||
except BaseException:
|
||||
if pool is not None:
|
||||
with registry.authority_guard():
|
||||
current = pool.get(worker_id)
|
||||
if current is live:
|
||||
current.registration_pending = was_pending
|
||||
raise
|
||||
if not revoked:
|
||||
if pool is not None:
|
||||
with registry.authority_guard():
|
||||
current = pool.get(worker_id)
|
||||
if current is live:
|
||||
current.registration_pending = was_pending
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
raise HTTPException(status_code=404, detail="No such worker.")
|
||||
registry.revoke(worker_id)
|
||||
if service.control_plane.running:
|
||||
service.control_plane.scheduler.on_disconnected(worker_id)
|
||||
service.control_plane.pool.breakers.forget_worker(worker_id)
|
||||
|
||||
# The tombstone committed before any egress/session mutation. Everything
|
||||
# below is loop-owned and published under the same scheduler authority read
|
||||
# used by next_assignment(), so no task can bind in the handoff window.
|
||||
with registry.authority_guard():
|
||||
if service.control_plane.running:
|
||||
if service.control_plane.servicer is not None:
|
||||
service.control_plane.servicer.revoke_worker_sessions(worker_id)
|
||||
service.control_plane.scheduler.on_disconnected(worker_id)
|
||||
service.control_plane.pool.breakers.forget_worker(worker_id)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return {"ok": True, "revoked": worker_id}
|
||||
|
||||
|
||||
@@ -375,7 +547,9 @@ async def submit_task(request: Request, body: SubmitTaskRequest) -> dict:
|
||||
|
||||
scheduler = service.control_plane.scheduler
|
||||
try:
|
||||
task = scheduler.submit(
|
||||
submit = getattr(scheduler, "submit_async", None)
|
||||
submit = submit if callable(submit) else scheduler.submit
|
||||
submitted = submit(
|
||||
operation=body.operation,
|
||||
engine=body.engine,
|
||||
model_id=body.model_id,
|
||||
@@ -384,6 +558,7 @@ async def submit_task(request: Request, body: SubmitTaskRequest) -> dict:
|
||||
deadline_seconds=body.deadline_seconds,
|
||||
pinned_worker_id=routing.decide().worker_id or None,
|
||||
)
|
||||
task = await submitted if asyncio.iscoroutine(submitted) else submitted
|
||||
except QueueFull as exc:
|
||||
raise HTTPException(status_code=429, detail=str(exc)) from exc
|
||||
|
||||
@@ -505,8 +680,33 @@ async def set_inbound_enabled(request: InboundEnableRequest) -> dict:
|
||||
"machine. Change that environment setting and restart VoiceStudio."
|
||||
),
|
||||
)
|
||||
|
||||
requested_bind = (
|
||||
inbound_service.normalise_bind_host(request.bind)
|
||||
if request.bind
|
||||
else inbound_service.bind_host()
|
||||
)
|
||||
requested_port = request.port or inbound_service.bind_port()
|
||||
if (
|
||||
request.enabled
|
||||
and inbound_service.node.running
|
||||
and (
|
||||
requested_bind != inbound_service.bind_host()
|
||||
or requested_port != inbound_service.node.port
|
||||
)
|
||||
):
|
||||
# start() is intentionally idempotent while a listener owns its
|
||||
# socket. Persisting a new endpoint here would make the UI report a
|
||||
# narrower/different bind while the original socket stayed live.
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(
|
||||
"Turn off Accept connections before changing its bind address "
|
||||
"or port."
|
||||
),
|
||||
)
|
||||
if request.bind:
|
||||
inbound_service.set_bind_host(request.bind)
|
||||
inbound_service.set_bind_host(requested_bind)
|
||||
if request.port:
|
||||
inbound_service.set_bind_port(request.port)
|
||||
inbound_service.set_enabled(request.enabled)
|
||||
@@ -535,6 +735,7 @@ def issue_inbound_key(request: IssueKeyRequest) -> dict:
|
||||
is stored, so it cannot be shown again, only replaced.
|
||||
"""
|
||||
from worker.inbound import service as inbound_service # noqa: PLC0415
|
||||
from worker.inbound.keys import KeyLimitExceeded # noqa: PLC0415
|
||||
|
||||
if not inbound_service.node.running:
|
||||
raise HTTPException(
|
||||
@@ -544,7 +745,10 @@ def issue_inbound_key(request: IssueKeyRequest) -> dict:
|
||||
"Settings → System → Remote workers → Accept connections first."
|
||||
),
|
||||
)
|
||||
issued = inbound_service.node.keys.issue(request.label)
|
||||
try:
|
||||
issued = inbound_service.node.keys.issue(request.label)
|
||||
except KeyLimitExceeded as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
return {
|
||||
"key_id": issued.key.key_id,
|
||||
"label": issued.key.label,
|
||||
@@ -555,12 +759,12 @@ def issue_inbound_key(request: IssueKeyRequest) -> dict:
|
||||
|
||||
|
||||
@router.delete("/inbound/keys/{key_id}")
|
||||
def revoke_inbound_key(key_id: str) -> dict:
|
||||
async def revoke_inbound_key(key_id: str) -> dict:
|
||||
"""Revoke one panel. Everyone else stays connected — the whole reason keys
|
||||
are per panel rather than one shared node key."""
|
||||
from worker.inbound import service as inbound_service # noqa: PLC0415
|
||||
|
||||
if not inbound_service.node.keys.revoke(key_id):
|
||||
if not await inbound_service.node.revoke_key(key_id):
|
||||
raise HTTPException(status_code=404, detail="No such key.")
|
||||
return inbound_service.node.snapshot()
|
||||
|
||||
@@ -579,6 +783,7 @@ async def add_inbound_connection(request: ConnectRequest) -> dict:
|
||||
"""Paste a connection string from a GPU machine and dial it."""
|
||||
from worker.inbound import service as inbound_service # noqa: PLC0415
|
||||
from worker.inbound.connection_string import InvalidConnectionString # noqa: PLC0415
|
||||
from worker.inbound.connector import InboundConnectionError # noqa: PLC0415
|
||||
|
||||
if not service.control_plane.running:
|
||||
raise HTTPException(
|
||||
@@ -597,12 +802,18 @@ async def add_inbound_connection(request: ConnectRequest) -> dict:
|
||||
# surfaces as "cannot connect", which is what a firewall, a wrong port
|
||||
# and a dead node all say too.
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except InboundConnectionError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
return {"endpoint": connection.endpoint, "connections": inbound_service.outbound.snapshot()}
|
||||
|
||||
|
||||
@router.delete("/inbound/connections/{endpoint}")
|
||||
async def remove_inbound_connection(endpoint: str) -> dict:
|
||||
from worker.inbound import service as inbound_service # noqa: PLC0415
|
||||
from worker.inbound.connector import InboundConnectionError # noqa: PLC0415
|
||||
|
||||
await inbound_service.outbound.remove(endpoint)
|
||||
try:
|
||||
await inbound_service.outbound.remove(endpoint)
|
||||
except InboundConnectionError as exc:
|
||||
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
||||
return {"connections": inbound_service.outbound.snapshot()}
|
||||
|
||||
@@ -0,0 +1,609 @@
|
||||
"""Nested subprocess ownership for desktop-managed backend operations.
|
||||
|
||||
The desktop owns the backend with an OS process group/Job. Engine and
|
||||
installer operations also need an independently terminable subtree: killing
|
||||
only their direct child on a timeout leaves uv/git/model workers holding pipes
|
||||
and mutating files. A small direct-child supervisor bridges both lifetimes.
|
||||
|
||||
On POSIX the supervisor is the unreaped leader of a nested process group. A
|
||||
control-pipe EOF (including kernel EOF when the backend dies) kills that group;
|
||||
the parent also drains the group before reaping its stable leader. On Windows
|
||||
the supervisor assigns the operation, while suspended, to a nested
|
||||
kill-on-close Job. The outer desktop Job still contains both levels.
|
||||
|
||||
Standalone/server launches use the same nested owner, preserving their
|
||||
independently terminable subtree without relying on ``taskkill`` or discovery.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import signal
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
_RESULT = struct.Struct("!i")
|
||||
_DESKTOP_MARKER = "OMNIVOICE_DESKTOP_CONTAINED"
|
||||
_DRAIN_FD_ENV = "OMNIVOICE_DESKTOP_DRAIN_FD"
|
||||
|
||||
|
||||
def backend_drain_fd(*, required: bool = False) -> Optional[int]:
|
||||
"""Validated Rust-owned drain writer inherited by the desktop backend."""
|
||||
if os.name != "posix" or os.environ.get(_DESKTOP_MARKER) != "1":
|
||||
return None
|
||||
try:
|
||||
fd = int(os.environ[_DRAIN_FD_ENV])
|
||||
os.fstat(fd)
|
||||
except (KeyError, ValueError, OSError) as exc:
|
||||
if required:
|
||||
raise RuntimeError(
|
||||
"desktop backend is missing its live nested-operation drain descriptor"
|
||||
) from exc
|
||||
return None
|
||||
return fd
|
||||
|
||||
|
||||
def secure_backend_drain_fd() -> None:
|
||||
"""Restore CLOEXEC after Rust's one intentional backend inheritance."""
|
||||
fd = backend_drain_fd(required=True)
|
||||
if fd is not None:
|
||||
os.set_inheritable(fd, False)
|
||||
|
||||
|
||||
class OwnedPopen:
|
||||
"""Popen-compatible handle for a desktop-owned nested operation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
proc: subprocess.Popen,
|
||||
control_fd: int,
|
||||
result_fd: int,
|
||||
) -> None:
|
||||
self._proc = proc
|
||||
self._control_fd: Optional[int] = control_fd
|
||||
self._result_fd: Optional[int] = result_fd
|
||||
self._returncode: Optional[int] = None
|
||||
self._lock = threading.RLock()
|
||||
|
||||
# Popen callers use these directly (protocol pipes and log drains).
|
||||
self.stdin = proc.stdin
|
||||
self.stdout = proc.stdout
|
||||
self.stderr = proc.stderr
|
||||
|
||||
@property
|
||||
def pid(self) -> int:
|
||||
return self._proc.pid
|
||||
|
||||
@property
|
||||
def args(self) -> Any:
|
||||
return self._proc.args
|
||||
|
||||
@property
|
||||
def returncode(self) -> Optional[int]:
|
||||
return self._returncode
|
||||
|
||||
def _close_control(self) -> None:
|
||||
fd, self._control_fd = self._control_fd, None
|
||||
if fd is not None:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# Cleanup is idempotent; another teardown path already closed it.
|
||||
pass
|
||||
|
||||
def _read_result(self, fallback: int) -> int:
|
||||
fd, self._result_fd = self._result_fd, None
|
||||
if fd is None:
|
||||
return fallback
|
||||
try:
|
||||
payload = b""
|
||||
while len(payload) < _RESULT.size:
|
||||
chunk = os.read(fd, _RESULT.size - len(payload))
|
||||
if not chunk:
|
||||
break
|
||||
payload += chunk
|
||||
return _RESULT.unpack(payload)[0] if len(payload) == _RESULT.size else fallback
|
||||
except OSError:
|
||||
return fallback
|
||||
finally:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# The descriptor may have been closed by cancellation cleanup.
|
||||
pass
|
||||
|
||||
def _posix_exited_unreaped(self) -> bool:
|
||||
flags = os.WEXITED | os.WNOHANG | os.WNOWAIT
|
||||
info = os.waitid(os.P_PID, self.pid, flags)
|
||||
return info is not None and info.si_pid != 0
|
||||
|
||||
def _posix_exited_reaping(self) -> Optional[int]:
|
||||
"""macOS fallback for :meth:`_posix_exited_unreaped` (#1656).
|
||||
|
||||
CPython on macOS does not expose ``os.waitid`` (HAVE_WAITID is not set
|
||||
in its build), so the WNOWAIT probe is unavailable there. This
|
||||
fallback *reaps* the wrapper with ``waitpid(WNOHANG)``: it returns
|
||||
the wrapper's exit code once it has exited, None while it is still
|
||||
running, and raises ``ChildProcessError`` when another owner already
|
||||
reaped it (the same refusal the waitid probe gives).
|
||||
|
||||
Reaping earlier than the WNOWAIT dance loses the pre-reap group kill
|
||||
in :meth:`poll`; that is safe because the supervisor's control-pipe
|
||||
EOF already terminates the whole nested group (#1635 design).
|
||||
"""
|
||||
pid, status = os.waitpid(self.pid, os.WNOHANG)
|
||||
if pid != self.pid:
|
||||
return None
|
||||
rc = os.waitstatus_to_exitcode(status)
|
||||
# Publish on the underlying Popen so its own wait()/poll() no-op.
|
||||
self._proc.returncode = rc
|
||||
return rc
|
||||
|
||||
def _posix_exit_state_reaping(self) -> Optional[int]:
|
||||
""":meth:`_posix_exited_reaping` plus one concession: if the leader
|
||||
was already reaped through *this* Popen (``_proc.returncode`` known),
|
||||
report that code rather than refusing — reaping by our own handle is
|
||||
not the foreign reaper the ECHILD refusal exists for."""
|
||||
try:
|
||||
return self._posix_exited_reaping()
|
||||
except ChildProcessError:
|
||||
return self._proc.returncode
|
||||
|
||||
def _signal_owned_group(self, sig: int) -> None:
|
||||
# The numeric group is safe only while its direct-child leader remains
|
||||
# ours and unreaped. ECHILD therefore refuses rather than guessing.
|
||||
try:
|
||||
os.waitid(os.P_PID, self.pid, os.WEXITED | os.WNOHANG | os.WNOWAIT)
|
||||
except ChildProcessError:
|
||||
return
|
||||
except AttributeError:
|
||||
# macOS CPython has no os.waitid (#1656). waitpid still proves
|
||||
# that this exact numeric pid is our live child: ECHILD refuses a
|
||||
# foreign-reaped/reused pid, while pid == self.pid records an exit
|
||||
# without ever signalling the now-unowned process-group number.
|
||||
try:
|
||||
pid, status = os.waitpid(self.pid, os.WNOHANG)
|
||||
except ChildProcessError:
|
||||
return
|
||||
if pid == self.pid:
|
||||
self._proc.returncode = os.waitstatus_to_exitcode(status)
|
||||
return
|
||||
try:
|
||||
os.killpg(self.pid, sig)
|
||||
except ProcessLookupError:
|
||||
# The owned group exited between the waitid probe and the signal.
|
||||
pass
|
||||
|
||||
def poll(self) -> Optional[int]:
|
||||
with self._lock:
|
||||
if self._returncode is not None:
|
||||
return self._returncode
|
||||
if os.name == "posix":
|
||||
try:
|
||||
if hasattr(os, "waitid"):
|
||||
if not self._posix_exited_unreaped():
|
||||
return None
|
||||
self._signal_owned_group(signal.SIGKILL)
|
||||
wrapper_rc = self._proc.wait()
|
||||
else:
|
||||
# macOS CPython: no os.waitid (#1656) — the reaping
|
||||
# probe already terminated/killed nothing; the group
|
||||
# is torn down by the control-pipe EOF in _close_control.
|
||||
wrapper_rc = self._posix_exit_state_reaping()
|
||||
if wrapper_rc is None:
|
||||
return None
|
||||
except ChildProcessError:
|
||||
# Never signal a potentially reused group after another
|
||||
# owner reaped the stable leader.
|
||||
return None
|
||||
else:
|
||||
wrapper_rc = self._proc.poll()
|
||||
if wrapper_rc is None:
|
||||
return None
|
||||
self._close_control()
|
||||
self._returncode = self._read_result(wrapper_rc)
|
||||
return self._returncode
|
||||
|
||||
def wait(self, timeout: Optional[float] = None) -> int:
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
while True:
|
||||
rc = self.poll()
|
||||
if rc is not None:
|
||||
return rc
|
||||
if deadline is not None and time.monotonic() >= deadline:
|
||||
raise subprocess.TimeoutExpired(self.args, timeout)
|
||||
time.sleep(0.01)
|
||||
|
||||
def terminate(self) -> None:
|
||||
with self._lock:
|
||||
if self._returncode is not None:
|
||||
return
|
||||
self._close_control()
|
||||
if os.name == "posix":
|
||||
self._signal_owned_group(signal.SIGTERM)
|
||||
else:
|
||||
# Closing the control pipe asks the supervisor to terminate
|
||||
# its nested Job. The stable wrapper handle is a fallback.
|
||||
try:
|
||||
self._proc.terminate()
|
||||
except OSError:
|
||||
# The wrapper exited after the return-code check.
|
||||
pass
|
||||
|
||||
def kill(self) -> None:
|
||||
with self._lock:
|
||||
if self._returncode is not None:
|
||||
return
|
||||
self._close_control()
|
||||
if os.name == "posix":
|
||||
self._signal_owned_group(signal.SIGKILL)
|
||||
else:
|
||||
try:
|
||||
self._proc.kill()
|
||||
except OSError:
|
||||
# The wrapper exited after the return-code check.
|
||||
pass
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._proc, name)
|
||||
|
||||
def __del__(self) -> None:
|
||||
self._close_control()
|
||||
fd, self._result_fd = self._result_fd, None
|
||||
if fd is not None:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# Finalization may race explicit wait or cancellation cleanup.
|
||||
pass
|
||||
|
||||
|
||||
def spawn_owned(argv: list[str], **kwargs: Any) -> "subprocess.Popen | OwnedPopen":
|
||||
"""Spawn an operation with a stable, independently terminable owner."""
|
||||
|
||||
drain_fd = backend_drain_fd(required=True) if os.name == "posix" else None
|
||||
control_read, control_write = os.pipe()
|
||||
result_read, result_write = os.pipe()
|
||||
control_token = control_read
|
||||
result_token = result_write
|
||||
if os.name == "nt":
|
||||
import msvcrt
|
||||
|
||||
control_token = msvcrt.get_osfhandle(control_read)
|
||||
result_token = msvcrt.get_osfhandle(result_write)
|
||||
wrapper_argv = _supervisor_argv(
|
||||
control_token,
|
||||
result_token,
|
||||
argv,
|
||||
)
|
||||
wrapper_kwargs = dict(kwargs)
|
||||
if os.name == "posix":
|
||||
wrapper_kwargs["start_new_session"] = True
|
||||
pass_fds = [control_read, result_write]
|
||||
if drain_fd is not None:
|
||||
pass_fds.append(drain_fd)
|
||||
if wrapper_kwargs.get("env") is not None:
|
||||
wrapper_env = dict(wrapper_kwargs["env"])
|
||||
wrapper_env[_DESKTOP_MARKER] = "1"
|
||||
wrapper_env[_DRAIN_FD_ENV] = str(drain_fd)
|
||||
wrapper_kwargs["env"] = wrapper_env
|
||||
wrapper_kwargs["pass_fds"] = tuple(pass_fds)
|
||||
else:
|
||||
# Python's Windows fd inheritance requires inheritable CRT handles.
|
||||
# All unrelated descriptors are non-inheritable by default (PEP 446).
|
||||
os.set_handle_inheritable(control_token, True)
|
||||
os.set_handle_inheritable(result_token, True)
|
||||
wrapper_kwargs["close_fds"] = False
|
||||
try:
|
||||
proc = subprocess.Popen(wrapper_argv, **wrapper_kwargs)
|
||||
except BaseException:
|
||||
# The finally block exclusively owns the child-side endpoints. Closing
|
||||
# them here as well risks closing a reused descriptor in another thread.
|
||||
for fd in (control_write, result_read):
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# A partial spawn may already have closed a parent-side endpoint.
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
for fd in (control_read, result_write):
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# Popen may have consumed an inherited child-side endpoint.
|
||||
pass
|
||||
return OwnedPopen(proc, control_write, result_read)
|
||||
|
||||
|
||||
def _supervisor_argv(
|
||||
control_token: int,
|
||||
result_token: int,
|
||||
argv: list[str],
|
||||
) -> list[str]:
|
||||
prefix = [sys.executable]
|
||||
if not getattr(sys, "frozen", False):
|
||||
prefix.append(str(Path(__file__).resolve().parents[1] / "main.py"))
|
||||
return [
|
||||
*prefix,
|
||||
"--supervise",
|
||||
str(control_token),
|
||||
str(result_token),
|
||||
"--",
|
||||
*map(str, argv),
|
||||
]
|
||||
|
||||
|
||||
def _write_result(fd: int, returncode: int) -> None:
|
||||
try:
|
||||
os.write(fd, _RESULT.pack(int(returncode)))
|
||||
except OSError:
|
||||
# The caller may have cancelled and closed its result reader.
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
# Writing or cancellation may already have closed the descriptor.
|
||||
pass
|
||||
|
||||
|
||||
def _operation_env() -> dict[str, str]:
|
||||
env = os.environ.copy()
|
||||
# The operation intentionally does not own the Rust drain writer. Avoid
|
||||
# exposing a stale numeric token which nested code could mistake as valid.
|
||||
env.pop(_DRAIN_FD_ENV, None)
|
||||
env.pop(_DESKTOP_MARKER, None)
|
||||
return env
|
||||
|
||||
|
||||
def _supervise_posix(control_fd: int, result_fd: int, argv: list[str]) -> int:
|
||||
def cancel_on_eof() -> None:
|
||||
try:
|
||||
while os.read(control_fd, 1):
|
||||
pass
|
||||
except OSError:
|
||||
# Closing the control descriptor is itself a cancellation signal.
|
||||
pass
|
||||
os.killpg(os.getpgrp(), signal.SIGKILL)
|
||||
|
||||
threading.Thread(target=cancel_on_eof, daemon=True).start()
|
||||
try:
|
||||
child = subprocess.Popen(argv, close_fds=True, env=_operation_env())
|
||||
rc = child.wait()
|
||||
except OSError:
|
||||
rc = 127
|
||||
_write_result(result_fd, rc)
|
||||
# Drain children which outlived the operation before the stable group
|
||||
# leader exits. SIGKILL intentionally includes this supervisor.
|
||||
os.killpg(os.getpgrp(), signal.SIGKILL)
|
||||
return rc # unreachable
|
||||
|
||||
|
||||
def _windows_job() -> tuple[Any, Any, Any]:
|
||||
import ctypes
|
||||
import ctypes.wintypes as wintypes
|
||||
|
||||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||||
kernel32.CloseHandle.argtypes = (wintypes.HANDLE,)
|
||||
kernel32.CloseHandle.restype = wintypes.BOOL
|
||||
kernel32.TerminateJobObject.argtypes = (wintypes.HANDLE, wintypes.UINT)
|
||||
kernel32.TerminateJobObject.restype = wintypes.BOOL
|
||||
kernel32.ReadFile.argtypes = (
|
||||
wintypes.HANDLE,
|
||||
ctypes.c_void_p,
|
||||
wintypes.DWORD,
|
||||
ctypes.POINTER(wintypes.DWORD),
|
||||
ctypes.c_void_p,
|
||||
)
|
||||
kernel32.ReadFile.restype = wintypes.BOOL
|
||||
kernel32.WriteFile.argtypes = (
|
||||
wintypes.HANDLE,
|
||||
ctypes.c_void_p,
|
||||
wintypes.DWORD,
|
||||
ctypes.POINTER(wintypes.DWORD),
|
||||
ctypes.c_void_p,
|
||||
)
|
||||
kernel32.WriteFile.restype = wintypes.BOOL
|
||||
create = kernel32.CreateJobObjectW
|
||||
create.argtypes = (ctypes.c_void_p, wintypes.LPCWSTR)
|
||||
create.restype = wintypes.HANDLE
|
||||
job = create(None, None)
|
||||
if not job:
|
||||
raise OSError(ctypes.get_last_error(), "CreateJobObjectW")
|
||||
|
||||
class BasicLimits(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("PerProcessUserTimeLimit", ctypes.c_longlong),
|
||||
("PerJobUserTimeLimit", ctypes.c_longlong),
|
||||
("LimitFlags", wintypes.DWORD),
|
||||
("MinimumWorkingSetSize", ctypes.c_size_t),
|
||||
("MaximumWorkingSetSize", ctypes.c_size_t),
|
||||
("ActiveProcessLimit", wintypes.DWORD),
|
||||
("Affinity", ctypes.c_size_t),
|
||||
("PriorityClass", wintypes.DWORD),
|
||||
("SchedulingClass", wintypes.DWORD),
|
||||
]
|
||||
|
||||
class IoCounters(ctypes.Structure):
|
||||
_fields_ = [(name, ctypes.c_ulonglong) for name in (
|
||||
"ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
|
||||
"ReadTransferCount", "WriteTransferCount", "OtherTransferCount",
|
||||
)]
|
||||
|
||||
class ExtendedLimits(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("BasicLimitInformation", BasicLimits),
|
||||
("IoInfo", IoCounters),
|
||||
("ProcessMemoryLimit", ctypes.c_size_t),
|
||||
("JobMemoryLimit", ctypes.c_size_t),
|
||||
("PeakProcessMemoryUsed", ctypes.c_size_t),
|
||||
("PeakJobMemoryUsed", ctypes.c_size_t),
|
||||
]
|
||||
|
||||
info = ExtendedLimits()
|
||||
info.BasicLimitInformation.LimitFlags = 0x00002000 # KILL_ON_JOB_CLOSE
|
||||
set_info = kernel32.SetInformationJobObject
|
||||
set_info.argtypes = (wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD)
|
||||
set_info.restype = wintypes.BOOL
|
||||
if not set_info(job, 9, ctypes.byref(info), ctypes.sizeof(info)):
|
||||
error = ctypes.get_last_error()
|
||||
kernel32.CloseHandle(job)
|
||||
raise OSError(error, "SetInformationJobObject")
|
||||
return job, kernel32, wintypes
|
||||
|
||||
|
||||
def _resume_windows_process(kernel32: Any, wintypes: Any, pid: int) -> None:
|
||||
import ctypes
|
||||
|
||||
class ThreadEntry(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("dwSize", wintypes.DWORD),
|
||||
("cntUsage", wintypes.DWORD),
|
||||
("th32ThreadID", wintypes.DWORD),
|
||||
("th32OwnerProcessID", wintypes.DWORD),
|
||||
("tpBasePri", wintypes.LONG),
|
||||
("tpDeltaPri", wintypes.LONG),
|
||||
("dwFlags", wintypes.DWORD),
|
||||
]
|
||||
|
||||
kernel32.CreateToolhelp32Snapshot.argtypes = (wintypes.DWORD, wintypes.DWORD)
|
||||
kernel32.CreateToolhelp32Snapshot.restype = wintypes.HANDLE
|
||||
kernel32.Thread32First.argtypes = (wintypes.HANDLE, ctypes.POINTER(ThreadEntry))
|
||||
kernel32.Thread32First.restype = wintypes.BOOL
|
||||
kernel32.Thread32Next.argtypes = (wintypes.HANDLE, ctypes.POINTER(ThreadEntry))
|
||||
kernel32.Thread32Next.restype = wintypes.BOOL
|
||||
kernel32.OpenThread.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD)
|
||||
kernel32.OpenThread.restype = wintypes.HANDLE
|
||||
kernel32.ResumeThread.argtypes = (wintypes.HANDLE,)
|
||||
kernel32.ResumeThread.restype = wintypes.DWORD
|
||||
|
||||
snapshot = kernel32.CreateToolhelp32Snapshot(0x00000004, 0)
|
||||
invalid = ctypes.c_void_p(-1).value
|
||||
if snapshot == invalid:
|
||||
raise OSError(ctypes.get_last_error(), "CreateToolhelp32Snapshot")
|
||||
try:
|
||||
entry = ThreadEntry(dwSize=ctypes.sizeof(ThreadEntry))
|
||||
found = kernel32.Thread32First(snapshot, ctypes.byref(entry))
|
||||
while found:
|
||||
if entry.th32OwnerProcessID == pid:
|
||||
thread = kernel32.OpenThread(0x0002, False, entry.th32ThreadID)
|
||||
if not thread:
|
||||
raise OSError(ctypes.get_last_error(), "OpenThread")
|
||||
try:
|
||||
if kernel32.ResumeThread(thread) == 0xFFFFFFFF:
|
||||
raise OSError(ctypes.get_last_error(), "ResumeThread")
|
||||
return
|
||||
finally:
|
||||
kernel32.CloseHandle(thread)
|
||||
found = kernel32.Thread32Next(snapshot, ctypes.byref(entry))
|
||||
finally:
|
||||
kernel32.CloseHandle(snapshot)
|
||||
raise OSError("suspended operation thread was not found")
|
||||
|
||||
|
||||
def _supervise_windows(control_fd: int, result_fd: int, argv: list[str]) -> int:
|
||||
import ctypes
|
||||
|
||||
job, kernel32, wintypes = _windows_job()
|
||||
cancelled = threading.Event()
|
||||
job_lock = threading.Lock()
|
||||
job_open = True
|
||||
|
||||
def terminate_job() -> None:
|
||||
with job_lock:
|
||||
if job_open:
|
||||
kernel32.TerminateJobObject(job, 1)
|
||||
|
||||
def cancel_on_eof() -> None:
|
||||
byte = ctypes.create_string_buffer(1)
|
||||
count = wintypes.DWORD()
|
||||
while kernel32.ReadFile(
|
||||
wintypes.HANDLE(control_fd), byte, 1, ctypes.byref(count), None
|
||||
) and count.value:
|
||||
pass
|
||||
kernel32.CloseHandle(wintypes.HANDLE(control_fd))
|
||||
cancelled.set()
|
||||
terminate_job()
|
||||
|
||||
threading.Thread(target=cancel_on_eof, daemon=True).start()
|
||||
child: Optional[subprocess.Popen] = None
|
||||
rc = 127
|
||||
try:
|
||||
child = subprocess.Popen(
|
||||
argv,
|
||||
close_fds=True,
|
||||
env=_operation_env(),
|
||||
creationflags=0x08000000 | 0x00000004, # NO_WINDOW | SUSPENDED
|
||||
)
|
||||
assign = kernel32.AssignProcessToJobObject
|
||||
assign.argtypes = (wintypes.HANDLE, wintypes.HANDLE)
|
||||
assign.restype = wintypes.BOOL
|
||||
if not assign(job, wintypes.HANDLE(child._handle)):
|
||||
raise OSError(ctypes.get_last_error(), "AssignProcessToJobObject")
|
||||
if cancelled.is_set():
|
||||
terminate_job()
|
||||
else:
|
||||
_resume_windows_process(kernel32, wintypes, child.pid)
|
||||
rc = child.wait()
|
||||
# A successful direct child may leave helpers behind; terminate the
|
||||
# nested stable Job before reporting completion.
|
||||
terminate_job()
|
||||
except OSError:
|
||||
terminate_job()
|
||||
if child is not None:
|
||||
try:
|
||||
# Assignment itself may have failed, leaving this suspended
|
||||
# process outside the nested Job. Terminate it through its
|
||||
# stable process handle before waiting; never strand an
|
||||
# unassigned operation or rely on the outer desktop Job.
|
||||
child.kill()
|
||||
except OSError:
|
||||
# The suspended child may have exited during Job teardown.
|
||||
pass
|
||||
try:
|
||||
child.wait(timeout=5)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
# The outer desktop Job remains the terminal containment fallback.
|
||||
pass
|
||||
finally:
|
||||
payload = _RESULT.pack(int(rc))
|
||||
payload_buffer = ctypes.create_string_buffer(payload)
|
||||
written = wintypes.DWORD()
|
||||
kernel32.WriteFile(
|
||||
wintypes.HANDLE(result_fd),
|
||||
payload_buffer,
|
||||
len(payload),
|
||||
ctypes.byref(written),
|
||||
None,
|
||||
)
|
||||
kernel32.CloseHandle(wintypes.HANDLE(result_fd))
|
||||
with job_lock:
|
||||
job_open = False
|
||||
kernel32.CloseHandle(job)
|
||||
return rc
|
||||
|
||||
|
||||
def supervisor_main(args: list[str]) -> int:
|
||||
if len(args) < 5 or args[0] != "--supervise" or args[3] != "--":
|
||||
return 2
|
||||
control_fd = int(args[1])
|
||||
result_fd = int(args[2])
|
||||
argv = args[4:]
|
||||
secure_backend_drain_fd()
|
||||
if os.name == "posix":
|
||||
return _supervise_posix(control_fd, result_fd, argv)
|
||||
return _supervise_windows(control_fd, result_fd, argv)
|
||||
|
||||
|
||||
def _main() -> int:
|
||||
return supervisor_main(sys.argv[1:])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(_main())
|
||||
@@ -177,6 +177,23 @@ def gfx_for_hsa_override(value: str) -> str | None:
|
||||
#: The ROCm kernel driver interface. Its absence, or its presence without
|
||||
#: permission, are the two commonest reasons a ROCm host silently runs on CPU.
|
||||
_KFD_DEVICE = "/dev/kfd"
|
||||
_DXG_DEVICE = "/dev/dxg"
|
||||
_DXG_RUNTIME_PATHS = (
|
||||
"/usr/lib/libdxcore.so",
|
||||
"/usr/lib/librocdxg.so",
|
||||
"/usr/share/rocdxg/dids.conf",
|
||||
)
|
||||
|
||||
|
||||
def _rocm_requires_dxg_detection(version: object) -> bool:
|
||||
"""Whether WSL's ROCDXG bridge still needs its explicit opt-in."""
|
||||
try:
|
||||
parts = str(version).split(".")
|
||||
return (int(parts[0]), int(parts[1])) < (7, 13)
|
||||
except (IndexError, TypeError, ValueError):
|
||||
# Unknown versions get the conservative advice. The variable is
|
||||
# harmless on newer runtimes and necessary on every older one.
|
||||
return True
|
||||
|
||||
|
||||
def why_no_gpu(torch) -> tuple[str, ...]:
|
||||
@@ -230,6 +247,40 @@ def why_no_gpu(torch) -> tuple[str, ...]:
|
||||
# /dev/kfd only exists on Linux; on any other platform its absence
|
||||
# says nothing, so don't invent a reason.
|
||||
if sys.platform.startswith("linux"):
|
||||
if not os.path.exists(_KFD_DEVICE) and os.path.exists(_DXG_DEVICE):
|
||||
if not os.access(_DXG_DEVICE, os.R_OK | os.W_OK):
|
||||
return (
|
||||
f"ROCm {hip} is installed and {_DXG_DEVICE} exists, "
|
||||
"but this process cannot open it — pass "
|
||||
"--device /dev/dxg to the WSL container",
|
||||
)
|
||||
dxg_detection = os.environ.get("HSA_ENABLE_DXG_DETECTION", "").strip()
|
||||
if dxg_detection == "0":
|
||||
return (
|
||||
f"ROCm {hip} is installed and {_DXG_DEVICE} is reachable, "
|
||||
"but HSA_ENABLE_DXG_DETECTION=0 explicitly disables the "
|
||||
"WSL GPU bridge; remove it or set it to 1",
|
||||
)
|
||||
if _rocm_requires_dxg_detection(hip) and dxg_detection != "1":
|
||||
return (
|
||||
f"ROCm {hip} is installed and {_DXG_DEVICE} is "
|
||||
"reachable, but this pre-7.13 runtime requires "
|
||||
"HSA_ENABLE_DXG_DETECTION=1 inside WSL containers",
|
||||
)
|
||||
missing = [
|
||||
path for path in _DXG_RUNTIME_PATHS if not os.path.exists(path)
|
||||
]
|
||||
if missing:
|
||||
return (
|
||||
f"ROCm {hip} can reach {_DXG_DEVICE}, but the WSL "
|
||||
"ROCDXG runtime mounts are incomplete; missing: "
|
||||
f"{', '.join(missing)}",
|
||||
)
|
||||
return (
|
||||
f"ROCm {hip} and the WSL ROCDXG bridge are reachable, "
|
||||
"but no GPU was enumerated — verify the AMD Windows "
|
||||
"driver, librocdxg/ROCm compatibility, and host `rocminfo`",
|
||||
)
|
||||
if not os.path.exists(_KFD_DEVICE):
|
||||
return (
|
||||
f"ROCm {hip} is installed but {_KFD_DEVICE} is not "
|
||||
|
||||
@@ -52,6 +52,7 @@ _REDACTED_VALUE = "***REDACTED***"
|
||||
# One-line "what to do" per docs-taxonomy key. Keys mirror error_docs_map's
|
||||
# taxonomy; the docs URL itself stays owned by error_docs_map.
|
||||
_HINTS: dict[str, str] = {
|
||||
"GPU_OOM": "Close other GPU-heavy apps or unload models, then retry. You can also choose CPU in Settings → Performance & Device or select a smaller TTS engine.",
|
||||
"WORKER_AT_CAPACITY": "Wait for a running job on that worker to finish, or choose another available worker and retry.",
|
||||
"MODEL_NOT_INSTALLED": "Install or enable this engine on the worker machine, then refresh its capabilities and retry.",
|
||||
"MODEL_NOT_DOWNLOADED": "Open Models, install this model on the selected worker, then retry when the download completes.",
|
||||
@@ -290,6 +291,9 @@ def append_hf_mirror_hint(text: str) -> str:
|
||||
# must NOT be added: its bare "timed out" trigger would stamp a "video server"
|
||||
# hint on a model-load timeout that leaks through the 500 handler.
|
||||
_CONTEXT_FREE_HINT_CLASSES = frozenset({
|
||||
# Device allocator signatures are specific enough to attach the shared
|
||||
# recovery without exposing CUDA's process table or filesystem paths.
|
||||
"GPU_OOM",
|
||||
"SOCKS_PROXY_SUPPORT_MISSING",
|
||||
"SSL_HANDSHAKE_FAILURE",
|
||||
# Its trigger is an exact OpenSSL string, so it cannot be confused with
|
||||
@@ -323,6 +327,38 @@ def append_hint(text: str) -> str:
|
||||
return f"{text} — {hint}" if hint else text
|
||||
|
||||
|
||||
_GPU_OOM_SIGNATURES = (
|
||||
"cuda out of memory",
|
||||
"cuda error: out of memory",
|
||||
"cuda_error_out_of_memory",
|
||||
"mps backend out of memory",
|
||||
"hip out of memory",
|
||||
"out of memory on device",
|
||||
)
|
||||
|
||||
|
||||
def is_gpu_oom(error: BaseException | str) -> bool:
|
||||
"""Recognize device OOMs through wrappers without importing torch."""
|
||||
pending: list[BaseException] = [error] if isinstance(error, BaseException) else []
|
||||
seen: set[int] = set()
|
||||
while pending:
|
||||
current = pending.pop()
|
||||
if id(current) in seen:
|
||||
continue
|
||||
seen.add(id(current))
|
||||
if type(current).__name__ == "OutOfMemoryError":
|
||||
return True
|
||||
if any(signature in str(current).lower() for signature in _GPU_OOM_SIGNATURES):
|
||||
return True
|
||||
if current.__cause__ is not None:
|
||||
pending.append(current.__cause__)
|
||||
if current.__context__ is not None:
|
||||
pending.append(current.__context__)
|
||||
if isinstance(error, str):
|
||||
return any(signature in error.lower() for signature in _GPU_OOM_SIGNATURES)
|
||||
return False
|
||||
|
||||
|
||||
def classify(reason: str) -> str:
|
||||
"""Map a failure reason to a docs-taxonomy key, or "" when unknown.
|
||||
|
||||
@@ -330,6 +366,8 @@ def classify(reason: str) -> str:
|
||||
backend log / diagnostic names the same class the UI deeplink will use.
|
||||
"""
|
||||
low = (reason or "").lower()
|
||||
if is_gpu_oom(low):
|
||||
return "GPU_OOM"
|
||||
if "pkg_resources" in low:
|
||||
return "PKG_RESOURCES_MISSING"
|
||||
if "quarantine" in low or "is damaged" in low or "gatekeeper" in low:
|
||||
|
||||
@@ -74,6 +74,48 @@ def stream_failure(code: str) -> dict[str, object]:
|
||||
return dict(failures.get(code, failures["generation_failed"]))
|
||||
|
||||
|
||||
def stream_generation_failure(error: BaseException | object) -> dict[str, object]:
|
||||
"""``generation_failed`` stream metadata, enriched with the actual cause.
|
||||
|
||||
The bare "Generation failed. Check the selected engine and try again." is
|
||||
the floor for an *unrecognized* failure. When the private exception DOES
|
||||
classify to a known failure class — a corrupt model cache, an unreachable
|
||||
Hugging Face mirror, a missing ffmpeg/ffprobe, a Windows paging-file limit,
|
||||
a SOCKS/TLS proxy problem, … — the stable VoiceStudio-owned remediation for
|
||||
that class is appended so the user can self-diagnose instead of guessing
|
||||
which engine or which failure. This is the same enrichment the classic
|
||||
(non-streaming) ``/generate`` 500 already gets via
|
||||
:func:`public_exception_response`; the in-band streaming error frame
|
||||
replaces the global 500 handler for a streaming request and used to bypass
|
||||
it entirely (#1607).
|
||||
|
||||
Only VoiceStudio-owned constants are copied — never a substring of
|
||||
``error`` (Constitution I). Never raises: a diagnosis failure must not
|
||||
replace the failure being diagnosed.
|
||||
"""
|
||||
payload = stream_failure("generation_failed")
|
||||
try:
|
||||
enriched = public_exception_response(error, fallback=str(payload["detail"]))
|
||||
except Exception:
|
||||
return payload
|
||||
hint = enriched.get("hint")
|
||||
if hint:
|
||||
payload["detail"] = enriched["detail"]
|
||||
payload["hint"] = hint
|
||||
topic = enriched.get("docs_topic")
|
||||
if topic:
|
||||
payload["docs_topic"] = topic
|
||||
try:
|
||||
from core import error_docs_map
|
||||
|
||||
url = error_docs_map.ERROR_DOCS.get(topic, "")
|
||||
except Exception:
|
||||
url = ""
|
||||
if url:
|
||||
payload["docs_url"] = url
|
||||
return payload
|
||||
|
||||
|
||||
def public_failure(
|
||||
logger: logging.Logger,
|
||||
log_message: str,
|
||||
|
||||
@@ -24,7 +24,7 @@ from pathlib import Path
|
||||
# tests/test_app_version.py::test_all_version_files_in_lockstep and bumped by
|
||||
# release.yml's version-bump job, so it stays equal to
|
||||
# pyproject/tauri.conf/Cargo/package.json.
|
||||
_FALLBACK_VERSION = "0.5.0"
|
||||
_FALLBACK_VERSION = "0.5.1"
|
||||
|
||||
|
||||
def _fallback_version() -> str:
|
||||
|
||||
@@ -98,6 +98,7 @@ def _platform_slug() -> str:
|
||||
darwin-x86_64
|
||||
windows-x86_64
|
||||
linux-x86_64
|
||||
linux-aarch64
|
||||
"""
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
@@ -107,6 +108,8 @@ def _platform_slug() -> str:
|
||||
return "darwin-x86_64"
|
||||
if system == "windows":
|
||||
return "windows-x86_64"
|
||||
if system == "linux" and machine in ("arm64", "aarch64"):
|
||||
return "linux-aarch64"
|
||||
# Linux + everything else falls into the linux slug.
|
||||
return "linux-x86_64"
|
||||
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
"""omnivoice-subprocess: the resident OmniVoice TTS engine in a crash-isolated
|
||||
sidecar process (#730/#1190).
|
||||
|
||||
The default ``omnivoice`` engine runs in-process on the GPU ``ThreadPoolExecutor``.
|
||||
The ``omnivoice`` engine runs in-process on CUDA, ROCm, and CPU. On MPS it is
|
||||
resolved to :class:`OmniVoiceMPSSubprocessBackend` so a fatal native allocator
|
||||
exit cannot take down the local API process.
|
||||
When a generate or load there exceeds its execution budget the pool is "reset"
|
||||
but the abandoned worker *thread* cannot be killed (Python cannot interrupt a
|
||||
native torch/MPS call), so it holds the MPS device until it finishes on its
|
||||
@@ -13,16 +15,11 @@ timeout the parent's watchdog calls ``proc.kill()``, reclaiming the child's
|
||||
VRAM/device, and the next request transparently respawns a fresh sidecar. That
|
||||
is the one thing the in-process engine structurally cannot do.
|
||||
|
||||
OPT-IN (Settings -> Engines, or ``OMNIVOICE_TTS_BACKEND=omnivoice-subprocess``);
|
||||
the in-process ``omnivoice`` stays the default so existing users see no change.
|
||||
The explicit ``omnivoice-subprocess`` id remains available on every host for
|
||||
operators who want the same containment elsewhere.
|
||||
|
||||
Tradeoff vs the in-process engine: identical model and quality, a little extra
|
||||
per-call overhead (one stdio round-trip), and it does not carry the native
|
||||
advanced-parameter surface (``t_shift`` / ``layer_penalty_factor`` /
|
||||
``position_temperature`` / ``class_temperature``) or parent-side seed
|
||||
determinism, because the generic ``backend.generate`` path does not forward
|
||||
those. Acceptable for unattended / reaction-triggered use where reliability
|
||||
matters more than those controls.
|
||||
Tradeoff vs the in-process engine: identical model, controls, seed behavior,
|
||||
and quality, with a little extra per-call overhead (one stdio round-trip).
|
||||
|
||||
Unlike IndexTTS / dots.tts / Supertonic-3, this sidecar runs under the PARENT
|
||||
interpreter (``venv_python() -> sys.executable``): the goal here is crash
|
||||
@@ -51,7 +48,7 @@ class OmniVoiceSubprocessBackend(SubprocessBackend):
|
||||
id = "omnivoice-subprocess"
|
||||
display_name = "OmniVoice (subprocess-isolated, killable on timeout)"
|
||||
_DEFAULT_SAMPLE_RATE = 24000
|
||||
gpu_compat = ("cuda", "mps", "cpu")
|
||||
gpu_compat = ("cuda", "rocm", "mps", "cpu")
|
||||
# Match OmniVoiceBackend: the measured floor below which a render that
|
||||
# should take seconds runs for minutes (the #1226/#1222 4 GB reports).
|
||||
min_vram_gb = 6.0
|
||||
@@ -102,4 +99,34 @@ class OmniVoiceSubprocessBackend(SubprocessBackend):
|
||||
return ["multi"]
|
||||
|
||||
|
||||
__all__ = ["OmniVoiceSubprocessBackend"]
|
||||
class OmniVoiceMPSSubprocessBackend(OmniVoiceSubprocessBackend):
|
||||
"""Effective ``omnivoice`` implementation on MPS.
|
||||
|
||||
Native torch/MPS allocator failures can terminate the process without a
|
||||
catchable Python exception. Keeping the same engine id and model surface in
|
||||
a child makes that failure recoverable while Settings, APIs, and saved
|
||||
projects continue to refer to ``omnivoice``.
|
||||
"""
|
||||
|
||||
id = "omnivoice"
|
||||
display_name = "VoiceStudio (k2-fsa/OmniVoice, 600+ languages)"
|
||||
supports_native_omnivoice_controls = True
|
||||
|
||||
def generate(self, text: str, **kw):
|
||||
from services.model_manager import make_room_before_generate
|
||||
|
||||
make_room_before_generate()
|
||||
try:
|
||||
return super().generate(text, **kw)
|
||||
except RuntimeError as exc:
|
||||
if "sidecar closed pipe mid-generate" not in str(exc):
|
||||
raise
|
||||
raise RuntimeError(
|
||||
"The isolated OmniVoice engine stopped during generation, "
|
||||
"usually because macOS reclaimed it under memory pressure. "
|
||||
"The VoiceStudio backend is still running. Close memory-heavy "
|
||||
"apps or select a smaller TTS engine, then retry."
|
||||
) from exc
|
||||
|
||||
|
||||
__all__ = ["OmniVoiceMPSSubprocessBackend", "OmniVoiceSubprocessBackend"]
|
||||
|
||||
@@ -50,6 +50,8 @@ OMNIVOICE_SAMPLE_RATE = 24000
|
||||
_GEN_KW_ALLOWLIST = (
|
||||
"language", "instruct", "duration", "num_step", "guidance_scale",
|
||||
"speed", "denoise", "postprocess_output", "preprocess_prompt",
|
||||
"t_shift", "layer_penalty_factor", "position_temperature",
|
||||
"class_temperature", "audio_chunk_duration", "audio_chunk_threshold",
|
||||
)
|
||||
|
||||
_model = None
|
||||
@@ -183,6 +185,12 @@ def _handle_synthesize(msg: dict, stdout) -> None:
|
||||
ref_text = msg.get("ref_text") or None
|
||||
gen_kw = {k: msg[k] for k in _GEN_KW_ALLOWLIST if k in msg}
|
||||
|
||||
seed = msg.get("seed")
|
||||
if seed is not None:
|
||||
import torch
|
||||
|
||||
torch.manual_seed(int(seed))
|
||||
|
||||
audios = model.generate(
|
||||
text=text, ref_audio=ref_audio, ref_text=ref_text, **gen_kw
|
||||
)
|
||||
|
||||
+23
-2
@@ -1,4 +1,3 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
@@ -10,6 +9,24 @@ _backend_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if _backend_dir not in sys.path:
|
||||
sys.path.insert(0, _backend_dir)
|
||||
|
||||
# PyInstaller re-executes this entry module when the frozen backend binary is
|
||||
# launched. Nested operation supervisors therefore dispatch here, before math,
|
||||
# logging, FastAPI, torch, or any application initialization. Source launches
|
||||
# use this same entry contract so frozen/source behavior cannot drift.
|
||||
if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == "--supervise":
|
||||
from core.contained_subprocess import supervisor_main
|
||||
|
||||
raise SystemExit(supervisor_main(sys.argv[1:]))
|
||||
|
||||
# Rust clears CLOEXEC only for the backend exec. Re-arm PEP 446 immediately:
|
||||
# nested supervisors receive this descriptor solely through explicit pass_fds,
|
||||
# so a third-party close_fds=False child cannot hold the desktop drain barrier.
|
||||
from core.contained_subprocess import secure_backend_drain_fd # noqa: E402
|
||||
|
||||
secure_backend_drain_fd()
|
||||
|
||||
import math # noqa: E402
|
||||
|
||||
# Windows: run every child process (ffmpeg, engine sidecars, yt-dlp, demucs, …)
|
||||
# WITHOUT popping a console window. The backend itself is spawned console-less by
|
||||
# the Tauri shell, so on Windows each console subprocess it launches would
|
||||
@@ -649,6 +666,7 @@ def _phase_a_build_inner() -> None:
|
||||
events,
|
||||
capture,
|
||||
capture_ws,
|
||||
speech_platform,
|
||||
dictation,
|
||||
openai_compat,
|
||||
tts_stream,
|
||||
@@ -668,7 +686,7 @@ def _phase_a_build_inner() -> None:
|
||||
system, profiles, exports, generation, dub_core, dub_generate,
|
||||
dub_export, dub_translate, projects, glossary, engines, tools,
|
||||
stories, setup, gallery, archetypes, describe_voice, community,
|
||||
batch, watermark, events, capture, capture_ws, dictation,
|
||||
batch, watermark, events, capture, capture_ws, speech_platform, dictation,
|
||||
openai_compat, tts_stream, marketplace, personas, sonitranslate,
|
||||
audiobook, longform_jobs, pronunciation, settings_router,
|
||||
media_tools_router, auth_router, _mcp_bindings_router, workers_router,
|
||||
@@ -991,6 +1009,9 @@ async def _phase_b(app: FastAPI) -> None:
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
from api.dependencies import validate_server_admin_key
|
||||
|
||||
validate_server_admin_key()
|
||||
# Startup watchdog (#632): a silent hang during startup (e.g. a model-load /
|
||||
# MCP deadlock on some platforms) means "Application startup complete" never
|
||||
# logs and the app sits forever with no error. If startup hasn't finished
|
||||
|
||||
@@ -190,6 +190,7 @@ class ParseSubtitleTextRequest(BaseModel):
|
||||
class DubIngestUrlRequest(BaseModel):
|
||||
url: str
|
||||
job_id: Optional[str] = None
|
||||
source_lang: Optional[str] = None
|
||||
# When true and the URL is a caption-bearing host (YouTube, Vimeo, TED…),
|
||||
# ask yt-dlp to also download the original-language + any additional
|
||||
# sub_langs as VTT. The UI uses this to seed a transcript without running
|
||||
|
||||
@@ -33,7 +33,9 @@ WS_TICKET_PREFIX = "ovs_ws_ticket_"
|
||||
_TOKEN_BYTES = 32
|
||||
_ENCODED_TOKEN_LENGTH = 43
|
||||
_TOKEN_BODY_RE = re.compile(rf"^[A-Za-z0-9_-]{{{_ENCODED_TOKEN_LENGTH}}}$")
|
||||
_ALLOWED_WS_PATHS = frozenset({"/ws/events", "/ws/transcribe"})
|
||||
_ALLOWED_WS_PATHS = frozenset(
|
||||
{"/ws/events", "/ws/transcribe", "/v1/audio/transcriptions/stream"}
|
||||
)
|
||||
_ADMIN_CAPABILITIES = frozenset({"consume", "admin"})
|
||||
_KEY_GENERATION_INFO = b"omnivoice-admin-key-generation-v1"
|
||||
|
||||
|
||||
@@ -88,10 +88,9 @@ def reset_pool_after_wedge(executor, *, what: str = "ASR") -> bool:
|
||||
|
||||
|
||||
# ── Consecutive-timeout streak → recommend the crash-isolated engine ────────
|
||||
# A pool reset restores *capacity*, but the wedged CTranslate2/whisperx thread
|
||||
# keeps its VRAM until the process exits. When guarded transcribes keep timing
|
||||
# out back-to-back in one session, resets clearly aren't recovering the
|
||||
# underlying hang — the durable fix is the crash-isolated sidecar engine
|
||||
# A timed-out CTranslate2/whisperx thread keeps its worker and VRAM until the
|
||||
# native call exits. When guarded transcribes keep timing out back-to-back in
|
||||
# one session, the durable fix is the crash-isolated sidecar engine
|
||||
# (services.subprocess_asr, #393), whose child process CAN be hard-killed to
|
||||
# reclaim the hung call and its VRAM. We only *recommend* it (log + error
|
||||
# message); we never switch engines automatically (owner rule: no silent
|
||||
@@ -146,23 +145,18 @@ def _isolated_engine_hint(streak: int) -> str:
|
||||
|
||||
async def run_transcribe_guarded(executor, fn, *, what: str = "ASR",
|
||||
timeout: float = ASR_TRANSCRIBE_TIMEOUT_S,
|
||||
timeout_env: str = "OMNIVOICE_ASR_TRANSCRIBE_TIMEOUT_S"):
|
||||
timeout_env: str = "OMNIVOICE_ASR_TRANSCRIBE_TIMEOUT_S",
|
||||
reset_on_timeout: bool = False):
|
||||
"""Run a blocking transcribe ``fn`` in ``executor`` with a hard wall-clock
|
||||
bound. On timeout, raise :class:`ASRTimeoutError` with guidance instead of
|
||||
letting the request hang forever.
|
||||
|
||||
``run_in_executor`` cannot cancel the underlying thread, so a wedged
|
||||
transcribe (a CTranslate2 / whisperx / VAD hang seen on some Windows + CUDA
|
||||
setups, #730) keeps occupying its GPU-pool worker. With a 1–2 worker pool
|
||||
that starves every *other* request — including TTS generate — and the next
|
||||
thing the user does surfaces as "Can't reach the local backend" even though
|
||||
the process is alive. So on timeout we also ``reset()`` the pool when it
|
||||
supports it (``_ResilientGpuPool``): the wedged thread is abandoned and the
|
||||
next submit gets a fresh worker, restoring capacity without an app restart.
|
||||
The orphaned thread still holds its VRAM until the process exits, which is
|
||||
why the message still recommends a smaller ASR model / Flush as the durable
|
||||
fix. Executors without ``reset`` (a plain ThreadPoolExecutor in tests) just
|
||||
get the bound + actionable error.
|
||||
``run_in_executor`` cannot cancel the underlying thread, so a timed-out
|
||||
in-process CTranslate2/whisperx call still owns its model and device. The
|
||||
default deliberately leaves that worker accounted for: swapping in a fresh
|
||||
pool and immediately retrying the same backend overlaps two native calls,
|
||||
which produced the Windows access violation in #1669. A caller backed by a
|
||||
genuinely killable process may opt into ``reset_on_timeout``.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
# Same SystemExit containment as the TTS pool (#1133 class): an ASR
|
||||
@@ -171,16 +165,16 @@ async def run_transcribe_guarded(executor, fn, *, what: str = "ASR",
|
||||
try:
|
||||
result = await asyncio.wait_for(fut, timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
# Free the poisoned pool so a hung transcribe can't keep starving TTS /
|
||||
# other ASR work (the "can't reach backend" symptom, #730).
|
||||
reset_pool_after_wedge(executor, what=what)
|
||||
if reset_on_timeout:
|
||||
reset_pool_after_wedge(executor, what=what)
|
||||
streak = _note_transcribe_timeout()
|
||||
msg = (
|
||||
f"{what} transcription exceeded {timeout:.0f}s and was abandoned — "
|
||||
"the backend is running, but the ASR model is too heavy for the "
|
||||
"available compute. Most often the GPU is VRAM-starved: the resident "
|
||||
"TTS model and a large ASR model (large-v3) contend for memory. "
|
||||
"Capacity was restored automatically, but for a durable fix Flush the "
|
||||
"The native call cannot be killed safely, so its capacity remains "
|
||||
"reserved until it exits. For a durable fix Flush the "
|
||||
"TTS model to free VRAM, pick a smaller ASR model in "
|
||||
f"Model Catalogue → Models, or set ASR to CPU. (Raise {timeout_env} "
|
||||
"for very long transcribes.)"
|
||||
|
||||
@@ -742,6 +742,74 @@ def _ensure_browser_playable_mp4(video_path: str) -> str:
|
||||
return video_path
|
||||
|
||||
|
||||
async def _ensure_browser_playable_mp4_for_job(job_id: str, video_path: str) -> str:
|
||||
"""Normalize an upload through the job's cancellable process registry."""
|
||||
is_mp4 = video_path.lower().endswith(".mp4")
|
||||
vcodec, acodec = await asyncio.to_thread(_probe_codecs, video_path)
|
||||
if is_mp4 and vcodec in _BROWSER_VIDEO_CODECS and acodec in _BROWSER_AUDIO_CODECS:
|
||||
return video_path
|
||||
|
||||
target = os.path.splitext(video_path)[0] + ".mp4"
|
||||
if target == video_path:
|
||||
target = os.path.splitext(video_path)[0] + ".browser.mp4"
|
||||
run_proc = run_proc_factory(job_id)
|
||||
ffmpeg_bin = find_ffmpeg()
|
||||
|
||||
async def attempt(cmd: list[str]) -> int:
|
||||
try:
|
||||
proc, _stdout, _stderr = await run_proc(cmd, timeout=1800.0)
|
||||
return proc.returncode
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Browser-media normalization process failed for %s: %s",
|
||||
log_safe(video_path),
|
||||
log_safe(exc),
|
||||
)
|
||||
return 1
|
||||
|
||||
rc = 1
|
||||
if not is_mp4:
|
||||
rc = await attempt(
|
||||
[
|
||||
ffmpeg_bin, "-y", "-i", video_path,
|
||||
"-c:v", "copy", "-c:a", "copy",
|
||||
"-movflags", "+faststart", target,
|
||||
]
|
||||
)
|
||||
if rc == 0 and os.path.exists(target):
|
||||
target_vcodec, target_acodec = await asyncio.to_thread(_probe_codecs, target)
|
||||
if (
|
||||
target_vcodec not in _BROWSER_VIDEO_CODECS
|
||||
or target_acodec not in _BROWSER_AUDIO_CODECS
|
||||
):
|
||||
rc = 1
|
||||
else:
|
||||
rc = 1
|
||||
if rc != 0:
|
||||
rc = await attempt(
|
||||
[
|
||||
ffmpeg_bin, "-y", "-i", video_path,
|
||||
"-c:v", "libx264", "-preset", "veryfast", "-crf", "23",
|
||||
"-pix_fmt", "yuv420p", "-c:a", "aac", "-b:a", "192k",
|
||||
"-movflags", "+faststart", target,
|
||||
]
|
||||
)
|
||||
if rc == 0 and os.path.exists(target) and target != video_path:
|
||||
try:
|
||||
os.remove(video_path)
|
||||
except OSError:
|
||||
pass # Best effort: the normalized target is already complete.
|
||||
return target
|
||||
logger.warning(
|
||||
"Could not transcode %s to browser-playable mp4 — the in-app "
|
||||
"video player may render this file as a black box.",
|
||||
log_safe(video_path),
|
||||
)
|
||||
return video_path
|
||||
|
||||
|
||||
# Bounded retry for transient download failures (#579/#598). yt-dlp's own
|
||||
# `retries`/`fragment_retries` cover per-fragment HTTP flakes, but a broken
|
||||
# pipe ([Errno 32]) raised while the write side of a pipe closes mid-stream
|
||||
@@ -1257,6 +1325,13 @@ async def ingest_pipeline(
|
||||
except Exception:
|
||||
dur = 0.0
|
||||
|
||||
# URL downloads already pass through this guard in yt_download_sync.
|
||||
# Uploaded videos did not, so a valid VP9/AV1/Opus upload could be
|
||||
# processed successfully but remain undecodable by the in-app WebView.
|
||||
# Codec probing/transcoding is blocking; keep it off the event loop.
|
||||
if source.get("kind") != "url" and input_type != "audio":
|
||||
video_path = await _ensure_browser_playable_mp4_for_job(job_id, video_path)
|
||||
|
||||
# Content-hash cache: reuse artifacts from previous matching jobs.
|
||||
content_hash = await asyncio.to_thread(compute_file_hash, audio_path)
|
||||
cached = find_cached_job(content_hash, job_id)
|
||||
@@ -1295,6 +1370,7 @@ async def ingest_pipeline(
|
||||
"scene_cuts": scene_cuts,
|
||||
"youtube_subs": youtube_subs_by_lang or None,
|
||||
"input_type": input_type,
|
||||
"source_lang_override": source.get("source_lang"),
|
||||
}
|
||||
if not put_and_save_job(
|
||||
job_id, full_job, filename=filename, duration=dur, content_hash=content_hash,
|
||||
@@ -1323,6 +1399,7 @@ async def ingest_pipeline(
|
||||
"scene_cuts": [],
|
||||
"youtube_subs": youtube_subs_by_lang or None,
|
||||
"input_type": input_type,
|
||||
"source_lang_override": source.get("source_lang"),
|
||||
}
|
||||
if not put_and_save_job(
|
||||
job_id, partial, filename=filename, duration=dur, content_hash=content_hash,
|
||||
|
||||
@@ -177,6 +177,9 @@ class LocalCall:
|
||||
queue_timeout: Optional[float] = None
|
||||
# The engine's declared VRAM floor; only shapes the timeout message.
|
||||
min_vram_gb: float = 0.0
|
||||
# Called once a local worker abandoned by its waiter can no longer touch
|
||||
# request-owned inputs. Normal completion does not call it (#1668).
|
||||
on_abandon: Optional[Callable[[], None]] = None
|
||||
# Some remote-first callers cannot construct the local callable without
|
||||
# loading the very model they are trying to offload. Prepare it only when
|
||||
# routing/fallback actually selects this machine.
|
||||
@@ -465,6 +468,7 @@ async def _run_local(call: LocalCall, *, admit: bool = False, executor=None) ->
|
||||
queue_timeout=call.queue_timeout,
|
||||
min_vram_gb=call.min_vram_gb,
|
||||
executor=executor,
|
||||
on_abandon=call.on_abandon,
|
||||
)
|
||||
|
||||
|
||||
@@ -499,7 +503,9 @@ async def _run_remote(
|
||||
deadline = _default_deadline(call.operation, params.get("text"))
|
||||
|
||||
try:
|
||||
task = scheduler.submit(
|
||||
submit = getattr(scheduler, "submit_async", None)
|
||||
submit = submit if callable(submit) else scheduler.submit
|
||||
submitted = submit(
|
||||
operation=call.operation,
|
||||
engine=call.engine,
|
||||
model_id=call.model_id,
|
||||
@@ -508,6 +514,7 @@ async def _run_remote(
|
||||
deadline_seconds=deadline,
|
||||
pinned_worker_id=decision.worker_id,
|
||||
)
|
||||
task = await submitted if asyncio.iscoroutine(submitted) else submitted
|
||||
except QueueFull as exc:
|
||||
raise _NotDispatched(str(exc)) from exc
|
||||
|
||||
|
||||
@@ -154,6 +154,17 @@ def bundled_dir() -> str:
|
||||
return os.path.join(media_tools_dir(), f"ffbin-{_FFBIN_COMMIT[:12]}", _platform_key())
|
||||
|
||||
|
||||
def _publish_bundled_on_path() -> None:
|
||||
"""Make a newly validated bundle visible to bare-name subprocess calls."""
|
||||
directory = os.path.abspath(bundled_dir())
|
||||
current = os.environ.get("PATH", "")
|
||||
entries = current.split(os.pathsep) if current else []
|
||||
if os.path.normcase(directory) in {os.path.normcase(entry) for entry in entries if entry}:
|
||||
return
|
||||
os.environ["PATH"] = os.pathsep.join([directory, *entries])
|
||||
logger.info("Published the acquired media-tool directory on PATH")
|
||||
|
||||
|
||||
def _exe(name: str) -> str:
|
||||
return f"{name}.exe" if sys.platform == "win32" else name
|
||||
|
||||
@@ -231,12 +242,21 @@ def acquire_bundled(wait: bool = False) -> dict:
|
||||
_ops["acquire"].update(state="running", progress=0.0, error=None)
|
||||
|
||||
if all(bundled_tool_path(t) and _binary_runs(bundled_tool_path(t)) for t in TOOLS):
|
||||
# The bundle may have arrived after startup's one-time PATH publish
|
||||
# (first-run acquisition is asynchronous). Make it visible to pydub
|
||||
# and other dependencies that launch ffmpeg/ffprobe by bare name now,
|
||||
# without requiring a backend restart (#1677).
|
||||
_publish_bundled_on_path()
|
||||
_set_op("acquire", state="done", progress=1.0)
|
||||
return _op_snapshot()["acquire"]
|
||||
|
||||
def _worker():
|
||||
try:
|
||||
_do_acquire()
|
||||
# Startup cannot publish binaries which do not exist yet. The
|
||||
# background worker must complete that second half atomically with
|
||||
# installation so the very next synthesis can use the tools.
|
||||
_publish_bundled_on_path()
|
||||
_set_op("acquire", state="done", progress=1.0, error=None)
|
||||
logger.info("media-tools: bundled ffmpeg/ffprobe installed at %s", bundled_dir())
|
||||
except Exception as e:
|
||||
|
||||
@@ -631,7 +631,8 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
timeout: "float | None" = None,
|
||||
executor=None,
|
||||
queue_timeout: "float | None" = None,
|
||||
min_vram_gb: float = 0.0):
|
||||
min_vram_gb: float = 0.0,
|
||||
on_abandon=None):
|
||||
"""Run blocking ``fn`` on the GPU pool, bounding **execution** — not the
|
||||
wait for a free worker.
|
||||
|
||||
@@ -659,6 +660,12 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
at 0 — the default, and correct for every non-TTS job on this pool
|
||||
(reference transcribe, watermarking, dub steps) — the under-provisioned-GPU
|
||||
wording is never used, because nothing measured says it applies (#1226).
|
||||
|
||||
``on_abandon`` is called once, after a job whose caller stopped waiting can
|
||||
no longer access its inputs. A queued job that is cancelled before it
|
||||
starts calls it immediately; a running thread calls it from ``_job``'s
|
||||
finalizer. Normal completion never calls it. This lets request-owned temp
|
||||
files outlive abandoned workers without delaying ordinary requests (#1668).
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
ex = executor if executor is not None else _get_gpu_pool()
|
||||
@@ -673,6 +680,24 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
# job's model-load heartbeats (#1367). A dict, not a nonlocal: the closure
|
||||
# runs on a pool thread while the waiter reads from the event loop.
|
||||
_ident_box: dict = {}
|
||||
_abandon_lock = threading.Lock()
|
||||
_abandon_state = {
|
||||
"requested": False,
|
||||
"finished": False,
|
||||
"callback_called": False,
|
||||
}
|
||||
|
||||
def _fire_abandon_callback() -> None:
|
||||
if on_abandon is None:
|
||||
return
|
||||
with _abandon_lock:
|
||||
if _abandon_state["callback_called"]:
|
||||
return
|
||||
_abandon_state["callback_called"] = True
|
||||
try:
|
||||
on_abandon()
|
||||
except Exception: # noqa: BLE001 — cleanup cannot hide the pool result
|
||||
logger.exception("%s abandon cleanup failed", _log_safe(what))
|
||||
|
||||
def _job():
|
||||
# First thing the worker does: tell the awaiting coroutine the
|
||||
@@ -689,8 +714,26 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
# Idents are reused by the OS; a stale heartbeat under this ident
|
||||
# must not vouch for some future job on the same thread.
|
||||
_MODEL_LOAD_ACTIVITY.pop(threading.get_ident(), None)
|
||||
with _abandon_lock:
|
||||
_abandon_state["finished"] = True
|
||||
abandoned = _abandon_state["requested"]
|
||||
if abandoned:
|
||||
_fire_abandon_callback()
|
||||
|
||||
concurrent_fut = ex.submit(_job)
|
||||
fut = asyncio.wrap_future(concurrent_fut, loop=loop)
|
||||
|
||||
def _abandon() -> None:
|
||||
# Keep the concurrent future so we can distinguish a job cancelled out
|
||||
# of the queue from a thread that Python cannot stop once it has begun.
|
||||
cancelled_before_start = concurrent_fut.cancel()
|
||||
with _abandon_lock:
|
||||
_abandon_state["requested"] = True
|
||||
finished = _abandon_state["finished"]
|
||||
fut.cancel()
|
||||
if cancelled_before_start or finished:
|
||||
_fire_abandon_callback()
|
||||
|
||||
fut = loop.run_in_executor(ex, _job)
|
||||
waiter = asyncio.ensure_future(started.wait())
|
||||
try:
|
||||
# Phase 1 — queue wait. Watch the future too, so a job that fails or is
|
||||
@@ -703,6 +746,7 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
# Caller went away (client disconnect). We stop awaiting the job, so
|
||||
# make sure its eventual result/exception is consumed rather than
|
||||
# logged as "Future exception was never retrieved".
|
||||
_abandon()
|
||||
fut.add_done_callback(_swallow_abandoned)
|
||||
raise
|
||||
finally:
|
||||
@@ -712,7 +756,7 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
# Never picked up: cancel it out of the queue (a not-yet-started
|
||||
# concurrent future cancels cleanly) and report saturation, NOT a
|
||||
# too-heavy job.
|
||||
fut.cancel()
|
||||
_abandon()
|
||||
fut.add_done_callback(_swallow_abandoned)
|
||||
stats = gpu_pool_stats(ex)
|
||||
logger.warning(
|
||||
@@ -783,14 +827,14 @@ async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
|
||||
# Caller went away mid-execution. The old wait_for cancelled the
|
||||
# wrapper itself; asyncio.wait does not, so do both halves here or the
|
||||
# eventual result is logged as "Future exception was never retrieved".
|
||||
fut.cancel()
|
||||
_abandon()
|
||||
fut.add_done_callback(_swallow_abandoned)
|
||||
raise
|
||||
except asyncio.TimeoutError as timeout_exc:
|
||||
# Parity with the old wait_for semantics: cancel the asyncio wrapper;
|
||||
# the worker thread keeps going regardless. Consume whatever it
|
||||
# eventually produces.
|
||||
fut.cancel()
|
||||
_abandon()
|
||||
fut.add_done_callback(_swallow_abandoned)
|
||||
# Capture the stacks BEFORE reset(): reset() replaces the executor, and
|
||||
# once the wedged thread is no longer a pool worker we can no longer
|
||||
@@ -2805,6 +2849,21 @@ async def preload_model():
|
||||
if model is not None:
|
||||
return # already loaded
|
||||
|
||||
# On MPS the configured ``omnivoice`` id resolves to a crash-isolated
|
||||
# sidecar. Warming the native singleton here would put the same fatal MPS
|
||||
# allocator risk back into the API process before the isolated engine is
|
||||
# ever asked to synthesize.
|
||||
try:
|
||||
from core.device_caps import detect_host_caps
|
||||
|
||||
if detect_host_caps().family == "mps":
|
||||
logger.info(
|
||||
"Native TTS preload skipped: OmniVoice uses crash isolation on this host."
|
||||
)
|
||||
return
|
||||
except Exception: # noqa: BLE001 -- preload selection must stay best-effort
|
||||
logger.debug("effective TTS preload selection failed", exc_info=True)
|
||||
|
||||
# A machine lending its GPU has no local user to warm the model FOR. This
|
||||
# preload exists to make the first /generate feel instant for the person
|
||||
# sitting in front of the app; on a headless node there is nobody sitting
|
||||
@@ -2894,7 +2953,10 @@ async def preload_model():
|
||||
"The TTS model could not be loaded. Settings → Logs → Backend "
|
||||
"has the full error."
|
||||
)
|
||||
_set_loading("failed", detail, error=detail)
|
||||
# `sub_stage` is a public API enum and the frontend keys failure state
|
||||
# off `error`. Keep the human-readable word "failed" in the detail,
|
||||
# not in the state machine (#1695).
|
||||
_set_loading("error", detail, error=detail)
|
||||
|
||||
def get_model_status():
|
||||
is_loaded = model is not None
|
||||
|
||||
@@ -79,6 +79,7 @@ class Segment:
|
||||
text: str
|
||||
speaker_id: str = "Speaker 1"
|
||||
id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
|
||||
extra: dict = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def duration(self) -> float:
|
||||
@@ -90,6 +91,7 @@ class Segment:
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
**self.extra,
|
||||
"id": self.id,
|
||||
"start": round(self.start, 2),
|
||||
"end": round(self.end, 2),
|
||||
@@ -98,6 +100,46 @@ class Segment:
|
||||
}
|
||||
|
||||
|
||||
def _merge_segment_extra(target: Segment, incoming: Segment, *, prepend: bool) -> None:
|
||||
"""Preserve editor metadata when cleanup folds ``incoming`` into ``target``."""
|
||||
for key, value in incoming.extra.items():
|
||||
target.extra.setdefault(key, value)
|
||||
|
||||
def joined(left: object, right: object) -> str:
|
||||
return _clean(f"{left or ''} {right or ''}")
|
||||
|
||||
target_original = target.extra.get("text_original")
|
||||
incoming_original = incoming.extra.get("text_original")
|
||||
if target_original is not None or incoming_original is not None:
|
||||
target.extra["text_original"] = (
|
||||
joined(incoming_original, target_original)
|
||||
if prepend
|
||||
else joined(target_original, incoming_original)
|
||||
)
|
||||
|
||||
raw_target_translations = target.extra.get("translations")
|
||||
raw_incoming_translations = incoming.extra.get("translations")
|
||||
target_translations = raw_target_translations if isinstance(raw_target_translations, dict) else {}
|
||||
incoming_translations = (
|
||||
raw_incoming_translations if isinstance(raw_incoming_translations, dict) else {}
|
||||
)
|
||||
if target_translations or incoming_translations:
|
||||
merged = {}
|
||||
languages = {
|
||||
*target_translations.keys(),
|
||||
*incoming_translations.keys(),
|
||||
}
|
||||
for language in languages:
|
||||
target_text = target_translations.get(language)
|
||||
incoming_text = incoming_translations.get(language)
|
||||
merged[language] = (
|
||||
joined(incoming_text, target_text)
|
||||
if prepend
|
||||
else joined(target_text, incoming_text)
|
||||
)
|
||||
target.extra["translations"] = merged
|
||||
|
||||
|
||||
def _clean(text: str) -> str:
|
||||
return _WS.sub(" ", (text or "").strip())
|
||||
|
||||
@@ -317,12 +359,14 @@ def _merge_short(segments: List[Segment]) -> List[Segment]:
|
||||
i += 1
|
||||
continue
|
||||
if target is prev:
|
||||
_merge_segment_extra(prev, s, prepend=False)
|
||||
prev.text = _clean(prev.text + " " + s.text)
|
||||
prev.end = max(prev.end, s.end)
|
||||
segments.pop(i)
|
||||
did_merge = True
|
||||
continue
|
||||
if target is nxt:
|
||||
_merge_segment_extra(nxt, s, prepend=True)
|
||||
nxt.text = _clean(s.text + " " + nxt.text)
|
||||
nxt.start = min(nxt.start, s.start)
|
||||
segments.pop(i)
|
||||
@@ -360,6 +404,7 @@ def _stitch_adjacent_shorts(segments: List[Segment]) -> List[Segment]:
|
||||
and b.duration <= STITCH_DUR
|
||||
and combined_dur <= MAX_DUR
|
||||
):
|
||||
_merge_segment_extra(a, b, prepend=False)
|
||||
a.text = _clean(a.text + " " + b.text)
|
||||
a.end = b.end
|
||||
segments.pop(i + 1)
|
||||
@@ -386,6 +431,11 @@ def clean_up_segments(segments: List[dict]) -> List[dict]:
|
||||
text=_clean(str(s.get("text", ""))),
|
||||
speaker_id=str(s.get("speaker_id") or "Speaker 1"),
|
||||
id=str(s.get("id") or uuid.uuid4().hex[:8]),
|
||||
extra={
|
||||
key: value
|
||||
for key, value in s.items()
|
||||
if key not in {"id", "start", "end", "text", "speaker_id"}
|
||||
},
|
||||
))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
|
||||
@@ -254,6 +254,31 @@ def get_text(key: str, default: Optional[str] = None) -> Optional[str]:
|
||||
return default
|
||||
|
||||
|
||||
def get_text_state(key: str) -> tuple[bool, str]:
|
||||
"""Return ``(is_present, value)`` without hiding storage failures.
|
||||
|
||||
Rollback snapshots must distinguish a missing row from an unreadable
|
||||
database. ``get_text`` deliberately collapses those cases for ordinary
|
||||
preference reads, so transactional callers use this strict variant.
|
||||
"""
|
||||
if key == _TOKEN_KEY or key.startswith(_SECRET_PREFIX):
|
||||
raise ValueError(
|
||||
"get_text_state refuses to read an encrypted secret row; "
|
||||
"use get_hf_token()/get_secret() for secrets"
|
||||
)
|
||||
from core.db import db_conn
|
||||
|
||||
with db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT value FROM settings WHERE key = ?", (key,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return False, ""
|
||||
if row[0] is None:
|
||||
return True, ""
|
||||
return True, str(row[0])
|
||||
|
||||
|
||||
def set_text(key: str, value: str) -> None:
|
||||
"""Persist a non-encrypted text value into the settings table.
|
||||
|
||||
@@ -274,6 +299,19 @@ def set_text(key: str, value: str) -> None:
|
||||
)
|
||||
|
||||
|
||||
def clear_text(key: str) -> None:
|
||||
"""Remove a non-encrypted text setting, preserving a missing-row default."""
|
||||
if key == _TOKEN_KEY or key.startswith(_SECRET_PREFIX):
|
||||
raise ValueError(
|
||||
"clear_text refuses to delete an encrypted secret row; "
|
||||
"use clear_hf_token()/clear_secret() for secrets"
|
||||
)
|
||||
from core.db import db_conn
|
||||
|
||||
with db_conn() as conn:
|
||||
conn.execute("DELETE FROM settings WHERE key = ?", (key,))
|
||||
|
||||
|
||||
# ── Phase 4 Plan 04-01 (GGUF-04): per-engine quant override ────────────────
|
||||
#
|
||||
# Settings row "gguf_quant_override" holds either:
|
||||
|
||||
@@ -60,6 +60,7 @@ from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
from core.config import DATA_DIR
|
||||
from core.contained_subprocess import OwnedPopen, spawn_owned
|
||||
|
||||
logger = logging.getLogger("omnivoice.sidecar_install")
|
||||
|
||||
@@ -998,6 +999,11 @@ def _step_persist(spec: SidecarSpec, job: dict) -> None:
|
||||
# ── Subprocess runner with live log capture ────────────────────────────────
|
||||
|
||||
|
||||
def _install_containment_kwargs() -> dict:
|
||||
"""Nested process-group/Job ownership is supplied by ``spawn_owned``."""
|
||||
return {}
|
||||
|
||||
|
||||
def _run_logged(job: dict, argv: list[str], *, timeout: float,
|
||||
env: "dict[str, str] | None" = None) -> int:
|
||||
"""Run *argv*, streaming combined stdout+stderr lines into the job log.
|
||||
@@ -1012,13 +1018,11 @@ def _run_logged(job: dict, argv: list[str], *, timeout: float,
|
||||
killed child — a blocking ``for line in proc.stdout`` on this thread
|
||||
would hang past the timeout waiting for pipe EOF.
|
||||
"""
|
||||
popen_kwargs: dict = {}
|
||||
if os.name == "posix":
|
||||
# New session → we can kill the whole process group on timeout
|
||||
# instead of only the direct child.
|
||||
popen_kwargs["start_new_session"] = True
|
||||
# ``spawn_owned`` creates the local timeout group/Job before the operation
|
||||
# starts and links it to backend death through its control pipe.
|
||||
popen_kwargs = _install_containment_kwargs()
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
proc = spawn_owned(
|
||||
argv,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
@@ -1053,29 +1057,18 @@ def _run_logged(job: dict, argv: list[str], *, timeout: float,
|
||||
|
||||
|
||||
def _kill_tree(proc: "subprocess.Popen") -> None:
|
||||
"""Kill the child and its whole process tree, on every platform.
|
||||
|
||||
POSIX: the child was started in its own session, so SIGKILL the group.
|
||||
Windows: ``proc.kill()`` only terminates the direct child — a git/uv
|
||||
helper it spawned would keep running (and writing into the checkout)
|
||||
past our timeout — so use ``taskkill /T`` to fell the tree.
|
||||
"""
|
||||
if os.name == "posix":
|
||||
import signal
|
||||
"""Kill an operation through its stable nested group/Job owner."""
|
||||
if isinstance(proc, OwnedPopen):
|
||||
# The retained supervisor/process-group or nested Job is the stable
|
||||
# per-operation owner. Do not fall back to a direct PID kill.
|
||||
proc.kill()
|
||||
try:
|
||||
os.killpg(proc.pid, signal.SIGKILL)
|
||||
return
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
pass # group already gone / not ours — fall through to plain kill
|
||||
else: # Windows
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/F", "/T", "/PID", str(proc.pid)],
|
||||
capture_output=True, timeout=15,
|
||||
)
|
||||
return
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
pass # taskkill unavailable/failed — fall through to plain kill
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
pass
|
||||
return
|
||||
# A test double or a legacy caller without the nested owner can only be
|
||||
# stopped through its stable direct-process handle.
|
||||
try:
|
||||
proc.kill()
|
||||
except OSError:
|
||||
|
||||
@@ -32,9 +32,9 @@ Threat-model summary (see Plan 02-01 frontmatter):
|
||||
AUTH-05 installed (``HFTokenRedactor``) on the root logger.
|
||||
T-02-04 — compromised sidecar emitting unexpected ops: parent allowlist
|
||||
``PARENT_INBOUND_OPS`` rejects everything else.
|
||||
T-02-05 — Tauri group-kill scope: ``start_new_session=True`` on Unix
|
||||
and ``CREATE_NEW_PROCESS_GROUP`` on Windows isolate the
|
||||
sidecar's process group.
|
||||
T-02-05 — nested containment: a retained supervisor process group/Job owns
|
||||
each engine operation and is linked to backend death by a control
|
||||
pipe, while still permitting independent timeout teardown.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -56,6 +56,7 @@ from typing import Optional
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from core.contained_subprocess import spawn_owned
|
||||
from services.tts_backend import TTSBackend
|
||||
|
||||
logger = logging.getLogger("omnivoice.subprocess_backend")
|
||||
@@ -470,13 +471,6 @@ class SubprocessBackend(TTSBackend):
|
||||
"env": env,
|
||||
"bufsize": 0, # unbuffered binary pipes
|
||||
}
|
||||
# Process-group isolation so the Tauri lib.rs group-kill in shutdown
|
||||
# doesn't escape into other children. See T-02-05.
|
||||
if sys.platform == "win32":
|
||||
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||||
else:
|
||||
kwargs["start_new_session"] = True
|
||||
|
||||
# `venv_python()` resolves the engine's interpreter, and on a cold
|
||||
# first run that is not cheap: it spawns each candidate to import the
|
||||
# engine (bounded, but tens of seconds on a slow disk), and if none is
|
||||
@@ -509,7 +503,7 @@ class SubprocessBackend(TTSBackend):
|
||||
self.id, Path(python_path).name, Path(script_path).name,
|
||||
)
|
||||
try:
|
||||
self._proc = subprocess.Popen([python_path, script_path], **kwargs)
|
||||
self._proc = spawn_owned([python_path, script_path], **kwargs)
|
||||
except OSError as exc:
|
||||
raise InvalidBinaryError(
|
||||
python_path,
|
||||
|
||||
@@ -669,7 +669,7 @@ class OmniVoiceBackend(TTSBackend):
|
||||
|
||||
id = "omnivoice"
|
||||
display_name = "VoiceStudio (k2-fsa/OmniVoice, 600+ languages)"
|
||||
gpu_compat = ("cuda", "mps", "cpu")
|
||||
gpu_compat = ("cuda", "rocm", "mps", "cpu")
|
||||
# Derived from the pool's own per-job budget (_GPU_VRAM_PER_JOB_GB = 5.0 in
|
||||
# model_manager, itself measured from the ~1.6 GB forward + autoregressive
|
||||
# decode and the co-loaded WhisperX on the clone path), plus room for the
|
||||
@@ -2398,6 +2398,7 @@ def list_backends() -> list[dict]:
|
||||
|
||||
out: list[dict] = []
|
||||
for bid, cls in _REGISTRY.items():
|
||||
cls = _effective_backend_class(bid, cls, caps.family)
|
||||
try:
|
||||
ok, msg = cls.is_available()
|
||||
except Exception:
|
||||
@@ -2471,10 +2472,29 @@ def list_backends() -> list[dict]:
|
||||
return out
|
||||
|
||||
|
||||
def _effective_backend_class(
|
||||
backend_id: str,
|
||||
backend_cls: type[TTSBackend],
|
||||
host_family: str | None = None,
|
||||
) -> type[TTSBackend]:
|
||||
"""Resolve host-specific containment without changing the configured id."""
|
||||
if backend_id != "omnivoice":
|
||||
return backend_cls
|
||||
if host_family is None:
|
||||
from core.device_caps import detect_host_caps
|
||||
|
||||
host_family = detect_host_caps().family
|
||||
if host_family != "mps":
|
||||
return backend_cls
|
||||
from engines.omnivoice_subprocess import OmniVoiceMPSSubprocessBackend
|
||||
|
||||
return OmniVoiceMPSSubprocessBackend
|
||||
|
||||
|
||||
def get_backend_class(backend_id: str) -> type[TTSBackend]:
|
||||
if backend_id not in _REGISTRY:
|
||||
raise ValueError(f"Unknown TTS backend: {backend_id!r}. Known: {list(_REGISTRY)}")
|
||||
return _REGISTRY[backend_id]
|
||||
return _effective_backend_class(backend_id, _REGISTRY[backend_id])
|
||||
|
||||
|
||||
def cloning_capable_engine_ids() -> list[str]:
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Dependency-free client for VoiceStudio's local speech platform."""
|
||||
@@ -0,0 +1,278 @@
|
||||
"""CLI/module bridge for terminals, editor extensions, and agent hooks.
|
||||
|
||||
The desktop app must be running for native dictation control. Batch
|
||||
transcription can also target a standalone or remote VoiceStudio backend.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ipaddress
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from pathlib import Path
|
||||
import secrets
|
||||
import sys
|
||||
from typing import Any
|
||||
from urllib import error, request
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
DEFAULT_CONTROL_URL = "http://127.0.0.1:3902"
|
||||
DEFAULT_ENGINE_URL = "http://127.0.0.1:3900"
|
||||
|
||||
|
||||
class SpeechClientError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class _RejectCredentialRedirect(request.HTTPRedirectHandler):
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ARG002
|
||||
raise SpeechClientError("VoiceStudio refused a credentialed redirect")
|
||||
|
||||
|
||||
def _join_url(base_url: str, path: str) -> str:
|
||||
return f"{base_url.rstrip('/')}/{path.lstrip('/')}"
|
||||
|
||||
|
||||
def _decode_error(exc: error.HTTPError) -> str:
|
||||
try:
|
||||
body = exc.read().decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
body = ""
|
||||
try:
|
||||
detail = json.loads(body)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
detail = body.strip()
|
||||
return f"HTTP {exc.code}: {detail or exc.reason}"
|
||||
|
||||
|
||||
def _is_loopback_host(host: str | None) -> bool:
|
||||
if not host:
|
||||
return False
|
||||
if host.lower() == "localhost":
|
||||
return True
|
||||
try:
|
||||
return ipaddress.ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _open(req: request.Request, timeout: float = 300.0) -> tuple[bytes, str]:
|
||||
target = urlsplit(req.full_url)
|
||||
scheme = target.scheme.lower()
|
||||
if scheme not in {"http", "https"}:
|
||||
raise SpeechClientError("VoiceStudio URLs must use http:// or https://")
|
||||
credentialed = bool(req.get_header("Authorization"))
|
||||
if credentialed and scheme != "https" and not _is_loopback_host(target.hostname):
|
||||
raise SpeechClientError("Remote VoiceStudio credentials require https://")
|
||||
try:
|
||||
opener = (
|
||||
request.build_opener(_RejectCredentialRedirect())
|
||||
if credentialed
|
||||
else request.build_opener()
|
||||
)
|
||||
with opener.open(req, timeout=timeout) as response: # noqa: S310
|
||||
return response.read(), response.headers.get("Content-Type", "")
|
||||
except error.HTTPError as exc:
|
||||
raise SpeechClientError(_decode_error(exc)) from exc
|
||||
except error.URLError as exc:
|
||||
raise SpeechClientError(f"VoiceStudio is unavailable: {exc.reason}") from exc
|
||||
|
||||
|
||||
def _json_request(method: str, url: str, payload: Any | None = None) -> Any:
|
||||
data = None if payload is None else json.dumps(payload).encode("utf-8")
|
||||
headers = {"Accept": "application/json"}
|
||||
if data is not None:
|
||||
headers["Content-Type"] = "application/json"
|
||||
body, _ = _open(request.Request(url, data=data, headers=headers, method=method), timeout=10.0)
|
||||
try:
|
||||
return json.loads(body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise SpeechClientError("VoiceStudio returned invalid JSON") from exc
|
||||
|
||||
|
||||
def _encode_multipart(
|
||||
*,
|
||||
filename: str,
|
||||
audio: bytes,
|
||||
fields: dict[str, str],
|
||||
boundary: str | None = None,
|
||||
) -> tuple[bytes, str]:
|
||||
boundary = boundary or f"voicestudio-{secrets.token_hex(16)}"
|
||||
marker = boundary.encode("ascii")
|
||||
parts: list[bytes] = []
|
||||
for name, value in fields.items():
|
||||
parts.extend(
|
||||
[
|
||||
b"--" + marker + b"\r\n",
|
||||
f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
|
||||
value.encode("utf-8"),
|
||||
b"\r\n",
|
||||
]
|
||||
)
|
||||
safe_filename = Path(filename).name.replace('"', "") or "audio.wav"
|
||||
content_type = mimetypes.guess_type(safe_filename)[0] or "application/octet-stream"
|
||||
if Path(safe_filename).suffix.lower() in {".wav", ".wave"}:
|
||||
content_type = "audio/wav"
|
||||
parts.extend(
|
||||
[
|
||||
b"--" + marker + b"\r\n",
|
||||
(
|
||||
'Content-Disposition: form-data; name="file"; '
|
||||
f'filename="{safe_filename}"\r\n'
|
||||
).encode(),
|
||||
f"Content-Type: {content_type}\r\n\r\n".encode(),
|
||||
audio,
|
||||
b"\r\n--" + marker + b"--\r\n",
|
||||
]
|
||||
)
|
||||
return b"".join(parts), f"multipart/form-data; boundary={boundary}"
|
||||
|
||||
|
||||
def _control(args: argparse.Namespace, action: str) -> int:
|
||||
method = "GET" if action in {"status", "capabilities"} else "POST"
|
||||
path = {
|
||||
"status": "/v1/status",
|
||||
"capabilities": "/v1/capabilities",
|
||||
"start": "/v1/dictation/start",
|
||||
"stop": "/v1/dictation/stop",
|
||||
"toggle": "/v1/dictation/toggle",
|
||||
}[action]
|
||||
result = _json_request(method, _join_url(args.control_url, path))
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
def _read_audio(path: str, stdin_filename: str) -> tuple[bytes, str]:
|
||||
if path == "-":
|
||||
return sys.stdin.buffer.read(), stdin_filename
|
||||
audio_path = Path(path)
|
||||
try:
|
||||
return audio_path.read_bytes(), audio_path.name
|
||||
except OSError as exc:
|
||||
display_name = path.replace("\\", "/").rsplit("/", 1)[-1] or "audio input"
|
||||
reason = exc.strerror or type(exc).__name__
|
||||
raise SpeechClientError(f"could not read '{display_name}': {reason}") from exc
|
||||
|
||||
|
||||
def _response_text(body: bytes, content_type: str) -> str:
|
||||
decoded = body.decode("utf-8", errors="replace")
|
||||
if "json" not in content_type.lower():
|
||||
return decoded
|
||||
try:
|
||||
payload = json.loads(decoded)
|
||||
except json.JSONDecodeError:
|
||||
return decoded
|
||||
if isinstance(payload, dict) and isinstance(payload.get("text"), str):
|
||||
return payload["text"]
|
||||
return decoded
|
||||
|
||||
|
||||
def _transcribe(args: argparse.Namespace) -> int:
|
||||
audio, filename = _read_audio(args.audio, args.stdin_filename)
|
||||
fields = {
|
||||
"model": args.model,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.language:
|
||||
fields["language"] = args.language
|
||||
body, content_type = _encode_multipart(filename=filename, audio=audio, fields=fields)
|
||||
headers = {"Content-Type": content_type, "Accept": "application/json, text/plain"}
|
||||
api_key = os.environ.get("OMNIVOICE_API_KEY", "").strip()
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
output_session_id = None
|
||||
if args.insert:
|
||||
session = _json_request(
|
||||
"POST", _join_url(args.control_url, "/v1/output/sessions")
|
||||
)
|
||||
output_session_id = session["session_id"]
|
||||
|
||||
session_needs_cleanup = output_session_id is not None
|
||||
try:
|
||||
response_body, response_type = _open(
|
||||
request.Request(
|
||||
_join_url(args.engine_url, "/v1/audio/transcriptions"),
|
||||
data=body,
|
||||
headers=headers,
|
||||
method="POST",
|
||||
)
|
||||
)
|
||||
if output_session_id is not None:
|
||||
_json_request(
|
||||
"POST",
|
||||
_join_url(
|
||||
args.control_url,
|
||||
f"/v1/output/sessions/{output_session_id}/insert",
|
||||
),
|
||||
{"text": _response_text(response_body, response_type)},
|
||||
)
|
||||
session_needs_cleanup = False
|
||||
finally:
|
||||
if session_needs_cleanup:
|
||||
try:
|
||||
_json_request(
|
||||
"DELETE",
|
||||
_join_url(args.control_url, f"/v1/output/sessions/{output_session_id}"),
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
# Best-effort cleanup must not replace the original failure or
|
||||
# KeyboardInterrupt that brought control into this finally.
|
||||
pass
|
||||
|
||||
sys.stdout.buffer.write(response_body)
|
||||
if response_body and not response_body.endswith(b"\n"):
|
||||
sys.stdout.buffer.write(b"\n")
|
||||
return 0
|
||||
|
||||
|
||||
def _parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="voicestudio-speech",
|
||||
description="Control and consume VoiceStudio's local speech platform.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control-url",
|
||||
default=os.environ.get("VOICESTUDIO_SPEECH_URL", DEFAULT_CONTROL_URL),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--engine-url",
|
||||
default=os.environ.get("VOICESTUDIO_URL", DEFAULT_ENGINE_URL),
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
for command in ("status", "capabilities", "start", "stop", "toggle"):
|
||||
subparsers.add_parser(command)
|
||||
|
||||
transcribe = subparsers.add_parser("transcribe")
|
||||
transcribe.add_argument("audio", help="audio file, or - for stdin")
|
||||
transcribe.add_argument("--stdin-filename", default="audio.wav")
|
||||
transcribe.add_argument("--model", default="whisper-1")
|
||||
transcribe.add_argument("--language")
|
||||
transcribe.add_argument(
|
||||
"--format",
|
||||
dest="response_format",
|
||||
choices=("json", "text", "verbose_json", "srt", "vtt"),
|
||||
default="text",
|
||||
)
|
||||
transcribe.add_argument(
|
||||
"--insert",
|
||||
action="store_true",
|
||||
help="insert the result into the app focused when this command starts",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = _parser().parse_args(argv)
|
||||
try:
|
||||
if args.command == "transcribe":
|
||||
return _transcribe(args)
|
||||
return _control(args, args.command)
|
||||
except (SpeechClientError, KeyError) as exc:
|
||||
print(f"voicestudio-speech: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -80,12 +80,10 @@ def test_timeout_error_is_a_timeouterror_subclass():
|
||||
assert issubclass(ASRTimeoutError, TimeoutError)
|
||||
|
||||
|
||||
def test_timeout_resets_a_resilient_pool_to_restore_capacity():
|
||||
# #730: a wedged transcribe holds its GPU-pool worker forever; with a 1-2
|
||||
# worker pool that starves TTS generate and surfaces as "can't reach
|
||||
# backend". On timeout, run_transcribe_guarded must reset() a pool that
|
||||
# supports it (the real _ResilientGpuPool) so the next submit gets a fresh
|
||||
# worker — capacity restored without an app restart.
|
||||
def test_timeout_does_not_overlap_an_in_process_native_worker():
|
||||
# #1669: reset() cannot kill the old native thread. A fresh pool let the
|
||||
# retry enter the same whisperx/CTranslate2 model concurrently and the
|
||||
# process died with 0xC0000005. Keep the old worker accounted for instead.
|
||||
class _FakePool(ThreadPoolExecutor):
|
||||
def __init__(self):
|
||||
super().__init__(max_workers=1)
|
||||
@@ -105,7 +103,7 @@ def test_timeout_resets_a_resilient_pool_to_restore_capacity():
|
||||
await run_transcribe_guarded(pool, _hang, what="Dub", timeout=0.2)
|
||||
|
||||
asyncio.run(_go())
|
||||
assert pool.reset_calls == 1
|
||||
assert pool.reset_calls == 0
|
||||
pool.shutdown(wait=False)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
"""Stable nested operation ownership (model-free, cross-platform seams)."""
|
||||
import ctypes
|
||||
import builtins
|
||||
import os
|
||||
import runpy
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from ctypes import wintypes
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from core import contained_subprocess as owned
|
||||
|
||||
|
||||
class _Call:
|
||||
def __init__(self, fn):
|
||||
self.fn = fn
|
||||
|
||||
def __call__(self, *args):
|
||||
return self.fn(*args)
|
||||
|
||||
|
||||
def test_supervisor_argv_uses_entry_module_for_source_and_frozen_binary(monkeypatch):
|
||||
monkeypatch.delattr(owned.sys, "frozen", raising=False)
|
||||
source = owned._supervisor_argv(3, 4, ["operation"])
|
||||
assert source[:2] == [sys.executable, str(Path(owned.__file__).parents[1] / "main.py")]
|
||||
assert source[2:] == ["--supervise", "3", "4", "--", "operation"]
|
||||
|
||||
monkeypatch.setattr(owned.sys, "frozen", True, raising=False)
|
||||
frozen = owned._supervisor_argv(3, 4, ["operation"])
|
||||
assert frozen == [sys.executable, "--supervise", "3", "4", "--", "operation"]
|
||||
|
||||
|
||||
def test_source_main_dispatches_supervisor_before_heavy_imports(monkeypatch):
|
||||
calls = []
|
||||
fake = types.ModuleType("core.contained_subprocess")
|
||||
fake.supervisor_main = lambda args: calls.append(args) or 23
|
||||
monkeypatch.setitem(sys.modules, "core.contained_subprocess", fake)
|
||||
main_path = Path(owned.__file__).parents[1] / "main.py"
|
||||
monkeypatch.setattr(
|
||||
sys,
|
||||
"argv",
|
||||
[str(main_path), "--supervise", "3", "4", "--", "operation"],
|
||||
)
|
||||
original_import = builtins.__import__
|
||||
|
||||
def guard_heavy_import(name, *args, **kwargs):
|
||||
if name == "math":
|
||||
raise AssertionError("supervisor dispatch reached application imports")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", guard_heavy_import)
|
||||
with pytest.raises(SystemExit, match="23"):
|
||||
runpy.run_path(str(main_path), run_name="__main__")
|
||||
assert calls == [["--supervise", "3", "4", "--", "operation"]]
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="Unix drain pipe contract")
|
||||
def test_drain_fd_is_explicitly_inherited_by_wrapper_but_not_operation(monkeypatch):
|
||||
drain_read, drain_write = os.pipe()
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_CONTAINED", "1")
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_DRAIN_FD", str(drain_write))
|
||||
owned.secure_backend_drain_fd()
|
||||
assert not os.get_inheritable(drain_write)
|
||||
implicit_probe = subprocess.check_output(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import os; "
|
||||
"fd=int(os.environ['OMNIVOICE_DESKTOP_DRAIN_FD']); "
|
||||
"\ntry: os.fstat(fd); print('leaked')"
|
||||
"\nexcept OSError: print('closed')",
|
||||
],
|
||||
close_fds=False,
|
||||
text=True,
|
||||
)
|
||||
assert implicit_probe.strip() == "closed"
|
||||
script = (
|
||||
"import os,time; token=os.environ.get('OMNIVOICE_DESKTOP_DRAIN_FD'); "
|
||||
"marker=os.environ.get('OMNIVOICE_DESKTOP_CONTAINED'); "
|
||||
"\nif token is None and marker is None: state='stripped'"
|
||||
"\nelse:"
|
||||
"\n try: os.fstat(int(token)); state='leaked'"
|
||||
"\n except OSError: state='closed'"
|
||||
"\nprint(state, flush=True); time.sleep(60)"
|
||||
)
|
||||
proc = owned.spawn_owned(
|
||||
[sys.executable, "-c", script],
|
||||
stdout=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
try:
|
||||
assert proc.stdout.readline().strip() == "stripped"
|
||||
os.close(drain_write)
|
||||
drain_write = -1
|
||||
os.set_blocking(drain_read, False)
|
||||
with pytest.raises(BlockingIOError):
|
||||
os.read(drain_read, 1) # wrapper still holds the only writer
|
||||
proc.kill()
|
||||
proc.wait(timeout=5)
|
||||
deadline = time.monotonic() + 2
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
if os.read(drain_read, 1) == b"":
|
||||
break
|
||||
except BlockingIOError:
|
||||
time.sleep(0.01)
|
||||
else:
|
||||
pytest.fail("wrapper exit did not close the desktop drain writer")
|
||||
finally:
|
||||
if drain_write >= 0:
|
||||
os.close(drain_write)
|
||||
os.close(drain_read)
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
|
||||
def test_invalid_or_missing_desktop_drain_fd_fails_safe(monkeypatch):
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_CONTAINED", "1")
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_DRAIN_FD", "not-an-fd")
|
||||
with pytest.raises(RuntimeError, match="missing its live.*drain descriptor"):
|
||||
owned.spawn_owned([sys.executable, "-c", "print('unsafe')"])
|
||||
|
||||
monkeypatch.delenv("OMNIVOICE_DESKTOP_DRAIN_FD")
|
||||
with pytest.raises(RuntimeError, match="missing its live.*drain descriptor"):
|
||||
owned.secure_backend_drain_fd()
|
||||
|
||||
monkeypatch.delenv("OMNIVOICE_DESKTOP_CONTAINED")
|
||||
assert owned.backend_drain_fd(required=True) is None
|
||||
proc = owned.spawn_owned(
|
||||
[sys.executable, "-c", "print('standalone')"],
|
||||
stdout=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
assert proc.stdout.readline().strip() == "standalone"
|
||||
assert proc.wait(timeout=5) == 0
|
||||
|
||||
|
||||
def test_windows_operation_is_in_kill_on_close_job_before_resume(monkeypatch):
|
||||
"""The child gets no instruction before stable nested Job assignment."""
|
||||
events = []
|
||||
job_closed = threading.Event()
|
||||
job = 99
|
||||
|
||||
def close_handle(handle):
|
||||
value = getattr(handle, "value", handle)
|
||||
events.append(("close", value))
|
||||
if value == job:
|
||||
job_closed.set()
|
||||
return True
|
||||
|
||||
kernel = type("Kernel", (), {})()
|
||||
kernel.AssignProcessToJobObject = _Call(
|
||||
lambda assigned_job, process: events.append(("assign", assigned_job, process)) or True
|
||||
)
|
||||
kernel.TerminateJobObject = _Call(
|
||||
lambda assigned_job, code: events.append(("terminate", assigned_job, code)) or True
|
||||
)
|
||||
kernel.WriteFile = _Call(
|
||||
lambda handle, payload, size, written, overlap: events.append(("write", size)) or True
|
||||
)
|
||||
kernel.CloseHandle = _Call(close_handle)
|
||||
|
||||
def read_control(*_args):
|
||||
job_closed.wait(2)
|
||||
return False
|
||||
|
||||
kernel.ReadFile = _Call(read_control)
|
||||
monkeypatch.setattr(owned, "_windows_job", lambda: (job, kernel, wintypes))
|
||||
monkeypatch.setattr(
|
||||
owned,
|
||||
"_resume_windows_process",
|
||||
lambda _kernel, _types, pid: events.append(("resume", pid)),
|
||||
)
|
||||
|
||||
class Child:
|
||||
_handle = 77
|
||||
pid = 123
|
||||
|
||||
def wait(self, timeout=None):
|
||||
events.append(("wait", timeout))
|
||||
return 0
|
||||
|
||||
monkeypatch.setattr(
|
||||
owned.subprocess,
|
||||
"Popen",
|
||||
lambda *args, **kwargs: events.append(("spawn", kwargs["creationflags"])) or Child(),
|
||||
)
|
||||
|
||||
assert owned._supervise_windows(11, 12, ["operation.exe"]) == 0
|
||||
assert job_closed.wait(1)
|
||||
|
||||
names = [event[0] for event in events]
|
||||
assert names.index("assign") < names.index("resume") < names.index("wait")
|
||||
assert names.index("wait") < names.index("terminate") < names.index("write")
|
||||
|
||||
|
||||
def test_windows_assignment_failure_kills_suspended_unowned_child(monkeypatch):
|
||||
"""A child outside the nested Job must be killed through its stable handle."""
|
||||
events = []
|
||||
job_closed = threading.Event()
|
||||
job = 99
|
||||
|
||||
def close_handle(handle):
|
||||
value = getattr(handle, "value", handle)
|
||||
events.append(("close", value))
|
||||
if value == job:
|
||||
job_closed.set()
|
||||
return True
|
||||
|
||||
kernel = type("Kernel", (), {})()
|
||||
kernel.AssignProcessToJobObject = _Call(
|
||||
lambda assigned_job, process: events.append(("assign", assigned_job, process))
|
||||
or False
|
||||
)
|
||||
kernel.TerminateJobObject = _Call(
|
||||
lambda assigned_job, code: events.append(("terminate", assigned_job, code)) or True
|
||||
)
|
||||
kernel.WriteFile = _Call(
|
||||
lambda handle, payload, size, written, overlap: events.append(("write", size)) or True
|
||||
)
|
||||
kernel.CloseHandle = _Call(close_handle)
|
||||
|
||||
def read_control(*_args):
|
||||
job_closed.wait(2)
|
||||
return False
|
||||
|
||||
kernel.ReadFile = _Call(read_control)
|
||||
monkeypatch.setattr(owned, "_windows_job", lambda: (job, kernel, wintypes))
|
||||
monkeypatch.setattr(ctypes, "get_last_error", lambda: 5, raising=False)
|
||||
|
||||
class Child:
|
||||
_handle = 77
|
||||
pid = 123
|
||||
|
||||
def kill(self):
|
||||
events.append(("kill",))
|
||||
|
||||
def wait(self, timeout=None):
|
||||
events.append(("wait", timeout))
|
||||
return 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
owned.subprocess,
|
||||
"Popen",
|
||||
lambda *args, **kwargs: events.append(("spawn", kwargs["creationflags"])) or Child(),
|
||||
)
|
||||
|
||||
assert owned._supervise_windows(11, 12, ["operation.exe"]) == 127
|
||||
assert job_closed.wait(1)
|
||||
|
||||
names = [event[0] for event in events]
|
||||
assert names.index("assign") < names.index("terminate") < names.index("kill")
|
||||
assert names.index("kill") < names.index("wait") < names.index("write")
|
||||
@@ -0,0 +1,126 @@
|
||||
"""macOS fallback for the os.waitid probe (#1656).
|
||||
|
||||
CPython on macOS does not expose os.waitid, so OwnedPopen's WNOWAIT dance
|
||||
crashed with AttributeError on every poll after the first spawn. These tests
|
||||
simulate that platform (monkeypatch os.waitid away) and pin the fallback:
|
||||
poll/wait/kill must work, exit codes must be real, and an already-reaped
|
||||
leader must be refused (ChildProcessError path), never signalled blind.
|
||||
"""
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from core import contained_subprocess as owned
|
||||
|
||||
|
||||
def _make_owned(argv):
|
||||
cr, cw = os.pipe()
|
||||
rr, rw = os.pipe()
|
||||
proc = subprocess.Popen(argv, start_new_session=True)
|
||||
os.close(cw)
|
||||
os.close(rw) # result writer gone: _read_result falls back to wrapper rc
|
||||
return owned.OwnedPopen(proc, cr, rr), proc
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def no_waitid(monkeypatch):
|
||||
monkeypatch.delattr(os, "waitid", raising=False)
|
||||
|
||||
|
||||
def test_poll_running_then_exited_without_waitid(no_waitid):
|
||||
h, _ = _make_owned([sys.executable, "-c", "import time; time.sleep(1.5)"])
|
||||
try:
|
||||
assert h.poll() is None, "running child must poll None"
|
||||
h._proc.wait()
|
||||
deadline = time.monotonic() + 5
|
||||
rc = None
|
||||
while rc is None and time.monotonic() < deadline:
|
||||
rc = h.poll()
|
||||
time.sleep(0.05)
|
||||
assert rc == 0
|
||||
assert h.poll() == 0
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
|
||||
|
||||
def test_poll_reports_real_exit_code_without_waitid(no_waitid):
|
||||
h, _ = _make_owned([sys.executable, "-c", "raise SystemExit(3)"])
|
||||
try:
|
||||
deadline = time.monotonic() + 5
|
||||
while h.poll() is None and time.monotonic() < deadline:
|
||||
time.sleep(0.05)
|
||||
assert h.poll() == 3
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
|
||||
|
||||
def test_wait_returns_after_kill_without_waitid(no_waitid):
|
||||
h, _ = _make_owned([sys.executable, "-c", "import time; time.sleep(30)"])
|
||||
try:
|
||||
h.kill()
|
||||
rc = h.wait(timeout=5)
|
||||
assert rc != 0
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
|
||||
|
||||
def test_reaped_by_own_popen_reports_code_without_waitid(no_waitid):
|
||||
h, proc = _make_owned([sys.executable, "-c", "pass"])
|
||||
try:
|
||||
proc.wait() # reaped through OUR handle: known code, not a refusal
|
||||
assert h.poll() == 0
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
|
||||
|
||||
def test_foreign_reaped_leader_is_refused_without_waitid(no_waitid):
|
||||
h, proc = _make_owned([sys.executable, "-c", "pass"])
|
||||
try:
|
||||
# Reap OUTSIDE this handle: Popen never learns the code, so poll must
|
||||
# refuse (None) rather than guess or signal a maybe-reused group.
|
||||
while True:
|
||||
pid, _ = os.waitpid(proc.pid, os.WNOHANG)
|
||||
if pid == proc.pid:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert h.poll() is None
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
|
||||
|
||||
def test_kill_after_pid_reuse_does_not_signal_without_waitid(no_waitid, monkeypatch):
|
||||
"""A foreign-reaped leader's reused numeric pid must not authorize killpg."""
|
||||
import signal as _signal
|
||||
|
||||
h, proc = _make_owned([sys.executable, "-c", "pass"])
|
||||
try:
|
||||
while True:
|
||||
pid, _ = os.waitpid(proc.pid, os.WNOHANG)
|
||||
if pid == proc.pid:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
# Model the numeric pid being reused: kill(pid, 0) would succeed even
|
||||
# though waitpid still reports that the original child is no longer
|
||||
# ours. The old guard therefore reached killpg and fails this test.
|
||||
monkeypatch.setattr(os, "kill", lambda _pid, _sig: None)
|
||||
signalled = []
|
||||
monkeypatch.setattr(os, "killpg", lambda pid, sig: signalled.append((pid, sig)))
|
||||
h._signal_owned_group(_signal.SIGKILL)
|
||||
assert signalled == []
|
||||
finally:
|
||||
h._close_control()
|
||||
if h._result_fd is not None:
|
||||
os.close(h._result_fd)
|
||||
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
from fastapi import UploadFile
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_ffmpeg_does_not_block_event_loop(monkeypatch, tmp_path):
|
||||
from api.routers import dub_core
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
started = asyncio.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def slow_ffmpeg(*_args, **_kwargs):
|
||||
loop.call_soon_threadsafe(started.set)
|
||||
assert release.wait(timeout=2)
|
||||
|
||||
monkeypatch.setattr(dub_core, "PREVIEW_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(dub_core, "find_ffmpeg", lambda: "ffmpeg")
|
||||
monkeypatch.setattr(dub_core.subprocess, "run", slow_ffmpeg)
|
||||
upload = UploadFile(filename="preview.mp4", file=io.BytesIO(b"video"))
|
||||
|
||||
before = loop.time()
|
||||
task = asyncio.create_task(dub_core.preview_upload(upload))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), timeout=1)
|
||||
assert loop.time() - before < 0.5
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
result = await task
|
||||
assert result["audioUrl"].endswith(".wav")
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abandoned_reader_keeps_adhoc_reference_until_worker_finishes(tmp_path):
|
||||
from api.routers.generation import (
|
||||
_TempReferenceLease,
|
||||
_run_with_reference_lease,
|
||||
)
|
||||
from services.model_manager import run_on_gpu_pool_guarded
|
||||
|
||||
reference = tmp_path / "reference.wav"
|
||||
reference.write_bytes(b"voice")
|
||||
lease = _TempReferenceLease(str(reference))
|
||||
started = threading.Event()
|
||||
release_worker = threading.Event()
|
||||
worker_read = threading.Event()
|
||||
|
||||
def read_reference():
|
||||
started.set()
|
||||
assert release_worker.wait(timeout=2)
|
||||
assert reference.read_bytes() == b"voice"
|
||||
worker_read.set()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
task = asyncio.create_task(
|
||||
_run_with_reference_lease(
|
||||
lease,
|
||||
lambda on_abandon: run_on_gpu_pool_guarded(
|
||||
read_reference,
|
||||
executor=executor,
|
||||
timeout=1,
|
||||
on_abandon=on_abandon,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
task.cancel()
|
||||
cancelled = await asyncio.gather(task, return_exceptions=True)
|
||||
assert isinstance(cancelled[0], asyncio.CancelledError)
|
||||
|
||||
lease.finish_request()
|
||||
assert reference.exists()
|
||||
release_worker.set()
|
||||
assert await asyncio.to_thread(worker_read.wait, 1)
|
||||
|
||||
for _ in range(100):
|
||||
if not reference.exists():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert not reference.exists()
|
||||
|
||||
|
||||
def test_normal_request_deletes_adhoc_reference_immediately(tmp_path):
|
||||
from api.routers.generation import _TempReferenceLease
|
||||
|
||||
reference = tmp_path / "reference.wav"
|
||||
reference.write_bytes(b"voice")
|
||||
lease = _TempReferenceLease(str(reference))
|
||||
|
||||
release = lease.acquire()
|
||||
release()
|
||||
lease.finish_request()
|
||||
|
||||
assert not reference.exists()
|
||||
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abandon_callback_waits_for_running_worker_to_finish():
|
||||
from services.model_manager import run_on_gpu_pool_guarded
|
||||
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
cleaned = threading.Event()
|
||||
|
||||
def job():
|
||||
started.set()
|
||||
assert release.wait(timeout=2)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
task = asyncio.create_task(
|
||||
run_on_gpu_pool_guarded(
|
||||
job,
|
||||
executor=executor,
|
||||
timeout=1,
|
||||
on_abandon=cleaned.set,
|
||||
)
|
||||
)
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
task.cancel()
|
||||
cancelled = await asyncio.gather(task, return_exceptions=True)
|
||||
assert isinstance(cancelled[0], asyncio.CancelledError)
|
||||
|
||||
assert not cleaned.is_set()
|
||||
release.set()
|
||||
assert await asyncio.to_thread(cleaned.wait, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queued_cancellation_releases_without_running_job():
|
||||
from services.model_manager import GpuPoolBusyError, run_on_gpu_pool_guarded
|
||||
|
||||
hog_started = threading.Event()
|
||||
release_hog = threading.Event()
|
||||
cleaned = threading.Event()
|
||||
queued_job_ran = threading.Event()
|
||||
|
||||
def hog():
|
||||
hog_started.set()
|
||||
assert release_hog.wait(timeout=2)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
hog_future = executor.submit(hog)
|
||||
assert hog_started.wait(timeout=1)
|
||||
try:
|
||||
with pytest.raises(GpuPoolBusyError):
|
||||
await run_on_gpu_pool_guarded(
|
||||
queued_job_ran.set,
|
||||
executor=executor,
|
||||
timeout=1,
|
||||
queue_timeout=0.05,
|
||||
on_abandon=cleaned.set,
|
||||
)
|
||||
assert cleaned.is_set()
|
||||
assert not queued_job_ran.is_set()
|
||||
finally:
|
||||
release_hog.set()
|
||||
hog_future.result(timeout=1)
|
||||
@@ -17,18 +17,31 @@ import json
|
||||
import math
|
||||
import array
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from services.subprocess_backend import SubprocessBackend, RECV_TIMEOUT_S
|
||||
from services.tts_backend import get_backend_class
|
||||
from engines.omnivoice_subprocess import OmniVoiceSubprocessBackend
|
||||
from services.subprocess_backend import (
|
||||
RECV_TIMEOUT_S,
|
||||
SubprocessBackend,
|
||||
)
|
||||
from services.tts_backend import OmniVoiceBackend, get_backend_class, list_backends
|
||||
from engines.omnivoice_subprocess import (
|
||||
OmniVoiceMPSSubprocessBackend,
|
||||
OmniVoiceSubprocessBackend,
|
||||
)
|
||||
|
||||
|
||||
# ── stub sidecar (model-free) ──────────────────────────────────────────────
|
||||
|
||||
STUB_SIDECAR = r'''
|
||||
import sys, json, struct, time, math, array, base64
|
||||
import sys, os, json, struct, time, math, array, base64, subprocess
|
||||
|
||||
def _send(o):
|
||||
b = json.dumps(o, separators=(",", ":")).encode()
|
||||
@@ -60,9 +73,20 @@ while True:
|
||||
sys.exit(0)
|
||||
elif op == "synthesize":
|
||||
t = m.get("text", "")
|
||||
if t == "CRASH":
|
||||
os._exit(137)
|
||||
if t == "HANG":
|
||||
while True: # wedge forever; the parent must hard-kill us
|
||||
time.sleep(1)
|
||||
if t == "HANG_CHILD":
|
||||
subprocess.Popen([
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import os,time; time.sleep(1); "
|
||||
"open(os.environ['OMNIVOICE_TIMEOUT_MARKER'], 'w').write('bad')",
|
||||
])
|
||||
while True:
|
||||
time.sleep(1)
|
||||
# Emit progress frames before the audio when asked, to exercise the
|
||||
# parent's progress-consuming recv loop (the cold-load fix).
|
||||
if t.startswith("PROG:"):
|
||||
@@ -98,6 +122,80 @@ def test_registry_resolves_to_subprocess_backend():
|
||||
assert get_backend_class("omnivoice-subprocess") is OmniVoiceSubprocessBackend
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("family", "expected_name"),
|
||||
[("mps", "OmniVoiceMPSSubprocessBackend"), ("cuda", "OmniVoiceBackend"),
|
||||
("cpu", "OmniVoiceBackend")],
|
||||
)
|
||||
def test_omnivoice_is_crash_isolated_only_on_mps(monkeypatch, family, expected_name):
|
||||
from core.device_caps import HostCaps
|
||||
|
||||
available = (family, "cpu") if family != "cpu" else ("cpu",)
|
||||
monkeypatch.setattr(
|
||||
"core.device_caps.detect_host_caps",
|
||||
lambda: HostCaps(family=family, available_families=available),
|
||||
)
|
||||
|
||||
resolved = get_backend_class("omnivoice")
|
||||
assert resolved.__name__ == expected_name
|
||||
if family != "mps":
|
||||
assert resolved is OmniVoiceBackend
|
||||
|
||||
|
||||
def test_engine_catalogue_reports_effective_mps_isolation(monkeypatch):
|
||||
from core.device_caps import HostCaps
|
||||
from services import tts_backend
|
||||
|
||||
monkeypatch.setattr(tts_backend, "_REGISTRY", {"omnivoice": OmniVoiceBackend})
|
||||
monkeypatch.setattr(
|
||||
"core.device_caps.detect_host_caps",
|
||||
lambda: HostCaps(family="mps", available_families=("mps", "cpu")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"engines.omnivoice_subprocess.OmniVoiceSubprocessBackend.is_available",
|
||||
classmethod(lambda cls: (True, "ready")),
|
||||
)
|
||||
|
||||
row = next(item for item in list_backends() if item["id"] == "omnivoice")
|
||||
assert row["isolation_mode"] == "subprocess"
|
||||
|
||||
|
||||
def test_mps_startup_does_not_preload_native_model(monkeypatch):
|
||||
from core.device_caps import HostCaps
|
||||
from services import model_manager
|
||||
|
||||
monkeypatch.setattr(
|
||||
"core.device_caps.detect_host_caps",
|
||||
lambda: HostCaps(family="mps", available_families=("mps", "cpu")),
|
||||
)
|
||||
monkeypatch.setenv("OMNIVOICE_TTS_BACKEND", "omnivoice")
|
||||
monkeypatch.setattr(model_manager, "model", None)
|
||||
|
||||
async def fail_load():
|
||||
raise AssertionError("native OmniVoice must not load in the API process on MPS")
|
||||
|
||||
monkeypatch.setattr(model_manager, "_load_model_with_timeout", fail_load)
|
||||
asyncio.run(model_manager.preload_model())
|
||||
|
||||
|
||||
def test_streaming_mps_path_does_not_load_native_model(monkeypatch):
|
||||
from api.routers.tts_stream import _resolve_stream_backend
|
||||
from services import model_manager, tts_backend
|
||||
|
||||
sentinel = object()
|
||||
monkeypatch.setattr(tts_backend, "active_backend_id", lambda: "omnivoice")
|
||||
monkeypatch.setattr(
|
||||
tts_backend, "get_backend_class", lambda _id: OmniVoiceMPSSubprocessBackend,
|
||||
)
|
||||
monkeypatch.setattr(tts_backend, "get_active_tts_backend", lambda: sentinel)
|
||||
|
||||
async def fail_load():
|
||||
raise AssertionError("streaming must not load native OmniVoice on MPS")
|
||||
|
||||
monkeypatch.setattr(model_manager, "get_model", fail_load)
|
||||
assert asyncio.run(_resolve_stream_backend(None)) is sentinel
|
||||
|
||||
|
||||
def test_is_marked_subprocess_isolated():
|
||||
# list_backends() detects isolation via this duck-typed marker, not issubclass.
|
||||
assert getattr(OmniVoiceSubprocessBackend, "_is_subprocess_isolated", False) is True
|
||||
@@ -136,6 +234,40 @@ def test_base_default_recv_timeout_is_60s():
|
||||
assert _PlainBackend().recv_timeout_s == 60.0
|
||||
|
||||
|
||||
def test_sidecar_spawn_delegates_all_containment_to_nested_owner(monkeypatch, tmp_path):
|
||||
from services import subprocess_backend as backend_module
|
||||
|
||||
captured = {}
|
||||
|
||||
class StubProcess:
|
||||
stderr = io.BytesIO()
|
||||
|
||||
@staticmethod
|
||||
def poll():
|
||||
return None
|
||||
|
||||
def fake_spawn(argv, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return StubProcess()
|
||||
|
||||
monkeypatch.setattr(_PlainBackend, "venv_python", classmethod(lambda cls: Path(sys.executable)))
|
||||
monkeypatch.setattr(
|
||||
_PlainBackend,
|
||||
"sidecar_script",
|
||||
classmethod(lambda cls: tmp_path / "stub.py"),
|
||||
)
|
||||
monkeypatch.setattr(backend_module, "spawn_owned", fake_spawn)
|
||||
monkeypatch.setattr(backend_module, "_ensure_reaper_running", lambda: None)
|
||||
backend = _PlainBackend()
|
||||
monkeypatch.setattr(backend, "_recv_with_timeout", lambda _timeout: {"op": "ready"})
|
||||
|
||||
try:
|
||||
backend._spawn()
|
||||
assert not ({"start_new_session", "creationflags", "preexec_fn"} & captured.keys())
|
||||
finally:
|
||||
backend._proc = None
|
||||
|
||||
|
||||
def test_omnivoice_subprocess_recv_timeout_overrides_default():
|
||||
b = OmniVoiceSubprocessBackend()
|
||||
assert b.recv_timeout_s == 300.0 # aligns with the generate budget
|
||||
@@ -205,6 +337,48 @@ def test_wedged_sidecar_is_hard_killed_and_recovers(stub_sidecar, monkeypatch):
|
||||
b.shutdown()
|
||||
|
||||
|
||||
def test_mps_proxy_survives_fatal_child_exit_and_recovers(stub_sidecar, monkeypatch):
|
||||
_use_stub(monkeypatch, stub_sidecar)
|
||||
monkeypatch.setattr(
|
||||
"services.model_manager.make_room_before_generate", lambda: None,
|
||||
)
|
||||
b = OmniVoiceMPSSubprocessBackend()
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="backend is still running"):
|
||||
b.generate("CRASH")
|
||||
assert b._proc is not None and b._proc.poll() is not None
|
||||
assert b.generate("ok").shape[1] == 24000
|
||||
finally:
|
||||
b.shutdown()
|
||||
|
||||
|
||||
def test_desktop_timeout_kills_engine_subtree_before_late_mutation(
|
||||
stub_sidecar, monkeypatch, tmp_path
|
||||
):
|
||||
marker = tmp_path / "late-engine-mutation"
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_CONTAINED", "1")
|
||||
drain_read, drain_write = os.pipe()
|
||||
monkeypatch.setenv("OMNIVOICE_DESKTOP_DRAIN_FD", str(drain_write))
|
||||
monkeypatch.setenv("OMNIVOICE_TIMEOUT_MARKER", str(marker))
|
||||
_use_stub(monkeypatch, stub_sidecar)
|
||||
monkeypatch.setattr(
|
||||
OmniVoiceSubprocessBackend,
|
||||
"recv_timeout_s",
|
||||
property(lambda self: 0.3),
|
||||
)
|
||||
b = OmniVoiceSubprocessBackend()
|
||||
try:
|
||||
with pytest.raises(RuntimeError):
|
||||
b.generate("HANG_CHILD")
|
||||
time.sleep(1.2)
|
||||
assert not marker.exists()
|
||||
assert b.generate("ok").shape[1] == 24000
|
||||
finally:
|
||||
b.shutdown()
|
||||
os.close(drain_write)
|
||||
os.close(drain_read)
|
||||
|
||||
|
||||
def test_generate_does_not_deadlock_when_called_on_gpu_pool_worker(stub_sidecar, monkeypatch):
|
||||
# Regression: /v1/audio/speech and /generate dispatch backend.generate() via
|
||||
# run_on_gpu_pool_guarded, i.e. ON a gpu-pool worker. generate() must NOT
|
||||
@@ -222,3 +396,94 @@ def test_generate_does_not_deadlock_when_called_on_gpu_pool_worker(stub_sidecar,
|
||||
assert tensor.shape[1] == 24000
|
||||
finally:
|
||||
b.shutdown()
|
||||
|
||||
|
||||
def test_sidecar_forwards_native_controls_and_applies_seed(monkeypatch):
|
||||
import torch
|
||||
from engines.omnivoice_subprocess import main as sidecar
|
||||
|
||||
calls = []
|
||||
seeds = []
|
||||
frames = []
|
||||
|
||||
class FakeModel:
|
||||
sampling_rate = 24000
|
||||
|
||||
def generate(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
return [torch.zeros(1, 16)]
|
||||
|
||||
monkeypatch.setattr(sidecar, "_load_model", lambda _stdout: FakeModel())
|
||||
monkeypatch.setattr(sidecar, "_send", lambda _stdout, frame: frames.append(frame))
|
||||
real_manual_seed = torch.manual_seed
|
||||
monkeypatch.setattr(
|
||||
torch, "manual_seed", lambda seed: (seeds.append(seed), real_manual_seed(seed))[1],
|
||||
)
|
||||
|
||||
sidecar._handle_synthesize({
|
||||
"text": "hello",
|
||||
"seed": 123,
|
||||
"t_shift": 0.4,
|
||||
"layer_penalty_factor": 0.2,
|
||||
"position_temperature": 0.7,
|
||||
"class_temperature": 0.8,
|
||||
"audio_chunk_duration": 10,
|
||||
"audio_chunk_threshold": 0.6,
|
||||
}, object())
|
||||
|
||||
assert seeds == [123]
|
||||
assert calls == [{
|
||||
"text": "hello",
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
"t_shift": 0.4,
|
||||
"layer_penalty_factor": 0.2,
|
||||
"position_temperature": 0.7,
|
||||
"class_temperature": 0.8,
|
||||
"audio_chunk_duration": 10,
|
||||
"audio_chunk_threshold": 0.6,
|
||||
}]
|
||||
assert frames[-1]["op"] == "audio"
|
||||
|
||||
|
||||
def test_generation_proxy_forwards_native_controls_and_seed():
|
||||
import torch
|
||||
from api.routers.generation import _run_backend_inference
|
||||
|
||||
calls = []
|
||||
|
||||
class Proxy:
|
||||
id = "omnivoice"
|
||||
display_name = "OmniVoice"
|
||||
sample_rate = 24000
|
||||
applies_own_mastering = True
|
||||
supports_native_omnivoice_controls = True
|
||||
|
||||
def generate(self, text, **kwargs):
|
||||
calls.append((text, kwargs))
|
||||
return torch.zeros(1, 240)
|
||||
|
||||
_run_backend_inference(
|
||||
Proxy(), "hello", "en", None, None, None, None,
|
||||
16, 2.0, 1.0, False, False, 321,
|
||||
t_shift=0.4, layer_penalty_factor=0.2,
|
||||
position_temperature=0.7, class_temperature=0.8,
|
||||
)
|
||||
|
||||
assert calls == [("hello", {
|
||||
"duration": None,
|
||||
"language": "en",
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
"instruct": None,
|
||||
"num_step": 16,
|
||||
"guidance_scale": 2.0,
|
||||
"speed": 1.0,
|
||||
"denoise": False,
|
||||
"postprocess_output": False,
|
||||
"t_shift": 0.4,
|
||||
"layer_penalty_factor": 0.2,
|
||||
"position_temperature": 0.7,
|
||||
"class_temperature": 0.8,
|
||||
"seed": 321,
|
||||
})]
|
||||
|
||||
+857
-86
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,61 @@
|
||||
"""Cancellation helpers for work that cannot be stopped mid-call."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeVar
|
||||
|
||||
_Result = TypeVar("_Result")
|
||||
|
||||
|
||||
async def drain_task(task: asyncio.Task[Any]) -> None:
|
||||
"""Wait for ``task`` even if the waiter is cancelled again."""
|
||||
while not task.done():
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except BaseException:
|
||||
break
|
||||
if task.done():
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
async def to_thread_and_drain_on_cancel(
|
||||
function: Callable[..., _Result], /, *args: Any
|
||||
) -> _Result:
|
||||
"""Run a blocking call without detaching it when its waiter is cancelled."""
|
||||
thread_task = asyncio.create_task(asyncio.to_thread(function, *args))
|
||||
try:
|
||||
return await asyncio.shield(thread_task)
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(thread_task)
|
||||
raise
|
||||
|
||||
|
||||
async def to_thread_and_defer_cancellation(
|
||||
function: Callable[..., _Result], /, *args: Any
|
||||
) -> tuple[_Result, bool]:
|
||||
"""Finish a durable call and report cancellation after its result is known.
|
||||
|
||||
Authority writes need their event-loop publication even when the HTTP
|
||||
caller disappears while SQLite is committing. Returning the cancellation
|
||||
flag lets the caller publish that result first, then propagate cancellation.
|
||||
"""
|
||||
thread_task = asyncio.create_task(asyncio.to_thread(function, *args))
|
||||
try:
|
||||
return await asyncio.shield(thread_task), False
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(thread_task)
|
||||
return thread_task.result(), True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"drain_task",
|
||||
"to_thread_and_defer_cancellation",
|
||||
"to_thread_and_drain_on_cancel",
|
||||
]
|
||||
@@ -59,10 +59,18 @@ _VRAM_PER_JOB_BYTES = 5 * 1024**3
|
||||
# cpu — oversubscription just thrashes
|
||||
_ALWAYS_SERIAL = frozenset({"mps", "mlx", "cpu", ""})
|
||||
|
||||
# Absolute ceiling regardless of how much memory a card reports. Beyond this
|
||||
# the bottleneck stops being VRAM and starts being scheduler overhead and
|
||||
# host-side I/O contention.
|
||||
_MAX_DERIVED = 4
|
||||
# Absolute protocol ceiling regardless of how much memory a peer reports.
|
||||
# Beyond this the bottleneck stops being VRAM and starts being scheduler
|
||||
# overhead and host-side I/O contention. It is public because every wire
|
||||
# boundary must clamp to the same number; a UINT32_MAX heartbeat must not grow
|
||||
# a scheduler queue that local derivation would never create.
|
||||
MAX_CONCURRENT_TASKS = 4
|
||||
|
||||
|
||||
def clamp_concurrency(value: int, *, allow_zero: bool = False) -> int:
|
||||
"""Bound an advertised concurrency value to the server's safe range."""
|
||||
minimum = 0 if allow_zero else 1
|
||||
return max(minimum, min(MAX_CONCURRENT_TASKS, int(value)))
|
||||
|
||||
# Bounds on how long a parked slot is held. The caller passes the timed-out
|
||||
# job's own execution budget — the longest its thread can still legitimately be
|
||||
@@ -104,7 +112,7 @@ def derive_concurrency(
|
||||
budget = max(min_model_bytes, _VRAM_PER_JOB_BYTES)
|
||||
if budget <= 0:
|
||||
return 1
|
||||
return max(1, min(_MAX_DERIVED, int(free_memory_bytes // budget)))
|
||||
return clamp_concurrency(int(free_memory_bytes // budget))
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -149,6 +157,9 @@ class WorkerCapacity:
|
||||
resident_models: set[str] = field(default_factory=set)
|
||||
slots: dict[str, ModelSlot] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.max_concurrent_tasks = clamp_concurrency(self.max_concurrent_tasks)
|
||||
|
||||
@staticmethod
|
||||
def slot_key(engine: str, model_id: str) -> str:
|
||||
return f"{engine}:{model_id}"
|
||||
@@ -198,6 +209,21 @@ class WorkerCapacity:
|
||||
)
|
||||
slot.active += 1
|
||||
|
||||
def reserve_unknown(self) -> None:
|
||||
"""Consume worker-wide capacity for claimed work we cannot classify.
|
||||
|
||||
Reconciliation will tell the peer to cancel a terminal or unknown
|
||||
attempt, but until that cancellation lands it is still using the GPU.
|
||||
"""
|
||||
self.active_tasks += 1
|
||||
|
||||
def release_unknown(self) -> bool:
|
||||
"""Release one exact reconciled claim with no model-slot identity."""
|
||||
if self.active_tasks <= 0:
|
||||
return False
|
||||
self.active_tasks -= 1
|
||||
return True
|
||||
|
||||
def release(
|
||||
self,
|
||||
engine: str,
|
||||
@@ -267,8 +293,12 @@ class WorkerCapacity:
|
||||
) -> None:
|
||||
"""Adopt a heartbeat snapshot. The worker is the source of truth for
|
||||
what it is actually running."""
|
||||
self.active_tasks = max(0, active_tasks)
|
||||
reported_ceiling = self.active_tasks + max(0, available_slots)
|
||||
self.active_tasks = clamp_concurrency(active_tasks, allow_zero=True)
|
||||
bounded_available = clamp_concurrency(available_slots, allow_zero=True)
|
||||
bounded_available = min(
|
||||
bounded_available, MAX_CONCURRENT_TASKS - self.active_tasks
|
||||
)
|
||||
reported_ceiling = self.active_tasks + bounded_available
|
||||
if reported_ceiling > 0:
|
||||
# Adopted, not merely grown. The worker computes this as its own
|
||||
# ``max_concurrent_tasks``, so a ceiling we refuse to lower is one
|
||||
@@ -305,4 +335,10 @@ class WorkerCapacity:
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["ModelSlot", "WorkerCapacity", "derive_concurrency"]
|
||||
__all__ = [
|
||||
"MAX_CONCURRENT_TASKS",
|
||||
"ModelSlot",
|
||||
"WorkerCapacity",
|
||||
"clamp_concurrency",
|
||||
"derive_concurrency",
|
||||
]
|
||||
|
||||
+520
-91
@@ -17,15 +17,19 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import errno
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import io
|
||||
import zipfile
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import zipfile
|
||||
from typing import Any, Awaitable, Callable, Optional
|
||||
|
||||
from worker.async_utils import drain_task
|
||||
from worker.errors import ErrorClass, WorkerError
|
||||
|
||||
logger = logging.getLogger("omnivoice.worker")
|
||||
@@ -44,6 +48,18 @@ INPUT_ERRORS_PARAM = "input_errors"
|
||||
# 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
|
||||
_STALE_INPUT_PARTIAL_SECONDS = 60 * 60.0
|
||||
|
||||
# Pruning runs in worker threads and every executor instance shares the same
|
||||
# on-disk cache, so active-path leases are process-wide and thread-safe.
|
||||
_INPUT_CACHE_LEASE_LOCK = threading.Lock()
|
||||
_INPUT_CACHE_LEASES: dict[str, int] = {}
|
||||
_INPUT_CACHE_FETCH_LEASES: dict[str, int] = {}
|
||||
_INPUT_CACHE_MUTATIONS: set[str] = set()
|
||||
# Concurrent fetches keep distinct partial files but serialize the instant a
|
||||
# verified generation is published at its content address.
|
||||
_INPUT_CACHE_FETCH_LOCKS: dict[str, asyncio.Lock] = {}
|
||||
_INPUT_CACHE_FETCH_USERS: dict[str, int] = {}
|
||||
|
||||
# on_progress(fraction: float, stage: str)
|
||||
# on_model_loading(fraction: float, detail: str)
|
||||
@@ -88,6 +104,7 @@ class TaskExecutor:
|
||||
self._on_model_loading = on_model_loading
|
||||
self._fetch_input = fetch_input
|
||||
self._input_dir = input_dir
|
||||
self._blocking_tasks: set[asyncio.Task] = set()
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
@@ -110,33 +127,35 @@ class TaskExecutor:
|
||||
"""
|
||||
operation = (assignment.operation or "tts").lower()
|
||||
params = _parse_params(assignment.params_json)
|
||||
params = await self._materialize_inputs(
|
||||
params, leased_inputs = await self._materialize_inputs(
|
||||
assignment, params, fetch_input or self._fetch_input
|
||||
)
|
||||
|
||||
handler = {
|
||||
"tts": self._run_tts,
|
||||
"clone": self._run_tts,
|
||||
"audiobook": self._run_audiobook,
|
||||
"dub_segments": self._run_dub_segments,
|
||||
}.get(operation)
|
||||
if handler is None:
|
||||
raise TaskFailure(
|
||||
WorkerError(
|
||||
error_class=ErrorClass.CAPABILITY,
|
||||
code="OPERATION_UNSUPPORTED",
|
||||
message=f"This worker cannot run '{operation}' tasks.",
|
||||
hint="Run this task locally, or use a worker that supports it.",
|
||||
try:
|
||||
handler = {
|
||||
"tts": self._run_tts,
|
||||
"clone": self._run_tts,
|
||||
"audiobook": self._run_audiobook,
|
||||
"dub_segments": self._run_dub_segments,
|
||||
}.get(operation)
|
||||
if handler is None:
|
||||
raise TaskFailure(
|
||||
WorkerError(
|
||||
error_class=ErrorClass.CAPABILITY,
|
||||
code="OPERATION_UNSUPPORTED",
|
||||
message=f"This worker cannot run '{operation}' tasks.",
|
||||
hint="Run this task locally, or use a worker that supports it.",
|
||||
)
|
||||
)
|
||||
return await handler(
|
||||
assignment,
|
||||
params,
|
||||
_Reporters(
|
||||
on_progress or self._on_progress,
|
||||
on_model_loading or self._on_model_loading,
|
||||
),
|
||||
)
|
||||
return await handler(
|
||||
assignment,
|
||||
params,
|
||||
_Reporters(
|
||||
on_progress or self._on_progress,
|
||||
on_model_loading or self._on_model_loading,
|
||||
),
|
||||
)
|
||||
finally:
|
||||
self._release_inputs_after_active_work(leased_inputs)
|
||||
|
||||
async def _run_dub_segments(self, assignment, params: dict, report: "_Reporters") -> dict:
|
||||
"""Render every requested dub line under one lease and return one bundle."""
|
||||
@@ -150,8 +169,9 @@ class TaskExecutor:
|
||||
))
|
||||
load_budget, run_budget = _budgets(assignment)
|
||||
await report.loading(0.0, f"preparing {assignment.engine}")
|
||||
backend = await self._bounded(
|
||||
asyncio.to_thread(self._load_backend, assignment.engine),
|
||||
backend = await self._bounded_thread(
|
||||
self._load_backend,
|
||||
assignment.engine,
|
||||
timeout=load_budget, code="MODEL_LOAD_TIMEOUT", what=f"Loading '{assignment.engine}'",
|
||||
)
|
||||
await report.loading(1.0, "model ready")
|
||||
@@ -159,11 +179,15 @@ class TaskExecutor:
|
||||
for index, row in enumerate(rows):
|
||||
row = dict(row)
|
||||
row["ref_audio"] = refs[index] if index < len(refs) else None
|
||||
audio = await self._bounded(
|
||||
asyncio.to_thread(self._synthesize_dub_segment, backend, row),
|
||||
audio = await self._bounded_thread(
|
||||
self._synthesize_dub_segment,
|
||||
backend,
|
||||
row,
|
||||
timeout=run_budget, code="EXECUTION_TIMEOUT", what=f"Dubbing segment {index + 1}",
|
||||
)
|
||||
payload, _meta = await asyncio.to_thread(self._encode, audio, row, backend)
|
||||
payload, _meta = await self._thread_call(
|
||||
self._encode, audio, row, backend
|
||||
)
|
||||
rendered.append((int(row.get("index", index)), payload))
|
||||
await report.progress((index + 1) / len(rows), f"segment {index + 1} of {len(rows)}")
|
||||
|
||||
@@ -184,9 +208,11 @@ class TaskExecutor:
|
||||
from services.text_normalization import normalize_for_tts
|
||||
|
||||
text = normalize_for_tts(row.get("text") or "", row.get("language"))
|
||||
seed = None
|
||||
if row.get("seed") is not None:
|
||||
import torch
|
||||
torch.manual_seed(int(row["seed"]))
|
||||
seed = int(row["seed"])
|
||||
torch.manual_seed(seed)
|
||||
kwargs = {
|
||||
"language": row.get("language") if row.get("language") != "Auto" else None,
|
||||
"ref_audio": row.get("ref_audio"), "ref_text": row.get("ref_text"),
|
||||
@@ -197,6 +223,11 @@ class TaskExecutor:
|
||||
"speed": float(row.get("speed") or 1.0), "denoise": True,
|
||||
"postprocess_output": True,
|
||||
}
|
||||
if (
|
||||
getattr(backend, "supports_native_omnivoice_controls", False)
|
||||
and seed is not None
|
||||
):
|
||||
kwargs["seed"] = seed
|
||||
audio = backend.generate(text=text, **kwargs)
|
||||
preset = row.get("effect_preset") or "broadcast"
|
||||
if preset != "raw":
|
||||
@@ -225,8 +256,9 @@ class TaskExecutor:
|
||||
load_budget, run_budget = _budgets(assignment)
|
||||
|
||||
await report.loading(0.0, f"preparing {assignment.engine}")
|
||||
backend = await self._bounded(
|
||||
asyncio.to_thread(self._load_backend, assignment.engine),
|
||||
backend = await self._bounded_thread(
|
||||
self._load_backend,
|
||||
assignment.engine,
|
||||
timeout=load_budget,
|
||||
code="MODEL_LOAD_TIMEOUT",
|
||||
what=f"Loading '{assignment.engine}'",
|
||||
@@ -234,16 +266,22 @@ class TaskExecutor:
|
||||
await report.loading(1.0, "model ready")
|
||||
|
||||
await report.progress(0.05, "synthesising")
|
||||
audio = await self._bounded(
|
||||
asyncio.to_thread(self._synthesize, backend, text, params),
|
||||
audio = await self._bounded_thread(
|
||||
self._synthesize,
|
||||
backend,
|
||||
text,
|
||||
params,
|
||||
timeout=run_budget,
|
||||
code="EXECUTION_TIMEOUT",
|
||||
what="Synthesis",
|
||||
)
|
||||
await report.progress(0.9, "encoding")
|
||||
|
||||
payload, meta = await self._bounded(
|
||||
asyncio.to_thread(self._encode, audio, params, backend),
|
||||
payload, meta = await self._bounded_thread(
|
||||
self._encode,
|
||||
audio,
|
||||
params,
|
||||
backend,
|
||||
timeout=run_budget,
|
||||
code="EXECUTION_TIMEOUT",
|
||||
what="Encoding",
|
||||
@@ -264,19 +302,26 @@ class TaskExecutor:
|
||||
))
|
||||
load_budget, run_budget = _budgets(assignment)
|
||||
await report.loading(0.0, f"preparing {assignment.engine}")
|
||||
backend = await self._bounded(
|
||||
asyncio.to_thread(self._load_backend, assignment.engine),
|
||||
backend = await self._bounded_thread(
|
||||
self._load_backend,
|
||||
assignment.engine,
|
||||
timeout=load_budget, code="MODEL_LOAD_TIMEOUT",
|
||||
what=f"Loading '{assignment.engine}'",
|
||||
)
|
||||
await report.loading(1.0, "model ready")
|
||||
await report.progress(0.05, "synthesising chapter")
|
||||
audio = await self._bounded(
|
||||
asyncio.to_thread(self._synthesize_audiobook, backend, spans, voices, params),
|
||||
audio = await self._bounded_thread(
|
||||
self._synthesize_audiobook,
|
||||
backend,
|
||||
spans,
|
||||
voices,
|
||||
params,
|
||||
timeout=run_budget, code="EXECUTION_TIMEOUT", what="Audiobook chapter",
|
||||
)
|
||||
await report.progress(0.9, "encoding")
|
||||
payload, meta = await asyncio.to_thread(self._encode, audio, params, backend)
|
||||
payload, meta = await self._thread_call(
|
||||
self._encode, audio, params, backend
|
||||
)
|
||||
await report.progress(1.0, "done")
|
||||
return {"meta": meta, "payload": payload}
|
||||
|
||||
@@ -294,7 +339,10 @@ class TaskExecutor:
|
||||
key: value for key, value in opts.to_manifest().items()
|
||||
if value is not None and key not in ("seed", "vary_repeats")
|
||||
}
|
||||
if isinstance(backend, OmniVoiceBackend):
|
||||
native_proxy = bool(
|
||||
getattr(backend, "supports_native_omnivoice_controls", False)
|
||||
)
|
||||
if isinstance(backend, OmniVoiceBackend) or native_proxy:
|
||||
extra.setdefault("num_step", 32)
|
||||
extra.setdefault("guidance_scale", 2.0)
|
||||
for key in ("emo_vector", "emo_text", "emo_alpha"):
|
||||
@@ -303,11 +351,13 @@ class TaskExecutor:
|
||||
def synth(text, index, speed=None):
|
||||
voice = voices[int(index)]
|
||||
base_seed = opts.seed if opts.seed is not None else voice.get("seed")
|
||||
seed = None
|
||||
if base_seed is not None:
|
||||
import torch
|
||||
nonce = occurrence["value"] if opts.vary_repeats else 0
|
||||
occurrence["value"] += 1
|
||||
torch.manual_seed(segment_seed(base_seed, text, nonce))
|
||||
seed = segment_seed(base_seed, text, nonce)
|
||||
torch.manual_seed(seed)
|
||||
kwargs = {
|
||||
"language": language,
|
||||
"ref_audio": voice.get("ref_audio"),
|
||||
@@ -316,6 +366,8 @@ class TaskExecutor:
|
||||
"speed": float(speed) if speed else 1.0,
|
||||
**extra,
|
||||
}
|
||||
if native_proxy and seed is not None:
|
||||
kwargs["seed"] = seed
|
||||
return backend.generate(text, **kwargs)
|
||||
|
||||
spans = [Span(voice_id=str(i), text=row.get("text", ""),
|
||||
@@ -329,7 +381,9 @@ class TaskExecutor:
|
||||
|
||||
# ── Inputs ────────────────────────────────────────────────────────────
|
||||
|
||||
async def _materialize_inputs(self, assignment, params: dict, fetch) -> dict:
|
||||
async def _materialize_inputs(
|
||||
self, assignment, params: dict, fetch
|
||||
) -> tuple[dict, list[str]]:
|
||||
"""Turn declared inputs into local files, then point the params at them.
|
||||
|
||||
The control plane sends artifact ids, never paths — its own paths mean
|
||||
@@ -352,7 +406,7 @@ class TaskExecutor:
|
||||
|
||||
refs = [ref for ref in (getattr(assignment, "inputs", None) or []) if ref.artifact_id]
|
||||
if not refs:
|
||||
return params
|
||||
return params, []
|
||||
if fetch is None:
|
||||
raise TaskFailure(
|
||||
WorkerError(
|
||||
@@ -365,16 +419,24 @@ class TaskExecutor:
|
||||
|
||||
_, 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)
|
||||
leased: list[str] = []
|
||||
try:
|
||||
for ref in refs:
|
||||
path = await self._bounded(
|
||||
self._fetch_one(ref, fetch, retain=True),
|
||||
timeout=min(run_budget, _FALLBACK_INPUT_FETCH_SECONDS),
|
||||
code="INPUT_FETCH_TIMEOUT",
|
||||
what=f"Fetching '{ref.filename or ref.artifact_id}'",
|
||||
)
|
||||
local[ref.artifact_id] = path
|
||||
leased.append(path)
|
||||
return _rewrite_params(params, local), leased
|
||||
except BaseException:
|
||||
for path in leased:
|
||||
_release_input_cache_path(path)
|
||||
raise
|
||||
|
||||
async def _fetch_one(self, ref, fetch) -> str:
|
||||
async def _fetch_one(self, ref, fetch, *, retain: bool = False) -> 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
|
||||
@@ -382,37 +444,127 @@ class TaskExecutor:
|
||||
worker — costs no transfer at all.
|
||||
"""
|
||||
directory = self._input_dir or default_input_dir()
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
await self._thread_call(_durable_makedirs, directory)
|
||||
destination = os.path.join(directory, _cache_name(ref))
|
||||
if _already_held(destination, ref):
|
||||
_touch(destination)
|
||||
return destination
|
||||
return await self._fetch_one_owned(
|
||||
ref,
|
||||
fetch,
|
||||
directory=directory,
|
||||
destination=destination,
|
||||
retain=retain,
|
||||
)
|
||||
|
||||
partial = f"{destination}.{uuid.uuid4().hex}.part"
|
||||
async def _fetch_one_owned(
|
||||
self,
|
||||
ref,
|
||||
fetch,
|
||||
*,
|
||||
directory: str,
|
||||
destination: str,
|
||||
retain: bool,
|
||||
) -> str:
|
||||
"""Validate, fetch, and safely publish one content address."""
|
||||
await _acquire_input_cache_path(destination, fetching=True)
|
||||
leased_result = destination
|
||||
succeeded = False
|
||||
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
|
||||
# Cache hits still hash the advertised content identity. Filename
|
||||
# plus size is not proof after disk corruption or external edits.
|
||||
if await self._thread_call(_already_held, destination, ref):
|
||||
await self._thread_call(_touch, destination)
|
||||
succeeded = True
|
||||
return destination
|
||||
|
||||
# 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
|
||||
partial = f"{destination}.{uuid.uuid4().hex}.part"
|
||||
_lease_input_cache_path(partial)
|
||||
finalized = False
|
||||
try:
|
||||
try:
|
||||
await fetch(ref, partial)
|
||||
except TaskFailure:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise _input_fetch_failure(ref, exc) from exc
|
||||
|
||||
# Hashing and durability barriers can both block on a large
|
||||
# source or slow disk. Keep them off the loop and drain before
|
||||
# cleanup so Windows never unlinks a file still in use.
|
||||
await self._thread_call(_verify, partial, ref)
|
||||
gate_key = _cache_path_key(destination)
|
||||
gate = _INPUT_CACHE_FETCH_LOCKS.setdefault(
|
||||
gate_key, asyncio.Lock()
|
||||
)
|
||||
_INPUT_CACHE_FETCH_USERS[gate_key] = (
|
||||
_INPUT_CACHE_FETCH_USERS.get(gate_key, 0) + 1
|
||||
)
|
||||
try:
|
||||
async with gate:
|
||||
# A concurrent fetch may have published these exact
|
||||
# bytes while this one was downloading its own partial.
|
||||
if await self._thread_call(
|
||||
_already_held, destination, ref
|
||||
):
|
||||
await self._thread_call(_discard, partial)
|
||||
finalized = True
|
||||
elif _claim_input_cache_mutation(destination):
|
||||
try:
|
||||
await self._thread_call(
|
||||
_durable_replace, partial, destination
|
||||
)
|
||||
finally:
|
||||
_finish_input_cache_mutation(destination)
|
||||
finalized = True
|
||||
else:
|
||||
# Another execution is actively reading the
|
||||
# canonical generation. Never unlink or replace
|
||||
# bytes underneath it; publish this verified fetch
|
||||
# under a leased sibling path and let a later
|
||||
# unshared fetch repair canonical.
|
||||
stem, suffix = os.path.splitext(destination)
|
||||
alternate = (
|
||||
f"{stem}.{uuid.uuid4().hex}.generation{suffix}"
|
||||
)
|
||||
await _acquire_input_cache_path(
|
||||
alternate, fetching=True
|
||||
)
|
||||
try:
|
||||
await self._thread_call(
|
||||
_durable_replace, partial, alternate
|
||||
)
|
||||
except BaseException:
|
||||
_release_input_cache_path(
|
||||
alternate, fetching=True
|
||||
)
|
||||
await self._thread_call(_discard, alternate)
|
||||
raise
|
||||
finalized = True
|
||||
_release_input_cache_path(
|
||||
destination, fetching=True
|
||||
)
|
||||
leased_result = alternate
|
||||
except OSError as exc:
|
||||
raise _input_fetch_failure(ref, exc) from exc
|
||||
finally:
|
||||
remaining = _INPUT_CACHE_FETCH_USERS[gate_key] - 1
|
||||
if remaining:
|
||||
_INPUT_CACHE_FETCH_USERS[gate_key] = remaining
|
||||
else:
|
||||
_INPUT_CACHE_FETCH_USERS.pop(gate_key, None)
|
||||
if _INPUT_CACHE_FETCH_LOCKS.get(gate_key) is gate:
|
||||
_INPUT_CACHE_FETCH_LOCKS.pop(gate_key, None)
|
||||
finally:
|
||||
if not finalized:
|
||||
await self._thread_call(_discard, partial)
|
||||
_release_input_cache_path(partial)
|
||||
|
||||
await self._thread_call(_prune_input_cache, directory)
|
||||
succeeded = True
|
||||
return leased_result
|
||||
finally:
|
||||
if retain and succeeded:
|
||||
_promote_input_cache_lease(leased_result)
|
||||
else:
|
||||
_release_input_cache_path(leased_result, fetching=True)
|
||||
|
||||
# ── Engine plumbing ───────────────────────────────────────────────────
|
||||
|
||||
@@ -508,6 +660,10 @@ class TaskExecutor:
|
||||
duration, num_step, guidance_scale, speed, denoise,
|
||||
postprocess_output, used_seed, effect_preset,
|
||||
max_chunk_chars, crossfade_ms,
|
||||
t_shift=params.get("t_shift"),
|
||||
layer_penalty_factor=params.get("layer_penalty_factor"),
|
||||
position_temperature=params.get("position_temperature"),
|
||||
class_temperature=params.get("class_temperature"),
|
||||
)
|
||||
except Exception as exc:
|
||||
from worker import errors as worker_errors # noqa: PLC0415
|
||||
@@ -551,6 +707,92 @@ class TaskExecutor:
|
||||
|
||||
# ── Bounding ──────────────────────────────────────────────────────────
|
||||
|
||||
def _release_inputs_after_active_work(self, paths: list[str]) -> None:
|
||||
"""Keep files leased while a timed-out engine thread still owns them."""
|
||||
pending = [task for task in self._blocking_tasks if not task.done()]
|
||||
if not pending:
|
||||
for path in paths:
|
||||
_release_input_cache_path(path)
|
||||
return
|
||||
remaining = {"count": len(pending)}
|
||||
|
||||
def finished(_task: asyncio.Task) -> None:
|
||||
remaining["count"] -= 1
|
||||
if remaining["count"] == 0:
|
||||
for path in paths:
|
||||
_release_input_cache_path(path)
|
||||
|
||||
for task in pending:
|
||||
task.add_done_callback(finished)
|
||||
|
||||
def _start_thread(self, function, /, *args) -> asyncio.Task:
|
||||
task = asyncio.create_task(asyncio.to_thread(function, *args))
|
||||
self._blocking_tasks.add(task)
|
||||
|
||||
def finished(completed: asyncio.Task) -> None:
|
||||
self._blocking_tasks.discard(completed)
|
||||
if not completed.cancelled():
|
||||
# Timed-out calls intentionally finish in the background. Read
|
||||
# their exception so asyncio never reports an unowned task.
|
||||
completed.exception()
|
||||
|
||||
task.add_done_callback(finished)
|
||||
return task
|
||||
|
||||
async def _thread_call(self, function, /, *args):
|
||||
task = self._start_thread(function, *args)
|
||||
try:
|
||||
return await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(task)
|
||||
raise
|
||||
|
||||
async def drain_active_work(self) -> None:
|
||||
"""Wait until every blocking engine call has relinquished the process."""
|
||||
cancelled = bool(
|
||||
(current := asyncio.current_task()) is not None and current.cancelling()
|
||||
)
|
||||
while self._blocking_tasks:
|
||||
for task in list(self._blocking_tasks):
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
# Cancellation cannot make a Python GPU thread stop. Hold
|
||||
# authority until it really exits, then propagate the
|
||||
# cancellation so callers never publish a false free slot.
|
||||
cancelled = True
|
||||
await drain_task(task)
|
||||
except BaseException:
|
||||
# The owner reports/classifies the engine exception. This
|
||||
# barrier only establishes that the thread has finished.
|
||||
pass
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def _bounded_thread(
|
||||
self, function, /, *args, timeout: float, code: str, what: str
|
||||
):
|
||||
"""Bound a blocking call without losing ownership of its live thread."""
|
||||
task = self._start_thread(function, *args)
|
||||
try:
|
||||
done, _pending = await asyncio.wait({task}, timeout=timeout)
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(task)
|
||||
raise
|
||||
if done:
|
||||
return task.result()
|
||||
# A GPU call cannot be killed. Return the timeout so the scheduler can
|
||||
# park its slot, but retain the task above so terminal authority loss
|
||||
# can drain it before claiming this worker has stopped.
|
||||
raise TaskFailure(
|
||||
WorkerError(
|
||||
error_class=ErrorClass.TIMEOUT,
|
||||
code=code,
|
||||
message=f"{what} exceeded the {timeout:g}s budget for this task.",
|
||||
hint="Try a shorter input, or a worker with more headroom.",
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _bounded(coro, *, timeout: float, code: str, what: str):
|
||||
"""Run ``coro`` under the server's budget for this phase.
|
||||
@@ -636,6 +878,158 @@ def default_input_dir() -> str:
|
||||
return os.path.join(tempfile.gettempdir(), "omnivoice-worker-inputs")
|
||||
|
||||
|
||||
def _cache_path_key(path: str) -> str:
|
||||
return os.path.normcase(os.path.abspath(path))
|
||||
|
||||
|
||||
def _lease_input_cache_path(path: str, *, fetching: bool = False) -> bool:
|
||||
key = _cache_path_key(path)
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
if key in _INPUT_CACHE_MUTATIONS:
|
||||
return False
|
||||
_INPUT_CACHE_LEASES[key] = _INPUT_CACHE_LEASES.get(key, 0) + 1
|
||||
if fetching:
|
||||
_INPUT_CACHE_FETCH_LEASES[key] = (
|
||||
_INPUT_CACHE_FETCH_LEASES.get(key, 0) + 1
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
async def _acquire_input_cache_path(
|
||||
path: str, *, fetching: bool = False
|
||||
) -> None:
|
||||
while not _lease_input_cache_path(path, fetching=fetching):
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
def _release_input_cache_path(path: str, *, fetching: bool = False) -> None:
|
||||
key = _cache_path_key(path)
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
if fetching:
|
||||
fetch_remaining = _INPUT_CACHE_FETCH_LEASES.get(key, 0) - 1
|
||||
if fetch_remaining > 0:
|
||||
_INPUT_CACHE_FETCH_LEASES[key] = fetch_remaining
|
||||
else:
|
||||
_INPUT_CACHE_FETCH_LEASES.pop(key, None)
|
||||
remaining = _INPUT_CACHE_LEASES.get(key, 0) - 1
|
||||
if remaining > 0:
|
||||
_INPUT_CACHE_LEASES[key] = remaining
|
||||
else:
|
||||
_INPUT_CACHE_LEASES.pop(key, None)
|
||||
|
||||
|
||||
def _promote_input_cache_lease(path: str) -> None:
|
||||
"""Turn a fetcher's lease into the active execution lease it returns."""
|
||||
key = _cache_path_key(path)
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
remaining = _INPUT_CACHE_FETCH_LEASES.get(key, 0) - 1
|
||||
if remaining > 0:
|
||||
_INPUT_CACHE_FETCH_LEASES[key] = remaining
|
||||
else:
|
||||
_INPUT_CACHE_FETCH_LEASES.pop(key, None)
|
||||
|
||||
|
||||
def _leased_input_cache_paths() -> set[str]:
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
return set(_INPUT_CACHE_LEASES)
|
||||
|
||||
|
||||
def _claim_input_cache_mutation(path: str) -> bool:
|
||||
key = _cache_path_key(path)
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
if key in _INPUT_CACHE_MUTATIONS:
|
||||
return False
|
||||
active_leases = _INPUT_CACHE_LEASES.get(
|
||||
key, 0
|
||||
) - _INPUT_CACHE_FETCH_LEASES.get(key, 0)
|
||||
# Fetchers can safely converge under the publication gate. A lease
|
||||
# already promoted to an execution may have this exact pathname open.
|
||||
if active_leases > 0:
|
||||
return False
|
||||
_INPUT_CACHE_MUTATIONS.add(key)
|
||||
return True
|
||||
|
||||
|
||||
def _finish_input_cache_mutation(path: str) -> None:
|
||||
key = _cache_path_key(path)
|
||||
with _INPUT_CACHE_LEASE_LOCK:
|
||||
_INPUT_CACHE_MUTATIONS.discard(key)
|
||||
|
||||
|
||||
def _fsync_file(path: str) -> None:
|
||||
with open(path, "r+b") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def _fsync_parent_directory(directory: str) -> None:
|
||||
directory_flag = getattr(os, "O_DIRECTORY", None)
|
||||
if directory_flag is None:
|
||||
return
|
||||
unsupported = {
|
||||
errno.EINVAL,
|
||||
getattr(errno, "ENOTSUP", errno.EINVAL),
|
||||
getattr(errno, "EOPNOTSUPP", errno.EINVAL),
|
||||
}
|
||||
try:
|
||||
descriptor = os.open(directory, os.O_RDONLY | directory_flag)
|
||||
except OSError as exc:
|
||||
if exc.errno in unsupported:
|
||||
return
|
||||
raise
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
except OSError as exc:
|
||||
if exc.errno not in unsupported:
|
||||
raise
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def _durable_makedirs(directory: str) -> None:
|
||||
target = os.path.abspath(directory)
|
||||
missing: list[str] = []
|
||||
current = target
|
||||
while not os.path.isdir(current):
|
||||
if os.path.exists(current):
|
||||
if os.path.isdir(current):
|
||||
break
|
||||
raise NotADirectoryError(current)
|
||||
missing.append(current)
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
for path in reversed(missing):
|
||||
try:
|
||||
os.mkdir(path)
|
||||
except FileExistsError:
|
||||
if not os.path.isdir(path):
|
||||
raise
|
||||
_fsync_parent_directory(os.path.dirname(path) or ".")
|
||||
if not missing:
|
||||
_fsync_parent_directory(os.path.dirname(target) or ".")
|
||||
|
||||
|
||||
def _durable_replace(source: str, destination: str) -> None:
|
||||
_fsync_file(source)
|
||||
os.replace(source, destination)
|
||||
_fsync_parent_directory(os.path.dirname(destination) or ".")
|
||||
|
||||
|
||||
def _input_fetch_failure(ref, error: BaseException) -> TaskFailure:
|
||||
return TaskFailure(
|
||||
WorkerError(
|
||||
# Transient on purpose: an id we cannot resolve now is far more
|
||||
# often a dropped stream/disk barrier 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}': {error}",
|
||||
hint="The control plane may have restarted; the task will be retried.",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _cache_name(ref) -> str:
|
||||
"""A safe, content-addressed local name for one input.
|
||||
|
||||
@@ -656,13 +1050,25 @@ def _cache_name(ref) -> str:
|
||||
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.
|
||||
The filename is content-addressed, but disks and external edits can still
|
||||
change bytes at that name. Re-hash the advertised identity before reuse.
|
||||
"""
|
||||
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
|
||||
if not os.path.isfile(path):
|
||||
return False
|
||||
if expected and os.path.getsize(path) != expected:
|
||||
return False
|
||||
expected_hash = (getattr(ref, "sha256", "") or "").strip().lower()
|
||||
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)
|
||||
if digest.hexdigest() != expected_hash:
|
||||
return False
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
@@ -723,22 +1129,45 @@ def _verify(path: str, ref) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _prune_input_cache(directory: str, limit_bytes: int = INPUT_CACHE_LIMIT_BYTES) -> None:
|
||||
"""Keep the input cache under its ceiling, oldest first."""
|
||||
def _prune_input_cache(
|
||||
directory: str,
|
||||
limit_bytes: int = INPUT_CACHE_LIMIT_BYTES,
|
||||
now: Optional[float] = None,
|
||||
) -> None:
|
||||
"""Keep the cache bounded without deleting inputs a task is still using."""
|
||||
try:
|
||||
entries = []
|
||||
total = 0
|
||||
stamp = time.time() if now is None else now
|
||||
leased = _leased_input_cache_paths()
|
||||
for name in os.listdir(directory):
|
||||
path = os.path.join(directory, name)
|
||||
if name.endswith(".part") or not os.path.isfile(path):
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
stat = os.stat(path)
|
||||
entries.append((stat.st_mtime, stat.st_size, path))
|
||||
key = _cache_path_key(path)
|
||||
is_partial = name.endswith(".part")
|
||||
if (
|
||||
is_partial
|
||||
and key not in leased
|
||||
and stamp - stat.st_mtime >= _STALE_INPUT_PARTIAL_SECONDS
|
||||
):
|
||||
os.remove(path)
|
||||
continue
|
||||
total += stat.st_size
|
||||
# Active finals and partial transfers count toward the ceiling but
|
||||
# cannot be evicted. Young unleased .part files may belong to a
|
||||
# process that has not yet rebuilt its in-memory lease after fork;
|
||||
# the age sweep will remove them if they are crash leftovers.
|
||||
if key not in leased and not is_partial:
|
||||
entries.append((stat.st_mtime, stat.st_size, path))
|
||||
for _mtime, size, path in sorted(entries):
|
||||
if total <= limit_bytes:
|
||||
break
|
||||
os.remove(path)
|
||||
try:
|
||||
os.remove(path)
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
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)
|
||||
|
||||
@@ -25,6 +25,7 @@ in the dialog that shows it once.
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import errno
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
@@ -301,10 +302,31 @@ def save_worker_key(path: str, keypair: WorkerKeypair) -> None:
|
||||
tmp = f"{path}.tmp"
|
||||
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
try:
|
||||
os.write(fd, keypair.private_bytes())
|
||||
finally:
|
||||
remaining = memoryview(keypair.private_bytes())
|
||||
while remaining:
|
||||
written = os.write(fd, remaining)
|
||||
if written <= 0:
|
||||
raise OSError("could not finish writing the worker identity key")
|
||||
remaining = remaining[written:]
|
||||
os.fsync(fd)
|
||||
except Exception:
|
||||
os.close(fd)
|
||||
os.replace(tmp, path)
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
else:
|
||||
os.close(fd)
|
||||
try:
|
||||
os.replace(tmp, path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
_fsync_parent_directory(directory)
|
||||
try:
|
||||
os.chmod(path, 0o600)
|
||||
except OSError:
|
||||
@@ -313,6 +335,30 @@ def save_worker_key(path: str, keypair: WorkerKeypair) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _fsync_parent_directory(directory: str) -> None:
|
||||
directory_flag = getattr(os, "O_DIRECTORY", None)
|
||||
if directory_flag is None:
|
||||
return
|
||||
unsupported = {
|
||||
errno.EINVAL,
|
||||
getattr(errno, "ENOTSUP", errno.EINVAL),
|
||||
getattr(errno, "EOPNOTSUPP", errno.EINVAL),
|
||||
}
|
||||
try:
|
||||
descriptor = os.open(directory, os.O_RDONLY | directory_flag)
|
||||
except OSError as exc:
|
||||
if exc.errno in unsupported:
|
||||
return
|
||||
raise
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
except OSError as exc:
|
||||
if exc.errno not in unsupported:
|
||||
raise
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def load_worker_key(path: str) -> Optional[WorkerKeypair]:
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -22,20 +22,106 @@ import logging
|
||||
import os
|
||||
import socket
|
||||
import ssl
|
||||
from typing import Optional
|
||||
from typing import BinaryIO, Optional
|
||||
|
||||
import grpc
|
||||
|
||||
from worker import identity, registry, tls
|
||||
from worker.async_utils import to_thread_and_drain_on_cancel
|
||||
from worker.inbound.connection_string import Connection
|
||||
from worker.inbound.listener import KEY_METADATA_KEY
|
||||
from worker.protocol.gen import worker_v1_pb2 as pb
|
||||
from worker.protocol.gen import worker_v1_pb2_grpc as pb_grpc
|
||||
from worker.transport.client import MAX_MESSAGE_BYTES, backoff_delay
|
||||
from worker.transport.client import (
|
||||
MAX_MESSAGE_BYTES,
|
||||
TerminalRegistrationError,
|
||||
backoff_delay,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_PUSH_CHUNK_BYTES = 1024 * 1024
|
||||
# Register remains provisional until the node confirms that it durably saved
|
||||
# the panel-assigned identity. Match the control plane's provisional-session
|
||||
# lifetime so a peer that stops after Register cannot strand this connector.
|
||||
_REGISTRATION_CONFIRMATION_TIMEOUT_SECONDS = 30.0
|
||||
_REMOTE_SHUTDOWN_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
_FileVersion = tuple[int, int, int, int, int]
|
||||
|
||||
|
||||
def _write_all(handle: BinaryIO, payload: bytes) -> None:
|
||||
remaining = memoryview(payload)
|
||||
while remaining:
|
||||
written = handle.write(remaining)
|
||||
if not written:
|
||||
raise OSError("result write made no progress")
|
||||
remaining = remaining[written:]
|
||||
|
||||
|
||||
def _remove_quietly(path: str) -> None:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(path)
|
||||
|
||||
|
||||
class InboundConnectionError(RuntimeError):
|
||||
"""A pasted inbound connection could not be validated or activated."""
|
||||
|
||||
|
||||
class InboundConnectionRollbackError(InboundConnectionError):
|
||||
"""A failed connection change could not restore its prior generation."""
|
||||
|
||||
|
||||
class RemoteShutdownUnavailable(InboundConnectionError):
|
||||
"""The node may retain work, but no live stream can revoke it safely."""
|
||||
|
||||
|
||||
def _file_version(stat: os.stat_result) -> _FileVersion:
|
||||
"""Fields that identify both a staged path and the bytes hashed from it."""
|
||||
return (
|
||||
int(stat.st_dev),
|
||||
int(stat.st_ino),
|
||||
int(stat.st_size),
|
||||
int(stat.st_mtime_ns),
|
||||
int(stat.st_ctime_ns),
|
||||
)
|
||||
|
||||
|
||||
def _hash_staged_input(path: str) -> tuple[int, str, _FileVersion]:
|
||||
"""Hash one stable generation without ever allocating the whole file."""
|
||||
digest = hashlib.sha256()
|
||||
received = 0
|
||||
with open(path, "rb") as handle:
|
||||
before = _file_version(os.fstat(handle.fileno()))
|
||||
while True:
|
||||
block = handle.read(_PUSH_CHUNK_BYTES)
|
||||
if not block:
|
||||
break
|
||||
received += len(block)
|
||||
digest.update(block)
|
||||
after = _file_version(os.fstat(handle.fileno()))
|
||||
if before != after or received != before[2]:
|
||||
raise RuntimeError("the staged task input changed while it was being hashed")
|
||||
return received, digest.hexdigest(), before
|
||||
|
||||
|
||||
def _validate_staged_input(path: str, expected: _FileVersion) -> None:
|
||||
"""Reject a replacement or in-place edit between hashing and streaming."""
|
||||
try:
|
||||
current = _file_version(os.stat(path))
|
||||
except OSError as exc:
|
||||
raise RuntimeError("the staged task input is no longer available") from exc
|
||||
if current != expected:
|
||||
raise RuntimeError("the staged task input changed before it could be sent")
|
||||
|
||||
|
||||
def _validate_open_staged_input(
|
||||
handle: BinaryIO, path: str, expected: _FileVersion
|
||||
) -> None:
|
||||
"""The open generation and its path must still be the bytes we hashed."""
|
||||
if _file_version(os.fstat(handle.fileno())) != expected:
|
||||
raise RuntimeError("the staged task input changed before it could be sent")
|
||||
_validate_staged_input(path, expected)
|
||||
|
||||
|
||||
def _fetch_pinned_certificate(
|
||||
@@ -70,9 +156,15 @@ class NodeConnection:
|
||||
self._connection = connection
|
||||
self._label = label or connection.host
|
||||
self._outbox: asyncio.Queue[pb.ServerMessage] = asyncio.Queue()
|
||||
self._active_session = None
|
||||
self._stub: Optional[pb_grpc.NodeServiceStub] = None
|
||||
self._worker_id = ""
|
||||
self._stop = asyncio.Event()
|
||||
self._session_closed = asyncio.Event()
|
||||
self._session_closed.set()
|
||||
self._shutdown_confirmed = asyncio.Event()
|
||||
self._registration_ready = asyncio.Event()
|
||||
self._remote_protocol_retained = False
|
||||
self._last_error = ""
|
||||
|
||||
@property
|
||||
@@ -105,6 +197,10 @@ class NodeConnection:
|
||||
attempt = 0
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except TerminalRegistrationError as exc:
|
||||
self._remote_protocol_retained = False
|
||||
self._last_error = str(exc)
|
||||
raise
|
||||
except Exception:
|
||||
attempt += 1
|
||||
self._last_error = "Connection failed; check the backend log for details."
|
||||
@@ -114,8 +210,143 @@ class NodeConnection:
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=delay)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._stop.is_set():
|
||||
return
|
||||
if self._shutdown_confirmed.is_set() and not self._remote_protocol_retained:
|
||||
self._stop.set()
|
||||
return
|
||||
if self._active_session is None:
|
||||
if self._remote_protocol_retained:
|
||||
raise RemoteShutdownUnavailable(
|
||||
"That GPU machine is offline and may still be running work. "
|
||||
"Reconnect it, then remove the connection again."
|
||||
)
|
||||
self._stop.set()
|
||||
return
|
||||
# EOF is indistinguishable from a network blip and deliberately keeps
|
||||
# node execution alive for reconnect. Send an explicit terminal frame
|
||||
# and wait for the node to drain before removal reports success.
|
||||
self._shutdown_confirmed.clear()
|
||||
await self._outbox.put(
|
||||
pb.ServerMessage(
|
||||
shutdown=pb.Shutdown(reason="This GPU-machine connection was removed.")
|
||||
)
|
||||
)
|
||||
confirmed = asyncio.create_task(self._shutdown_confirmed.wait())
|
||||
disconnected = asyncio.create_task(self._session_closed.wait())
|
||||
try:
|
||||
done, _pending = await asyncio.wait(
|
||||
{confirmed, disconnected},
|
||||
timeout=_REMOTE_SHUTDOWN_TIMEOUT_SECONDS,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
if confirmed not in done and not self._shutdown_confirmed.is_set():
|
||||
raise RemoteShutdownUnavailable(
|
||||
"The GPU machine disconnected before it confirmed shutdown. "
|
||||
"Reconnect it, then remove the connection again."
|
||||
)
|
||||
finally:
|
||||
confirmed.cancel()
|
||||
disconnected.cancel()
|
||||
await asyncio.gather(confirmed, disconnected, return_exceptions=True)
|
||||
self._remote_protocol_retained = False
|
||||
self._stop.set()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""End this process without revoking reconnectable remote work."""
|
||||
self._stop.set()
|
||||
|
||||
def confirm_remote_shutdown(self, session) -> None:
|
||||
if self._active_session is not session:
|
||||
return
|
||||
self._remote_protocol_retained = False
|
||||
self._shutdown_confirmed.set()
|
||||
|
||||
def confirm_registration(self, session) -> None:
|
||||
"""Publish readiness only after the shared servicer activated the session."""
|
||||
if self._active_session is session:
|
||||
self._registration_ready.set()
|
||||
|
||||
async def wait_until_registered(
|
||||
self, task: asyncio.Task, *, timeout: float = 30.0
|
||||
) -> None:
|
||||
"""Wait for activation or surface a terminal/background dial failure."""
|
||||
ready = asyncio.create_task(self._registration_ready.wait())
|
||||
try:
|
||||
done, _pending = await asyncio.wait(
|
||||
{ready, task}, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if ready in done:
|
||||
return
|
||||
if task in done:
|
||||
if task.cancelled():
|
||||
raise InboundConnectionError(
|
||||
"The GPU-machine connection stopped before it became ready."
|
||||
)
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
raise InboundConnectionError(str(exc)) from exc
|
||||
raise InboundConnectionError(
|
||||
"That GPU machine did not finish connecting in time."
|
||||
)
|
||||
finally:
|
||||
ready.cancel()
|
||||
await asyncio.gather(ready, return_exceptions=True)
|
||||
|
||||
async def probe(self) -> None:
|
||||
"""Authenticate a replacement paste without publishing a worker session."""
|
||||
try:
|
||||
certificate_pem = await asyncio.to_thread(
|
||||
_fetch_pinned_certificate, self._connection
|
||||
)
|
||||
async with self._channel(certificate_pem) as channel:
|
||||
stub = pb_grpc.NodeServiceStub(channel)
|
||||
metadata = ((KEY_METADATA_KEY, self._connection.secret),)
|
||||
stream = stub.Attach(self._outbound(), metadata=metadata)
|
||||
try:
|
||||
first = await asyncio.wait_for(
|
||||
stream.read(),
|
||||
timeout=_REGISTRATION_CONFIRMATION_TIMEOUT_SECONDS,
|
||||
)
|
||||
finally:
|
||||
stream.cancel()
|
||||
except InboundConnectionError:
|
||||
raise
|
||||
except grpc.aio.AioRpcError as exc:
|
||||
detail = exc.details() or "The GPU machine rejected this connection."
|
||||
raise InboundConnectionError(detail) from exc
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise InboundConnectionError(
|
||||
"That GPU machine did not answer in time."
|
||||
) from exc
|
||||
except Exception as exc:
|
||||
raise InboundConnectionError(str(exc)) from exc
|
||||
|
||||
if first == grpc.aio.EOF or first.WhichOneof("payload") != "register":
|
||||
raise InboundConnectionError(
|
||||
"That machine answered, but not as a VoiceStudio GPU node."
|
||||
)
|
||||
request = first.register
|
||||
validate = getattr(self._servicer, "validate_inbound_request", None)
|
||||
refusal = validate(request) if callable(validate) else None
|
||||
if refusal is not None and refusal.error.code:
|
||||
raise InboundConnectionError(
|
||||
f"{refusal.error.code}: {refusal.error.message}"
|
||||
)
|
||||
public_key = bytes(request.public_key)
|
||||
if len(public_key) != 32:
|
||||
raise InboundConnectionError("That machine sent no usable identity.")
|
||||
key_id = identity.key_id_for(public_key)
|
||||
if registry.is_revoked(key_id):
|
||||
raise InboundConnectionError(
|
||||
"This GPU machine was removed from this app. Add it again to use it."
|
||||
)
|
||||
known = registry.get_by_key_id(key_id)
|
||||
if known is not None and not self._proves_key_possession(request, known):
|
||||
raise InboundConnectionError(
|
||||
"That machine could not prove its saved identity."
|
||||
)
|
||||
|
||||
async def _connect_once(self) -> None:
|
||||
# A fresh outbox per attempt. The queue used to be built once and
|
||||
# reused, so anything a dying session left behind became the NEXT
|
||||
@@ -123,7 +354,10 @@ class NodeConnection:
|
||||
# registration it requires first, aborted the call, and the pair span
|
||||
# at full speed: on hardware this reached session epoch 2445 inside a
|
||||
# second, with the log reading "Locally aborted" over and over.
|
||||
self._worker_id = ""
|
||||
self._stub = None
|
||||
self._outbox = asyncio.Queue()
|
||||
self._active_session = None
|
||||
certificate_pem = await asyncio.to_thread(
|
||||
_fetch_pinned_certificate, self._connection
|
||||
)
|
||||
@@ -140,33 +374,98 @@ class NodeConnection:
|
||||
"That machine answered, but not as a VoiceStudio GPU node."
|
||||
)
|
||||
|
||||
response = self._register(first.register)
|
||||
response = await self._register(first.register)
|
||||
if response.error.code:
|
||||
# A refusal here is a decision, not a blip: the node is a
|
||||
# different machine than the one this key was trusted for, or
|
||||
# its version cannot work with ours. Reconnecting cannot fix
|
||||
# either, so surface it rather than looping.
|
||||
raise RuntimeError(f"{response.error.code}: {response.error.message}")
|
||||
# either. Deliver the verdict before surfacing it locally so
|
||||
# the node can retire work retained across the dead stream;
|
||||
# closing first strands that executor with nobody left able to
|
||||
# cancel it.
|
||||
await self._outbox.put(pb.ServerMessage(registered=response))
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
stream.read(),
|
||||
timeout=_REGISTRATION_CONFIRMATION_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (asyncio.TimeoutError, grpc.aio.AioRpcError):
|
||||
pass
|
||||
raise TerminalRegistrationError(
|
||||
f"{response.error.code}: {response.error.message}"
|
||||
)
|
||||
|
||||
self._worker_id = response.worker_id
|
||||
self._stub = stub
|
||||
self._last_error = ""
|
||||
await self._outbox.put(pb.ServerMessage(registered=response))
|
||||
|
||||
session = self._servicer.session_for(self._worker_id)
|
||||
if session is None:
|
||||
raise RuntimeError("the session went away before the stream opened")
|
||||
|
||||
pump = asyncio.create_task(self._pump_outbound(session))
|
||||
try:
|
||||
await self._servicer.run_inbound_stream(session, _Frames(stream), self)
|
||||
await self._complete_registration(stream, response, stub)
|
||||
finally:
|
||||
pump.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await pump
|
||||
self._stub = None
|
||||
# Idempotent after activation; essential before it. A user can
|
||||
# remove this connection while the node is still persisting
|
||||
# identity, and cancellation must release the old worker's
|
||||
# scheduling gate immediately rather than wait for expiry.
|
||||
self._servicer.discard_unopened_session(
|
||||
response.worker_id, session_token=response.session_token
|
||||
)
|
||||
|
||||
def _register(self, request: pb.RegisterRequest) -> pb.RegisterResponse:
|
||||
async def _complete_registration(self, stream, response, stub) -> None:
|
||||
"""Validate durable acceptance, then run the exact issued session."""
|
||||
try:
|
||||
confirmation = await asyncio.wait_for(
|
||||
stream.read(), timeout=_REGISTRATION_CONFIRMATION_TIMEOUT_SECONDS
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(
|
||||
"That GPU machine did not confirm registration in time."
|
||||
) from exc
|
||||
except grpc.aio.AioRpcError as exc:
|
||||
detail = exc.details() or ""
|
||||
error_code = detail.partition(":")[0].strip()
|
||||
if exc.code() == grpc.StatusCode.FAILED_PRECONDITION and error_code in {
|
||||
"AUTH_FAILED",
|
||||
"LOCAL_STATE",
|
||||
"UPGRADE_REQUIRED",
|
||||
}:
|
||||
raise TerminalRegistrationError(detail) from exc
|
||||
raise
|
||||
if confirmation == grpc.aio.EOF:
|
||||
raise RuntimeError(
|
||||
"That GPU machine disconnected before confirming registration."
|
||||
)
|
||||
if confirmation.WhichOneof("payload") != "heartbeat":
|
||||
raise RuntimeError(
|
||||
"That GPU machine sent an invalid registration confirmation."
|
||||
)
|
||||
|
||||
session = self._servicer.session_for(
|
||||
response.worker_id, session_token=response.session_token
|
||||
)
|
||||
if session is None:
|
||||
raise RuntimeError("the session went away before the stream opened")
|
||||
|
||||
self._worker_id = response.worker_id
|
||||
self._stub = stub
|
||||
self._last_error = ""
|
||||
self._active_session = session
|
||||
self._remote_protocol_retained = True
|
||||
self._shutdown_confirmed.clear()
|
||||
self._session_closed.clear()
|
||||
|
||||
pump = asyncio.create_task(self._pump_outbound(session))
|
||||
try:
|
||||
await self._servicer.run_inbound_stream(
|
||||
session, _Frames(stream, first=confirmation), self
|
||||
)
|
||||
finally:
|
||||
pump.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await pump
|
||||
self._stub = None
|
||||
self._worker_id = ""
|
||||
if self._active_session is session:
|
||||
self._active_session = None
|
||||
self._session_closed.set()
|
||||
|
||||
async def _register(self, request: pb.RegisterRequest) -> pb.RegisterResponse:
|
||||
"""Trust on first sight, then require the same key forever after.
|
||||
|
||||
Pasting the connection string is the consent — the user went to the
|
||||
@@ -175,14 +474,28 @@ class NodeConnection:
|
||||
is a licence for a different machine to answer at that address later,
|
||||
which is why the key is bound on first contact.
|
||||
"""
|
||||
refusal = self._servicer.validate_inbound_request(request)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
worker, refusal = await to_thread_and_drain_on_cancel(
|
||||
self._authenticate_registration, request
|
||||
)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
return await self._servicer.establish_session(
|
||||
worker, request, address=self._connection.endpoint
|
||||
)
|
||||
|
||||
def _authenticate_registration(self, request: pb.RegisterRequest):
|
||||
"""Resolve inbound identity without running SQLite on the app loop."""
|
||||
public_key = bytes(request.public_key)
|
||||
if len(public_key) != 32:
|
||||
return self._servicer._refuse(
|
||||
return None, self._servicer._refuse(
|
||||
"AUTH_FAILED", "That machine sent no usable identity."
|
||||
)
|
||||
key_id = identity.key_id_for(public_key)
|
||||
if registry.is_revoked(key_id):
|
||||
return self._servicer._refuse(
|
||||
return None, self._servicer._refuse(
|
||||
"AUTH_FAILED",
|
||||
"This GPU machine was removed from this app. Add it again to use it.",
|
||||
)
|
||||
@@ -219,14 +532,12 @@ class NodeConnection:
|
||||
)
|
||||
worker = known
|
||||
if worker is None:
|
||||
return self._servicer._refuse(
|
||||
return None, self._servicer._refuse(
|
||||
"AUTH_FAILED",
|
||||
"That machine could not prove it is the one this key was added for.",
|
||||
)
|
||||
|
||||
return self._servicer.register_inbound(
|
||||
worker, request, address=self._connection.endpoint
|
||||
)
|
||||
return worker, None
|
||||
|
||||
@staticmethod
|
||||
def _proves_key_possession(request: pb.RegisterRequest, known) -> bool:
|
||||
@@ -249,9 +560,34 @@ class NodeConnection:
|
||||
public_key, message, bytes(request.challenge_signature)
|
||||
)
|
||||
|
||||
async def _outbound(self):
|
||||
def _outbound(self):
|
||||
# grpc closes request iterators itself when the peer ends a stream. An
|
||||
# async generator can still be suspended in ``Queue.get`` at that
|
||||
# point, making its concurrent ``aclose`` fail and leak teardown into
|
||||
# the next channel. A plain async iterator has no generator-finalizer
|
||||
# race and keeps the same one-frame-at-a-time backpressure.
|
||||
return _OutboundFrames(self)
|
||||
|
||||
def fence_session_egress(self, session) -> None:
|
||||
"""Drop frames copied before a replacement generation activated."""
|
||||
if self._active_session is not session:
|
||||
return
|
||||
while True:
|
||||
yield await self._outbox.get()
|
||||
try:
|
||||
self._outbox.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
def revoke_session(self, session) -> None:
|
||||
"""Synchronously fence frames already copied into the request queue."""
|
||||
if self._active_session is not session:
|
||||
return
|
||||
while True:
|
||||
try:
|
||||
self._outbox.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
self._outbox.put_nowait(None)
|
||||
|
||||
async def _pump_outbound(self, session) -> None:
|
||||
"""Move the servicer's per-session outbox onto the dialled stream.
|
||||
@@ -260,8 +596,18 @@ class NodeConnection:
|
||||
to cross into the request generator instead, because this side is the
|
||||
caller.
|
||||
"""
|
||||
while True:
|
||||
await self._outbox.put(await session.outbox.get())
|
||||
task = asyncio.current_task()
|
||||
if task is not None:
|
||||
session.egress_tasks.add(task)
|
||||
try:
|
||||
while not session.revoked and not getattr(session, "egress_fenced", False):
|
||||
message = await session.outbox.get()
|
||||
if session.revoked or getattr(session, "egress_fenced", False):
|
||||
return
|
||||
await self._outbox.put(message)
|
||||
finally:
|
||||
if task is not None:
|
||||
session.egress_tasks.discard(task)
|
||||
|
||||
# ── Artifacts ─────────────────────────────────────────────────────────
|
||||
|
||||
@@ -277,32 +623,52 @@ class NodeConnection:
|
||||
if stub is None:
|
||||
raise RuntimeError("that GPU machine is not connected")
|
||||
|
||||
size = os.path.getsize(path)
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as handle:
|
||||
digest.update(handle.read())
|
||||
size, digest, version = await to_thread_and_drain_on_cancel(
|
||||
_hash_staged_input, path
|
||||
)
|
||||
# Hashing and the gRPC request are separate operations. Re-resolve the
|
||||
# path immediately before handing the iterator to gRPC so a replaced
|
||||
# staging file is never described by the old generation's digest.
|
||||
await to_thread_and_drain_on_cancel(_validate_staged_input, path, version)
|
||||
|
||||
declared = pb.ArtifactRef()
|
||||
declared.CopyFrom(ref)
|
||||
declared.size_bytes = size
|
||||
declared.sha256 = digest.hexdigest()
|
||||
declared.sha256 = digest
|
||||
if not declared.filename:
|
||||
declared.filename = os.path.basename(path)
|
||||
|
||||
async def chunks():
|
||||
offset = 0
|
||||
with open(path, "rb") as handle:
|
||||
while True:
|
||||
data = handle.read(_PUSH_CHUNK_BYTES)
|
||||
handle = await to_thread_and_drain_on_cancel(open, path, "rb")
|
||||
try:
|
||||
await to_thread_and_drain_on_cancel(
|
||||
_validate_open_staged_input, handle, path, version
|
||||
)
|
||||
while offset < size:
|
||||
data = await to_thread_and_drain_on_cancel(
|
||||
handle.read, min(_PUSH_CHUNK_BYTES, size - offset)
|
||||
)
|
||||
if not data:
|
||||
break
|
||||
raise RuntimeError(
|
||||
"the staged task input changed before it could be sent"
|
||||
)
|
||||
offset += len(data)
|
||||
last = offset == size
|
||||
if last:
|
||||
# Do not publish the terminal frame until both the open
|
||||
# generation and its path still match what was hashed.
|
||||
await to_thread_and_drain_on_cancel(
|
||||
_validate_open_staged_input, handle, path, version
|
||||
)
|
||||
yield pb.ArtifactChunk(
|
||||
ref=declared,
|
||||
offset=offset - len(data),
|
||||
data=data,
|
||||
last=offset >= size,
|
||||
last=last,
|
||||
)
|
||||
finally:
|
||||
await to_thread_and_drain_on_cancel(handle.close)
|
||||
|
||||
ack = await stub.PushInput(
|
||||
chunks(), metadata=((KEY_METADATA_KEY, self._connection.secret),)
|
||||
@@ -327,7 +693,11 @@ class NodeConnection:
|
||||
offset = 0
|
||||
complete = False
|
||||
try:
|
||||
with open(destination, "wb") as handle:
|
||||
handle = None
|
||||
try:
|
||||
handle = await to_thread_and_drain_on_cancel(
|
||||
open, destination, "wb"
|
||||
)
|
||||
async for chunk in stub.FetchResult(
|
||||
request, metadata=((KEY_METADATA_KEY, self._connection.secret),)
|
||||
):
|
||||
@@ -339,27 +709,31 @@ class NodeConnection:
|
||||
raise RuntimeError(
|
||||
"the result is larger than the control plane accepts"
|
||||
)
|
||||
handle.write(chunk.data)
|
||||
digest.update(chunk.data)
|
||||
offset += len(chunk.data)
|
||||
data = bytes(chunk.data)
|
||||
await to_thread_and_drain_on_cancel(_write_all, handle, data)
|
||||
digest.update(data)
|
||||
offset += len(data)
|
||||
if chunk.last:
|
||||
complete = True
|
||||
break
|
||||
finally:
|
||||
if handle is not None:
|
||||
await to_thread_and_drain_on_cancel(handle.close)
|
||||
except asyncio.CancelledError:
|
||||
await to_thread_and_drain_on_cancel(_remove_quietly, destination)
|
||||
raise
|
||||
except Exception:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
await to_thread_and_drain_on_cancel(_remove_quietly, destination)
|
||||
raise
|
||||
|
||||
# A truncated file that is renamed into place and called done is the
|
||||
# exact failure the upload path was hardened against; the pull
|
||||
# direction gets the same treatment.
|
||||
if not complete:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
await to_thread_and_drain_on_cancel(_remove_quietly, destination)
|
||||
raise RuntimeError("the result ended before its final chunk")
|
||||
if ref.sha256 and digest.hexdigest() != ref.sha256:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
await to_thread_and_drain_on_cancel(_remove_quietly, destination)
|
||||
raise RuntimeError(
|
||||
"the result did not match the checksum that machine declared"
|
||||
)
|
||||
@@ -368,14 +742,46 @@ class NodeConnection:
|
||||
class _Frames:
|
||||
"""Adapts a gRPC client stream to the ``async for`` the read loop expects."""
|
||||
|
||||
def __init__(self, stream) -> None:
|
||||
def __init__(self, stream, *, first=None) -> None:
|
||||
self._stream = stream
|
||||
self._first = first
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._first is not None:
|
||||
message = self._first
|
||||
self._first = None
|
||||
return message
|
||||
message = await self._stream.read()
|
||||
if message == grpc.aio.EOF:
|
||||
raise StopAsyncIteration
|
||||
return message
|
||||
|
||||
|
||||
class _OutboundFrames:
|
||||
"""Cancellation-safe request iterator for the inverted Attach stream."""
|
||||
|
||||
def __init__(self, connection: NodeConnection) -> None:
|
||||
self._connection = connection
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
connection = self._connection
|
||||
while True:
|
||||
message = await connection._outbox.get()
|
||||
if message is None:
|
||||
raise StopAsyncIteration
|
||||
session = connection._active_session
|
||||
if session is not None:
|
||||
if session.revoked:
|
||||
raise StopAsyncIteration
|
||||
if (
|
||||
getattr(session, "egress_fenced", False)
|
||||
and message.WhichOneof("payload") != "shutdown"
|
||||
):
|
||||
continue
|
||||
return message
|
||||
|
||||
+192
-17
@@ -15,6 +15,7 @@ settings store.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -45,6 +46,20 @@ _MAX_FAILURES = 5
|
||||
_LOCKOUT_SECONDS = 60.0
|
||||
_FAILURE_WINDOW_SECONDS = 300.0
|
||||
|
||||
# ``Attach`` is the only RPC that records presence. Persisting on every
|
||||
# reconnect lets an authenticated peer turn harmless telemetry into an fsync
|
||||
# storm on the gRPC event loop, so coalesce it to a useful reporting cadence.
|
||||
_LAST_SEEN_PERSIST_INTERVAL_SECONDS = 60.0
|
||||
|
||||
# Authentication deliberately scans every stored hash in constant time. Keep
|
||||
# that work and the JSON credential file bounded even if an administrator
|
||||
# repeatedly issues replacements.
|
||||
MAX_PANEL_KEYS = 256
|
||||
|
||||
|
||||
class KeyLimitExceeded(RuntimeError):
|
||||
"""No additional panel credential can be retained safely."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class PanelKey:
|
||||
@@ -89,6 +104,31 @@ def _peer_host(peer: str) -> str:
|
||||
return peer
|
||||
|
||||
|
||||
def _fsync_parent_directory(directory: str) -> None:
|
||||
"""Make a preceding directory-entry replacement durable when supported."""
|
||||
directory_flag = getattr(os, "O_DIRECTORY", None)
|
||||
if directory_flag is None:
|
||||
return
|
||||
unsupported = {
|
||||
errno.EINVAL,
|
||||
getattr(errno, "ENOTSUP", errno.EINVAL),
|
||||
getattr(errno, "EOPNOTSUPP", errno.EINVAL),
|
||||
}
|
||||
try:
|
||||
descriptor = os.open(directory, os.O_RDONLY | directory_flag)
|
||||
except OSError as exc:
|
||||
if exc.errno in unsupported:
|
||||
return
|
||||
raise
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
except OSError as exc:
|
||||
if exc.errno not in unsupported:
|
||||
raise
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
@dataclass
|
||||
class IssuedKey:
|
||||
"""The one and only time the plaintext exists outside the caller's hands."""
|
||||
@@ -113,6 +153,10 @@ class KeyStore:
|
||||
self._connection_secrets: dict[str, str] = {}
|
||||
self._connection_fingerprints: dict[str, str] = {}
|
||||
self._failures: dict[str, _Failures] = {}
|
||||
# A failed persistence attempt must remain denied in this process but
|
||||
# still be retryable. Keeping this separate from PanelKey.revoked lets
|
||||
# the next DELETE attempt write the durable transition again.
|
||||
self._pending_revocations: set[str] = set()
|
||||
self._load()
|
||||
|
||||
# ── Persistence ───────────────────────────────────────────────────────
|
||||
@@ -170,10 +214,31 @@ class KeyStore:
|
||||
# `identity.save_worker_key` uses for the Ed25519 private key.
|
||||
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
try:
|
||||
os.write(fd, payload)
|
||||
finally:
|
||||
remaining = memoryview(payload)
|
||||
while remaining:
|
||||
written = os.write(fd, remaining)
|
||||
if written <= 0:
|
||||
raise OSError("could not finish writing the inbound key file")
|
||||
remaining = remaining[written:]
|
||||
os.fsync(fd)
|
||||
except Exception:
|
||||
os.close(fd)
|
||||
os.replace(tmp, self._path)
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
else:
|
||||
os.close(fd)
|
||||
try:
|
||||
os.replace(tmp, self._path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
_fsync_parent_directory(directory)
|
||||
try:
|
||||
os.chmod(self._path, 0o600)
|
||||
except OSError:
|
||||
@@ -196,8 +261,37 @@ class KeyStore:
|
||||
created_at=now,
|
||||
)
|
||||
with self._lock:
|
||||
previous = self._keys.get(key.key_id)
|
||||
pruned: dict[str, PanelKey] = {}
|
||||
if previous is None and len(self._keys) >= MAX_PANEL_KEYS:
|
||||
revoked = sorted(
|
||||
(
|
||||
stored
|
||||
for stored in self._keys.values()
|
||||
if stored.revoked
|
||||
and stored.key_id not in self._pending_revocations
|
||||
),
|
||||
key=lambda stored: stored.created_at,
|
||||
)
|
||||
while len(self._keys) >= MAX_PANEL_KEYS and revoked:
|
||||
stale = revoked.pop(0)
|
||||
pruned[stale.key_id] = self._keys.pop(stale.key_id)
|
||||
if len(self._keys) >= MAX_PANEL_KEYS:
|
||||
self._keys.update(pruned)
|
||||
raise KeyLimitExceeded(
|
||||
"This GPU machine already has as many panel keys as it accepts. "
|
||||
"Revoke an unused key, then try again."
|
||||
)
|
||||
self._keys[key.key_id] = key
|
||||
self._save_locked()
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
if previous is None:
|
||||
self._keys.pop(key.key_id, None)
|
||||
else:
|
||||
self._keys[key.key_id] = previous
|
||||
self._keys.update(pruned)
|
||||
raise
|
||||
return IssuedKey(key=key, secret=secret)
|
||||
|
||||
def revoke(self, key_id: str) -> bool:
|
||||
@@ -206,8 +300,14 @@ class KeyStore:
|
||||
key = self._keys.get(key_id)
|
||||
if key is None or key.revoked:
|
||||
return False
|
||||
self._pending_revocations.add(key_id)
|
||||
key.revoked = True
|
||||
self._save_locked()
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
key.revoked = False
|
||||
raise
|
||||
self._pending_revocations.discard(key_id)
|
||||
return True
|
||||
|
||||
def remember_worker_id(self, key_id: str, worker_id: str) -> None:
|
||||
@@ -216,15 +316,46 @@ class KeyStore:
|
||||
return
|
||||
with self._lock:
|
||||
key = self._keys.get(key_id)
|
||||
if key is None or key.worker_id == worker_id:
|
||||
if (
|
||||
key is None
|
||||
or key.revoked
|
||||
or key_id in self._pending_revocations
|
||||
):
|
||||
raise PermissionError("the panel key was revoked during registration")
|
||||
if key.worker_id == worker_id:
|
||||
return
|
||||
previous_worker_id = key.worker_id
|
||||
key.worker_id = worker_id
|
||||
self._save_locked()
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
# A callback retry must attempt the durable write again. If
|
||||
# the failed value remains in memory, the equality fast path
|
||||
# above accepts it as saved and the node reconnects with an id
|
||||
# that disappears on process restart.
|
||||
key.worker_id = previous_worker_id
|
||||
raise
|
||||
|
||||
def worker_id_for(self, key_id: str) -> str:
|
||||
with self._lock:
|
||||
key = self._keys.get(key_id)
|
||||
return key.worker_id if key is not None else ""
|
||||
return (
|
||||
key.worker_id
|
||||
if key is not None
|
||||
and not key.revoked
|
||||
and key_id not in self._pending_revocations
|
||||
else ""
|
||||
)
|
||||
|
||||
def is_active(self, key_id: str) -> bool:
|
||||
"""Whether this key still has authority to use an existing session."""
|
||||
with self._lock:
|
||||
key = self._keys.get(key_id)
|
||||
return (
|
||||
key is not None
|
||||
and not key.revoked
|
||||
and key_id not in self._pending_revocations
|
||||
)
|
||||
|
||||
def list_keys(self) -> list[dict]:
|
||||
with self._lock:
|
||||
@@ -232,7 +363,10 @@ class KeyStore:
|
||||
|
||||
def any_active(self) -> bool:
|
||||
with self._lock:
|
||||
return any(not k.revoked for k in self._keys.values())
|
||||
return any(
|
||||
not key.revoked and key.key_id not in self._pending_revocations
|
||||
for key in self._keys.values()
|
||||
)
|
||||
|
||||
# ── Panel-side connection credentials ───────────────────────────────
|
||||
|
||||
@@ -241,10 +375,23 @@ class KeyStore:
|
||||
) -> None:
|
||||
"""Persist a pasted node secret outside the UI-readable settings store."""
|
||||
with self._lock:
|
||||
previous_secret = self._connection_secrets.get(endpoint)
|
||||
previous_fingerprint = self._connection_fingerprints.get(endpoint)
|
||||
self._connection_secrets[endpoint] = secret
|
||||
if fingerprint:
|
||||
self._connection_fingerprints[endpoint] = fingerprint
|
||||
self._save_locked()
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
if previous_secret is None:
|
||||
self._connection_secrets.pop(endpoint, None)
|
||||
else:
|
||||
self._connection_secrets[endpoint] = previous_secret
|
||||
if previous_fingerprint is None:
|
||||
self._connection_fingerprints.pop(endpoint, None)
|
||||
else:
|
||||
self._connection_fingerprints[endpoint] = previous_fingerprint
|
||||
raise
|
||||
|
||||
def connection_secret(self, endpoint: str) -> str:
|
||||
with self._lock:
|
||||
@@ -256,9 +403,19 @@ class KeyStore:
|
||||
|
||||
def forget_connection_secret(self, endpoint: str) -> None:
|
||||
with self._lock:
|
||||
if self._connection_secrets.pop(endpoint, None) is not None:
|
||||
self._connection_fingerprints.pop(endpoint, None)
|
||||
previous_secret = self._connection_secrets.get(endpoint)
|
||||
if previous_secret is None:
|
||||
return
|
||||
previous_fingerprint = self._connection_fingerprints.get(endpoint)
|
||||
self._connection_secrets.pop(endpoint, None)
|
||||
self._connection_fingerprints.pop(endpoint, None)
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
self._connection_secrets[endpoint] = previous_secret
|
||||
if previous_fingerprint is not None:
|
||||
self._connection_fingerprints[endpoint] = previous_fingerprint
|
||||
raise
|
||||
|
||||
# ── Authentication ────────────────────────────────────────────────────
|
||||
|
||||
@@ -268,7 +425,9 @@ class KeyStore:
|
||||
record = self._failures.get(peer)
|
||||
return record is not None and record.locked_until > self._now()
|
||||
|
||||
def authenticate(self, secret: str, *, peer: str = "") -> Optional[PanelKey]:
|
||||
def authenticate(
|
||||
self, secret: str, *, peer: str = "", record_seen: bool = True
|
||||
) -> Optional[PanelKey]:
|
||||
"""Return the matching live key, or None.
|
||||
|
||||
Compares against every stored key in constant time and does not stop at
|
||||
@@ -286,7 +445,11 @@ class KeyStore:
|
||||
candidate = hash_secret(secret) if secret else ""
|
||||
matched: Optional[PanelKey] = None
|
||||
for key in self._keys.values():
|
||||
if key.revoked or not candidate:
|
||||
if (
|
||||
key.revoked
|
||||
or key.key_id in self._pending_revocations
|
||||
or not candidate
|
||||
):
|
||||
continue
|
||||
if constant_time_equals(key.secret_hash, candidate):
|
||||
matched = key
|
||||
@@ -296,9 +459,21 @@ class KeyStore:
|
||||
return None
|
||||
|
||||
self._failures.pop(peer_host, None)
|
||||
matched.last_seen_at = now
|
||||
matched.last_seen_peer = peer
|
||||
self._save_locked()
|
||||
if record_seen and (
|
||||
matched.last_seen_at <= 0.0
|
||||
or now - matched.last_seen_at
|
||||
>= _LAST_SEEN_PERSIST_INTERVAL_SECONDS
|
||||
):
|
||||
previous_at = matched.last_seen_at
|
||||
previous_peer = matched.last_seen_peer
|
||||
matched.last_seen_at = now
|
||||
matched.last_seen_peer = peer
|
||||
try:
|
||||
self._save_locked()
|
||||
except Exception:
|
||||
matched.last_seen_at = previous_at
|
||||
matched.last_seen_peer = previous_peer
|
||||
raise
|
||||
return matched
|
||||
|
||||
def _record_failure_locked(self, peer: str, now: float) -> None:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -14,11 +14,13 @@ second box — which is why neither implies the other.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from worker.async_utils import drain_task, to_thread_and_drain_on_cancel
|
||||
from worker.inbound.artifacts import ArtifactStore, KeyedArtifactTransport
|
||||
from worker.inbound.connection_log import ConnectionLog
|
||||
from worker.inbound.connection_string import (
|
||||
@@ -27,6 +29,7 @@ from worker.inbound.connection_string import (
|
||||
format_connection,
|
||||
parse_connection,
|
||||
)
|
||||
from worker.inbound.connector import InboundConnectionRollbackError
|
||||
from worker.inbound.keys import KeyStore
|
||||
from worker.inbound.listener import DEFAULT_BIND, DEFAULT_PORT, NodeListener
|
||||
|
||||
@@ -38,6 +41,41 @@ _PORT_KEY = "inbound_node_port"
|
||||
_SAVED_KEY = "inbound_saved_nodes"
|
||||
|
||||
|
||||
async def _finish_rollback(rollback, *, description: str) -> None:
|
||||
"""Finish a lifecycle rollback even if its caller is cancelled again."""
|
||||
task = asyncio.create_task(rollback, name="inbound-connection-rollback")
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except BaseException:
|
||||
await drain_task(task)
|
||||
if task.cancelled():
|
||||
raise InboundConnectionRollbackError(
|
||||
f"Could not {description}: rollback was cancelled."
|
||||
)
|
||||
return task.result()
|
||||
|
||||
|
||||
def _normalise_listener_host(value: str) -> str:
|
||||
"""Return the bare identity gRPC and X.509 expect for an IP literal."""
|
||||
candidate = (value or "").strip()
|
||||
inner = (
|
||||
candidate[1:-1]
|
||||
if len(candidate) >= 2
|
||||
and candidate.startswith("[")
|
||||
and candidate.endswith("]")
|
||||
else candidate
|
||||
)
|
||||
try:
|
||||
return str(ipaddress.ip_address(inner))
|
||||
except ValueError:
|
||||
return candidate
|
||||
|
||||
|
||||
def normalise_bind_host(value: str) -> str:
|
||||
"""Canonicalise a requested listener host before comparing or saving it."""
|
||||
return _normalise_listener_host(value)
|
||||
|
||||
|
||||
def _setting(name: str, default: str = "") -> str:
|
||||
try:
|
||||
from services import settings_store # noqa: PLC0415
|
||||
@@ -86,13 +124,15 @@ def bind_host() -> str:
|
||||
point. Reaching this node from another machine should be a decision
|
||||
somebody made, not a side effect of turning the feature on.
|
||||
"""
|
||||
return (
|
||||
os.environ.get("OMNIVOICE_INBOUND_BIND") or _setting(_BIND_KEY) or DEFAULT_BIND
|
||||
return _normalise_listener_host(
|
||||
os.environ.get("OMNIVOICE_INBOUND_BIND")
|
||||
or _setting(_BIND_KEY)
|
||||
or DEFAULT_BIND
|
||||
)
|
||||
|
||||
|
||||
def set_bind_host(value: str) -> None:
|
||||
_set_setting(_BIND_KEY, (value or "").strip() or DEFAULT_BIND)
|
||||
_set_setting(_BIND_KEY, _normalise_listener_host(value) or DEFAULT_BIND)
|
||||
|
||||
|
||||
def bind_port() -> int:
|
||||
@@ -148,7 +188,9 @@ def is_exposed(host: Optional[str] = None) -> bool:
|
||||
connection string then admits clients beyond this machine. Transport
|
||||
remains pinned TLS (docs/adr/inbound-node-mode.md).
|
||||
"""
|
||||
return (host if host is not None else bind_host()) not in (
|
||||
return _normalise_listener_host(
|
||||
host if host is not None else bind_host()
|
||||
).lower() not in (
|
||||
"127.0.0.1",
|
||||
"localhost",
|
||||
"::1",
|
||||
@@ -189,6 +231,7 @@ class InboundNode:
|
||||
self._keys: Optional[KeyStore] = None
|
||||
self._log = ConnectionLog()
|
||||
self._idle_sweep: Optional[asyncio.Task] = None
|
||||
self._lifecycle_lock = asyncio.Lock()
|
||||
self.startup_error: Optional[str] = None
|
||||
|
||||
@property
|
||||
@@ -209,18 +252,12 @@ class InboundNode:
|
||||
def port(self) -> int:
|
||||
return self._listener.port if self._listener else 0
|
||||
|
||||
def _client_factory(self, artifacts: KeyedArtifactTransport, key_id: str):
|
||||
# Imported here so a machine that never accepts connections does not
|
||||
# pay for the executor or grpc at startup.
|
||||
def _prepare_client(self, key_id: str) -> dict:
|
||||
"""Probe keys, host and accelerators away from the listener loop."""
|
||||
from worker import capabilities # noqa: PLC0415
|
||||
from worker.agent import _paths as agent_paths # noqa: PLC0415
|
||||
from worker.executor import TaskExecutor # noqa: PLC0415
|
||||
from worker.identity import load_or_create_worker_key # noqa: PLC0415
|
||||
from worker.transport.client import ( # noqa: PLC0415
|
||||
WorkerClient,
|
||||
WorkerConfig,
|
||||
describe_host,
|
||||
)
|
||||
from worker.transport.client import describe_host # noqa: PLC0415
|
||||
|
||||
locations = agent_paths()
|
||||
os.makedirs(locations["root"], exist_ok=True)
|
||||
@@ -228,6 +265,29 @@ class InboundNode:
|
||||
discovered = capabilities.discover(include_unavailable=True)
|
||||
host = describe_host()
|
||||
host["gpus"] = capabilities.describe_gpus()
|
||||
return {
|
||||
"keypair": keypair,
|
||||
"discovered": discovered,
|
||||
"host": host,
|
||||
"worker_id": self.keys.worker_id_for(key_id),
|
||||
"max_concurrent_tasks": capabilities.max_concurrent_tasks(discovered),
|
||||
}
|
||||
|
||||
async def _client_factory(
|
||||
self, artifacts: KeyedArtifactTransport, key_id: str
|
||||
):
|
||||
# Imported here so a machine that never accepts connections does not
|
||||
# pay for the executor or grpc at startup.
|
||||
from worker import capabilities # noqa: PLC0415
|
||||
from worker.executor import TaskExecutor # noqa: PLC0415
|
||||
from worker.transport.client import ( # noqa: PLC0415
|
||||
WorkerClient,
|
||||
WorkerConfig,
|
||||
)
|
||||
|
||||
prepared = await to_thread_and_drain_on_cancel(
|
||||
self._prepare_client, key_id
|
||||
)
|
||||
|
||||
executor = TaskExecutor()
|
||||
return WorkerClient(
|
||||
@@ -235,23 +295,28 @@ class InboundNode:
|
||||
endpoint="",
|
||||
cert_fingerprint="",
|
||||
certificate_pem=b"",
|
||||
keypair=keypair,
|
||||
keypair=prepared["keypair"],
|
||||
# Per panel key, not per node: each panel keeps its own
|
||||
# registry, so the same machine is a different worker id to
|
||||
# each of them, and the node signs its challenge over that id.
|
||||
worker_id=self.keys.worker_id_for(key_id),
|
||||
worker_id=prepared["worker_id"],
|
||||
enrollment_token="",
|
||||
max_concurrent_tasks=capabilities.max_concurrent_tasks(discovered),
|
||||
capabilities=discovered,
|
||||
host=host,
|
||||
max_concurrent_tasks=prepared["max_concurrent_tasks"],
|
||||
capabilities=prepared["discovered"],
|
||||
host=prepared["host"],
|
||||
),
|
||||
execute=executor.execute,
|
||||
capability_probe=lambda: capabilities.discover(include_unavailable=True),
|
||||
on_registered=lambda wid: self.keys.remember_worker_id(key_id, wid),
|
||||
artifacts=artifacts,
|
||||
drain_active_work=executor.drain_active_work,
|
||||
)
|
||||
|
||||
async def start(self) -> None:
|
||||
async with self._lifecycle_lock:
|
||||
await self._start()
|
||||
|
||||
async def _start(self) -> None:
|
||||
if self._listener is not None:
|
||||
return
|
||||
self.startup_error = None
|
||||
@@ -264,9 +329,18 @@ class InboundNode:
|
||||
)
|
||||
try:
|
||||
await listener.start(host=bind_host(), port=bind_port())
|
||||
except asyncio.CancelledError:
|
||||
# NodeListener cleans a partially bound server before returning.
|
||||
# If that cleanup itself failed it retains the handle; publish it
|
||||
# here so a later stop can retry rather than losing a live socket.
|
||||
if listener.running:
|
||||
self._listener = listener
|
||||
raise
|
||||
except Exception as exc:
|
||||
# A node that cannot listen must say so in the UI rather than look
|
||||
# enabled and quietly accept nothing.
|
||||
if listener.running:
|
||||
self._listener = listener
|
||||
self.startup_error = str(exc)
|
||||
logger.error("Could not start the inbound listener: %s", exc)
|
||||
return
|
||||
@@ -282,12 +356,48 @@ class InboundNode:
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
sweep, self._idle_sweep = self._idle_sweep, None
|
||||
if sweep is not None:
|
||||
sweep.cancel()
|
||||
listener, self._listener = self._listener, None
|
||||
if listener is not None:
|
||||
await listener.stop()
|
||||
async with self._lifecycle_lock:
|
||||
await self._stop()
|
||||
|
||||
async def _stop(self) -> None:
|
||||
sweep = self._idle_sweep
|
||||
listener = self._listener
|
||||
|
||||
async def shutdown() -> None:
|
||||
if sweep is not None:
|
||||
sweep.cancel()
|
||||
await asyncio.gather(sweep, return_exceptions=True)
|
||||
if listener is not None:
|
||||
await listener.stop()
|
||||
|
||||
stopping = asyncio.create_task(shutdown(), name="inbound-node-stop")
|
||||
try:
|
||||
await asyncio.shield(stopping)
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(stopping)
|
||||
if stopping.cancelled():
|
||||
raise
|
||||
failure = stopping.exception()
|
||||
if failure is not None:
|
||||
raise failure
|
||||
if self._idle_sweep is sweep:
|
||||
self._idle_sweep = None
|
||||
if self._listener is listener:
|
||||
self._listener = None
|
||||
raise
|
||||
except BaseException:
|
||||
await drain_task(stopping)
|
||||
raise
|
||||
if self._idle_sweep is sweep:
|
||||
self._idle_sweep = None
|
||||
if self._listener is listener:
|
||||
self._listener = None
|
||||
|
||||
async def revoke_key(self, key_id: str) -> bool:
|
||||
"""Durably revoke one panel and withdraw all of its live sessions."""
|
||||
if self._listener is not None:
|
||||
return await self._listener.revoke_key_and_wait(key_id)
|
||||
return self.keys.revoke(key_id)
|
||||
|
||||
def connection_string(self, secret: str, *, host: Optional[str] = None) -> str:
|
||||
"""The one artifact a user copies to another machine.
|
||||
@@ -330,6 +440,8 @@ class OutboundNodes:
|
||||
self._connections: dict[str, object] = {}
|
||||
self._tasks: dict[str, asyncio.Task] = {}
|
||||
self._credentials = credentials
|
||||
self._lifecycle_lock = asyncio.Lock()
|
||||
self._servicer = None
|
||||
|
||||
@property
|
||||
def credentials(self) -> KeyStore:
|
||||
@@ -367,67 +479,337 @@ class OutboundNodes:
|
||||
async def add(self, text: str, servicer) -> Connection:
|
||||
"""Parse, save and dial. Raises InvalidConnectionString on a bad paste."""
|
||||
connection = parse_connection(text)
|
||||
entries = self.saved()
|
||||
# Keyed by endpoint: re-pasting a rotated key for the same machine
|
||||
# replaces it rather than leaving a dead entry that retries forever.
|
||||
entries = [e for e in entries if _endpoint_of(e) != connection.endpoint]
|
||||
entries.append(connection.endpoint)
|
||||
self.credentials.remember_connection_secret(
|
||||
connection.endpoint, connection.secret, connection.fingerprint
|
||||
)
|
||||
self._save(entries)
|
||||
async with self._lifecycle_lock:
|
||||
return await self._add(connection, servicer)
|
||||
|
||||
# Tear down any live session to this machine BEFORE dialling. Without
|
||||
# this, re-pasting for an already-connected machine saved the new key
|
||||
# and then short-circuited on the existing connection — so a wrong key
|
||||
# reported success, kept working on the old session, and only failed
|
||||
# after a restart, by which time nothing pointed at the paste that
|
||||
# caused it. Verified on hardware.
|
||||
await self._drop(connection.endpoint)
|
||||
await self._dial(connection, servicer)
|
||||
return connection
|
||||
|
||||
async def _drop(self, endpoint: str) -> None:
|
||||
connection = self._connections.pop(endpoint, None)
|
||||
task = self._tasks.pop(endpoint, None)
|
||||
if connection is not None:
|
||||
await connection.stop()
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
|
||||
async def remove(self, endpoint: str) -> bool:
|
||||
entries = [e for e in self.saved() if _endpoint_of(e) != endpoint]
|
||||
self._save(entries)
|
||||
self.credentials.forget_connection_secret(endpoint)
|
||||
existed = endpoint in self._connections
|
||||
await self._drop(endpoint)
|
||||
return existed
|
||||
|
||||
async def start_all(self, servicer) -> None:
|
||||
for entry in self.saved():
|
||||
try:
|
||||
await self._dial(self._connection_for(entry), servicer)
|
||||
except InvalidConnectionString as exc:
|
||||
logger.warning(
|
||||
"Ignoring a saved connection that no longer parses: %s", exc
|
||||
)
|
||||
|
||||
async def _dial(self, connection: Connection, servicer) -> None:
|
||||
async def _add(self, connection: Connection, servicer) -> Connection:
|
||||
from worker.inbound.connector import NodeConnection # noqa: PLC0415
|
||||
|
||||
if connection.endpoint in self._connections:
|
||||
original_entries = self.saved()
|
||||
old_secret = self.credentials.connection_secret(connection.endpoint)
|
||||
old_fingerprint = self.credentials.connection_fingerprint(connection.endpoint)
|
||||
existing = self._connections.get(connection.endpoint)
|
||||
existing_task = self._tasks.get(connection.endpoint)
|
||||
|
||||
# Pasting the already-running key is idempotent. Probing it would be a
|
||||
# duplicate Attach and could disturb state deliberately retained by
|
||||
# that same key.
|
||||
if (
|
||||
existing is not None
|
||||
and old_secret == connection.secret
|
||||
and old_fingerprint == connection.fingerprint
|
||||
):
|
||||
if existing_task is not None and not existing_task.done():
|
||||
return connection
|
||||
# Terminal registration failures leave their diagnostic connector
|
||||
# in the snapshot. Re-pasting after an upgrade/repair must really
|
||||
# redial, not mistake that dead object for a healthy connection.
|
||||
if self._connections.get(connection.endpoint) is existing:
|
||||
self._connections.pop(connection.endpoint, None)
|
||||
if self._tasks.get(connection.endpoint) is existing_task:
|
||||
self._tasks.pop(connection.endpoint, None)
|
||||
try:
|
||||
await self._dial(connection, servicer, wait_until_ready=True)
|
||||
except BaseException as operation:
|
||||
try:
|
||||
await _finish_rollback(
|
||||
self._restore_failed_redial(
|
||||
connection.endpoint, existing, existing_task
|
||||
),
|
||||
description="restore the previous inbound connector",
|
||||
)
|
||||
except InboundConnectionRollbackError as rollback:
|
||||
raise rollback from operation
|
||||
raise
|
||||
return connection
|
||||
|
||||
if existing is not None:
|
||||
# Authenticate and apply identity/version policy before touching
|
||||
# the only working connector or its durable credential.
|
||||
await NodeConnection(servicer, connection).probe()
|
||||
|
||||
entries = [
|
||||
entry
|
||||
for entry in original_entries
|
||||
if _endpoint_of(entry) != connection.endpoint
|
||||
]
|
||||
# Keyed by endpoint: re-pasting a rotated key for the same machine
|
||||
# replaces it rather than leaving a dead entry that retries forever.
|
||||
entries.append(connection.endpoint)
|
||||
try:
|
||||
if existing is not None:
|
||||
# An offline connector can own work retained on the node. Its
|
||||
# shutdown guard must run before replacement state is persisted.
|
||||
await self._drop(connection.endpoint)
|
||||
self.credentials.remember_connection_secret(
|
||||
connection.endpoint, connection.secret, connection.fingerprint
|
||||
)
|
||||
self._save(entries)
|
||||
await self._dial(
|
||||
connection, servicer, wait_until_ready=existing is not None
|
||||
)
|
||||
except BaseException as operation:
|
||||
try:
|
||||
await _finish_rollback(
|
||||
self._rollback_add(
|
||||
connection,
|
||||
servicer,
|
||||
existing,
|
||||
original_entries,
|
||||
old_secret,
|
||||
old_fingerprint,
|
||||
),
|
||||
description="restore the previous inbound connection",
|
||||
)
|
||||
except InboundConnectionRollbackError as rollback:
|
||||
raise rollback from operation
|
||||
raise
|
||||
return connection
|
||||
|
||||
async def _restore_failed_redial(
|
||||
self, endpoint: str, existing, existing_task: Optional[asyncio.Task]
|
||||
) -> None:
|
||||
failure = None
|
||||
candidate = self._connections.get(endpoint)
|
||||
candidate_task = self._tasks.get(endpoint)
|
||||
if candidate is not None and candidate is not existing:
|
||||
try:
|
||||
await self._close_candidate(endpoint, candidate, candidate_task)
|
||||
except BaseException as exc:
|
||||
failure = exc
|
||||
if endpoint not in self._connections:
|
||||
self._connections[endpoint] = existing
|
||||
if endpoint not in self._tasks and existing_task is not None:
|
||||
self._tasks[endpoint] = existing_task
|
||||
if failure is not None:
|
||||
raise InboundConnectionRollbackError(
|
||||
"The previous GPU-machine connector could not be restored safely."
|
||||
) from failure
|
||||
|
||||
async def _close_candidate(self, endpoint: str, candidate, candidate_task) -> None:
|
||||
try:
|
||||
close = getattr(candidate, "close", None)
|
||||
if callable(close):
|
||||
await close()
|
||||
finally:
|
||||
if candidate_task is not None:
|
||||
candidate_task.cancel()
|
||||
await asyncio.gather(candidate_task, return_exceptions=True)
|
||||
if self._connections.get(endpoint) is candidate:
|
||||
self._connections.pop(endpoint, None)
|
||||
if self._tasks.get(endpoint) is candidate_task:
|
||||
self._tasks.pop(endpoint, None)
|
||||
|
||||
async def _rollback_add(
|
||||
self,
|
||||
connection: Connection,
|
||||
servicer,
|
||||
existing,
|
||||
original_entries: list[str],
|
||||
old_secret: str,
|
||||
old_fingerprint: str,
|
||||
) -> None:
|
||||
"""Restore both live and durable generations after a failed replacement."""
|
||||
endpoint = connection.endpoint
|
||||
failures = []
|
||||
candidate = self._connections.get(endpoint)
|
||||
candidate_task = self._tasks.get(endpoint)
|
||||
if candidate is not None and candidate is not existing:
|
||||
try:
|
||||
await self._close_candidate(endpoint, candidate, candidate_task)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
try:
|
||||
if old_secret:
|
||||
self.credentials.remember_connection_secret(
|
||||
endpoint, old_secret, old_fingerprint
|
||||
)
|
||||
else:
|
||||
self.credentials.forget_connection_secret(endpoint)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
try:
|
||||
self._save(original_entries)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
if (
|
||||
not failures
|
||||
and existing is not None
|
||||
and old_secret
|
||||
and endpoint not in self._connections
|
||||
):
|
||||
try:
|
||||
await self._dial(
|
||||
Connection(
|
||||
host=connection.host,
|
||||
port=connection.port,
|
||||
secret=old_secret,
|
||||
fingerprint=old_fingerprint,
|
||||
),
|
||||
servicer,
|
||||
)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
if failures:
|
||||
raise InboundConnectionRollbackError(
|
||||
"The previous GPU-machine connection could not be restored safely. "
|
||||
"It remains stopped; fix its connection/settings storage, then "
|
||||
"paste the original connection again."
|
||||
) from failures[0]
|
||||
|
||||
async def _drop(self, endpoint: str) -> None:
|
||||
connection = self._connections.get(endpoint)
|
||||
task = self._tasks.get(endpoint)
|
||||
if connection is not None:
|
||||
await connection.stop()
|
||||
if self._connections.get(endpoint) is connection:
|
||||
self._connections.pop(endpoint, None)
|
||||
if self._tasks.get(endpoint) is task:
|
||||
self._tasks.pop(endpoint, None)
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
async def remove(self, endpoint: str) -> bool:
|
||||
async with self._lifecycle_lock:
|
||||
return await self._remove(endpoint)
|
||||
|
||||
async def _remove(self, endpoint: str) -> bool:
|
||||
existed = endpoint in self._connections
|
||||
original_entries = self.saved()
|
||||
secret = self.credentials.connection_secret(endpoint)
|
||||
fingerprint = self.credentials.connection_fingerprint(endpoint)
|
||||
previous_connection = None
|
||||
if secret and fingerprint:
|
||||
try:
|
||||
previous_connection = self._connection_for(endpoint)
|
||||
except InvalidConnectionString:
|
||||
pass
|
||||
# A disconnected node deliberately retains work for reconnect. Do not
|
||||
# erase the only connector/key capable of delivering terminal shutdown.
|
||||
try:
|
||||
await self._drop(endpoint)
|
||||
entries = [e for e in original_entries if _endpoint_of(e) != endpoint]
|
||||
# Remove the protected credential first. If that durable write
|
||||
# fails, restore both durable generations and the live connector.
|
||||
self.credentials.forget_connection_secret(endpoint)
|
||||
self._save(entries)
|
||||
except BaseException as operation:
|
||||
try:
|
||||
await _finish_rollback(
|
||||
self._rollback_remove(
|
||||
endpoint,
|
||||
original_entries,
|
||||
secret,
|
||||
fingerprint,
|
||||
previous_connection,
|
||||
),
|
||||
description="restore removed inbound connection state",
|
||||
)
|
||||
except InboundConnectionRollbackError as rollback:
|
||||
raise rollback from operation
|
||||
raise
|
||||
return existed
|
||||
|
||||
async def _rollback_remove(
|
||||
self,
|
||||
endpoint: str,
|
||||
original_entries: list[str],
|
||||
secret: str,
|
||||
fingerprint: str,
|
||||
previous_connection: Optional[Connection],
|
||||
) -> None:
|
||||
failures = []
|
||||
try:
|
||||
if secret:
|
||||
self.credentials.remember_connection_secret(
|
||||
endpoint, secret, fingerprint
|
||||
)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
try:
|
||||
self._save(original_entries)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
if (
|
||||
not failures
|
||||
and previous_connection is not None
|
||||
and self._servicer is not None
|
||||
and endpoint not in self._connections
|
||||
):
|
||||
try:
|
||||
await self._dial(previous_connection, self._servicer)
|
||||
except BaseException as exc:
|
||||
failures.append(exc)
|
||||
if failures:
|
||||
raise InboundConnectionRollbackError(
|
||||
"The removed GPU-machine connection could not be restored safely. "
|
||||
"It remains stopped; fix its connection/settings storage, then retry."
|
||||
) from failures[0]
|
||||
|
||||
async def start_all(self, servicer) -> None:
|
||||
async with self._lifecycle_lock:
|
||||
for entry in self.saved():
|
||||
try:
|
||||
await self._dial(self._connection_for(entry), servicer)
|
||||
except InvalidConnectionString as exc:
|
||||
logger.warning(
|
||||
"Ignoring a saved connection that no longer parses: %s", exc
|
||||
)
|
||||
|
||||
async def _dial(
|
||||
self, connection: Connection, servicer, *, wait_until_ready: bool = False
|
||||
) -> None:
|
||||
from worker.inbound.connector import NodeConnection # noqa: PLC0415
|
||||
|
||||
self._servicer = servicer
|
||||
existing = self._connections.get(connection.endpoint)
|
||||
if existing is not None:
|
||||
if wait_until_ready and getattr(existing, "_connection", None) != connection:
|
||||
from worker.inbound.connector import ( # noqa: PLC0415
|
||||
InboundConnectionError,
|
||||
)
|
||||
|
||||
raise InboundConnectionError(
|
||||
"A different connection to that GPU machine is already active."
|
||||
)
|
||||
return
|
||||
node = NodeConnection(servicer, connection)
|
||||
self._connections[connection.endpoint] = node
|
||||
self._tasks[connection.endpoint] = asyncio.create_task(
|
||||
task = asyncio.create_task(
|
||||
node.run_forever(), name=f"inbound-node-{connection.endpoint}"
|
||||
)
|
||||
task.add_done_callback(self._observe_connection_result)
|
||||
self._tasks[connection.endpoint] = task
|
||||
if wait_until_ready:
|
||||
await node.wait_until_registered(task)
|
||||
|
||||
@staticmethod
|
||||
def _observe_connection_result(task: asyncio.Task) -> None:
|
||||
"""Retrieve terminal dial errors; NodeConnection retains the UI detail."""
|
||||
if task.cancelled():
|
||||
return
|
||||
try:
|
||||
task.exception()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def stop(self) -> None:
|
||||
async with self._lifecycle_lock:
|
||||
await self._stop_all()
|
||||
|
||||
async def _stop_all(self) -> None:
|
||||
for connection in list(self._connections.values()):
|
||||
await connection.stop()
|
||||
for task in list(self._tasks.values()):
|
||||
close = getattr(connection, "close", None)
|
||||
if callable(close):
|
||||
await close()
|
||||
else:
|
||||
await connection.stop()
|
||||
tasks = list(self._tasks.values())
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._connections.clear()
|
||||
self._tasks.clear()
|
||||
|
||||
|
||||
+20
-5
@@ -17,7 +17,7 @@ from dataclasses import dataclass, field
|
||||
from typing import Iterator, Optional
|
||||
|
||||
from worker.breaker import BreakerRegistry
|
||||
from worker.capacity import WorkerCapacity, derive_concurrency
|
||||
from worker.capacity import WorkerCapacity, clamp_concurrency, derive_concurrency
|
||||
from worker.clock import resolve
|
||||
from worker.identity import Session
|
||||
from worker.registry import RemoteWorker
|
||||
@@ -58,6 +58,9 @@ class ConnectedWorker:
|
||||
# The address this worker connected FROM, as the control plane saw it.
|
||||
address: str = ""
|
||||
draining: bool = False
|
||||
# Registration handoff temporarily stops new assignments without
|
||||
# conflating that transport state with a user-requested drain/shutdown.
|
||||
registration_pending: bool = False
|
||||
# Attempt ids this worker claims to be running. Rebuilt on every reconnect
|
||||
# from its own report, never inferred.
|
||||
in_flight: set[str] = field(default_factory=set)
|
||||
@@ -82,7 +85,11 @@ class ConnectedWorker:
|
||||
"""
|
||||
if self.stale():
|
||||
return "offline"
|
||||
if self.draining or self.capacity.available_slots <= 0:
|
||||
if (
|
||||
self.draining
|
||||
or self.registration_pending
|
||||
or self.capacity.available_slots <= 0
|
||||
):
|
||||
return "busy"
|
||||
return "ready"
|
||||
|
||||
@@ -175,7 +182,7 @@ class WorkerPool:
|
||||
epoch=epoch,
|
||||
capacity=WorkerCapacity(
|
||||
worker_id=record.id,
|
||||
max_concurrent_tasks=max(1, max_concurrent_tasks),
|
||||
max_concurrent_tasks=clamp_concurrency(max_concurrent_tasks),
|
||||
backend=backend,
|
||||
),
|
||||
connected_at=stamp,
|
||||
@@ -231,6 +238,10 @@ class WorkerPool:
|
||||
def disconnect(self, worker_id: str) -> Optional[ConnectedWorker]:
|
||||
return self._connected.pop(worker_id, None)
|
||||
|
||||
def restore_connection(self, worker: ConnectedWorker) -> None:
|
||||
"""Restore an exact live snapshot after replacement activation fails."""
|
||||
self._connected[worker.worker_id] = worker
|
||||
|
||||
def get(self, worker_id: str) -> Optional[ConnectedWorker]:
|
||||
return self._connected.get(worker_id)
|
||||
|
||||
@@ -300,10 +311,14 @@ class WorkerPool:
|
||||
worker.capacity.slots[key] = ModelSlot(
|
||||
engine=cap.get("engine", ""),
|
||||
model_id=cap.get("model_id", ""),
|
||||
derived_concurrency=max(0, declared),
|
||||
derived_concurrency=clamp_concurrency(
|
||||
declared, allow_zero=True
|
||||
),
|
||||
)
|
||||
else:
|
||||
slot.derived_concurrency = max(0, declared)
|
||||
slot.derived_concurrency = clamp_concurrency(
|
||||
declared, allow_zero=True
|
||||
)
|
||||
|
||||
def stale_workers(self, *, now: Optional[float] = None) -> list[ConnectedWorker]:
|
||||
return [w for w in self if w.stale(now=now)]
|
||||
|
||||
@@ -7,7 +7,9 @@
|
||||
// Rules of the road (goal_v2.md A5):
|
||||
// * Additive-only within v1. Never renumber, never reuse a field number.
|
||||
// * Version negotiation happens at Register; the server may refuse with
|
||||
// UPGRADE_REQUIRED. Supported skew window is N-2.
|
||||
// UPGRADE_REQUIRED. Semantic protocol v2 is intentionally incompatible
|
||||
// with v1 because enrollment became a durable two-phase handshake; never
|
||||
// infer a release-based skew window across that boundary.
|
||||
// * The Control stream carries SMALL messages only. Artifacts (reference
|
||||
// audio in, rendered audio/video out) move through UploadResult /
|
||||
// DownloadArtifact. A large payload on the control stream head-of-line
|
||||
|
||||
+254
-53
@@ -19,12 +19,15 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from core.db import db_conn
|
||||
from worker.capacity import clamp_concurrency
|
||||
from worker.clock import resolve
|
||||
from worker import identity
|
||||
|
||||
@@ -33,6 +36,14 @@ logger = logging.getLogger("omnivoice.worker")
|
||||
# Default scheduling preference. Higher wins; equal priorities fall through to
|
||||
# least-busy, which is the actual default behaviour for a homogeneous setup.
|
||||
_DEFAULT_PRIORITY = 50
|
||||
_AUTHORITY_LOCK = threading.RLock()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def authority_guard():
|
||||
"""Serialize durable authority changes with live-session publication."""
|
||||
with _AUTHORITY_LOCK:
|
||||
yield
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -180,6 +191,54 @@ def purge_expired_enrollments(*, now: Optional[float] = None) -> int:
|
||||
# ── Workers ────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _new_worker(
|
||||
*,
|
||||
name: str,
|
||||
public_key: bytes,
|
||||
endpoint: str,
|
||||
host: Optional[dict],
|
||||
capabilities: Optional[list[dict]],
|
||||
max_concurrent_tasks: int,
|
||||
consent_granted: bool,
|
||||
stamp: float,
|
||||
) -> RemoteWorker:
|
||||
key_id = identity.key_id_for(public_key)
|
||||
return RemoteWorker(
|
||||
id=uuid.uuid4().hex[:12],
|
||||
name=name or key_id,
|
||||
key_id=key_id,
|
||||
public_key=public_key,
|
||||
endpoint=endpoint,
|
||||
host=host or {},
|
||||
capabilities=capabilities or [],
|
||||
max_concurrent_tasks=clamp_concurrency(max_concurrent_tasks),
|
||||
consent_granted_at=stamp if consent_granted else None,
|
||||
created_at=stamp,
|
||||
)
|
||||
|
||||
|
||||
def _insert_worker(conn, worker: RemoteWorker) -> None:
|
||||
conn.execute(
|
||||
"INSERT INTO remote_workers "
|
||||
"(id, name, key_id, public_key, enabled, revoked, priority, endpoint, host_json, "
|
||||
" capabilities_json, max_concurrent_tasks, session_epoch, consent_granted_at, created_at) "
|
||||
"VALUES (?, ?, ?, ?, 1, 0, ?, ?, ?, ?, ?, 0, ?, ?)",
|
||||
(
|
||||
worker.id,
|
||||
worker.name,
|
||||
worker.key_id,
|
||||
worker.public_key,
|
||||
worker.priority,
|
||||
worker.endpoint,
|
||||
json.dumps(worker.host),
|
||||
json.dumps(worker.capabilities),
|
||||
worker.max_concurrent_tasks,
|
||||
worker.consent_granted_at,
|
||||
worker.created_at,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def enroll_worker(
|
||||
*,
|
||||
name: str,
|
||||
@@ -198,49 +257,128 @@ def enroll_worker(
|
||||
your own desktop is not agreeing to use someone else's.
|
||||
"""
|
||||
stamp = resolve(now)
|
||||
key_id = identity.key_id_for(public_key)
|
||||
existing = get_by_key_id(key_id)
|
||||
if existing is not None:
|
||||
return existing
|
||||
worker = RemoteWorker(
|
||||
id=uuid.uuid4().hex[:12],
|
||||
name=name or key_id,
|
||||
key_id=key_id,
|
||||
worker = _new_worker(
|
||||
name=name,
|
||||
public_key=public_key,
|
||||
endpoint=endpoint,
|
||||
host=host or {},
|
||||
capabilities=capabilities or [],
|
||||
max_concurrent_tasks=max(1, int(max_concurrent_tasks)),
|
||||
consent_granted_at=stamp if consent_granted else None,
|
||||
created_at=stamp,
|
||||
host=host,
|
||||
capabilities=capabilities,
|
||||
max_concurrent_tasks=max_concurrent_tasks,
|
||||
consent_granted=consent_granted,
|
||||
stamp=stamp,
|
||||
)
|
||||
with db_conn() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO remote_workers "
|
||||
"(id, name, key_id, public_key, enabled, revoked, priority, endpoint, host_json, "
|
||||
" capabilities_json, max_concurrent_tasks, session_epoch, consent_granted_at, created_at) "
|
||||
"VALUES (?, ?, ?, ?, 1, 0, ?, ?, ?, ?, ?, 0, ?, ?)",
|
||||
(
|
||||
worker.id,
|
||||
worker.name,
|
||||
worker.key_id,
|
||||
worker.public_key,
|
||||
worker.priority,
|
||||
worker.endpoint,
|
||||
json.dumps(worker.host),
|
||||
json.dumps(worker.capabilities),
|
||||
worker.max_concurrent_tasks,
|
||||
worker.consent_granted_at,
|
||||
worker.created_at,
|
||||
),
|
||||
)
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE key_id = ?", (worker.key_id,)
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
return _row_to_worker(row)
|
||||
_insert_worker(conn, worker)
|
||||
logger.info("Enrolled remote worker %s (%s)", worker.name, worker.key_id)
|
||||
return worker
|
||||
|
||||
|
||||
def enroll_with_token(
|
||||
token: identity.EnrollmentToken,
|
||||
*,
|
||||
name: str,
|
||||
public_key: bytes,
|
||||
endpoint: str = "",
|
||||
host: Optional[dict] = None,
|
||||
capabilities: Optional[list[dict]] = None,
|
||||
max_concurrent_tasks: int = 1,
|
||||
consent_granted: bool = True,
|
||||
now: Optional[float] = None,
|
||||
) -> Optional[RemoteWorker]:
|
||||
"""Consume a token and bind its worker identity in one transaction."""
|
||||
stamp = resolve(now)
|
||||
worker = _new_worker(
|
||||
name=name,
|
||||
public_key=public_key,
|
||||
endpoint=endpoint,
|
||||
host=host,
|
||||
capabilities=capabilities,
|
||||
max_concurrent_tasks=max_concurrent_tasks,
|
||||
consent_granted=consent_granted,
|
||||
stamp=stamp,
|
||||
)
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
enrollment = conn.execute(
|
||||
"SELECT secret_hash, expires_at, used_at FROM remote_worker_enrollments "
|
||||
"WHERE token_id = ?",
|
||||
(token.token_id,),
|
||||
).fetchone()
|
||||
if (
|
||||
enrollment is None
|
||||
or enrollment["used_at"] is not None
|
||||
or stamp > float(enrollment["expires_at"])
|
||||
or not identity.constant_time_equals(
|
||||
enrollment["secret_hash"], token.secret_hash
|
||||
)
|
||||
):
|
||||
return None
|
||||
existing_row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE key_id = ?", (worker.key_id,)
|
||||
).fetchone()
|
||||
if existing_row is not None and bool(existing_row["revoked"]):
|
||||
return None
|
||||
enrolled = _row_to_worker(existing_row) if existing_row is not None else worker
|
||||
consumed = conn.execute(
|
||||
"UPDATE remote_worker_enrollments SET used_at = ?, used_by_worker = ? "
|
||||
"WHERE token_id = ? AND used_at IS NULL",
|
||||
(stamp, enrolled.id, token.token_id),
|
||||
)
|
||||
if consumed.rowcount != 1:
|
||||
return None
|
||||
if existing_row is None:
|
||||
_insert_worker(conn, worker)
|
||||
if existing_row is None:
|
||||
logger.info("Enrolled remote worker %s (%s)", worker.name, worker.key_id)
|
||||
return enrolled
|
||||
|
||||
|
||||
def recover_enrollment_with_token(
|
||||
token: identity.EnrollmentToken, *, public_key: bytes
|
||||
) -> Optional[RemoteWorker]:
|
||||
"""Resolve a spent token only to the exact key it originally enrolled.
|
||||
|
||||
This is only the durable lookup half of recovery. The transport must also
|
||||
verify a fresh signature from this key before issuing another session; a
|
||||
spent token and an observed public key are not proof of private-key
|
||||
possession.
|
||||
"""
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
enrollment = conn.execute(
|
||||
"SELECT secret_hash, used_at, used_by_worker "
|
||||
"FROM remote_worker_enrollments WHERE token_id = ?",
|
||||
(token.token_id,),
|
||||
).fetchone()
|
||||
if (
|
||||
enrollment is None
|
||||
or enrollment["used_at"] is None
|
||||
or not enrollment["used_by_worker"]
|
||||
or not identity.constant_time_equals(
|
||||
enrollment["secret_hash"], token.secret_hash
|
||||
)
|
||||
):
|
||||
return None
|
||||
row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE id = ?",
|
||||
(enrollment["used_by_worker"],),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
worker = _row_to_worker(row)
|
||||
if worker.revoked or worker.public_key != public_key:
|
||||
return None
|
||||
return worker
|
||||
|
||||
|
||||
def get(worker_id: str) -> Optional[RemoteWorker]:
|
||||
with db_conn() as conn:
|
||||
row = conn.execute("SELECT * FROM remote_workers WHERE id = ?", (worker_id,)).fetchone()
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE id = ?", (worker_id,)
|
||||
).fetchone()
|
||||
return _row_to_worker(row) if row else None
|
||||
|
||||
|
||||
@@ -268,7 +406,7 @@ def begin_session(worker_id: str, *, now: Optional[float] = None) -> int:
|
||||
otherwise delivers two accepts for one assignment.
|
||||
"""
|
||||
stamp = resolve(now)
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE remote_workers SET session_epoch = session_epoch + 1, last_seen_at = ? WHERE id = ?",
|
||||
(stamp, worker_id),
|
||||
@@ -292,29 +430,79 @@ def update_capabilities(
|
||||
capabilities: list[dict],
|
||||
host: Optional[dict] = None,
|
||||
max_concurrent_tasks: Optional[int] = None,
|
||||
_conn=None,
|
||||
) -> None:
|
||||
sets = ["capabilities_json = ?"]
|
||||
params: list = [json.dumps(capabilities)]
|
||||
if host is not None:
|
||||
sets.append("host_json = ?")
|
||||
params.append(json.dumps(host))
|
||||
if max_concurrent_tasks is not None:
|
||||
sets.append("max_concurrent_tasks = ?")
|
||||
params.append(max(1, int(max_concurrent_tasks)))
|
||||
params.append(worker_id)
|
||||
with db_conn() as conn:
|
||||
conn.execute(f"UPDATE remote_workers SET {', '.join(sets)} WHERE id = ?", params)
|
||||
params = (
|
||||
json.dumps(capabilities),
|
||||
json.dumps(host) if host is not None else None,
|
||||
clamp_concurrency(max_concurrent_tasks)
|
||||
if max_concurrent_tasks is not None
|
||||
else None,
|
||||
worker_id,
|
||||
)
|
||||
sql = """
|
||||
UPDATE remote_workers
|
||||
SET capabilities_json = ?,
|
||||
host_json = COALESCE(?, host_json),
|
||||
max_concurrent_tasks = COALESCE(?, max_concurrent_tasks)
|
||||
WHERE id = ?
|
||||
"""
|
||||
if _conn is not None:
|
||||
# The caller owns the surrounding SQLite transaction. Acquiring the
|
||||
# authority lock inside it can deadlock against revoke, which takes
|
||||
# that lock before waiting for the same database write lock.
|
||||
_conn.execute(sql, params)
|
||||
else:
|
||||
with _AUTHORITY_LOCK:
|
||||
with db_conn() as conn:
|
||||
conn.execute(sql, params)
|
||||
|
||||
|
||||
def update_policy(
|
||||
worker_id: str,
|
||||
*,
|
||||
name: Optional[str] = None,
|
||||
enabled: Optional[bool] = None,
|
||||
priority: Optional[int] = None,
|
||||
) -> Optional[RemoteWorker]:
|
||||
"""Atomically update user-controlled policy and return the committed row."""
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE id = ?", (worker_id,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
if name is not None or enabled is not None or priority is not None:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE remote_workers
|
||||
SET name = COALESCE(?, name),
|
||||
enabled = COALESCE(?, enabled),
|
||||
priority = COALESCE(?, priority)
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
name,
|
||||
1 if enabled is True else 0 if enabled is False else None,
|
||||
max(0, min(100, int(priority))) if priority is not None else None,
|
||||
worker_id,
|
||||
),
|
||||
)
|
||||
row = conn.execute(
|
||||
"SELECT * FROM remote_workers WHERE id = ?", (worker_id,)
|
||||
).fetchone()
|
||||
return _row_to_worker(row) if row is not None else None
|
||||
|
||||
|
||||
def set_enabled(worker_id: str, enabled: bool) -> None:
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE remote_workers SET enabled = ? WHERE id = ?", (1 if enabled else 0, worker_id)
|
||||
)
|
||||
|
||||
|
||||
def set_priority(worker_id: str, priority: int) -> None:
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE remote_workers SET priority = ? WHERE id = ?",
|
||||
(max(0, min(100, int(priority))), worker_id),
|
||||
@@ -322,7 +510,7 @@ def set_priority(worker_id: str, priority: int) -> None:
|
||||
|
||||
|
||||
def rename(worker_id: str, name: str) -> None:
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
conn.execute("UPDATE remote_workers SET name = ? WHERE id = ?", (name, worker_id))
|
||||
|
||||
|
||||
@@ -334,7 +522,7 @@ def revoke(worker_id: str, *, now: Optional[float] = None) -> bool:
|
||||
who could simply enroll again.
|
||||
"""
|
||||
stamp = resolve(now)
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
cur = conn.execute(
|
||||
"UPDATE remote_workers SET revoked = 1, revoked_at = ?, enabled = 0 WHERE id = ?",
|
||||
(stamp, worker_id),
|
||||
@@ -345,15 +533,23 @@ def revoke(worker_id: str, *, now: Optional[float] = None) -> bool:
|
||||
|
||||
|
||||
def is_revoked(key_id: str) -> bool:
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT revoked FROM remote_workers WHERE key_id = ?", (key_id,)
|
||||
).fetchone()
|
||||
return bool(row and row["revoked"])
|
||||
|
||||
|
||||
def is_enabled(worker_id: str) -> bool:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT enabled FROM remote_workers WHERE id = ?", (worker_id,)
|
||||
).fetchone()
|
||||
return bool(row and row["enabled"])
|
||||
|
||||
|
||||
def grant_consent(worker_id: str, *, now: Optional[float] = None) -> None:
|
||||
with db_conn() as conn:
|
||||
with _AUTHORITY_LOCK, db_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE remote_workers SET consent_granted_at = ? WHERE id = ?",
|
||||
(resolve(now), worker_id),
|
||||
@@ -397,16 +593,20 @@ def authenticate(
|
||||
|
||||
__all__ = [
|
||||
"RemoteWorker",
|
||||
"authority_guard",
|
||||
"authenticate",
|
||||
"begin_session",
|
||||
"create_enrollment",
|
||||
"enroll_with_token",
|
||||
"enroll_worker",
|
||||
"get",
|
||||
"get_by_key_id",
|
||||
"grant_consent",
|
||||
"is_revoked",
|
||||
"is_enabled",
|
||||
"list_workers",
|
||||
"purge_expired_enrollments",
|
||||
"recover_enrollment_with_token",
|
||||
"redeem_enrollment",
|
||||
"rename",
|
||||
"revoke",
|
||||
@@ -414,4 +614,5 @@ __all__ = [
|
||||
"set_priority",
|
||||
"touch",
|
||||
"update_capabilities",
|
||||
"update_policy",
|
||||
]
|
||||
|
||||
+325
-24
@@ -27,14 +27,20 @@ be minutes.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import enum
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, fields
|
||||
from functools import partial
|
||||
from typing import Callable, Optional
|
||||
|
||||
from worker import deadlines as deadline_policy
|
||||
from worker import task_store
|
||||
from worker import registry, task_store
|
||||
from worker.async_utils import (
|
||||
to_thread_and_defer_cancellation,
|
||||
to_thread_and_drain_on_cancel,
|
||||
)
|
||||
from worker.breaker import Attribution
|
||||
from worker.capacity import WorkerCapacity
|
||||
from worker.clock import resolve
|
||||
@@ -100,6 +106,29 @@ class SchedulerStopped(RuntimeError):
|
||||
"""
|
||||
|
||||
|
||||
def _adopt_task_state(target: Task, source: Task) -> None:
|
||||
"""Commit a reconciled copy while preserving public task/attempt identities."""
|
||||
current_attempts = {attempt.attempt_id: attempt for attempt in target.attempts}
|
||||
adopted_attempts = []
|
||||
for source_attempt in source.attempts:
|
||||
target_attempt = current_attempts.get(source_attempt.attempt_id)
|
||||
if target_attempt is None:
|
||||
target_attempt = copy.deepcopy(source_attempt)
|
||||
else:
|
||||
for item in fields(Attempt):
|
||||
setattr(
|
||||
target_attempt,
|
||||
item.name,
|
||||
copy.deepcopy(getattr(source_attempt, item.name)),
|
||||
)
|
||||
adopted_attempts.append(target_attempt)
|
||||
|
||||
for item in fields(Task):
|
||||
if item.name != "attempts":
|
||||
setattr(target, item.name, copy.deepcopy(getattr(source, item.name)))
|
||||
target.attempts[:] = adopted_attempts
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Assignment:
|
||||
"""One task bound to one worker, ready to send."""
|
||||
@@ -110,6 +139,18 @@ class Assignment:
|
||||
deadlines: deadline_policy.Deadlines
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ReconciliationGeneration:
|
||||
"""Immutable task copies persisted before their live generation is adopted."""
|
||||
|
||||
worker_id: str
|
||||
stamp: float
|
||||
originals: tuple[Task, ...]
|
||||
snapshots: tuple[Task, ...]
|
||||
candidates: tuple[Task, ...]
|
||||
zombies: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""Owns the queue, the selection pipeline, and the deadline sweeper."""
|
||||
|
||||
@@ -137,6 +178,10 @@ class Scheduler:
|
||||
# task_id → when its progress was last written through. Cleared with
|
||||
# the task, so it cannot outlive what it describes.
|
||||
self._progress_saved_at: dict[str, float] = {}
|
||||
# Async producers serialize admission while durable input staging is
|
||||
# off-loop. This keeps queue/idempotency decisions atomic without
|
||||
# holding the event loop hostage to a multi-gigabyte copy and hash.
|
||||
self._submission_lock = asyncio.Lock()
|
||||
|
||||
# ── Persistence seam ──────────────────────────────────────────────────
|
||||
|
||||
@@ -281,12 +326,99 @@ class Scheduler:
|
||||
)
|
||||
if deadline_seconds:
|
||||
task.deadline_at = stamp + deadline_seconds
|
||||
self._tasks[task.task_id] = task
|
||||
if self._persist:
|
||||
task_store.create(task, now=stamp)
|
||||
persisted = task_store.create(task, now=stamp)
|
||||
if persisted.task_id != task.task_id:
|
||||
self._tasks.setdefault(persisted.task_id, persisted)
|
||||
return persisted
|
||||
task = persisted
|
||||
# Publish to the dispatcher only after durable admission succeeds. A
|
||||
# failed DB/input-stage write must not send work the API rejected.
|
||||
self._tasks[task.task_id] = task
|
||||
self._emit("queued", task)
|
||||
return task
|
||||
|
||||
async def submit_async(
|
||||
self,
|
||||
*,
|
||||
operation: str,
|
||||
engine: str,
|
||||
model_id: str,
|
||||
params: Optional[dict] = None,
|
||||
priority: PriorityClass = PriorityClass.INTERACTIVE,
|
||||
idempotency_key: Optional[str] = None,
|
||||
max_attempts: int = 3,
|
||||
deadline_seconds: Optional[float] = None,
|
||||
pinned_worker_id: Optional[str] = None,
|
||||
now: Optional[float] = None,
|
||||
) -> Task:
|
||||
"""Admit durably without hashing/copying task inputs on the app loop."""
|
||||
async with self._submission_lock:
|
||||
stamp = resolve(now)
|
||||
if idempotency_key:
|
||||
for existing in self._tasks.values():
|
||||
if existing.idempotency_key == idempotency_key:
|
||||
return existing
|
||||
if self._persist:
|
||||
stored, cancelled = await to_thread_and_defer_cancellation(
|
||||
task_store.get_by_idempotency_key, idempotency_key
|
||||
)
|
||||
if stored is not None:
|
||||
self._tasks.setdefault(stored.task_id, stored)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return stored
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
if self.queue_depth >= self.max_queue_depth:
|
||||
raise QueueFull(
|
||||
f"The remote task queue is full ({self.max_queue_depth} waiting). "
|
||||
"Wait for current work to finish, or add another worker."
|
||||
)
|
||||
|
||||
task = Task(
|
||||
task_id=uuid.uuid4().hex[:16],
|
||||
operation=operation,
|
||||
engine=engine,
|
||||
model_id=model_id,
|
||||
params=params or {},
|
||||
priority=priority,
|
||||
idempotency_key=idempotency_key,
|
||||
max_attempts=max_attempts,
|
||||
created_at=stamp,
|
||||
pinned_worker_id=pinned_worker_id,
|
||||
)
|
||||
if deadline_seconds:
|
||||
task.deadline_at = stamp + deadline_seconds
|
||||
if self._persist:
|
||||
persisted, cancelled = await to_thread_and_defer_cancellation(
|
||||
partial(task_store.create, task, now=stamp)
|
||||
)
|
||||
if persisted.task_id != task.task_id:
|
||||
self._tasks.setdefault(persisted.task_id, persisted)
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return persisted
|
||||
task = persisted
|
||||
if cancelled:
|
||||
# The durable copy cannot be interrupted midway. Settle it
|
||||
# before publication so cancellation can never dispatch a
|
||||
# task whose caller did not finish submitting it.
|
||||
task.cancel(reason="submission was cancelled", now=resolve())
|
||||
await to_thread_and_drain_on_cancel(
|
||||
partial(task_store.save, task, now=resolve())
|
||||
)
|
||||
self._tasks[task.task_id] = task
|
||||
self._emit("cancelled", task)
|
||||
raise asyncio.CancelledError
|
||||
|
||||
# The dispatcher only learns the task after input bytes and its DB
|
||||
# row are durable. A failed stage cannot escape as runnable work.
|
||||
self._tasks[task.task_id] = task
|
||||
self._emit("queued", task)
|
||||
return task
|
||||
|
||||
def adopt(self, task: Task) -> None:
|
||||
"""Take ownership of a task loaded from disk after a restart."""
|
||||
self._tasks[task.task_id] = task
|
||||
@@ -378,7 +510,11 @@ class Scheduler:
|
||||
for worker in self.pool:
|
||||
if task.pinned_worker_id and worker.worker_id != task.pinned_worker_id:
|
||||
continue
|
||||
if not worker.record.schedulable or worker.draining:
|
||||
if (
|
||||
not worker.record.schedulable
|
||||
or worker.draining
|
||||
or worker.registration_pending
|
||||
):
|
||||
continue
|
||||
if worker.stale(now=stamp):
|
||||
continue
|
||||
@@ -469,6 +605,13 @@ class Scheduler:
|
||||
empty or every waiting task is blocked on capacity. A task with no
|
||||
capable worker at all is failed here rather than left to age out.
|
||||
"""
|
||||
# Selection and binding are one authority read. A revoke/disable route
|
||||
# holds the same lock until its durable row and live pool record agree,
|
||||
# so another thread cannot bind work in that publication window.
|
||||
with registry.authority_guard():
|
||||
return self._next_assignment(now=now)
|
||||
|
||||
def _next_assignment(self, *, now: Optional[float] = None) -> Optional[Assignment]:
|
||||
stamp = resolve(now)
|
||||
for task in self._queued_order():
|
||||
try:
|
||||
@@ -700,9 +843,25 @@ class Scheduler:
|
||||
if epoch is not None and attempt.session_epoch != epoch:
|
||||
return False, task
|
||||
|
||||
committed, attempt = task.commit_result(
|
||||
attempt_id, result_ref=result_ref, session_epoch=epoch, now=stamp
|
||||
)
|
||||
if self._persist:
|
||||
# Prepare and persist the terminal generation before publishing it
|
||||
# to the live graph. A failed database write must leave the
|
||||
# redelivery path looking exactly as it did before this frame:
|
||||
# still in flight, still consuming capacity, and not yet credited
|
||||
# as a breaker success.
|
||||
candidate = copy.deepcopy(task)
|
||||
committed, _ = candidate.commit_result(
|
||||
attempt_id, result_ref=result_ref, session_epoch=epoch, now=stamp
|
||||
)
|
||||
if committed:
|
||||
task_store.commit_result(candidate, result_json=result, now=stamp)
|
||||
else:
|
||||
task_store.save(candidate, now=stamp)
|
||||
_adopt_task_state(task, candidate)
|
||||
else:
|
||||
committed, attempt = task.commit_result(
|
||||
attempt_id, result_ref=result_ref, session_epoch=epoch, now=stamp
|
||||
)
|
||||
self._release_slot(task, attempt, now=stamp)
|
||||
worker = self.pool.get(attempt.worker_id)
|
||||
if worker is not None and committed:
|
||||
@@ -711,10 +870,6 @@ class Scheduler:
|
||||
WorkerCapacity.slot_key(task.engine, task.model_id),
|
||||
now=stamp,
|
||||
)
|
||||
if committed and self._persist:
|
||||
task_store.commit_result(task, result_json=result, now=stamp)
|
||||
elif self._persist:
|
||||
self._save(task, now=stamp)
|
||||
self._emit("completed" if committed else "duplicate", task)
|
||||
return committed, task
|
||||
|
||||
@@ -779,28 +934,85 @@ class Scheduler:
|
||||
This is where the duplicate-execution bug would live if a disconnect
|
||||
were treated as a failure: the worker may be seconds from delivering.
|
||||
"""
|
||||
generation = self.prepare_disconnected(worker_id, now=now)
|
||||
try:
|
||||
# A disconnect is one generation, just like reconnect. If any
|
||||
# write fails, neither the database nor the live task graph may
|
||||
# contain a half-marked set of attempts.
|
||||
self.persist_reconciliation(generation)
|
||||
return self.apply_disconnected(generation)
|
||||
finally:
|
||||
# A failed persistence write must not leave a dead stream's pool
|
||||
# entry schedulable.
|
||||
self.pool.disconnect(worker_id)
|
||||
|
||||
def prepare_disconnected(
|
||||
self,
|
||||
worker_id: str,
|
||||
*,
|
||||
now: Optional[float] = None,
|
||||
include_task_ids: Optional[set[str]] = None,
|
||||
) -> _ReconciliationGeneration:
|
||||
"""Build, but do not publish, one disconnected task generation."""
|
||||
stamp = resolve(now)
|
||||
affected: list[Task] = []
|
||||
for task in self._tasks.values():
|
||||
included = include_task_ids or set()
|
||||
originals = tuple(
|
||||
task
|
||||
for task in self._tasks.values()
|
||||
if task.task_id in included
|
||||
or (
|
||||
task.active_attempt is not None
|
||||
and task.active_attempt.worker_id == worker_id
|
||||
)
|
||||
)
|
||||
snapshots = tuple(copy.deepcopy(task) for task in originals)
|
||||
candidates = tuple(copy.deepcopy(task) for task in snapshots)
|
||||
for task in candidates:
|
||||
attempt = task.active_attempt
|
||||
if attempt is None or attempt.worker_id != worker_id:
|
||||
continue
|
||||
grace = deadline_policy.default_grace_seconds(task.operation)
|
||||
task.mark_disconnected(attempt.attempt_id, grace_seconds=grace, now=stamp)
|
||||
affected.append(task)
|
||||
self._save(task, now=stamp)
|
||||
self._emit("worker_lost", task)
|
||||
self.pool.disconnect(worker_id)
|
||||
return affected
|
||||
task.mark_disconnected(
|
||||
attempt.attempt_id, grace_seconds=grace, now=stamp
|
||||
)
|
||||
return _ReconciliationGeneration(
|
||||
worker_id=worker_id,
|
||||
stamp=stamp,
|
||||
originals=originals,
|
||||
snapshots=snapshots,
|
||||
candidates=candidates,
|
||||
)
|
||||
|
||||
def on_reconnected(
|
||||
self, worker_id: str, *, in_flight: set[str], now: Optional[float] = None
|
||||
self,
|
||||
worker_id: str,
|
||||
*,
|
||||
in_flight: set[str],
|
||||
now: Optional[float] = None,
|
||||
before_persist: Optional[Callable[[object], None]] = None,
|
||||
) -> list[str]:
|
||||
"""Reconcile against what the worker says it is running.
|
||||
|
||||
Returns attempt ids the worker should cancel — work we have already
|
||||
written off, which it must stop burning a GPU on.
|
||||
"""
|
||||
generation = self.prepare_reconnected(
|
||||
worker_id, in_flight=in_flight, now=now
|
||||
)
|
||||
self.persist_reconciliation(
|
||||
generation, before_persist=before_persist
|
||||
)
|
||||
return self.apply_reconnected(generation)
|
||||
|
||||
def prepare_reconnected(
|
||||
self,
|
||||
worker_id: str,
|
||||
*,
|
||||
in_flight: set[str],
|
||||
now: Optional[float] = None,
|
||||
include_task_ids: Optional[set[str]] = None,
|
||||
) -> _ReconciliationGeneration:
|
||||
"""Build, but do not publish, one reconnect reconciliation."""
|
||||
stamp = resolve(now)
|
||||
known = {
|
||||
attempt.attempt_id: attempt
|
||||
@@ -813,7 +1025,20 @@ class Scheduler:
|
||||
for attempt_id in in_flight
|
||||
if attempt_id not in known or known[attempt_id].state.terminal
|
||||
]
|
||||
for task in list(self._tasks.values()):
|
||||
included = include_task_ids or set()
|
||||
originals = tuple(
|
||||
task
|
||||
for task in self._tasks.values()
|
||||
if task.task_id in included
|
||||
or (
|
||||
not task.state.terminal
|
||||
and task.active_attempt is not None
|
||||
and task.active_attempt.worker_id == worker_id
|
||||
)
|
||||
)
|
||||
snapshots = tuple(copy.deepcopy(task) for task in originals)
|
||||
candidates = tuple(copy.deepcopy(task) for task in snapshots)
|
||||
for task in candidates:
|
||||
if task.state.terminal:
|
||||
continue
|
||||
reconcile(
|
||||
@@ -826,8 +1051,84 @@ class Scheduler:
|
||||
resume_lease_seconds=self._budget_for(task).progress_lease_seconds,
|
||||
now=stamp,
|
||||
)
|
||||
self._save(task, now=stamp)
|
||||
return zombies
|
||||
return _ReconciliationGeneration(
|
||||
worker_id=worker_id,
|
||||
stamp=stamp,
|
||||
originals=originals,
|
||||
snapshots=snapshots,
|
||||
candidates=candidates,
|
||||
zombies=tuple(zombies),
|
||||
)
|
||||
|
||||
def persist_reconciliation(
|
||||
self,
|
||||
generation: _ReconciliationGeneration,
|
||||
*,
|
||||
before_persist: Optional[Callable[[object], None]] = None,
|
||||
) -> None:
|
||||
"""Persist a prepared generation without touching the live task graph."""
|
||||
# Nothing in the live graph changes until the complete generation is
|
||||
# durable. A failed write therefore cannot queue duplicate execution
|
||||
# while the old worker is still finishing the original attempt.
|
||||
if self._persist or before_persist is not None:
|
||||
task_store.save_many(
|
||||
generation.candidates if self._persist else [],
|
||||
now=generation.stamp,
|
||||
before_save=before_persist,
|
||||
)
|
||||
|
||||
def reconciliation_is_current(
|
||||
self, generation: _ReconciliationGeneration
|
||||
) -> bool:
|
||||
"""Did the live graph remain unchanged while persistence was off-loop?"""
|
||||
return all(
|
||||
self._tasks.get(original.task_id) is original
|
||||
and original == snapshot
|
||||
for original, snapshot in zip(
|
||||
generation.originals, generation.snapshots
|
||||
)
|
||||
)
|
||||
|
||||
def apply_reconnected(
|
||||
self, generation: _ReconciliationGeneration
|
||||
) -> list[str]:
|
||||
"""Publish a durable reconnect generation on the scheduler's loop."""
|
||||
if not self.reconciliation_is_current(generation):
|
||||
raise RuntimeError("reconnect reconciliation generation changed")
|
||||
changed: list[Task] = []
|
||||
for original, candidate in zip(
|
||||
generation.originals, generation.candidates
|
||||
):
|
||||
if original != candidate:
|
||||
changed.append(original)
|
||||
_adopt_task_state(original, candidate)
|
||||
for task in changed:
|
||||
self._emit(
|
||||
"failed"
|
||||
if task.state.terminal
|
||||
else "requeued"
|
||||
if task.state is TaskState.QUEUED
|
||||
else "resumed",
|
||||
task,
|
||||
)
|
||||
return list(generation.zombies)
|
||||
|
||||
def apply_disconnected(
|
||||
self, generation: _ReconciliationGeneration
|
||||
) -> list[Task]:
|
||||
"""Publish a durable disconnect generation on the scheduler's loop."""
|
||||
if not self.reconciliation_is_current(generation):
|
||||
raise RuntimeError("disconnect reconciliation generation changed")
|
||||
changed: list[Task] = []
|
||||
for original, candidate in zip(
|
||||
generation.originals, generation.candidates
|
||||
):
|
||||
if original != candidate:
|
||||
changed.append(original)
|
||||
_adopt_task_state(original, candidate)
|
||||
for task in changed:
|
||||
self._emit("worker_lost", task)
|
||||
return list(generation.originals)
|
||||
|
||||
def cancel(self, task_id: str, *, reason: str = "cancelled", now: Optional[float] = None) -> bool:
|
||||
task = self._tasks.get(task_id)
|
||||
|
||||
+124
-19
@@ -15,6 +15,7 @@ import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from worker.clock import resolve
|
||||
|
||||
@@ -24,10 +25,52 @@ logger = logging.getLogger("omnivoice.worker")
|
||||
_SWEEP_INTERVAL_SECONDS = 5.0
|
||||
# How often the dispatcher looks for queued work it can place.
|
||||
_DISPATCH_INTERVAL_SECONDS = 1.0
|
||||
# Finished rows and their artifacts remain inspectable for a week, then leave
|
||||
# in bounded batches. The loop keeps a long-lived control plane from growing
|
||||
# forever while the startup pass covers apps that are rarely left open.
|
||||
_ARTIFACT_GC_INTERVAL_SECONDS = 60 * 60.0
|
||||
_ARTIFACT_GC_BATCH_SIZE = 250
|
||||
|
||||
DEFAULT_PORT = 7443
|
||||
|
||||
|
||||
class EndpointCertificateError(ValueError):
|
||||
"""An advertised endpoint is not usable with the live TLS certificate."""
|
||||
|
||||
|
||||
def _endpoint_host(endpoint: str) -> str:
|
||||
"""Extract the hostname from the gRPC ``host:port`` enrollment target."""
|
||||
target = (endpoint or "").strip()
|
||||
try:
|
||||
parsed = urlsplit(f"//{target}")
|
||||
host = parsed.hostname
|
||||
port = parsed.port
|
||||
except ValueError as exc:
|
||||
raise EndpointCertificateError(
|
||||
f"Worker endpoint must be host:port; got {endpoint!r}."
|
||||
) from exc
|
||||
if (
|
||||
not host
|
||||
or port is None
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
or parsed.path
|
||||
or parsed.query
|
||||
or parsed.fragment
|
||||
):
|
||||
raise EndpointCertificateError(
|
||||
f"Worker endpoint must be host:port; got {endpoint!r}."
|
||||
)
|
||||
return host
|
||||
|
||||
|
||||
def _format_endpoint(host: str, port: int) -> str:
|
||||
clean_host = (host or "").strip().strip("[]")
|
||||
if ":" in clean_host:
|
||||
clean_host = f"[{clean_host}]"
|
||||
return f"{clean_host}:{port}"
|
||||
|
||||
|
||||
def remote_workers_enabled() -> bool:
|
||||
"""Opt-in gate. Off unless the user turned it on.
|
||||
|
||||
@@ -93,6 +136,7 @@ class ControlPlane:
|
||||
self._server = None
|
||||
self._tasks: list[asyncio.Task] = []
|
||||
self._started = False
|
||||
self._lifecycle_lock = asyncio.Lock()
|
||||
self.startup_error: Optional[str] = None
|
||||
# The port we actually bound, which is not necessarily the configured
|
||||
# one — an enrollment token carries this, so advertising the config
|
||||
@@ -108,6 +152,10 @@ class ControlPlane:
|
||||
return self.credentials.fingerprint if self.credentials else ""
|
||||
|
||||
async def start(self, *, port: Optional[int] = None) -> None:
|
||||
async with self._lifecycle_lock:
|
||||
await self._start(port=port)
|
||||
|
||||
async def _start(self, *, port: Optional[int] = None) -> None:
|
||||
if self._started:
|
||||
return
|
||||
# Imported here, not at module scope: a user who never enables remote
|
||||
@@ -117,25 +165,39 @@ class ControlPlane:
|
||||
from worker.scheduler import Scheduler # noqa: PLC0415
|
||||
from worker.transport.server import WorkerServicer, serve # noqa: PLC0415
|
||||
|
||||
locations = paths()
|
||||
os.makedirs(locations["root"], exist_ok=True)
|
||||
self.credentials = tls.load_or_create(
|
||||
locations["certificate"], locations["private_key"]
|
||||
)
|
||||
self.pool = WorkerPool()
|
||||
self.scheduler = Scheduler(self.pool)
|
||||
# Recover anything that was in flight when the app last quit. The
|
||||
# workers holding those tasks may still be rendering.
|
||||
self.scheduler.restore()
|
||||
|
||||
self.servicer = WorkerServicer(
|
||||
self.scheduler,
|
||||
self.pool,
|
||||
artifact_dir=locations["artifacts"],
|
||||
cert_fingerprint=self.credentials.fingerprint,
|
||||
)
|
||||
self._port = port or control_port()
|
||||
try:
|
||||
locations = paths()
|
||||
os.makedirs(locations["root"], exist_ok=True)
|
||||
wanted_hostnames = tls.default_hostnames()
|
||||
advertised_host = _endpoint_host(self.default_endpoint())
|
||||
if advertised_host not in wanted_hostnames:
|
||||
wanted_hostnames.append(advertised_host)
|
||||
self.credentials = tls.load_or_create(
|
||||
locations["certificate"],
|
||||
locations["private_key"],
|
||||
hostnames=wanted_hostnames,
|
||||
)
|
||||
self.pool = WorkerPool()
|
||||
self.scheduler = Scheduler(self.pool)
|
||||
try:
|
||||
await self._artifact_gc_once(
|
||||
locations["artifacts"], retry_upload_parts=False
|
||||
)
|
||||
except Exception:
|
||||
# Retention is best-effort. A locked file or temporarily busy
|
||||
# database must not make remote compute unavailable.
|
||||
logger.exception("Remote worker startup artifact sweep failed")
|
||||
# Recover anything that was in flight when the app last quit. The
|
||||
# workers holding those tasks may still be rendering.
|
||||
self.scheduler.restore()
|
||||
|
||||
self.servicer = WorkerServicer(
|
||||
self.scheduler,
|
||||
self.pool,
|
||||
artifact_dir=locations["artifacts"],
|
||||
cert_fingerprint=self.credentials.fingerprint,
|
||||
)
|
||||
self._server = await serve(
|
||||
self.servicer,
|
||||
port=self._port,
|
||||
@@ -148,12 +210,20 @@ class ControlPlane:
|
||||
self._tasks = [
|
||||
asyncio.create_task(self._sweep_loop(), name="worker-sweep"),
|
||||
asyncio.create_task(self._dispatch_loop(), name="worker-dispatch"),
|
||||
asyncio.create_task(
|
||||
self._artifact_gc_loop(locations["artifacts"]),
|
||||
name="worker-artifact-gc",
|
||||
),
|
||||
]
|
||||
self._started = True
|
||||
self.startup_error = None
|
||||
logger.info("Remote worker control plane started on port %d", self._port)
|
||||
|
||||
async def stop(self) -> None:
|
||||
async with self._lifecycle_lock:
|
||||
await self._stop()
|
||||
|
||||
async def _stop(self) -> None:
|
||||
for task in self._tasks:
|
||||
task.cancel()
|
||||
if self._tasks:
|
||||
@@ -199,6 +269,29 @@ class ControlPlane:
|
||||
except Exception:
|
||||
logger.exception("Worker sweep failed")
|
||||
|
||||
async def _artifact_gc_once(
|
||||
self, artifact_root: str, *, retry_upload_parts: bool = True
|
||||
) -> int:
|
||||
"""Purge one bounded retention batch without blocking the event loop."""
|
||||
from worker import task_store # noqa: PLC0415
|
||||
|
||||
removed = await asyncio.to_thread(
|
||||
task_store.purge_finished,
|
||||
root=artifact_root,
|
||||
limit=_ARTIFACT_GC_BATCH_SIZE,
|
||||
)
|
||||
if retry_upload_parts and self.servicer is not None:
|
||||
self.servicer.sweep_orphaned_upload_parts()
|
||||
return removed
|
||||
|
||||
async def _artifact_gc_loop(self, artifact_root: str) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(_ARTIFACT_GC_INTERVAL_SECONDS)
|
||||
try:
|
||||
await self._artifact_gc_once(artifact_root)
|
||||
except Exception:
|
||||
logger.exception("Remote worker artifact sweep failed")
|
||||
|
||||
async def _dispatch_loop(self) -> None:
|
||||
"""Place queued work on eligible workers.
|
||||
|
||||
@@ -233,10 +326,21 @@ class ControlPlane:
|
||||
|
||||
def create_enrollment(self, *, endpoint: str = "", label: str = "", ttl_seconds: int = 900):
|
||||
"""Mint a join token carrying this control plane's fingerprint."""
|
||||
from worker import tls # noqa: PLC0415
|
||||
from worker import registry # noqa: PLC0415
|
||||
|
||||
advertised_endpoint = endpoint.strip() or self.default_endpoint()
|
||||
advertised_host = _endpoint_host(advertised_endpoint)
|
||||
if self.credentials is not None and not tls.covers(
|
||||
self.credentials, advertised_host
|
||||
):
|
||||
raise EndpointCertificateError(
|
||||
f"The running certificate does not cover {advertised_host!r}. "
|
||||
"Set OMNIVOICE_WORKER_ENDPOINT_HOST to that hostname and restart "
|
||||
"VoiceStudio before creating this enrollment."
|
||||
)
|
||||
return registry.create_enrollment(
|
||||
endpoint=endpoint or self.default_endpoint(),
|
||||
endpoint=advertised_endpoint,
|
||||
cert_fingerprint=self.fingerprint,
|
||||
label=label,
|
||||
ttl_seconds=ttl_seconds,
|
||||
@@ -262,7 +366,7 @@ class ControlPlane:
|
||||
or tls.primary_ip()
|
||||
or "127.0.0.1"
|
||||
)
|
||||
return f"{host}:{self._port or control_port()}"
|
||||
return _format_endpoint(host, self._port or control_port())
|
||||
|
||||
def snapshot(self, *, now: Optional[float] = None) -> dict:
|
||||
"""Everything the workers UI needs in one call."""
|
||||
@@ -380,6 +484,7 @@ async def stop() -> None:
|
||||
__all__ = [
|
||||
"ControlPlane",
|
||||
"DEFAULT_PORT",
|
||||
"EndpointCertificateError",
|
||||
"control_plane",
|
||||
"control_port",
|
||||
"paths",
|
||||
|
||||
+305
-66
@@ -17,6 +17,7 @@ flips the task to completed, so the ack can only follow a durable fact.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
@@ -25,7 +26,8 @@ import os
|
||||
import re
|
||||
import shutil
|
||||
import time
|
||||
from typing import Iterable, Iterator, Optional
|
||||
import uuid
|
||||
from typing import Callable, Iterable, Iterator, Optional
|
||||
|
||||
from core.db import db_conn
|
||||
from core.path_security import UnsafePath, resolve_within, safe_filename
|
||||
@@ -133,6 +135,7 @@ INPUTS_PARAM_KEY = "inputs"
|
||||
|
||||
_HASH_CHUNK_BYTES = 1024 * 1024
|
||||
_SAFE_EXTENSION = re.compile(r"^\.[A-Za-z0-9]{1,8}$")
|
||||
_CONTENT_ARTIFACT = re.compile(r"^([0-9a-f]{64})(?:\.[A-Za-z0-9]{1,8})?$")
|
||||
|
||||
|
||||
class InputStagingError(RuntimeError):
|
||||
@@ -143,6 +146,60 @@ class InputStagingError(RuntimeError):
|
||||
"""
|
||||
|
||||
|
||||
def _fsync_parent_directory(directory: str) -> None:
|
||||
"""Persist directory entry changes where the platform supports it."""
|
||||
directory_flag = getattr(os, "O_DIRECTORY", None)
|
||||
if directory_flag is None:
|
||||
return
|
||||
unsupported = {
|
||||
errno.EINVAL,
|
||||
getattr(errno, "ENOTSUP", errno.EINVAL),
|
||||
getattr(errno, "EOPNOTSUPP", errno.EINVAL),
|
||||
}
|
||||
try:
|
||||
descriptor = os.open(directory, os.O_RDONLY | directory_flag)
|
||||
except OSError as exc:
|
||||
if exc.errno in unsupported:
|
||||
return
|
||||
raise
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
except OSError as exc:
|
||||
if exc.errno not in unsupported:
|
||||
raise
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def _fsync_file(path: str) -> None:
|
||||
with open(path, "r+b") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def _durable_makedirs(directory: str) -> None:
|
||||
"""Create a directory hierarchy and persist each parent entry."""
|
||||
target = os.path.abspath(directory)
|
||||
missing: list[str] = []
|
||||
current = target
|
||||
while not os.path.isdir(current):
|
||||
if os.path.exists(current):
|
||||
raise NotADirectoryError(current)
|
||||
missing.append(current)
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
for path in reversed(missing):
|
||||
try:
|
||||
os.mkdir(path)
|
||||
except FileExistsError:
|
||||
if not os.path.isdir(path):
|
||||
raise
|
||||
_fsync_parent_directory(os.path.dirname(path) or ".")
|
||||
if not missing:
|
||||
_fsync_parent_directory(os.path.dirname(target) or ".")
|
||||
|
||||
|
||||
def artifact_root(*, create_dir: bool = True) -> str:
|
||||
"""The directory the control plane serves artifacts from.
|
||||
|
||||
@@ -154,7 +211,7 @@ def artifact_root(*, create_dir: bool = True) -> str:
|
||||
|
||||
root = paths()["artifacts"]
|
||||
if create_dir:
|
||||
os.makedirs(os.path.join(root, INPUTS_DIRNAME), exist_ok=True)
|
||||
_durable_makedirs(os.path.join(root, INPUTS_DIRNAME))
|
||||
return root
|
||||
|
||||
|
||||
@@ -183,7 +240,40 @@ def _digest(path: str) -> tuple[str, int]:
|
||||
return digest.hexdigest(), size
|
||||
|
||||
|
||||
def stage_input(source: str, *, root: Optional[str] = None, now: Optional[float] = None) -> dict:
|
||||
def _staged_entry_matches(path: str, entry: dict) -> bool:
|
||||
"""Verify staged bytes against both metadata and their content address."""
|
||||
if not os.path.isfile(path):
|
||||
return False
|
||||
artifact_id = str(entry.get("artifact_id") or "")
|
||||
portable_name = artifact_id.replace("\\", "/").rsplit("/", 1)[-1]
|
||||
named = _CONTENT_ARTIFACT.fullmatch(portable_name)
|
||||
if named is None:
|
||||
return False
|
||||
try:
|
||||
actual_digest, actual_size = _digest(path)
|
||||
except OSError:
|
||||
return False
|
||||
recorded_digest = str(entry.get("sha256") or "").strip().lower()
|
||||
recorded_size = entry.get("size_bytes")
|
||||
if recorded_digest and actual_digest != recorded_digest:
|
||||
return False
|
||||
if recorded_size is not None:
|
||||
try:
|
||||
if actual_size != int(recorded_size):
|
||||
return False
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if actual_digest != named.group(1):
|
||||
return False
|
||||
# Backfill metadata on a legacy row once its content address proves it.
|
||||
entry["sha256"] = actual_digest
|
||||
entry["size_bytes"] = actual_size
|
||||
return True
|
||||
|
||||
|
||||
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
|
||||
@@ -195,27 +285,58 @@ def stage_input(source: str, *, root: Optional[str] = None, now: Optional[float]
|
||||
try:
|
||||
digest, size = _digest(source)
|
||||
except OSError as exc:
|
||||
raise InputStagingError(f"Could not read the task input {source!r}: {exc}") from 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
|
||||
raise InputStagingError(
|
||||
f"Refusing to stage {source!r} outside the artifact store"
|
||||
) from exc
|
||||
|
||||
partial = destination.with_name(
|
||||
f".{destination.name}.{uuid.uuid4().hex}.part"
|
||||
)
|
||||
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")
|
||||
_durable_makedirs(str(destination.parent))
|
||||
expected = {
|
||||
"artifact_id": artifact_id,
|
||||
"sha256": digest,
|
||||
"size_bytes": size,
|
||||
}
|
||||
if not _staged_entry_matches(str(destination), expected):
|
||||
shutil.copyfile(source, partial)
|
||||
copied_digest, copied_size = _digest(str(partial))
|
||||
if copied_digest != digest or copied_size != size:
|
||||
raise InputStagingError(
|
||||
f"The task input {source!r} changed while it was being staged."
|
||||
)
|
||||
_fsync_file(str(partial))
|
||||
os.replace(partial, destination)
|
||||
_fsync_parent_directory(str(destination.parent))
|
||||
# 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))
|
||||
_fsync_file(str(destination))
|
||||
_fsync_parent_directory(str(destination.parent))
|
||||
except OSError as exc:
|
||||
raise InputStagingError(f"Could not stage the task input {source!r}: {exc}") from exc
|
||||
raise InputStagingError(
|
||||
f"Could not stage the task input {source!r}: {exc}"
|
||||
) from exc
|
||||
finally:
|
||||
try:
|
||||
os.remove(partial)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError:
|
||||
logger.debug(
|
||||
"Could not remove the staged-input partial %s",
|
||||
partial,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
filename = os.path.basename(str(source)) or f"{digest}{_extension(source)}"
|
||||
return {
|
||||
@@ -254,8 +375,13 @@ def ensure_staged(
|
||||
"""
|
||||
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:
|
||||
entries: list[dict] = (
|
||||
[e for e in recorded if isinstance(e, dict)]
|
||||
if isinstance(recorded, list)
|
||||
else []
|
||||
)
|
||||
base = root or artifact_root()
|
||||
if entries:
|
||||
# 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.
|
||||
@@ -263,7 +389,10 @@ def ensure_staged(
|
||||
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())
|
||||
path = resolve_within(base, artifact_id)
|
||||
available = bool(
|
||||
artifact_id and _staged_entry_matches(str(path), entry)
|
||||
)
|
||||
except UnsafePath:
|
||||
available = False
|
||||
if available:
|
||||
@@ -271,7 +400,7 @@ def ensure_staged(
|
||||
continue
|
||||
source = str(entry.get("source") or "")
|
||||
if source and os.path.isfile(source):
|
||||
replacement = stage_input(source, root=root, now=now)
|
||||
replacement = stage_input(source, root=base, now=now)
|
||||
replacement.update(key=entry.get("key"), index=entry.get("index"))
|
||||
refreshed.append(replacement)
|
||||
else:
|
||||
@@ -289,7 +418,7 @@ def ensure_staged(
|
||||
# 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 = stage_input(value, root=base, now=now)
|
||||
entry["key"] = key
|
||||
entry["index"] = index
|
||||
entries.append(entry)
|
||||
@@ -345,6 +474,63 @@ def _referenced_artifacts(conn) -> set[str]:
|
||||
return referenced
|
||||
|
||||
|
||||
def _purge_result_directories(
|
||||
task_ids: Iterable[str], *, root: Optional[str] = None
|
||||
) -> tuple[list[str], int]:
|
||||
"""Delete result directories before their owning rows become unreachable."""
|
||||
task_ids = list(task_ids)
|
||||
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 task_ids, 0
|
||||
|
||||
cleaned: list[str] = []
|
||||
removed = 0
|
||||
for task_id in task_ids:
|
||||
try:
|
||||
path = resolve_within(base, safe_filename(task_id))
|
||||
except UnsafePath:
|
||||
continue
|
||||
if not os.path.exists(path):
|
||||
try:
|
||||
_fsync_parent_directory(base)
|
||||
except OSError:
|
||||
logger.debug(
|
||||
"Could not persist task artifact cleanup at %s",
|
||||
base,
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
cleaned.append(task_id)
|
||||
continue
|
||||
if not os.path.isdir(path):
|
||||
logger.warning("Refusing to purge non-directory task artifact %s", path)
|
||||
continue
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
except OSError:
|
||||
logger.debug("Could not purge task artifacts at %s", path, exc_info=True)
|
||||
continue
|
||||
try:
|
||||
_fsync_parent_directory(base)
|
||||
except OSError:
|
||||
# The bytes are gone from this process's view, but the directory
|
||||
# deletion is not a crash-durable fact yet. Keep the DB row as the
|
||||
# retry index until a later sweep can establish that barrier.
|
||||
logger.debug(
|
||||
"Could not persist task artifact cleanup at %s",
|
||||
base,
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
cleaned.append(task_id)
|
||||
removed += 1
|
||||
return cleaned, removed
|
||||
|
||||
|
||||
def purge_artifacts(
|
||||
task_ids: Iterable[str], referenced: set[str], *, cutoff: float, root: Optional[str] = None
|
||||
) -> int:
|
||||
@@ -356,7 +542,7 @@ def purge_artifacts(
|
||||
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
|
||||
_cleaned, removed = _purge_result_directories(task_ids, root=root)
|
||||
try:
|
||||
base = root or artifact_root(create_dir=False)
|
||||
except Exception: # pragma: no cover — no data dir at all
|
||||
@@ -365,15 +551,6 @@ def purge_artifacts(
|
||||
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)
|
||||
@@ -484,6 +661,25 @@ def _upsert_attempts(conn, task: Task) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _save_with_conn(conn, task: Task, *, stamp: float) -> None:
|
||||
conn.execute(
|
||||
"UPDATE remote_tasks SET state=?, excluded_json=?, error_json=?, result_ref=?, "
|
||||
"updated_at=?, deadline_at=?, finished_at=?, pinned_worker_id=? WHERE id=?",
|
||||
(
|
||||
task.state.value,
|
||||
json.dumps(sorted(task.excluded_workers)),
|
||||
_dump_error(task.error),
|
||||
task.result_ref,
|
||||
stamp,
|
||||
task.deadline_at,
|
||||
task.finished_at,
|
||||
task.pinned_worker_id,
|
||||
task.task_id,
|
||||
),
|
||||
)
|
||||
_upsert_attempts(conn, task)
|
||||
|
||||
|
||||
def save(task: Task, *, now: Optional[float] = None) -> None:
|
||||
"""Write the whole task + attempt graph.
|
||||
|
||||
@@ -492,22 +688,22 @@ def save(task: Task, *, now: Optional[float] = None) -> None:
|
||||
"""
|
||||
stamp = resolve(now)
|
||||
with db_conn() as conn:
|
||||
conn.execute(
|
||||
"UPDATE remote_tasks SET state=?, excluded_json=?, error_json=?, result_ref=?, "
|
||||
"updated_at=?, deadline_at=?, finished_at=?, pinned_worker_id=? WHERE id=?",
|
||||
(
|
||||
task.state.value,
|
||||
json.dumps(sorted(task.excluded_workers)),
|
||||
_dump_error(task.error),
|
||||
task.result_ref,
|
||||
stamp,
|
||||
task.deadline_at,
|
||||
task.finished_at,
|
||||
task.pinned_worker_id,
|
||||
task.task_id,
|
||||
),
|
||||
)
|
||||
_upsert_attempts(conn, task)
|
||||
_save_with_conn(conn, task, stamp=stamp)
|
||||
|
||||
|
||||
def save_many(
|
||||
tasks: Iterable[Task],
|
||||
*,
|
||||
now: Optional[float] = None,
|
||||
before_save: Optional[Callable[[object], None]] = None,
|
||||
) -> None:
|
||||
"""Persist one reconciliation generation atomically."""
|
||||
stamp = resolve(now)
|
||||
with db_conn() as conn:
|
||||
if before_save is not None:
|
||||
before_save(conn)
|
||||
for task in tasks:
|
||||
_save_with_conn(conn, task, stamp=stamp)
|
||||
|
||||
|
||||
def commit_result(
|
||||
@@ -584,23 +780,30 @@ def load_unfinished() -> list[Task]:
|
||||
still be rendering, and reconciliation decides each one's fate once the
|
||||
workers reconnect.
|
||||
"""
|
||||
live = ", ".join(f"'{s.value}'" for s in TaskState if not s.terminal)
|
||||
live = json.dumps([s.value for s in TaskState if not s.terminal])
|
||||
with db_conn() as conn:
|
||||
rows = conn.execute(
|
||||
f"SELECT * FROM remote_tasks WHERE state IN ({live}) ORDER BY priority ASC, created_at ASC"
|
||||
"""
|
||||
SELECT * FROM remote_tasks
|
||||
WHERE state IN (SELECT value FROM json_each(?))
|
||||
ORDER BY priority ASC, created_at ASC
|
||||
""",
|
||||
(live,),
|
||||
).fetchall()
|
||||
return [_row_to_task(r, _attempts_for(conn, r["id"])) for r in rows]
|
||||
|
||||
|
||||
def list_tasks(*, states: Optional[Iterable[TaskState]] = None, limit: int = 100) -> list[Task]:
|
||||
sql = "SELECT * FROM remote_tasks"
|
||||
params: list = []
|
||||
if states:
|
||||
placeholders = ", ".join("?" for _ in states)
|
||||
sql += f" WHERE state IN ({placeholders})"
|
||||
params.extend(s.value for s in states)
|
||||
sql += " ORDER BY created_at DESC LIMIT ?"
|
||||
params.append(limit)
|
||||
sql = """
|
||||
SELECT * FROM remote_tasks
|
||||
WHERE state IN (SELECT value FROM json_each(?))
|
||||
ORDER BY created_at DESC LIMIT ?
|
||||
"""
|
||||
params = (json.dumps([s.value for s in states]), limit)
|
||||
else:
|
||||
sql = "SELECT * FROM remote_tasks ORDER BY created_at DESC LIMIT ?"
|
||||
params = (limit,)
|
||||
with db_conn() as conn:
|
||||
rows = conn.execute(sql, params).fetchall()
|
||||
return [_row_to_task(r, _attempts_for(conn, r["id"])) for r in rows]
|
||||
@@ -611,6 +814,7 @@ def purge_finished(
|
||||
older_than_seconds: float = 7 * 24 * 3600,
|
||||
now: Optional[float] = None,
|
||||
root: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> int:
|
||||
"""Drop old finished tasks — rows *and* the bytes they own.
|
||||
|
||||
@@ -619,31 +823,66 @@ def purge_finished(
|
||||
audio. Neither was ever deleted, so the feature grew the user's disk for
|
||||
as long as they used it.
|
||||
"""
|
||||
if limit is not None and limit <= 0:
|
||||
return 0
|
||||
cutoff = resolve(now) - older_than_seconds
|
||||
terminal = ", ".join(f"'{s.value}'" for s in TaskState if s.terminal)
|
||||
terminal = json.dumps([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,),
|
||||
"""
|
||||
SELECT id FROM remote_tasks
|
||||
WHERE state IN (SELECT value FROM json_each(?))
|
||||
AND finished_at < ?
|
||||
ORDER BY finished_at ASC, id ASC
|
||||
LIMIT ?
|
||||
""",
|
||||
(terminal, cutoff, -1 if limit is None else int(limit)),
|
||||
).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 < ?)",
|
||||
(cutoff,),
|
||||
)
|
||||
cur = conn.execute(
|
||||
f"DELETE FROM remote_tasks WHERE state IN ({terminal}) AND finished_at < ?", (cutoff,)
|
||||
)
|
||||
removed = cur.rowcount
|
||||
# Results are attempt-scoped. Delete them before their task rows, so a
|
||||
# crash or transient Windows lock cannot erase the only index from which a
|
||||
# future sweep could find those bytes.
|
||||
cleaned, _artifacts_removed = _purge_result_directories(doomed, root=root)
|
||||
with db_conn() as conn:
|
||||
eligible: list[str] = []
|
||||
if cleaned:
|
||||
eligible = [
|
||||
row["id"]
|
||||
for row in conn.execute(
|
||||
"""
|
||||
SELECT id FROM remote_tasks
|
||||
WHERE id IN (SELECT value FROM json_each(?))
|
||||
AND state IN (SELECT value FROM json_each(?))
|
||||
AND finished_at < ?
|
||||
""",
|
||||
(json.dumps(cleaned), terminal, cutoff),
|
||||
).fetchall()
|
||||
]
|
||||
removed = 0
|
||||
if eligible:
|
||||
conn.execute(
|
||||
"""
|
||||
DELETE FROM remote_task_attempts
|
||||
WHERE task_id IN (SELECT value FROM json_each(?))
|
||||
""",
|
||||
(json.dumps(eligible),),
|
||||
)
|
||||
cur = conn.execute(
|
||||
"""
|
||||
DELETE FROM remote_tasks
|
||||
WHERE id IN (SELECT value FROM json_each(?))
|
||||
""",
|
||||
(json.dumps(eligible),),
|
||||
)
|
||||
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)
|
||||
# Shared content-addressed inputs remain discoverable without their old
|
||||
# task row, so they can be swept after the transaction.
|
||||
purge_artifacts((), referenced, cutoff=cutoff, root=root)
|
||||
return removed
|
||||
|
||||
|
||||
|
||||
+296
-34
@@ -20,15 +20,21 @@ There is deliberately no way to disable verification.
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as _dt
|
||||
import errno
|
||||
import ipaddress
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import ssl
|
||||
import tempfile
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
from cryptography import x509
|
||||
from cryptography.exceptions import InvalidSignature, UnsupportedAlgorithm
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
from cryptography.x509.oid import NameOID
|
||||
@@ -39,6 +45,8 @@ logger = logging.getLogger("omnivoice.worker")
|
||||
|
||||
_CERT_VALID_DAYS = 825 # the CA/Browser Forum maximum; long enough to be quiet
|
||||
_RENEW_WITHIN_DAYS = 30
|
||||
_PENDING_PAIR_MAGIC = b"omnivoice-tls-pair-v1\n"
|
||||
_PROCESS_CREDENTIAL_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def unverified_client_context() -> ssl.SSLContext:
|
||||
@@ -72,7 +80,7 @@ def _san_entries(hostnames: list[str]) -> list[x509.GeneralName]:
|
||||
"""
|
||||
entries: list[x509.GeneralName] = []
|
||||
for host in hostnames:
|
||||
host = (host or "").strip()
|
||||
host = _normalise_host(host)
|
||||
if not host:
|
||||
continue
|
||||
try:
|
||||
@@ -118,9 +126,23 @@ def covers(credentials: "ServerCredentials", host: str) -> bool:
|
||||
).value
|
||||
except Exception:
|
||||
return False
|
||||
names = set(san.get_values_for_type(x509.DNSName))
|
||||
names |= {str(ip) for ip in san.get_values_for_type(x509.IPAddress)}
|
||||
return host in names
|
||||
names = {_normalise_host(name) for name in san.get_values_for_type(x509.DNSName)}
|
||||
names |= {
|
||||
_normalise_host(str(ip)) for ip in san.get_values_for_type(x509.IPAddress)
|
||||
}
|
||||
return _normalise_host(host) in names
|
||||
|
||||
|
||||
def _normalise_host(host: str) -> str:
|
||||
"""Canonicalise a SAN identity for comparison, not for DNS resolution."""
|
||||
candidate = (host or "").strip().strip("[]")
|
||||
try:
|
||||
return str(ipaddress.ip_address(candidate))
|
||||
except ValueError:
|
||||
# DNS names are case-insensitive, and a trailing root dot does not
|
||||
# name a different host. cryptography intentionally does no matching
|
||||
# for us because the certificate is its own trust root.
|
||||
return candidate.rstrip(".").lower()
|
||||
|
||||
|
||||
def default_hostnames() -> list[str]:
|
||||
@@ -206,25 +228,50 @@ def load_or_create(
|
||||
desktop control plane presents as "my workers all went offline" with
|
||||
nothing in the UI explaining why.
|
||||
"""
|
||||
existing = _load(cert_path, key_path)
|
||||
wanted = hostnames or default_hostnames()
|
||||
# Every explicitly requested identity must be present. This matters for an
|
||||
# inbound listener bound to a user-entered LAN address: keeping a stable
|
||||
# certificate that does not name that address makes mandatory hostname
|
||||
# verification fail even though its fingerprint is correct.
|
||||
missing = [
|
||||
host for host in wanted if existing is not None and not covers(existing, host)
|
||||
]
|
||||
if existing is not None and not _expiring_soon(existing) and not missing:
|
||||
return existing
|
||||
if existing is not None:
|
||||
reason = (
|
||||
"expiring" if _expiring_soon(existing) else "missing a requested hostname"
|
||||
)
|
||||
logger.info("Control-plane certificate is %s — regenerating.", reason)
|
||||
credentials = generate_self_signed(hostnames=wanted)
|
||||
_save(cert_path, key_path, credentials)
|
||||
return credentials
|
||||
with _credential_lock(cert_path):
|
||||
existing = _load(cert_path, key_path)
|
||||
pending_present, pending = _load_pending_pair(cert_path, key_path)
|
||||
if existing is not None and pending_present:
|
||||
# A valid final pair is either the old generation (the transaction
|
||||
# never began replacing it) or the fully committed new one. Both
|
||||
# are safer than rotating a pin merely because cleanup was cut
|
||||
# short. The lock makes it safe to remove the stale journal here.
|
||||
_remove_pending_pair(cert_path, key_path)
|
||||
elif existing is None and pending is not None:
|
||||
logger.warning(
|
||||
"Recovering an interrupted control-plane TLS credential update."
|
||||
)
|
||||
_install_pair(cert_path, key_path, pending)
|
||||
_remove_pending_pair(cert_path, key_path)
|
||||
existing = pending
|
||||
|
||||
wanted = hostnames or default_hostnames()
|
||||
# Every explicitly requested identity must be present. This matters for
|
||||
# a listener bound to a user-entered address: keeping a stable
|
||||
# certificate that does not name it makes mandatory hostname
|
||||
# verification fail even though its fingerprint is correct.
|
||||
missing = [
|
||||
host
|
||||
for host in wanted
|
||||
if existing is not None and not covers(existing, host)
|
||||
]
|
||||
if existing is not None and not _expiring_soon(existing) and not missing:
|
||||
return existing
|
||||
if existing is not None:
|
||||
reason = (
|
||||
"expiring"
|
||||
if _expiring_soon(existing)
|
||||
else "missing a requested hostname"
|
||||
)
|
||||
logger.info("Control-plane certificate is %s — regenerating.", reason)
|
||||
elif pending_present and pending is None:
|
||||
logger.warning(
|
||||
"Ignoring a corrupt interrupted TLS credential update and "
|
||||
"generating a new pair."
|
||||
)
|
||||
credentials = generate_self_signed(hostnames=wanted)
|
||||
_save(cert_path, key_path, credentials)
|
||||
return credentials
|
||||
|
||||
|
||||
def _load(cert_path: str, key_path: str) -> Optional[ServerCredentials]:
|
||||
@@ -233,11 +280,40 @@ def _load(cert_path: str, key_path: str) -> Optional[ServerCredentials]:
|
||||
cert_pem = fh.read()
|
||||
with open(key_path, "rb") as fh:
|
||||
key_pem = fh.read()
|
||||
except (FileNotFoundError, PermissionError):
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
return _credentials_from_pem(cert_pem, key_pem)
|
||||
|
||||
|
||||
def _credentials_from_pem(
|
||||
cert_pem: bytes, key_pem: bytes
|
||||
) -> Optional[ServerCredentials]:
|
||||
"""Parse a credential pair only when the private key belongs to the cert."""
|
||||
try:
|
||||
certificate = x509.load_pem_x509_certificate(cert_pem)
|
||||
except ValueError:
|
||||
private_key = serialization.load_pem_private_key(key_pem, password=None)
|
||||
certificate_public_key = certificate.public_key().public_bytes(
|
||||
serialization.Encoding.DER,
|
||||
serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
private_public_key = private_key.public_key().public_bytes(
|
||||
serialization.Encoding.DER,
|
||||
serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
certificate.public_key().verify(
|
||||
certificate.signature,
|
||||
certificate.tbs_certificate_bytes,
|
||||
ec.ECDSA(certificate.signature_hash_algorithm),
|
||||
)
|
||||
except (
|
||||
AttributeError,
|
||||
InvalidSignature,
|
||||
TypeError,
|
||||
UnsupportedAlgorithm,
|
||||
ValueError,
|
||||
):
|
||||
return None
|
||||
if certificate_public_key != private_public_key:
|
||||
return None
|
||||
return ServerCredentials(
|
||||
certificate_pem=cert_pem,
|
||||
@@ -258,19 +334,205 @@ def _expiring_soon(
|
||||
|
||||
|
||||
def _save(cert_path: str, key_path: str, credentials: ServerCredentials) -> None:
|
||||
os.makedirs(os.path.dirname(os.path.abspath(cert_path)), exist_ok=True)
|
||||
with open(cert_path, "wb") as fh:
|
||||
fh.write(credentials.certificate_pem)
|
||||
# The private key gets the same 0600 treatment as worker keys.
|
||||
fd = os.open(key_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
try:
|
||||
os.write(fd, credentials.private_key_pem)
|
||||
finally:
|
||||
os.close(fd)
|
||||
"""Durably replace a pair, leaving enough state to finish after a crash.
|
||||
|
||||
No filesystem atomically renames two independent paths. A mode-0600
|
||||
journal is therefore made durable first; the loader uses it only when the
|
||||
final paths are torn or corrupt. Each final path is itself an atomic
|
||||
sibling rename, and the journal is removed only after both directory
|
||||
entries and their contents are durable.
|
||||
"""
|
||||
pending_path = _pending_pair_path(cert_path, key_path)
|
||||
_atomic_write(pending_path, _encode_pending_pair(credentials))
|
||||
_install_pair(cert_path, key_path, credentials)
|
||||
_remove_pending_pair(cert_path, key_path)
|
||||
|
||||
|
||||
def _install_pair(
|
||||
cert_path: str, key_path: str, credentials: ServerCredentials
|
||||
) -> None:
|
||||
# Key first, certificate second: the public certificate is the commit
|
||||
# marker. A crash between them is detected by _load and recovered from the
|
||||
# already-durable pending pair.
|
||||
_atomic_write(key_path, credentials.private_key_pem)
|
||||
try:
|
||||
os.chmod(key_path, 0o600)
|
||||
except OSError:
|
||||
# Windows and some network filesystems do not honour POSIX modes; the
|
||||
# key remains in the app's per-user data directory there.
|
||||
pass
|
||||
_atomic_write(cert_path, credentials.certificate_pem)
|
||||
installed = _load(cert_path, key_path)
|
||||
if installed is None or installed.fingerprint != credentials.fingerprint:
|
||||
raise OSError("TLS credential pair failed verification after persistence")
|
||||
|
||||
|
||||
def _atomic_write(path: str, payload: bytes) -> None:
|
||||
"""Fsync and atomically replace one mode-restricted credential file."""
|
||||
directory = os.path.dirname(os.path.abspath(path))
|
||||
_durable_makedirs(directory)
|
||||
fd, temporary = tempfile.mkstemp(
|
||||
prefix=f".{os.path.basename(path)}.", suffix=".tmp", dir=directory
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as fh:
|
||||
fh.write(payload)
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
os.replace(temporary, path)
|
||||
_fsync_parent_directory(directory)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(temporary)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def _pending_pair_path(cert_path: str, key_path: str) -> str:
|
||||
cert_name = os.path.basename(cert_path)
|
||||
key_name = os.path.basename(key_path)
|
||||
directory = os.path.dirname(os.path.abspath(cert_path))
|
||||
return os.path.join(directory, f".{cert_name}.{key_name}.pending")
|
||||
|
||||
|
||||
def _encode_pending_pair(credentials: ServerCredentials) -> bytes:
|
||||
certificate = credentials.certificate_pem
|
||||
return (
|
||||
_PENDING_PAIR_MAGIC
|
||||
+ str(len(certificate)).encode("ascii")
|
||||
+ b"\n"
|
||||
+ certificate
|
||||
+ credentials.private_key_pem
|
||||
)
|
||||
|
||||
|
||||
def _load_pending_pair(
|
||||
cert_path: str, key_path: str
|
||||
) -> tuple[bool, Optional[ServerCredentials]]:
|
||||
path = _pending_pair_path(cert_path, key_path)
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
raw = fh.read()
|
||||
except FileNotFoundError:
|
||||
return False, None
|
||||
if not raw.startswith(_PENDING_PAIR_MAGIC):
|
||||
return True, None
|
||||
size_line, separator, payload = raw[len(_PENDING_PAIR_MAGIC) :].partition(b"\n")
|
||||
if not separator:
|
||||
return True, None
|
||||
try:
|
||||
certificate_size = int(size_line)
|
||||
except ValueError:
|
||||
return True, None
|
||||
if certificate_size <= 0 or certificate_size >= len(payload):
|
||||
return True, None
|
||||
return True, _credentials_from_pem(
|
||||
payload[:certificate_size], payload[certificate_size:]
|
||||
)
|
||||
|
||||
|
||||
def _remove_pending_pair(cert_path: str, key_path: str) -> None:
|
||||
path = _pending_pair_path(cert_path, key_path)
|
||||
try:
|
||||
os.unlink(path)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
_fsync_parent_directory(os.path.dirname(path) or ".")
|
||||
|
||||
|
||||
def _fsync_parent_directory(directory: str) -> None:
|
||||
"""Make a preceding directory-entry change durable where supported."""
|
||||
directory_flag = getattr(os, "O_DIRECTORY", None)
|
||||
if directory_flag is None:
|
||||
return
|
||||
unsupported = {
|
||||
errno.EINVAL,
|
||||
getattr(errno, "ENOTSUP", errno.EINVAL),
|
||||
getattr(errno, "EOPNOTSUPP", errno.EINVAL),
|
||||
}
|
||||
try:
|
||||
descriptor = os.open(directory, os.O_RDONLY | directory_flag)
|
||||
except OSError as exc:
|
||||
if exc.errno in unsupported:
|
||||
return
|
||||
raise
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
except OSError as exc:
|
||||
if exc.errno not in unsupported:
|
||||
raise
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def _durable_makedirs(directory: str) -> None:
|
||||
"""Persist every newly-created credential-directory entry."""
|
||||
target = os.path.abspath(directory)
|
||||
missing: list[str] = []
|
||||
current = target
|
||||
while not os.path.isdir(current):
|
||||
if os.path.exists(current):
|
||||
if os.path.isdir(current):
|
||||
break
|
||||
raise NotADirectoryError(current)
|
||||
missing.append(current)
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
|
||||
for path in reversed(missing):
|
||||
try:
|
||||
os.mkdir(path)
|
||||
except FileExistsError:
|
||||
if not os.path.isdir(path):
|
||||
raise
|
||||
_fsync_parent_directory(os.path.dirname(path) or ".")
|
||||
|
||||
if not missing:
|
||||
# Retry the exact barrier that may have failed after a prior mkdir.
|
||||
_fsync_parent_directory(os.path.dirname(target) or ".")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _credential_lock(cert_path: str) -> Iterator[None]:
|
||||
"""Serialise first-start/renewal across threads and app processes."""
|
||||
directory = os.path.dirname(os.path.abspath(cert_path))
|
||||
_durable_makedirs(directory)
|
||||
lock_path = os.path.join(directory, f".{os.path.basename(cert_path)}.lock")
|
||||
with _PROCESS_CREDENTIAL_LOCK:
|
||||
descriptor = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
|
||||
is_windows = os.name == "nt"
|
||||
acquired = False
|
||||
try:
|
||||
if is_windows:
|
||||
import msvcrt # noqa: PLC0415
|
||||
|
||||
if os.fstat(descriptor).st_size == 0:
|
||||
os.write(descriptor, b"\0")
|
||||
os.fsync(descriptor)
|
||||
os.lseek(descriptor, 0, os.SEEK_SET)
|
||||
msvcrt.locking(descriptor, msvcrt.LK_LOCK, 1)
|
||||
else:
|
||||
import fcntl # noqa: PLC0415
|
||||
|
||||
fcntl.flock(descriptor, fcntl.LOCK_EX)
|
||||
acquired = True
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
if acquired and is_windows:
|
||||
import msvcrt # noqa: PLC0415
|
||||
|
||||
os.lseek(descriptor, 0, os.SEEK_SET)
|
||||
msvcrt.locking(descriptor, msvcrt.LK_UNLCK, 1)
|
||||
elif acquired:
|
||||
import fcntl # noqa: PLC0415
|
||||
|
||||
fcntl.flock(descriptor, fcntl.LOCK_UN)
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def pin_matches(certificate_der: bytes, expected_fingerprint: str) -> bool:
|
||||
|
||||
@@ -48,8 +48,10 @@ from typing import Awaitable, Callable, Optional, Protocol
|
||||
|
||||
import grpc
|
||||
|
||||
from worker.async_utils import drain_task, to_thread_and_drain_on_cancel
|
||||
from worker import errors as worker_errors
|
||||
from worker import identity
|
||||
from worker.capacity import clamp_concurrency
|
||||
from worker.errors import ErrorClass, WorkerError
|
||||
from worker.identity import EnrollmentToken, WorkerKeypair
|
||||
from worker.protocol.gen import worker_v1_pb2 as pb
|
||||
@@ -103,6 +105,8 @@ class ArtifactTransport(Protocol):
|
||||
|
||||
async def stage_in(self, ref: pb.ArtifactRef, destination: str) -> None: ...
|
||||
|
||||
def result_acked(self, artifacts: list[pb.ArtifactRef]) -> None: ...
|
||||
|
||||
# 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
|
||||
@@ -115,6 +119,29 @@ _MIN_KEEPALIVE_INTERVAL_SECONDS = 0.05
|
||||
_EXECUTOR_KWARGS = frozenset({"on_progress", "on_model_loading", "fetch_input"})
|
||||
|
||||
|
||||
def _write_all(handle, payload: bytes) -> None:
|
||||
"""Write a complete chunk, including through short-writing file wrappers."""
|
||||
remaining = memoryview(payload)
|
||||
while remaining:
|
||||
written = handle.write(remaining)
|
||||
if written is None or written <= 0:
|
||||
raise OSError("input destination made no write progress")
|
||||
remaining = remaining[written:]
|
||||
|
||||
|
||||
def _close_and_remove(handle, destination: str) -> None:
|
||||
"""Finish file cleanup as one blocking operation after cancellation."""
|
||||
if handle is not None:
|
||||
with contextlib.suppress(OSError):
|
||||
handle.close()
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
|
||||
|
||||
class TerminalRegistrationError(RuntimeError):
|
||||
"""A registration failure that reconnecting cannot repair."""
|
||||
|
||||
|
||||
def keepalive_interval(lease_seconds: float) -> float:
|
||||
"""How often a running task must renew its progress lease.
|
||||
|
||||
@@ -267,28 +294,44 @@ class WorkerClient:
|
||||
cancel: Optional[Callable[[str], Awaitable[None]]] = None,
|
||||
capability_probe: Optional[Callable[[], list[dict]]] = None,
|
||||
on_registered: Optional[Callable[[str], None]] = None,
|
||||
on_activated: Optional[Callable[[str], None]] = None,
|
||||
artifacts: Optional["ArtifactTransport"] = None,
|
||||
drain_active_work: Optional[Callable[[], Awaitable[None]]] = None,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.config.max_concurrent_tasks = clamp_concurrency(
|
||||
self.config.max_concurrent_tasks
|
||||
)
|
||||
self._execute = execute
|
||||
self._cancel = cancel
|
||||
self._capability_probe = capability_probe
|
||||
# Discovery imports engine adapters and inspects model storage. Share
|
||||
# one off-loop probe across reconnect/task/prewarm/idle refresh races;
|
||||
# a cancelled waiter drains it before returning so no detached probe
|
||||
# can mutate global engine state after authority is gone.
|
||||
self._capability_probe_task: Optional[asyncio.Task] = None
|
||||
# Outbound mode moves artifacts with RPCs this side initiates
|
||||
# (UploadResult / DownloadArtifact), which is only possible because
|
||||
# this side dialled. In inbound mode the node cannot call the panel at
|
||||
# all, so both directions are driven from the panel and this hook
|
||||
# swaps in the staging that makes that work. None means outbound.
|
||||
self._artifacts = artifacts
|
||||
self._drain_active_work = drain_active_work
|
||||
# Lets the agent persist the server-assigned id. Without it a restarted
|
||||
# worker signs its challenge with an empty worker_id, the signature
|
||||
# never matches, and reconnecting needs a fresh enrollment token —
|
||||
# which would make key-based identity pointless.
|
||||
self._on_registered = on_registered
|
||||
# Register only reserves a provisional server generation. Readiness is
|
||||
# published separately, after ConfigUpdate proves Control activated it.
|
||||
self._on_activated = on_activated
|
||||
self._activation_confirmed = False
|
||||
self._reporter_kwargs = _accepted_reporter_kwargs(execute)
|
||||
self._outbox = _Outbox()
|
||||
self._pending: dict[str, PendingResult] = {}
|
||||
self._running: dict[str, asyncio.Task] = {}
|
||||
self._keepalives: dict[str, asyncio.Task] = {}
|
||||
self._maintenance: set[asyncio.Task] = set()
|
||||
self._epoch = 0
|
||||
self._session_token = ""
|
||||
# Negotiated by ConfigUpdate; None means "use the executor's own
|
||||
@@ -298,6 +341,12 @@ class WorkerClient:
|
||||
# same channel rather than the control stream.
|
||||
self._stub = None
|
||||
self._stop = asyncio.Event()
|
||||
# Drain is a graceful reconnect, not terminal shutdown. This event
|
||||
# half-closes only the current stream after every active result is ACKed
|
||||
# while ``_stop`` remains reserved for cancelling the agent itself.
|
||||
self._reconnect_requested = asyncio.Event()
|
||||
self._draining = False
|
||||
self._accepting_assignments = True
|
||||
|
||||
# ── Connection ────────────────────────────────────────────────────────
|
||||
|
||||
@@ -318,6 +367,8 @@ class WorkerClient:
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
"""Connect, serve, and reconnect until stopped."""
|
||||
if not self._stop.is_set():
|
||||
self._accepting_assignments = True
|
||||
attempt = 0
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
@@ -325,6 +376,13 @@ class WorkerClient:
|
||||
attempt = 0
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except TerminalRegistrationError:
|
||||
# The control plane has made a durable decision that this
|
||||
# identity may not reconnect. Work deliberately survives an
|
||||
# ordinary network drop, but must not survive revocation and
|
||||
# keep using the GPU with no authority able to cancel it.
|
||||
await self._cancel_active_work()
|
||||
raise
|
||||
except Exception as exc:
|
||||
attempt += 1
|
||||
delay = backoff_delay(attempt)
|
||||
@@ -338,6 +396,52 @@ class WorkerClient:
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._stop.set()
|
||||
self._reconnect_requested.set()
|
||||
await self._cancel_active_work()
|
||||
|
||||
async def _cancel_active_work(self) -> None:
|
||||
"""Cancel retained tasks without permanently disabling reconnect."""
|
||||
# Close admission before taking any snapshots. Attach/Control may still
|
||||
# have a frame ready while revocation drains an uninterruptible GPU
|
||||
# call; accepting that frame here lets it escape the snapshot entirely.
|
||||
self._accepting_assignments = False
|
||||
maintenance = list(self._maintenance)
|
||||
for task in maintenance:
|
||||
task.cancel()
|
||||
running = list(self._running.items())
|
||||
keepalives = list(self._keepalives.values())
|
||||
for key, task in running:
|
||||
self._stop_keepalive(key)
|
||||
task.cancel()
|
||||
for keepalive in keepalives:
|
||||
keepalive.cancel()
|
||||
cancel_callbacks = []
|
||||
if self._cancel is not None:
|
||||
cancel_callbacks = [
|
||||
asyncio.create_task(self._cancel(key.split("/")[0]))
|
||||
for key, _task in running
|
||||
]
|
||||
# Cancel every maintenance, task, and keepalive wrapper before awaiting
|
||||
# any uninterruptible one. A prewarm stuck in a model-load thread must
|
||||
# not delay revocation of active user renders.
|
||||
draining = [
|
||||
*maintenance,
|
||||
*(task for _key, task in running),
|
||||
*keepalives,
|
||||
*cancel_callbacks,
|
||||
]
|
||||
if draining:
|
||||
await asyncio.gather(
|
||||
*draining, return_exceptions=True
|
||||
)
|
||||
self._maintenance.clear()
|
||||
for key, task in running:
|
||||
if self._running.get(key) is task:
|
||||
self._running.pop(key, None)
|
||||
if self._drain_active_work is not None:
|
||||
await self._drain_active_work()
|
||||
self._keepalives.clear()
|
||||
self._pending.clear()
|
||||
|
||||
async def _connect_once(self) -> None:
|
||||
async with self._channel() as channel:
|
||||
@@ -359,6 +463,8 @@ class WorkerClient:
|
||||
try:
|
||||
async for message in stream:
|
||||
await self._on_server_message(message)
|
||||
if self._stop.is_set():
|
||||
break
|
||||
finally:
|
||||
heartbeat.cancel()
|
||||
# The channel closes with this block, so a stub kept past it
|
||||
@@ -375,7 +481,27 @@ class WorkerClient:
|
||||
# before; inbound calls them from the Attach handler. Neither mode gets its
|
||||
# own copy of registration, zombie reconciliation or redelivery.
|
||||
|
||||
def build_register_request(self) -> pb.RegisterRequest:
|
||||
async def _probe_capabilities(self) -> list[dict]:
|
||||
if self._capability_probe is None:
|
||||
return list(self.config.capabilities or [])
|
||||
task = self._capability_probe_task
|
||||
if task is None or task.done():
|
||||
task = asyncio.create_task(
|
||||
to_thread_and_drain_on_cancel(self._capability_probe),
|
||||
name="worker-capability-probe",
|
||||
)
|
||||
self._capability_probe_task = task
|
||||
try:
|
||||
capabilities = await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
await drain_task(task)
|
||||
raise
|
||||
finally:
|
||||
if task.done() and self._capability_probe_task is task:
|
||||
self._capability_probe_task = None
|
||||
return list(capabilities or [])
|
||||
|
||||
async def build_register_request(self) -> pb.RegisterRequest:
|
||||
"""This worker's self-description. Identical in both modes."""
|
||||
challenge = identity.new_challenge()
|
||||
nonce = identity.new_challenge()
|
||||
@@ -387,9 +513,8 @@ class WorkerClient:
|
||||
nonce=nonce,
|
||||
)
|
||||
)
|
||||
capabilities = (
|
||||
self._capability_probe() if self._capability_probe else self.config.capabilities
|
||||
)
|
||||
capabilities = await self._probe_capabilities()
|
||||
self.config.capabilities = capabilities
|
||||
return pb.RegisterRequest(
|
||||
envelope=pb.Envelope(sequence=self._epoch),
|
||||
protocol_version_min=PROTOCOL_VERSION,
|
||||
@@ -403,7 +528,9 @@ class WorkerClient:
|
||||
key_id=self.config.keypair.key_id,
|
||||
host=codec.host_to_pb(self.config.host or describe_host()),
|
||||
capabilities=[codec.capability_to_pb(c) for c in capabilities],
|
||||
max_concurrent_tasks=self.config.max_concurrent_tasks,
|
||||
max_concurrent_tasks=clamp_concurrency(
|
||||
self.config.max_concurrent_tasks
|
||||
),
|
||||
in_flight=[
|
||||
codec.task_ref(t.split("/")[0], t.split("/")[1], self._epoch)
|
||||
for t in self._running
|
||||
@@ -415,27 +542,80 @@ class WorkerClient:
|
||||
async def accept_registration(self, response: pb.RegisterResponse) -> None:
|
||||
"""Adopt the control plane's answer and recover in-flight state."""
|
||||
if response.error.code:
|
||||
raise RuntimeError(f"{response.error.code}: {response.error.message}")
|
||||
raise TerminalRegistrationError(
|
||||
f"{response.error.code}: {response.error.message}"
|
||||
)
|
||||
# An enrollment token is already spent when this response arrives.
|
||||
# Commit the reconnect identity before adopting the live session; if
|
||||
# local durable state cannot be written, retrying the spent token can
|
||||
# never repair the worker and must reach the caller immediately.
|
||||
if self._on_registered is not None:
|
||||
try:
|
||||
await to_thread_and_drain_on_cancel(
|
||||
self._on_registered, response.worker_id
|
||||
)
|
||||
except Exception as exc:
|
||||
raise TerminalRegistrationError(
|
||||
"LOCAL_STATE: accepted enrollment could not be persisted"
|
||||
) from exc
|
||||
|
||||
self._draining = False
|
||||
self._reconnect_requested.clear()
|
||||
if not self._stop.is_set():
|
||||
self._accepting_assignments = True
|
||||
self._activation_confirmed = False
|
||||
self._epoch = response.session_epoch
|
||||
self._session_token = response.session_token
|
||||
self.config.worker_id = response.worker_id
|
||||
# The token is spent; every later connection proves key possession.
|
||||
self.config.enrollment_token = ""
|
||||
if self._on_registered is not None:
|
||||
try:
|
||||
self._on_registered(response.worker_id)
|
||||
except Exception:
|
||||
logger.warning("Could not persist the worker id", exc_info=True)
|
||||
|
||||
authoritative = {ref.attempt_id for ref in response.authoritative_in_flight}
|
||||
await self._cancel_zombies(authoritative)
|
||||
await self._redeliver_pending()
|
||||
|
||||
def confirm_activation(self) -> None:
|
||||
"""Publish readiness once Control proves the provisional session live."""
|
||||
if self._activation_confirmed:
|
||||
return
|
||||
if self._on_activated is not None:
|
||||
try:
|
||||
self._on_activated(self.config.worker_id)
|
||||
except Exception as exc:
|
||||
raise TerminalRegistrationError(
|
||||
"LOCAL_STATE: activated enrollment could not be published"
|
||||
) from exc
|
||||
self._activation_confirmed = True
|
||||
|
||||
async def next_outbound(self) -> pb.WorkerMessage:
|
||||
"""The next frame this worker wants to send."""
|
||||
return await self._outbox.get()
|
||||
|
||||
def prepare_inbound_session(self) -> None:
|
||||
"""Start a fresh stream while retaining running and pending work.
|
||||
|
||||
An inbound listener creates the protocol owner once per panel key, not
|
||||
once per transport generation. Frames queued for the dead stream are
|
||||
stale, but `_running` and `_pending` are precisely the state the next
|
||||
Register must reconcile and redeliver.
|
||||
"""
|
||||
self._outbox = _Outbox()
|
||||
self._session_token = ""
|
||||
self._stub = None
|
||||
self._draining = False
|
||||
self._reconnect_requested.clear()
|
||||
if not self._stop.is_set():
|
||||
self._accepting_assignments = True
|
||||
|
||||
@property
|
||||
def reconnect_requested(self) -> bool:
|
||||
return self._reconnect_requested.is_set()
|
||||
|
||||
@property
|
||||
def outbound_pending(self) -> bool:
|
||||
"""Whether a terminal/control frame still needs transport delivery."""
|
||||
return not self._outbox.empty()
|
||||
|
||||
def start_heartbeat(self, response: pb.RegisterResponse) -> asyncio.Task:
|
||||
"""Begin the heartbeat this session's liveness depends on.
|
||||
|
||||
@@ -458,14 +638,27 @@ class WorkerClient:
|
||||
await self._on_server_message(message)
|
||||
|
||||
async def _register(self, stub) -> pb.RegisterResponse:
|
||||
return await stub.Register(self.build_register_request())
|
||||
return await stub.Register(await self.build_register_request())
|
||||
|
||||
# ── Outbound ──────────────────────────────────────────────────────────
|
||||
|
||||
async def _outbound(self):
|
||||
while True:
|
||||
message = await self._outbox.get()
|
||||
yield message
|
||||
while not self._reconnect_requested.is_set():
|
||||
message = asyncio.create_task(self._outbox.get())
|
||||
reconnect = asyncio.create_task(self._reconnect_requested.wait())
|
||||
try:
|
||||
done, _pending = await asyncio.wait(
|
||||
{message, reconnect}, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if message not in done:
|
||||
return
|
||||
yield message.result()
|
||||
self._maybe_finish_drain()
|
||||
finally:
|
||||
for task in (message, reconnect):
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(message, reconnect, return_exceptions=True)
|
||||
|
||||
async def _send(self, message: pb.WorkerMessage, *, bulk: bool = False) -> None:
|
||||
"""Enqueue a frame. ``bulk`` is for anything that can be large.
|
||||
@@ -479,17 +672,19 @@ class WorkerClient:
|
||||
async def _heartbeat_loop(self, interval: float) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
await self._send(
|
||||
pb.WorkerMessage(
|
||||
heartbeat=pb.Heartbeat(
|
||||
active_tasks=len(self._running),
|
||||
available_slots=max(
|
||||
0, self.config.max_concurrent_tasks - len(self._running)
|
||||
),
|
||||
resident_models=self._resident_models(),
|
||||
)
|
||||
)
|
||||
await self._send(self.heartbeat_message())
|
||||
|
||||
def heartbeat_message(self) -> pb.WorkerMessage:
|
||||
"""Build the worker's current liveness/capacity frame."""
|
||||
return pb.WorkerMessage(
|
||||
heartbeat=pb.Heartbeat(
|
||||
active_tasks=len(self._running),
|
||||
available_slots=max(
|
||||
0, self.config.max_concurrent_tasks - len(self._running)
|
||||
),
|
||||
resident_models=self._resident_models(),
|
||||
)
|
||||
)
|
||||
|
||||
def _resident_models(self) -> list[str]:
|
||||
return [
|
||||
@@ -503,7 +698,7 @@ class WorkerClient:
|
||||
if self._capability_probe is None:
|
||||
return
|
||||
try:
|
||||
capabilities = self._capability_probe()
|
||||
capabilities = await self._probe_capabilities()
|
||||
self.config.capabilities = capabilities
|
||||
await self._send(pb.WorkerMessage(capabilities=pb.CapabilityUpdate(
|
||||
capabilities=[codec.capability_to_pb(c) for c in capabilities]
|
||||
@@ -531,9 +726,14 @@ class WorkerClient:
|
||||
# keepalive for an attempt the server has disowned is exactly the
|
||||
# frame that resurrects a cancelled task.
|
||||
self._stop_keepalive(key)
|
||||
task = self._running.pop(key, None)
|
||||
task = self._running.get(key)
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
# CancelAck releases server capacity. Do not send it while an
|
||||
# uninterruptible synthesis/load thread still owns the GPU, or the
|
||||
# replacement assignment can overlap and corrupt or OOM the worker.
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
self._running.pop(key, None)
|
||||
if self._cancel is not None:
|
||||
await self._cancel(key.split("/")[0])
|
||||
|
||||
@@ -550,24 +750,76 @@ class WorkerClient:
|
||||
)
|
||||
elif kind == "result_ack":
|
||||
# Only now is it safe to forget the result.
|
||||
self._pending.pop(self._key(message.result_ack.ref), None)
|
||||
pending = self._pending.pop(self._key(message.result_ack.ref), None)
|
||||
if pending is not None and self._artifacts is not None:
|
||||
result_acked_async = getattr(
|
||||
self._artifacts, "result_acked_async", None
|
||||
)
|
||||
if callable(result_acked_async):
|
||||
await result_acked_async(pending.artifacts)
|
||||
else:
|
||||
result_acked = getattr(self._artifacts, "result_acked", None)
|
||||
if callable(result_acked):
|
||||
result_acked(pending.artifacts)
|
||||
self._maybe_finish_drain()
|
||||
elif kind == "config":
|
||||
if message.config.max_concurrent_tasks:
|
||||
self.config.max_concurrent_tasks = message.config.max_concurrent_tasks
|
||||
self.config.max_concurrent_tasks = clamp_concurrency(
|
||||
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)
|
||||
self.confirm_activation()
|
||||
elif kind == "ping":
|
||||
# Answer immediately; the server times the round trip.
|
||||
await self._send(pb.WorkerMessage(pong=pb.Pong(nonce=message.ping.nonce)))
|
||||
elif kind == "drain":
|
||||
self._stop.set()
|
||||
self._accepting_assignments = False
|
||||
self._draining = True
|
||||
if message.drain.reconnect_to:
|
||||
self.config.endpoint = message.drain.reconnect_to
|
||||
self._maybe_finish_drain()
|
||||
elif kind == "shutdown":
|
||||
await self._cancel_active_work()
|
||||
self._stop.set()
|
||||
if self._artifacts is not None:
|
||||
await self._send(
|
||||
pb.WorkerMessage(
|
||||
goodbye=pb.WorkerGoodbye(
|
||||
reason="The control-plane connection was removed."
|
||||
)
|
||||
)
|
||||
)
|
||||
# Publish the terminal acknowledgement before asking the inbound
|
||||
# Attach loop to leave. Reversing these lets that loop observe
|
||||
# reconnect_requested and close the stream with Goodbye still
|
||||
# queued, so the control plane cannot prove remote work drained.
|
||||
self._reconnect_requested.set()
|
||||
elif kind == "prewarm":
|
||||
asyncio.create_task(self._on_prewarm(message.prewarm), name="worker-prewarm")
|
||||
if not self._accepting_assignments:
|
||||
return
|
||||
task = asyncio.create_task(
|
||||
self._on_prewarm(message.prewarm), name="worker-prewarm"
|
||||
)
|
||||
self._maintenance.add(task)
|
||||
task.add_done_callback(self._maintenance_finished)
|
||||
|
||||
def _maintenance_finished(self, task: asyncio.Task) -> None:
|
||||
self._maintenance.discard(task)
|
||||
self._maybe_finish_drain()
|
||||
|
||||
def _maybe_finish_drain(self) -> None:
|
||||
if (
|
||||
self._draining
|
||||
and not self._running
|
||||
and not self._pending
|
||||
and not self._maintenance
|
||||
and self._outbox.empty()
|
||||
):
|
||||
self._reconnect_requested.set()
|
||||
|
||||
async def _on_prewarm(self, request: pb.PrewarmRequest) -> None:
|
||||
"""Load/download a catalog model, then report the resulting capability."""
|
||||
@@ -591,15 +843,20 @@ class WorkerClient:
|
||||
await self._install_catalog_repo(repo_ids[0])
|
||||
from worker.executor import TaskExecutor # noqa: PLC0415
|
||||
|
||||
await asyncio.to_thread(TaskExecutor._load_backend, engine)
|
||||
await to_thread_and_drain_on_cancel(TaskExecutor._load_backend, engine)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.warning("Prewarm failed for %s", engine or request.model_id, exc_info=True)
|
||||
finally:
|
||||
await self.refresh_capabilities()
|
||||
await self.refresh_capabilities()
|
||||
|
||||
async def _install_catalog_repo(self, repo_id: str) -> None:
|
||||
"""Run the existing setup installer and pipe its hf_progress upstream."""
|
||||
from api.routers.setup.download import InstallModelRequest, install_model # noqa: PLC0415
|
||||
from api.routers.setup.download import ( # noqa: PLC0415
|
||||
InstallModelRequest,
|
||||
cancel_install_and_wait,
|
||||
install_model,
|
||||
)
|
||||
from utils import download_aggregator, hf_progress # noqa: PLC0415
|
||||
|
||||
hf_progress.install()
|
||||
@@ -627,12 +884,20 @@ class WorkerClient:
|
||||
loop.call_soon_threadsafe(_finish)
|
||||
|
||||
listener_id = hf_progress.register_listener(listener)
|
||||
started_install = False
|
||||
install_completed = False
|
||||
try:
|
||||
await install_model(InstallModelRequest(repo_id=repo_id, target="local"))
|
||||
response = await install_model(
|
||||
InstallModelRequest(repo_id=repo_id, target="local")
|
||||
)
|
||||
started_install = response.get("status") == "install_started"
|
||||
event = await asyncio.wait_for(terminal, timeout=_FALLBACK_MODEL_LOAD_SECONDS)
|
||||
if event.get("phase") != "install_done":
|
||||
raise RuntimeError(event.get("error") or "model install did not complete")
|
||||
install_completed = True
|
||||
finally:
|
||||
if started_install and not install_completed:
|
||||
await cancel_install_and_wait(repo_id)
|
||||
hf_progress.unregister_listener(listener_id)
|
||||
|
||||
@staticmethod
|
||||
@@ -641,6 +906,20 @@ class WorkerClient:
|
||||
|
||||
async def _on_assignment(self, assignment: pb.TaskAssignment) -> None:
|
||||
key = self._key(assignment.ref)
|
||||
if not self._accepting_assignments or self._stop.is_set():
|
||||
await self._send(
|
||||
pb.WorkerMessage(
|
||||
rejected=pb.TaskRejected(
|
||||
ref=assignment.ref,
|
||||
error=pb.Error(
|
||||
error_class=pb.ERROR_CLASS_TRANSIENT,
|
||||
code="WORKER_STOPPING",
|
||||
message="The worker is relinquishing this control plane.",
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
return
|
||||
if len(self._running) >= self.config.max_concurrent_tasks:
|
||||
# Declining because we are full is normal and penalty-free; the
|
||||
# scheduler's view of our capacity is only ever advisory.
|
||||
@@ -672,9 +951,11 @@ class WorkerClient:
|
||||
# release the slot too, or the reserved task keeps running work
|
||||
# the scheduler never saw accepted — and double-executes after
|
||||
# reassignment. The stream-death case lands here as well.
|
||||
task = self._running.pop(key, None)
|
||||
task = self._running.get(key)
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
self._running.pop(key, None)
|
||||
raise
|
||||
|
||||
async def _run(self, assignment: pb.TaskAssignment) -> None:
|
||||
@@ -745,14 +1026,32 @@ class WorkerClient:
|
||||
failure: WorkerError = (
|
||||
exc.error if isinstance(exc, TaskFailure) else worker_errors.from_exception(exc)
|
||||
)
|
||||
if (
|
||||
failure.error_class is ErrorClass.TIMEOUT
|
||||
and self._drain_active_work is not None
|
||||
):
|
||||
# A timeout only ends the coroutine's wait; Python cannot
|
||||
# cancel the GPU thread underneath it. Do not send FAILED or
|
||||
# release this admission slot until the executor proves that
|
||||
# work relinquished the device, otherwise a capacity-1 worker
|
||||
# can accept a replacement on top of the timed-out render.
|
||||
await self._drain_active_work()
|
||||
await self._fail(assignment.ref, failure)
|
||||
finally:
|
||||
# Also covers the abnormal exits — cancellation, a crash between
|
||||
# the two _stop_keepalive calls above — so the timer can never
|
||||
# outlive the task that owns it.
|
||||
self._stop_keepalive(key)
|
||||
self._running.pop(key, None)
|
||||
await self.refresh_capabilities()
|
||||
# Discovery may inspect the just-used backend. Keep the admission
|
||||
# reservation until that off-loop probe has drained, otherwise a
|
||||
# heartbeat can advertise the slot while this generation still
|
||||
# owns task-finalization work. The inner finally also releases it
|
||||
# when shutdown cancels the refresh waiter.
|
||||
try:
|
||||
await self.refresh_capabilities()
|
||||
finally:
|
||||
self._running.pop(key, None)
|
||||
self._maybe_finish_drain()
|
||||
|
||||
async def _fail(self, ref: pb.TaskRef, error: WorkerError) -> None:
|
||||
await self._send(
|
||||
@@ -855,8 +1154,12 @@ class WorkerClient:
|
||||
await self._report_upload(ref, 0.0)
|
||||
|
||||
offset = 0
|
||||
metadata = ((SESSION_METADATA_KEY, self._session_token),)
|
||||
for _ in range(_MAX_UPLOAD_RESUMES):
|
||||
ack = await stub.UploadResult(self._result_chunks(ref, artifact, payload, offset))
|
||||
ack = await stub.UploadResult(
|
||||
self._result_chunks(ref, artifact, payload, offset),
|
||||
metadata=metadata,
|
||||
)
|
||||
if ack.committed:
|
||||
break
|
||||
resumed = int(ack.bytes_received)
|
||||
@@ -967,25 +1270,29 @@ class WorkerClient:
|
||||
request.session_token = self._session_token
|
||||
offset = 0
|
||||
complete = False
|
||||
handle = None
|
||||
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
|
||||
handle = await to_thread_and_drain_on_cancel(open, destination, "wb")
|
||||
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"
|
||||
)
|
||||
await to_thread_and_drain_on_cancel(_write_all, handle, chunk.data)
|
||||
offset += len(chunk.data)
|
||||
if chunk.last:
|
||||
complete = True
|
||||
break
|
||||
await to_thread_and_drain_on_cancel(handle.close)
|
||||
handle = None
|
||||
except asyncio.CancelledError:
|
||||
await to_thread_and_drain_on_cancel(_close_and_remove, handle, destination)
|
||||
raise
|
||||
except Exception:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
await to_thread_and_drain_on_cancel(_close_and_remove, handle, destination)
|
||||
raise
|
||||
if not complete:
|
||||
with contextlib.suppress(OSError):
|
||||
os.remove(destination)
|
||||
await to_thread_and_drain_on_cancel(_close_and_remove, None, destination)
|
||||
raise RuntimeError("input download ended before its final chunk")
|
||||
|
||||
def _executor_kwargs(self, assignment: pb.TaskAssignment) -> dict[str, Callable]:
|
||||
@@ -1120,6 +1427,7 @@ __all__ = [
|
||||
"MAX_MESSAGE_BYTES",
|
||||
"UPLOAD_STAGE",
|
||||
"PendingResult",
|
||||
"TerminalRegistrationError",
|
||||
"WorkerClient",
|
||||
"WorkerConfig",
|
||||
"backoff_delay",
|
||||
|
||||
@@ -10,7 +10,7 @@ import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from worker.capacity import derive_concurrency
|
||||
from worker.capacity import clamp_concurrency, derive_concurrency
|
||||
from worker.deadlines import Deadlines
|
||||
from worker.errors import ErrorClass, WorkerError
|
||||
from worker.lifecycle import Attempt, PriorityClass, Task
|
||||
@@ -225,7 +225,7 @@ def capability_to_pb(cap: dict) -> pb.ModelCapability:
|
||||
resident=bool(cap.get("resident")),
|
||||
min_memory_bytes=int(cap.get("min_memory_bytes") or 0),
|
||||
precision=str(cap.get("precision") or ""),
|
||||
derived_concurrency=max(0, declared),
|
||||
derived_concurrency=clamp_concurrency(declared, allow_zero=True),
|
||||
cpu_fallback=bool(cap.get("cpu_fallback")),
|
||||
repo_ids=list(cap.get("repo_ids") or []),
|
||||
display_name=str(cap.get("display_name") or ""),
|
||||
@@ -243,7 +243,9 @@ def capability_from_pb(message: pb.ModelCapability) -> dict:
|
||||
"resident": message.resident,
|
||||
"min_memory_bytes": message.min_memory_bytes,
|
||||
"precision": message.precision,
|
||||
"derived_concurrency": message.derived_concurrency,
|
||||
"derived_concurrency": clamp_concurrency(
|
||||
message.derived_concurrency, allow_zero=True
|
||||
),
|
||||
"cpu_fallback": message.cpu_fallback,
|
||||
"repo_ids": list(message.repo_ids),
|
||||
"display_name": message.display_name,
|
||||
|
||||
+2349
-186
File diff suppressed because it is too large
Load Diff
+2
-1
@@ -7,7 +7,8 @@ VoiceStudio installer. Today:
|
||||
|------|------------|---------|
|
||||
| `omnivoice-tts-darwin-arm64` | `ServeurpersoCom/omnivoice.cpp` @ pinned SHA | GGUF inference runtime — Apple Silicon |
|
||||
| `omnivoice-tts-darwin-x86_64` | same | Intel Mac |
|
||||
| `omnivoice-tts-linux-x86_64` | same | Linux |
|
||||
| `omnivoice-tts-linux-x86_64` | same | Linux (x86_64) |
|
||||
| `omnivoice-tts-linux-aarch64` | same | Linux ARM64 — Apple Silicon under Asahi; built with GGML Vulkan where the toolchain supports it, so the Honeykrisp driver can accelerate generation |
|
||||
| `omnivoice-tts-windows-x86_64.exe` | same | Windows |
|
||||
| `checksums.sha256` | computed by `scripts/build-omnivoice-tts.sh` | SHA-256 manifest — verified by `VoiceStudioGGUFBackend.is_available()` |
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
},
|
||||
"frontend": {
|
||||
"name": "omnivoice-studio",
|
||||
"version": "0.5.0",
|
||||
"version": "0.5.1",
|
||||
"dependencies": {
|
||||
"@fontsource-variable/inter": "^5.3.0",
|
||||
"@fontsource-variable/source-serif-4": "^5.3.0",
|
||||
|
||||
@@ -5,9 +5,11 @@
|
||||
# docker compose -f deploy/docker-compose.yml --profile cpu up # CPU mode
|
||||
# docker compose -f deploy/docker-compose.yml --profile gpu up # NVIDIA GPU
|
||||
# docker compose -f deploy/docker-compose.yml --profile rocm up # AMD GPU (ROCm)
|
||||
# docker compose -f deploy/docker-compose.yml --profile worker-gpu up # NVIDIA worker
|
||||
# docker compose -f deploy/docker-compose.yml --profile worker-rocm up # AMD worker
|
||||
#
|
||||
# All services bind to port 3900, so they MUST be opt-in via profiles —
|
||||
# otherwise `compose up` would race them and one would fail to bind.
|
||||
# The Studio services publish the UI on loopback and the authenticated TLS
|
||||
# worker control plane on 7443. Worker-only services publish no host port.
|
||||
#
|
||||
# First run downloads ~4 GB of models. Progress is shown in logs.
|
||||
# Open http://localhost:3900 once the health check passes.
|
||||
@@ -15,9 +17,9 @@
|
||||
# SECURITY: The port is bound to 127.0.0.1 by default — only this
|
||||
# machine can reach the API. To expose VoiceStudio on your LAN (or
|
||||
# through a reverse proxy / tunnel), change the port mapping to
|
||||
# "0.0.0.0:3900:3900" or "3900:3900". VoiceStudio itself ships no
|
||||
# authentication — if you expose it, put it behind a reverse proxy
|
||||
# with auth (Caddy basic_auth, nginx + htpasswd, Tailscale, etc.).
|
||||
# "0.0.0.0:3900:3900" or "3900:3900". Export a long random
|
||||
# OMNIVOICE_API_KEY before starting a Studio profile; use HTTPS or an
|
||||
# encrypted private overlay whenever traffic leaves a fully trusted LAN.
|
||||
# ──────────────────────────────────────────────────────────────
|
||||
|
||||
services:
|
||||
@@ -33,6 +35,7 @@ services:
|
||||
profiles: ["cpu"]
|
||||
ports:
|
||||
- "127.0.0.1:3900:3900"
|
||||
- "${OMNIVOICE_WORKER_PUBLISH_HOST:-127.0.0.1}:${OMNIVOICE_WORKER_PORT:-7443}:${OMNIVOICE_WORKER_PORT:-7443}"
|
||||
volumes:
|
||||
- omnivoice-data:/app/omnivoice_data
|
||||
environment:
|
||||
@@ -53,6 +56,11 @@ services:
|
||||
# discoverable. If you front the container with your own auth proxy on
|
||||
# loopback, set this to 0 to re-enable the strict gate.
|
||||
- OMNIVOICE_SERVER_MODE=1
|
||||
# Required for server-mode settings, diagnostics, and other admin
|
||||
# mutations because Docker NAT cannot prove a browser is loopback.
|
||||
- OMNIVOICE_API_KEY=${OMNIVOICE_API_KEY:?export a long random OMNIVOICE_API_KEY before starting a Studio profile}
|
||||
- OMNIVOICE_WORKER_PORT=${OMNIVOICE_WORKER_PORT:-7443}
|
||||
- OMNIVOICE_WORKER_ENDPOINT_HOST=${OMNIVOICE_WORKER_ENDPOINT_HOST:-}
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-sf", "http://localhost:3900/health"]
|
||||
interval: 30s
|
||||
@@ -71,6 +79,7 @@ services:
|
||||
profiles: ["gpu"]
|
||||
ports:
|
||||
- "127.0.0.1:3900:3900"
|
||||
- "${OMNIVOICE_WORKER_PUBLISH_HOST:-127.0.0.1}:${OMNIVOICE_WORKER_PORT:-7443}:${OMNIVOICE_WORKER_PORT:-7443}"
|
||||
volumes:
|
||||
- omnivoice-data:/app/omnivoice_data
|
||||
environment:
|
||||
@@ -86,6 +95,9 @@ services:
|
||||
# See the CPU service above — relaxes the loopback origin gate for the
|
||||
# headless Docker deployment (issue #261). Set to 0 to re-enable it.
|
||||
- OMNIVOICE_SERVER_MODE=1
|
||||
- OMNIVOICE_API_KEY=${OMNIVOICE_API_KEY:?export a long random OMNIVOICE_API_KEY before starting a Studio profile}
|
||||
- OMNIVOICE_WORKER_PORT=${OMNIVOICE_WORKER_PORT:-7443}
|
||||
- OMNIVOICE_WORKER_ENDPOINT_HOST=${OMNIVOICE_WORKER_ENDPOINT_HOST:-}
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-sf", "http://localhost:3900/health"]
|
||||
interval: 30s
|
||||
@@ -121,6 +133,7 @@ services:
|
||||
profiles: ["rocm"]
|
||||
ports:
|
||||
- "127.0.0.1:3900:3900"
|
||||
- "${OMNIVOICE_WORKER_PUBLISH_HOST:-127.0.0.1}:${OMNIVOICE_WORKER_PORT:-7443}:${OMNIVOICE_WORKER_PORT:-7443}"
|
||||
devices:
|
||||
- /dev/kfd
|
||||
- /dev/dri
|
||||
@@ -136,6 +149,9 @@ services:
|
||||
# loopback origin gate for the headless Docker deployment.
|
||||
- OMNIVOICE_BIND_HOST=0.0.0.0
|
||||
- OMNIVOICE_SERVER_MODE=1
|
||||
- OMNIVOICE_API_KEY=${OMNIVOICE_API_KEY:?export a long random OMNIVOICE_API_KEY before starting a Studio profile}
|
||||
- OMNIVOICE_WORKER_PORT=${OMNIVOICE_WORKER_PORT:-7443}
|
||||
- OMNIVOICE_WORKER_ENDPOINT_HOST=${OMNIVOICE_WORKER_ENDPOINT_HOST:-}
|
||||
# RDNA3 consumer cards (RX 7900 XTX/XT and friends, gfx1100): if the
|
||||
# GPU is not detected, uncomment the override below. The backend
|
||||
# auto-sets it for known consumer GFX IDs, so try without it first.
|
||||
@@ -148,5 +164,79 @@ services:
|
||||
start_period: 180s
|
||||
restart: unless-stopped
|
||||
|
||||
# ── Worker-only NVIDIA GPU mode
|
||||
# Generate a join code on the control plane, then start with:
|
||||
# OMNIVOICE_WORKER_TOKEN='ovw_…' docker compose \
|
||||
# -f deploy/docker-compose.yml --profile worker-gpu up -d
|
||||
# No port is published: uvicorn only hosts the application lifespan that
|
||||
# owns the outbound worker agent. The browser UI is not needed or exposed.
|
||||
omnivoice-worker-gpu:
|
||||
image: ghcr.io/debpalash/omnivoice-studio:latest
|
||||
build:
|
||||
context: ..
|
||||
dockerfile: deploy/Dockerfile
|
||||
container_name: omnivoice-worker-gpu
|
||||
profiles: ["worker-gpu"]
|
||||
entrypoint: ["python3", "-m", "uvicorn"]
|
||||
command: ["backend.main:app", "--host", "127.0.0.1", "--port", "3900"]
|
||||
volumes:
|
||||
- omnivoice-worker-data:/app/omnivoice_data
|
||||
environment:
|
||||
- HF_HOME=/app/omnivoice_data/huggingface
|
||||
- HF_TOKEN=${HF_TOKEN:-}
|
||||
- OMNIVOICE_DATA_DIR=/app/omnivoice_data
|
||||
- PYTHONPATH=/app/backend
|
||||
- PYTHONUNBUFFERED=1
|
||||
- OMNIVOICE_SERVER_MODE=0
|
||||
- OMNIVOICE_WORKER_MODE=1
|
||||
- OMNIVOICE_WORKER_TOKEN=${OMNIVOICE_WORKER_TOKEN:-}
|
||||
- OMNIVOICE_WORKER_ENDPOINT=${OMNIVOICE_WORKER_ENDPOINT:-}
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: 1
|
||||
capabilities: [gpu]
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-fsS", "http://127.0.0.1:3900/workers/agent/readiness"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 6
|
||||
start_period: 180s
|
||||
restart: unless-stopped
|
||||
|
||||
# ── Worker-only AMD GPU (ROCm) mode
|
||||
omnivoice-worker-rocm:
|
||||
image: ghcr.io/debpalash/omnivoice-studio:rocm
|
||||
container_name: omnivoice-worker-rocm
|
||||
profiles: ["worker-rocm"]
|
||||
entrypoint: ["python3", "-m", "uvicorn"]
|
||||
command: ["backend.main:app", "--host", "127.0.0.1", "--port", "3900"]
|
||||
devices:
|
||||
- /dev/kfd
|
||||
- /dev/dri
|
||||
volumes:
|
||||
- omnivoice-worker-rocm-data:/app/omnivoice_data
|
||||
environment:
|
||||
- HF_HOME=/app/omnivoice_data/huggingface
|
||||
- HF_TOKEN=${HF_TOKEN:-}
|
||||
- OMNIVOICE_DATA_DIR=/app/omnivoice_data
|
||||
- PYTHONPATH=/app/backend
|
||||
- PYTHONUNBUFFERED=1
|
||||
- OMNIVOICE_SERVER_MODE=0
|
||||
- OMNIVOICE_WORKER_MODE=1
|
||||
- OMNIVOICE_WORKER_TOKEN=${OMNIVOICE_WORKER_TOKEN:-}
|
||||
- OMNIVOICE_WORKER_ENDPOINT=${OMNIVOICE_WORKER_ENDPOINT:-}
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-fsS", "http://127.0.0.1:3900/workers/agent/readiness"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 6
|
||||
start_period: 180s
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
omnivoice-data:
|
||||
omnivoice-worker-data:
|
||||
omnivoice-worker-rocm-data:
|
||||
|
||||
@@ -43,21 +43,30 @@ the entire pipeline runs on CPU, just slower. Pull size: ~5 GB compressed
|
||||
## Quick start (CPU)
|
||||
|
||||
```bash
|
||||
export OMNIVOICE_API_KEY="$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')"
|
||||
|
||||
docker run -d --name omnivoice \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
palashdeb/omnivoice-studio:latest
|
||||
```
|
||||
|
||||
Open <http://localhost:3900>. The first run downloads a few GB of model weights —
|
||||
follow `docker logs -f omnivoice` to watch progress.
|
||||
follow `docker logs -f omnivoice` to watch progress. When the UI asks for an
|
||||
API key, paste the generated value; settings and diagnostic actions require
|
||||
this administrator session because Docker NAT hides the browser's true
|
||||
loopback origin.
|
||||
|
||||
## Quick start (NVIDIA GPU)
|
||||
|
||||
```bash
|
||||
export OMNIVOICE_API_KEY="$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')"
|
||||
|
||||
docker run -d --name omnivoice --gpus all \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
palashdeb/omnivoice-studio:latest
|
||||
@@ -74,9 +83,12 @@ CUDA-only and runs on CPU on AMD hardware). No toolkit needed — pass the GPU
|
||||
through as device nodes; the host only needs the `amdgpu` kernel driver:
|
||||
|
||||
```bash
|
||||
export OMNIVOICE_API_KEY="$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')"
|
||||
|
||||
docker run -d --name omnivoice \
|
||||
--device /dev/kfd --device /dev/dri \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
palashdeb/omnivoice-studio:rocm
|
||||
@@ -87,8 +99,9 @@ Podman users: same two `--device` flags (Quadlet: `AddDevice=/dev/kfd` +
|
||||
`-e HSA_OVERRIDE_GFX_VERSION=11.0.0` if the GPU isn't detected — details in
|
||||
the [Docker install guide](https://github.com/debpalash/VoiceStudio/blob/main/docs/install/docker.md).
|
||||
|
||||
There's also a Compose file in the repo with `cpu` / `gpu` / `rocm` profiles
|
||||
— see the [Docker install guide](https://github.com/debpalash/VoiceStudio/blob/main/docs/install/docker.md).
|
||||
There's also a Compose file in the repo with `cpu` / `gpu` / `rocm` profiles,
|
||||
plus `worker-gpu` / `worker-rocm` profiles that lend a headless GPU without
|
||||
publishing the web UI — see the [Docker install guide](https://github.com/debpalash/VoiceStudio/blob/main/docs/install/docker.md).
|
||||
|
||||
---
|
||||
|
||||
@@ -98,12 +111,12 @@ There's also a Compose file in the repo with `cpu` / `gpu` / `rocm` profiles
|
||||
|-----|--------------|
|
||||
| `:latest` | **Rolling preview** — latest commit on `main`, at or ahead of the last release. This is the preview channel; pin `:stable` for production. |
|
||||
| `:stable` | Most recent versioned release (updated on every `v*` git tag) |
|
||||
| `:0.5.0` | Exact release version |
|
||||
| `:0.5.1` | Exact release version |
|
||||
| `:0.5` | Latest patch within the `0.5` minor |
|
||||
| `:main` | Alias of the same rolling `main` build as `:latest` |
|
||||
| `:sha-xxxxxxx` | A specific commit (produced by manual workflow dispatch) |
|
||||
| `:rocm` | **AMD GPU (ROCm) build** of the rolling preview — the ROCm analogue of `:latest` |
|
||||
| `:stable-rocm`, `:0.5.0-rocm`, `:0.5-rocm`, `:sha-xxxxxxx-rocm` | ROCm builds of the corresponding tags above |
|
||||
| `:stable-rocm`, `:0.5.1-rocm`, `:0.5-rocm`, `:sha-xxxxxxx-rocm` | ROCm builds of the corresponding tags above |
|
||||
|
||||
Preview builds always come from `main` and never version-sort below `:stable`,
|
||||
so upgrades flow naturally. The same images and tags
|
||||
|
||||
+1
-1
@@ -201,7 +201,7 @@ None on the critical path to world-class. All are answers to real demand.
|
||||
| Kill per-segment disk round-trip | 🟡 | Long-video assembly stays disk-backed to bound RAM. Unchanged same-rate natural segments now skip the redundant decode → scratch encode → decode cycle; fresh segments still persist once and reload for assembly. |
|
||||
| Cold start ≤1.5 s to first audible sample | 🟡 | Installed models preload in the background and `scripts/bench_pipeline.py` measures cold/warm synthesis; target is not yet verified. |
|
||||
| Speculative regeneration on hover | ⏳ | — |
|
||||
| Crash-sandbox engines (subprocess isolation) | 🟡 | Killable sidecar engines and opt-in `omnivoice-subprocess` are live; the default in-process engine can still take down the server on a native crash. |
|
||||
| Crash-sandbox engines (subprocess isolation) | 🟡 | Killable sidecar engines and opt-in `omnivoice-subprocess` are live; the default OmniVoice engine now routes through the crash sandbox on MPS, while CUDA/ROCm/CPU retain the lower-overhead in-process path. |
|
||||
| Interaction budgets (<50 ms UI, <200 ms preview, <4 s first seg) | 🟡 | `/ws/tts` reports real TTFA, total generation time and RTF; frontend responsiveness instrumentation exists, but no cross-surface budget gate yet. |
|
||||
| Dedicated dev-week per quarter | ⏳ | Cadence not yet booked. |
|
||||
|
||||
|
||||
@@ -77,7 +77,7 @@ Both spike URLs are **real, live, and the intended artifacts**. The "VoiceStudio
|
||||
|
||||
# GGUF variant — adds a bundled binary, not a Python dep
|
||||
# Built once per platform in CI and stored alongside the Tauri installer:
|
||||
# bin/omnivoice-tts-{darwin-arm64,darwin-x86_64,windows-x86_64,linux-x86_64}
|
||||
# bin/omnivoice-tts-{darwin-arm64,darwin-x86_64,windows-x86_64,linux-x86_64,linux-aarch64}
|
||||
# Quants pulled at first use via huggingface_hub
|
||||
```
|
||||
|
||||
@@ -149,7 +149,8 @@ User opens app / first run │ Settings → Engines → Default
|
||||
┌─────────────────────────────────────────────────────────────────────────┐
|
||||
│ SubprocessBackend (Phase 2 primitive) │
|
||||
│ │
|
||||
│ spawn: bin/omnivoice-tts-{darwin-arm64|x86_64|win|linux} │
|
||||
│ spawn: bin/omnivoice-tts-{darwin-arm64|x86_64|win|linux| │
|
||||
│ linux-aarch64} │
|
||||
│ --model $HF_HUB_CACHE/.../omnivoice-base-{Q}.gguf │
|
||||
│ --codec $HF_HUB_CACHE/.../omnivoice-tokenizer-{Q}.gguf │
|
||||
│ --lang <user lang> │
|
||||
@@ -229,7 +230,8 @@ bin/ # bundled in installer (per platform)
|
||||
├── omnivoice-tts-darwin-arm64
|
||||
├── omnivoice-tts-darwin-x86_64
|
||||
├── omnivoice-tts-windows-x86_64.exe
|
||||
└── omnivoice-tts-linux-x86_64
|
||||
├── omnivoice-tts-linux-x86_64
|
||||
└── omnivoice-tts-linux-aarch64
|
||||
|
||||
.planning/decisions/
|
||||
├── SPIKE-01-gguf.md # ADR (this research → human review → ratified)
|
||||
|
||||
+16
-7
@@ -1,10 +1,14 @@
|
||||
# Agentic voice: VoiceStudio as a TTS/STT provider
|
||||
|
||||
VoiceStudio exposes an **OpenAI-compatible API**, so any agent framework that
|
||||
VoiceStudio exposes a **local speech platform**—OpenAI-compatible batch audio,
|
||||
a versioned transcription WebSocket, native dictation control, and MCP—so any agent framework that
|
||||
speaks to OpenAI's audio endpoints can use your local VoiceStudio for speech —
|
||||
in your own cloned voice, with nothing leaving your machine. You bring the
|
||||
agent runtime; VoiceStudio is the voice.
|
||||
|
||||
For dictating directly into Claude Code, Codex, Pi, Antigravity CLI, Herdr, or
|
||||
another focused prompt, use the [Rust control sidecar](speech-platform.md).
|
||||
|
||||
This is "agentic v1": VoiceStudio is a provider, not the orchestrator. You wire
|
||||
your own agent (a support line, a desk assistant, a Discord persona) and point
|
||||
its TTS/STT at VoiceStudio.
|
||||
@@ -16,13 +20,17 @@ its TTS/STT at VoiceStudio.
|
||||
|
||||
## The endpoints
|
||||
|
||||
VoiceStudio serves these on `http://localhost:3900/v1` (or your
|
||||
[remote backend URL](remote-gpu.md)):
|
||||
VoiceStudio's service root is `http://localhost:3900` (or your
|
||||
[remote backend URL](remote-gpu.md)). OpenAI-compatible clients use
|
||||
`http://localhost:3900/v1` as their base URL, while discovery stays at the
|
||||
service root: `http://localhost:3900/.well-known/voicestudio-speech`.
|
||||
|
||||
| OpenAI route | VoiceStudio support |
|
||||
|---|---|
|
||||
| `POST /v1/audio/speech` | TTS. `model` = engine id, `voice` = a voice-profile id (your clone) or preset, `response_format` incl. `pcm` and `wav`, `speed`. Default output is 24 kHz. |
|
||||
| `POST /v1/audio/transcriptions` | STT (Whisper-family). |
|
||||
| `WS /v1/audio/transcriptions/stream` | Live partial/final STT from PCM or WebM. |
|
||||
| `GET /.well-known/voicestudio-speech` | Machine-readable transport discovery. |
|
||||
| `GET /v1/audio/voices` | list available voices (VoiceStudio extension). |
|
||||
|
||||
A contract test (`tests/test_agentic_provider_contract.py`) pins this request
|
||||
@@ -74,10 +82,11 @@ single local agent, pipecat is lighter.
|
||||
|
||||
## Remote backend
|
||||
|
||||
Running VoiceStudio on a [remote GPU box](remote-gpu.md)? Use that backend's URL
|
||||
as `base_url` and pass its `OMNIVOICE_API_KEY` as the `api_key` — the same
|
||||
bearer the rest of the app uses. Keep it on your tailnet, not the open
|
||||
internet.
|
||||
Running VoiceStudio on a [remote GPU box](remote-gpu.md)? Append `/v1` to that
|
||||
backend's service-root URL for the OpenAI client's `base_url`, and pass its
|
||||
`OMNIVOICE_API_KEY` as the `api_key` — the same bearer the rest of the app uses.
|
||||
Keep the unmodified service root for `/.well-known/voicestudio-speech`
|
||||
discovery, and keep the backend on your tailnet, not the open internet.
|
||||
|
||||
## Use your own voice responsibly
|
||||
|
||||
|
||||
@@ -46,6 +46,15 @@ failing at spawn time; build one with
|
||||
`scripts/build-omnivoice-tts.sh --platform <slug>` or use the default
|
||||
in-process engine.
|
||||
|
||||
**Linux ARM64 (Asahi Apple Silicon):** the `linux-aarch64` binary prefers
|
||||
GGML's Vulkan backend when built on a host with `glslc` and the Khronos
|
||||
SPIRV headers installed (Arch: `pacman -S shaderc spirv-headers`; Debian:
|
||||
`apt install glslc libvulkan-dev spirv-headers`), so Apple GPUs accelerate
|
||||
generation through the open-source Honeykrisp driver. Without those deps the
|
||||
build falls back to CPU. Expect roughly 2–4x slower generation than macOS
|
||||
Metal while upstream Mesa and llama.cpp Vulkan optimizations mature; still
|
||||
well ahead of CPU-only.
|
||||
|
||||
## Integrity and self-healing
|
||||
|
||||
Before reporting ready, the engine:
|
||||
|
||||
@@ -6,12 +6,12 @@ wedged generation can be hard-killed and its VRAM/device reclaimed.
|
||||
|
||||
## Why this engine exists
|
||||
|
||||
The default `omnivoice` engine runs in-process on the GPU worker pool. On
|
||||
VRAM-tight machines (Apple Silicon MPS especially) a heavy generation or model
|
||||
load can exceed its execution budget. When that happens the worker is
|
||||
"abandoned" but **cannot be killed** (Python cannot interrupt a native torch /
|
||||
MPS call), so it keeps holding the GPU device until it finishes on its own, and
|
||||
every later synth queues behind it and hangs (#730 / #1190).
|
||||
An in-process `omnivoice` engine runs on the GPU worker pool. On VRAM-tight
|
||||
machines a heavy generation or model load can exceed its execution budget.
|
||||
When that happens the worker is "abandoned" but **cannot be killed** (Python
|
||||
cannot interrupt a native torch call), so it keeps holding the GPU device until
|
||||
it finishes on its own, and every later synth queues behind it and hangs
|
||||
(#730 / #1190).
|
||||
|
||||
`omnivoice-subprocess` runs the model in a child process spawned via the same
|
||||
`SubprocessBackend` primitive used by IndexTTS, Supertonic-3, and dots.tts. A
|
||||
@@ -26,21 +26,23 @@ structurally cannot do.
|
||||
must recover on its own instead of hanging until a manual restart.
|
||||
- **VRAM-starved MPS hosts** that hit the abandoned-worker cascade.
|
||||
|
||||
For interactive single-shot use on a machine with comfortable VRAM, the default
|
||||
in-process `omnivoice` engine is faster (no stdio round-trip) and remains the
|
||||
default.
|
||||
On Apple Silicon, the default `omnivoice` id automatically uses this isolated
|
||||
implementation. CUDA, ROCm, and CPU keep the in-process implementation and its
|
||||
lower call overhead.
|
||||
|
||||
## Selecting it
|
||||
|
||||
- **Settings -> Engines**, or
|
||||
- **Model Catalogue → Engines**, or
|
||||
- `OMNIVOICE_TTS_BACKEND=omnivoice-subprocess`
|
||||
|
||||
It is **opt-in**; the in-process engine stays the default, so existing setups
|
||||
see no change.
|
||||
The explicit engine is opt-in on CUDA, ROCm, and CPU. Apple Silicon gets the
|
||||
same isolation automatically while keeping the default `omnivoice` id in APIs,
|
||||
Settings, and saved projects.
|
||||
|
||||
## Platform support
|
||||
|
||||
- **CUDA, MPS, and CPU** (same as the in-process VoiceStudio engine).
|
||||
- **CUDA, AMD ROCm on Linux, MPS, and CPU** (same as the in-process
|
||||
VoiceStudio engine).
|
||||
- **No extra install.** Unlike IndexTTS / dots.tts / Supertonic-3, this sidecar
|
||||
runs under VoiceStudio's own interpreter, because the goal here is crash
|
||||
isolation, not dependency isolation. If the default `omnivoice` engine works
|
||||
@@ -53,11 +55,9 @@ see no change.
|
||||
- A wedged generation is **killed and recovered** at the recv-timeout deadline
|
||||
(`OMNIVOICE_SIDECAR_RECV_TIMEOUT_S`, default 300s, aligned with the generate
|
||||
budget) instead of hanging indefinitely.
|
||||
- It does **not** carry the native advanced-parameter surface
|
||||
(`t_shift` / `layer_penalty_factor` / `position_temperature` /
|
||||
`class_temperature`) or parent-side seed determinism, because the generic
|
||||
engine path does not forward those. For plain voice-clone and design
|
||||
synthesis this is a non-issue.
|
||||
- The default Apple Silicon proxy preserves native advanced parameters,
|
||||
deterministic seeds, and longform quality settings across the process
|
||||
boundary.
|
||||
- The recv-timeout deadline is per call and assumes the route's text chunking:
|
||||
`/generate` and `/v1/audio/speech` split long text into pieces of at most
|
||||
`max_chunk_chars` before calling the engine, so each call stays short. A
|
||||
|
||||
@@ -9,7 +9,8 @@ and dictation all run on it out of the box.
|
||||
|
||||
- You want cloning plus the broadest language coverage (see
|
||||
[languages.md](../languages.md)).
|
||||
- You have a GPU (CUDA or Apple Silicon MPS) with ~6 GB VRAM or more.
|
||||
- You have a GPU (CUDA, AMD ROCm on Linux, or Apple Silicon MPS) with
|
||||
~6 GB VRAM or more.
|
||||
- You just installed VoiceStudio — it's already selected.
|
||||
|
||||
For low-VRAM or CPU-only machines, the
|
||||
@@ -18,7 +19,7 @@ quantized native binary with a much smaller memory footprint.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Runs on CUDA, MPS (Apple Silicon), or CPU — auto-detected.
|
||||
- Runs on CUDA, AMD ROCm (Linux), MPS (Apple Silicon), or CPU — auto-detected.
|
||||
- Recommended VRAM floor: **6 GB** on a dedicated GPU. This is the only
|
||||
engine with a measured floor: on 4 GB cards (GTX 1650 Ti, Quadro P2000 —
|
||||
issues [#1226](https://github.com/debpalash/VoiceStudio/issues/1226) /
|
||||
@@ -43,8 +44,10 @@ The env var overrides the persisted UI choice.
|
||||
|
||||
- Weights load lazily on first use and are shared with the rest of the app
|
||||
(dubbing, dictation) — the model is never double-loaded.
|
||||
- On CUDA the model runs fp16 with `torch.compile`; a speech recognizer is
|
||||
co-loaded for the cloning path.
|
||||
- On CUDA and ROCm the model runs fp16 with `torch.compile`; PyTorch exposes
|
||||
ROCm/HIP devices through its `cuda` API, while VoiceStudio's engine matrix
|
||||
reports the hardware as ROCm. A speech recognizer is co-loaded for the
|
||||
cloning path.
|
||||
- Output is 24 kHz mono; the shared mastering chain (highpass + compressor)
|
||||
is tuned for this rate and applied automatically.
|
||||
- Cloning takes a short reference clip (`ref_audio`); 3–10 seconds is the
|
||||
|
||||
@@ -4,6 +4,12 @@ VoiceStudio dictation records from the system-wide shortcut, transcribes
|
||||
locally, and—where the desktop permits it—inserts the result into the app where
|
||||
the shortcut was pressed. The pill never needs keyboard focus.
|
||||
|
||||
The same flow is available to other applications through the bundled Rust
|
||||
control sidecar. Herdr actions, editor extensions, agent hooks, and scripts can
|
||||
start or stop VoiceStudio's capture over loopback HTTP/JSON-RPC or stream their
|
||||
own microphone audio to the versioned WebSocket API. See the
|
||||
[local speech platform](../speech-platform.md) for the protocol and examples.
|
||||
|
||||
## Use it
|
||||
|
||||
1. Choose an installed dictation model in the Model Catalogue.
|
||||
|
||||
+96
-9
@@ -13,18 +13,31 @@ and [`palashdeb/omnivoice-studio` on Docker Hub](https://hub.docker.com/r/palash
|
||||
> |-----|--------------|
|
||||
> | `:latest` | **Rolling preview** — latest commit on `main`, at or ahead of the last release. This is the preview channel; pin `:stable` for production. |
|
||||
> | `:stable` | Most recent versioned release (updated on every `v*` git tag) |
|
||||
> | `:0.5.0` | Exact release version |
|
||||
> | `:0.5.1` | Exact release version |
|
||||
> | `:0.5` | Latest patch within the 0.5 minor |
|
||||
> | `:main` | Alias of the same rolling `main` build as `:latest` |
|
||||
> | `:sha-xxxxxxx` | Specific commit (produced by manual workflow dispatch) |
|
||||
> | `:rocm` | **AMD GPU (ROCm) build** of the rolling preview — the ROCm analogue of `:latest` |
|
||||
> | `:stable-rocm`, `:0.5.0-rocm`, `:0.5-rocm`, `:sha-xxxxxxx-rocm` | ROCm builds of the corresponding CUDA tags above |
|
||||
> | `:stable-rocm`, `:0.5.1-rocm`, `:0.5-rocm`, `:sha-xxxxxxx-rocm` | ROCm builds of the corresponding CUDA tags above |
|
||||
>
|
||||
> Versioning rule: preview builds always come from `main` and never
|
||||
> version-sort below `:stable` — upgrades flow naturally.
|
||||
>
|
||||
> **Note on the update-channel toggle:** The update-channel UI (Settings → About → Update channel) is part of the Tauri desktop app's built-in auto-updater. It does **not** apply to the Docker image — the Docker image is the headless web-server build. To update your Docker deployment, pull the new image tag and recreate the container (`docker compose pull && docker compose up -d`).
|
||||
|
||||
Docker's NAT prevents the backend from proving that a browser is on the host,
|
||||
so server-mode settings and diagnostics require an administrator API key even
|
||||
when the published port is loopback-only. Generate one before using any Studio
|
||||
profile or `docker run` command below:
|
||||
|
||||
```bash
|
||||
export OMNIVOICE_API_KEY="$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')"
|
||||
```
|
||||
|
||||
Keep that shell open until the container starts. The web UI asks for this key
|
||||
and exchanges it for a short-lived browser session; it does not persist the
|
||||
master key.
|
||||
|
||||
## Pull and run (CPU)
|
||||
|
||||
```bash
|
||||
@@ -32,6 +45,7 @@ docker pull ghcr.io/debpalash/omnivoice-studio:latest
|
||||
|
||||
docker run -d --name omnivoice \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
ghcr.io/debpalash/omnivoice-studio:latest
|
||||
@@ -51,6 +65,7 @@ Open [http://localhost:3900](http://localhost:3900). The first run downloads
|
||||
```bash
|
||||
docker run -d --name omnivoice --gpus all \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
ghcr.io/debpalash/omnivoice-studio:latest
|
||||
@@ -71,11 +86,43 @@ as plain device nodes (no container toolkit needed):
|
||||
docker run -d --name omnivoice \
|
||||
--device /dev/kfd --device /dev/dri \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
ghcr.io/debpalash/omnivoice-studio:rocm
|
||||
```
|
||||
|
||||
### AMD GPU on WSL2
|
||||
|
||||
WSL exposes AMD compute through `/dev/dxg`, not native Linux's `/dev/kfd` and
|
||||
`/dev/dri`. First install ROCm and `librocdxg` in the WSL distribution and
|
||||
confirm the host-side `rocminfo` lists the GPU. Then use the WSL-specific
|
||||
bridge flags from AMD's `librocdxg` container contract:
|
||||
|
||||
```bash
|
||||
docker run -d --name omnivoice \
|
||||
--device /dev/dxg \
|
||||
-v /usr/lib/wsl/lib/libdxcore.so:/usr/lib/libdxcore.so \
|
||||
-v /opt/rocm/lib/librocdxg.so:/usr/lib/librocdxg.so \
|
||||
-v /opt/rocm/share/rocdxg/dids.conf:/usr/share/rocdxg/dids.conf \
|
||||
-e HSA_ENABLE_DXG_DETECTION=1 \
|
||||
--cap-add SYS_PTRACE \
|
||||
--security-opt seccomp=unconfined \
|
||||
--ipc=host --shm-size 8G \
|
||||
-p 127.0.0.1:3900:3900 \
|
||||
-e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-v omnivoice-data:/app/omnivoice_data \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
ghcr.io/debpalash/omnivoice-studio:rocm
|
||||
```
|
||||
|
||||
The image currently uses ROCm 7.2.x, so `HSA_ENABLE_DXG_DETECTION=1` is
|
||||
required; AMD removed that requirement only in ROCk 7.13. The ptrace and
|
||||
unconfined-seccomp flags weaken container isolation, so keep the published port
|
||||
on `127.0.0.1` and do not run untrusted workloads in this container. See AMD's
|
||||
[`librocdxg` WSL container instructions](https://github.com/ROCm/librocdxg#4-container-launch--wsl-specific-flags)
|
||||
for the driver/runtime compatibility matrix.
|
||||
|
||||
The same flags work with **Podman** (`podman run --device /dev/kfd
|
||||
--device /dev/dri …`); in a **Quadlet** unit that's two `AddDevice=` lines:
|
||||
|
||||
@@ -87,9 +134,10 @@ AddDevice=/dev/kfd
|
||||
AddDevice=/dev/dri
|
||||
PublishPort=127.0.0.1:3900:3900
|
||||
Volume=omnivoice-data:/app/omnivoice_data
|
||||
Environment=OMNIVOICE_API_KEY=replace-with-a-long-random-key
|
||||
```
|
||||
|
||||
Release pins exist too: `:stable-rocm`, `:0.5.0-rocm`, `:0.5-rocm` mirror
|
||||
Release pins exist too: `:stable-rocm`, `:0.5.1-rocm`, `:0.5-rocm` mirror
|
||||
the CUDA tags exactly.
|
||||
|
||||
> **Consumer cards and APUs (RX 6000/7000, Strix Point/Halo):** the backend
|
||||
@@ -119,7 +167,11 @@ ROCm container `omnivoice-studio-rocm` (CPU: `omnivoice-studio`, NVIDIA:
|
||||
|
||||
(ROCm-built PyTorch reports through `torch.cuda.*` — `True` plus your card's
|
||||
name means torch can see the GPU.) That check alone isn't proof the app is
|
||||
using it: **Settings → System** shows the device VoiceStudio actually resolved.
|
||||
using it: **Settings → Performance & Device** shows the device VoiceStudio
|
||||
actually resolved.
|
||||
**Model Catalogue → Engines** should report both `omnivoice` and
|
||||
`omnivoice-subprocess` as accelerated on ROCm, rather than a CPU-fallback
|
||||
warning.
|
||||
If it reads `cpu` while the command above prints `True`, the backend log line
|
||||
starting `Falling back to CPU:` names the architecture mismatch it hit.
|
||||
|
||||
@@ -129,8 +181,8 @@ interpreter. To verify this invariant on an older or custom image, compare
|
||||
torch.version.hip)"` with `docker exec <container> sh -c 'tr "\\0" " "
|
||||
</proc/1/cmdline'`; PID 1 must begin with `python3 -m uvicorn`.
|
||||
|
||||
If the command prints `False`, **Settings → System** now says why, and the
|
||||
three answers need different fixes:
|
||||
If the command prints `False`, run **Settings → About → Run self-check**;
|
||||
the GPU row says why. Native Linux has three common answers:
|
||||
|
||||
| What it says | What to do |
|
||||
|---|---|
|
||||
@@ -138,9 +190,16 @@ three answers need different fixes:
|
||||
| `this process cannot open it` | A group problem. Run `ls -l /dev/kfd /dev/dri/render*` **on the host**, and pass those GIDs with `--group-add`. The numbers differ between machines — a `--group-add 39` copied from someone else's command grants nothing. |
|
||||
| `no GPU was enumerated` | The device nodes are fine and the runtime still found nothing — usually a card newer than the image's ROCm. Check `rocminfo` on the host, and see the `HSA_OVERRIDE_GFX_VERSION` note above. |
|
||||
|
||||
On WSL, the self-check instead distinguishes a missing `/dev/dxg` permission,
|
||||
the pre-7.13 `HSA_ENABLE_DXG_DETECTION` opt-in, and incomplete ROCDXG runtime
|
||||
mounts.
|
||||
|
||||
## Docker Compose (recommended)
|
||||
|
||||
```bash
|
||||
# Generate this once in the shell that runs Compose.
|
||||
export OMNIVOICE_API_KEY="$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')"
|
||||
|
||||
# CPU
|
||||
docker compose -f deploy/docker-compose.yml --profile cpu up -d
|
||||
|
||||
@@ -156,6 +215,31 @@ on the host. The backend inside the container binds to `0.0.0.0` so the
|
||||
host port mapping can forward — the host-side `127.0.0.1` binding is what
|
||||
enforces loopback-only.
|
||||
|
||||
### Worker-only GPU container
|
||||
|
||||
To lend a headless GPU to VoiceStudio running on another machine, generate a
|
||||
join code on that control plane and start one of the worker profiles:
|
||||
|
||||
```bash
|
||||
# NVIDIA
|
||||
OMNIVOICE_WORKER_TOKEN='ovw_…' docker compose \
|
||||
-f deploy/docker-compose.yml --profile worker-gpu up -d
|
||||
|
||||
# AMD / ROCm
|
||||
OMNIVOICE_WORKER_TOKEN='ovw_…' docker compose \
|
||||
-f deploy/docker-compose.yml --profile worker-rocm up -d
|
||||
```
|
||||
|
||||
These profiles publish no HTTP port and require no browser UI. The join code
|
||||
must advertise a LAN or private-overlay address the container can reach, not
|
||||
the control plane's `127.0.0.1`. Enrollment state persists in a dedicated
|
||||
volume, so the container reconnects after a restart even though the join code
|
||||
is single-use. Container health becomes green only after the control plane
|
||||
accepts that registration; a missing or invalid token stays unhealthy instead
|
||||
of reporting the generic web backend as ready. See [Remote GPU
|
||||
workers](../remote-workers.md) for enrollment, approval, routing, and security
|
||||
details.
|
||||
|
||||
## LAN access
|
||||
|
||||
<a id="lan-access"></a>
|
||||
@@ -182,7 +266,8 @@ prebuilt image via `docker run -e` (the older `VITE_OMNIVOICE_API` is inlined at
|
||||
*build* time and cannot be set on a prebuilt image):
|
||||
|
||||
```bash
|
||||
docker run -e OMNIVOICE_PUBLIC_API_BASE=https://api.your-host.example \
|
||||
docker run -e OMNIVOICE_API_KEY="$OMNIVOICE_API_KEY" \
|
||||
-e OMNIVOICE_PUBLIC_API_BASE=https://api.your-host.example \
|
||||
-p 0.0.0.0:3900:3900 \
|
||||
ghcr.io/debpalash/omnivoice-studio:latest
|
||||
```
|
||||
@@ -230,8 +315,10 @@ Two paths are worth persisting across container restarts:
|
||||
origin, but Docker's NAT makes every request look non-loopback, so the gate
|
||||
used to 403 the whole admin UI (issue #261). The image now ships with
|
||||
`OMNIVOICE_SERVER_MODE=1`, which relaxes that gate for the headless
|
||||
deployment — exposure is instead governed by your `-p` port mapping (keep the
|
||||
`127.0.0.1:` prefix to stay local) plus the optional share PIN. If you front
|
||||
deployment. Admin mutations still require `OMNIVOICE_API_KEY`; all commands
|
||||
above pass it into the container, and the UI prompts for it on first use.
|
||||
Exposure is governed by your `-p` port mapping (keep the `127.0.0.1:` prefix
|
||||
to stay local) plus authentication. If you front
|
||||
the container with your own auth proxy on loopback, set `OMNIVOICE_SERVER_MODE=0`
|
||||
to re-enable the strict gate.
|
||||
- **Media-preview 404 in LAN mode:** see the [LAN access](#lan-access) section
|
||||
|
||||
+17
-4
@@ -38,13 +38,14 @@ Everything above, plus the toolchain:
|
||||
libwebkit2gtk-4.1-dev libgtk-3-dev libpango1.0-dev libcairo2-dev \
|
||||
libsoup-3.0-dev libgdk-pixbuf-2.0-dev \
|
||||
libayatana-appindicator3-dev librsvg2-dev libssl-dev libxdo-dev \
|
||||
gstreamer1.0-plugins-good \
|
||||
libasound2-dev build-essential curl wget file
|
||||
|
||||
# Fedora
|
||||
sudo dnf install webkit2gtk4.1-devel libappindicator-gtk3-devel librsvg2-devel openssl-devel
|
||||
sudo dnf install webkit2gtk4.1-devel libappindicator-gtk3-devel librsvg2-devel openssl-devel libxdo-devel gstreamer1-plugins-good
|
||||
|
||||
# Arch
|
||||
sudo pacman -S --needed base-devel webkit2gtk-4.1 libayatana-appindicator librsvg openssl xdotool
|
||||
sudo pacman -S --needed base-devel webkit2gtk-4.1 libayatana-appindicator librsvg openssl xdotool gst-plugins-good
|
||||
```
|
||||
|
||||
- Optional: a **Hugging Face token** for diarization + the larger TTS engines
|
||||
@@ -86,6 +87,13 @@ pkg-config --exists \
|
||||
&& echo "Tauri system libraries are ready"
|
||||
```
|
||||
|
||||
`bun desktop` also checks the native `libxdo` linker input and GStreamer's
|
||||
`autoaudiosink` before starting. The latter is required even if you do not plan
|
||||
to record: WebKitGTK 2.52 aborts its renderer when a page creates an audio
|
||||
element without that plugin, which otherwise turns a running app blank. The
|
||||
launcher prints one distro-specific install command when either dependency is
|
||||
missing.
|
||||
|
||||
The first app launch downloads model weights on demand. Subsequent launches
|
||||
reuse the Rust build, Python environment, and installed models.
|
||||
|
||||
@@ -378,9 +386,14 @@ dependency sync — matched to the app's pinned `torch==2.8.0` (the rocm6.2
|
||||
index only ever published up to torch 2.5.1, so it silently failed the
|
||||
reinstall and left the CPU-only CUDA build in place).
|
||||
|
||||
**2. Environment variable (existing installs / headless).** Set
|
||||
**2. Environment variable (existing installs / headless / source).** Set
|
||||
`OMNIVOICE_TORCH_VARIANT=rocm` before launching — the next bootstrap performs
|
||||
the same ROCm reinstall. `OMNIVOICE_TORCH_INDEX=<url>` overrides the wheel
|
||||
the same ROCm reinstall. Source installs honour it too:
|
||||
`OMNIVOICE_TORCH_VARIANT=rocm bun run desktop` swaps torch right after
|
||||
`uv sync` and launches the backend without re-syncing, so the wheel is not
|
||||
reverted on the next start (#1665). Without the variable, `bun run desktop`
|
||||
restores the lockfile's CUDA build — a hand-swapped ROCm wheel does not
|
||||
survive it. `OMNIVOICE_TORCH_INDEX=<url>` overrides the wheel
|
||||
index when you need a different ROCm version — e.g. AMD publishes newer
|
||||
driver-matched builds (7.2.x) at `repo.radeon.com` as a `--find-links` page
|
||||
rather than a PyPI-style index:
|
||||
|
||||
@@ -632,14 +632,28 @@ the error persistently on a current build, that's section **14** (a wedged GPU
|
||||
job), section **14d** (the backend never started), or the crash notice above —
|
||||
not this window.
|
||||
|
||||
Desktop startup, **Retry**, storage reset, setup re-entry, in-app uninstall,
|
||||
app shutdown, and automatic crash recovery also share one backend lifecycle
|
||||
owner. Overlapping start attempts wait and attach to the healthy process,
|
||||
while reset/setup/uninstall keep exclusive ownership through teardown and
|
||||
disk changes, and shutdown joins any in-progress launch before stopping it.
|
||||
They no longer launch a second backend that fails on port 3900, leave a
|
||||
misleading crash notice, delete an environment from under a starting child,
|
||||
or orphan one on exit. Quitting also interrupts a first-run `uv` install
|
||||
instead of waiting for a long download to finish. Unix builds give backend
|
||||
lifespan cleanup a bounded SIGTERM grace period; Windows' hidden backend has no
|
||||
console, so its initial stop is best-effort and may proceed directly to the
|
||||
bounded forced tree cleanup. Surviving subprocess engines cannot retain ports
|
||||
or files (#1635).
|
||||
|
||||
## 14c. "Can't reach the backend" in a browser — `bun run dev`, Docker, or LAN share
|
||||
|
||||
**Symptom:** you're using VoiceStudio **outside the desktop app** — the dev
|
||||
stack (`bun run dev`), a Docker deployment, or a shared/remote backend — and
|
||||
requests fail with a "can't reach the backend" error.
|
||||
|
||||
These deployments have no desktop shell to supervise the backend, so newer
|
||||
builds make the backend **self-forensicate** instead:
|
||||
These deployments have no desktop shell to supervise the backend. Newer builds
|
||||
make the backend **self-forensicate**, and the dev runner adds bounded recovery:
|
||||
|
||||
- **The error tells you what it knows.** It now says whether the backend *was
|
||||
answering and stopped* ("it was answering 12 s ago … likely crashed or was
|
||||
@@ -649,11 +663,15 @@ builds make the backend **self-forensicate** instead:
|
||||
dev, `docker logs <container>` / `journalctl` on a server. (If Docker
|
||||
serves the page itself, the page can go down together with the backend —
|
||||
check the container first.)
|
||||
- **Dev exit banner.** `bun run dev`'s backend runs through
|
||||
`scripts/dev-backend.mjs`: when uvicorn dies with a non-zero exit, a boxed
|
||||
- **Dev crash recovery and exit banner.** `bun run dev`'s backend runs through
|
||||
`scripts/dev-backend.mjs`, which owns Python-file reloads directly so a dead
|
||||
server cannot hide behind a still-running uvicorn reload parent. When the
|
||||
backend dies with a non-zero exit, a boxed
|
||||
banner prints the exit code/signal, the last 20 lines of `omnivoice.log`,
|
||||
and an OOM-check hint (`journalctl -k | grep -i oom` on Linux) before
|
||||
`concurrently` tears the stack down.
|
||||
and an OOM-check hint (`journalctl -k | grep -i oom` on Linux), then restarts
|
||||
it after one second. Three restarts are allowed in a rolling 60-second
|
||||
window; a fourth crash exits and lets `concurrently` tear down the broken
|
||||
stack. Ctrl+C and other deliberate shutdowns never restart it.
|
||||
- **Crash notice on the next start.** The backend keeps a **run sentinel**
|
||||
(`run_sentinel.json` in its data folder) while running and clears it on a
|
||||
clean shutdown. If a start finds a stale sentinel whose process is gone,
|
||||
@@ -732,7 +750,7 @@ IPC custom protocol failed, Tauri will now use the postMessage interface instead
|
||||
TypeError: Failed to fetch
|
||||
```
|
||||
|
||||
**Cause:** the crash corrupted the WebView2 profile cache at
|
||||
**Cause:** the crash corrupted cache directories inside the WebView2 profile at
|
||||
`%LOCALAPPDATA%\com.debpalash.omnivoice-studio\EBWebView`. Both the IPC custom
|
||||
protocol *and* its postMessage fallback break, so the splash never hears the
|
||||
"ready" signal from the app shell (issue #879).
|
||||
@@ -740,16 +758,25 @@ protocol *and* its postMessage fallback break, so the splash never hears the
|
||||
**Fix:** current builds handle this automatically — if the splash gets no IPC
|
||||
signal within ~10 s it checks the backend over plain HTTP and proceeds on its
|
||||
own; if the backend isn't up either, after ~45 s a recovery panel appears with
|
||||
**Repair and restart** (Windows), which clears the WebView cache and relaunches.
|
||||
Your voices, projects, and settings are not touched — only browser display data
|
||||
is cleared.
|
||||
**Repair and restart** (Windows), which clears cache-only directories and
|
||||
relaunches. It deliberately preserves `Default\Local Storage` and
|
||||
`Default\IndexedDB`, where browser-owned settings and long-form projects live.
|
||||
|
||||
On older builds (≤ 0.3.8), or if the automatic repair fails, do it manually:
|
||||
quit VoiceStudio, delete the folder below, then start the app again.
|
||||
quit VoiceStudio, delete only the cache directories below, then start the app
|
||||
again. Do not delete the whole `EBWebView` profile; doing so also deletes
|
||||
browser-owned projects and settings.
|
||||
|
||||
<!-- validate: skip -->
|
||||
```powershell
|
||||
Remove-Item -Recurse -Force "$env:LOCALAPPDATA\com.debpalash.omnivoice-studio\EBWebView"
|
||||
$voiceStudioWebView = "$env:LOCALAPPDATA\com.debpalash.omnivoice-studio\EBWebView"
|
||||
@(
|
||||
"Default\Cache", "Default\Code Cache", "Default\GPUCache", "Default\DawnCache",
|
||||
"Default\Service Worker\CacheStorage", "Default\Service Worker\ScriptCache",
|
||||
"GPUCache", "DawnCache", "ShaderCache", "GrShaderCache", "GraphiteDawnCache"
|
||||
) | ForEach-Object {
|
||||
Remove-Item -Recurse -Force -ErrorAction SilentlyContinue (Join-Path $voiceStudioWebView $_)
|
||||
}
|
||||
```
|
||||
|
||||
## 16. macOS: microphone permission never prompts, VoiceStudio never appears in System Settings
|
||||
|
||||
+61
-13
@@ -23,7 +23,9 @@ administration.
|
||||
|
||||
## What you need
|
||||
|
||||
* VoiceStudio on both machines, on versions no more than two releases apart.
|
||||
* VoiceStudio builds with a compatible worker protocol on both machines. The
|
||||
durable-enrollment v2 boundary requires updating both sides; the app refuses
|
||||
an unsafe pairing with an update instruction before any task runs.
|
||||
* The worker machine must be able to **reach** this one over the network. Same
|
||||
LAN is enough at home; across networks, a VPN such as
|
||||
[Tailscale](https://tailscale.com/) is the reliable answer. The worker dials
|
||||
@@ -41,6 +43,23 @@ Settings → System → Remote workers → turn on **Use remote workers**.
|
||||
The panel shows the address workers should connect to, and a **Generate token**
|
||||
button.
|
||||
|
||||
For a Docker Compose Studio, start it with the host address workers can reach;
|
||||
Compose publishes the TLS worker port (`7443`) separately from the loopback-only
|
||||
web UI:
|
||||
|
||||
```bash
|
||||
OMNIVOICE_WORKER_ENDPOINT_HOST=192.168.1.20 \
|
||||
OMNIVOICE_WORKER_PUBLISH_HOST=0.0.0.0 docker compose \
|
||||
-f deploy/docker-compose.yml --profile gpu up -d
|
||||
```
|
||||
|
||||
Use the host's LAN or private-overlay address, not the container's bridge IP.
|
||||
The worker port is published on loopback by default; setting
|
||||
`OMNIVOICE_WORKER_PUBLISH_HOST=0.0.0.0` is the explicit opt-in that makes it
|
||||
reachable from the LAN. Keep the default when a host-side tunnel or proxy
|
||||
provides reachability. Until remote workers are enabled in VoiceStudio, the
|
||||
container has no process listening on the published control-plane port.
|
||||
|
||||
**2. Generate a join code.**
|
||||
|
||||
The panel shows it as text **and as a QR code**, with a countdown. Copy it, or
|
||||
@@ -64,9 +83,36 @@ another code.
|
||||
Headless machines still take the environment route:
|
||||
|
||||
```bash
|
||||
OMNIVOICE_WORKER_TOKEN='ovw_…' OMNIVOICE_WORKER_MODE=1 omnivoice
|
||||
OMNIVOICE_WORKER_TOKEN='ovw_…' OMNIVOICE_WORKER_MODE=1 \
|
||||
uv run uvicorn backend.main:app --host 127.0.0.1 --port 3900
|
||||
```
|
||||
|
||||
Run that command from the repository root. Uvicorn hosts the application
|
||||
lifespan that owns the worker agent; binding it to loopback means no Studio UI
|
||||
is exposed, and no browser interaction is required.
|
||||
|
||||
For a worker-only NVIDIA Docker container, use the included Compose profile:
|
||||
|
||||
```bash
|
||||
OMNIVOICE_WORKER_TOKEN='ovw_…' docker compose \
|
||||
-f deploy/docker-compose.yml --profile worker-gpu up -d
|
||||
```
|
||||
|
||||
Use `worker-rocm` instead for AMD GPUs. Neither profile publishes an HTTP
|
||||
port. The control-plane address inside the join code must be reachable from
|
||||
the container, so use its LAN or private-overlay address rather than
|
||||
`127.0.0.1`. Worker identity, pinned certificate, and endpoint persist in the
|
||||
profile's data volume. After the first successful enrollment, restarts ignore
|
||||
that same now-spent environment token and reconnect by proving possession of
|
||||
the identity key. Replacing it with a fresh join code can move a non-revoked
|
||||
worker to another control plane. A revoked identity remains revoked; start
|
||||
with a fresh worker data volume to generate a new identity.
|
||||
|
||||
The container reports healthy only after the control plane accepts its initial
|
||||
registration. A missing, malformed, expired, or rejected join code leaves the
|
||||
worker service running for diagnosis but unhealthy; inspect its logs, correct
|
||||
the token, and recreate the container.
|
||||
|
||||
`OMNIVOICE_WORKER_MODE` wins over the in-app switch when it is set, so a
|
||||
deployment that pins worker mode cannot be turned off from the UI — the panel
|
||||
says so instead of showing a switch that springs back.
|
||||
@@ -216,12 +262,13 @@ in 45s" — and **Resume** clears it immediately when you've fixed the machine.
|
||||
launch VoiceStudio recovers those tasks and reconciles with each worker about
|
||||
what is genuinely still in flight.
|
||||
|
||||
**Version or feature mismatch.** The protocol keeps a two-release compatibility
|
||||
window, but release numbers alone do not prove that a worker understands every
|
||||
additive command. Registration therefore also declares named features for task
|
||||
inputs, progress leases, remote model downloads, and the voice-identity render
|
||||
pipeline. A worker outside the
|
||||
version window, or one missing a required feature, is refused with
|
||||
**Version or feature mismatch.** Release numbers alone do not prove that a
|
||||
worker understands every additive command. Registration negotiates an explicit
|
||||
protocol range and declares named features for task inputs, progress leases,
|
||||
remote model downloads, and the voice-identity render pipeline. Durable
|
||||
enrollment changed the handshake from protocol v1 to v2, so that boundary is
|
||||
intentionally incompatible in either direction. A worker outside the supported
|
||||
protocol range, or one missing a required feature, is refused with
|
||||
`UPGRADE_REQUIRED` and an update instruction before any task runs. It can never
|
||||
silently render without reference audio, substitute a different voice, or leave
|
||||
a download stuck at 0%.
|
||||
@@ -280,13 +327,14 @@ new code.
|
||||
| `OMNIVOICE_REMOTE_WORKERS` | `1`/`0` — enable without the UI (headless, Docker) |
|
||||
| `OMNIVOICE_WORKER_PORT` | Control-plane port (default `7443`) |
|
||||
| `OMNIVOICE_WORKER_ENDPOINT_HOST` | Override the address shown to workers |
|
||||
| `OMNIVOICE_WORKER_PUBLISH_HOST` | Compose-only host address for publishing the control-plane port (default `127.0.0.1`; set `0.0.0.0` to opt into LAN reachability) |
|
||||
| `OMNIVOICE_INBOUND_NODE` | `1`/`0` — accept connections from other panels |
|
||||
| `OMNIVOICE_INBOUND_BIND` | Address to accept them on (default `127.0.0.1`) |
|
||||
| `OMNIVOICE_INBOUND_PORT` | Port to accept them on (default `7444`) |
|
||||
| `OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` | How long a model may sit unused before its VRAM is handed back (default `600`, minimum `5`) |
|
||||
| `OMNIVOICE_IDLE_SWEEP_SECONDS` | How often that check runs (default `60`, minimum `1`) |
|
||||
| `OMNIVOICE_WORKER_MODE` | `1` on the worker machine — overrides the in-app switch |
|
||||
| `OMNIVOICE_WORKER_TOKEN` | Join code, first run only (the in-app Join box is the usual route) |
|
||||
| `OMNIVOICE_WORKER_TOKEN` | Join code, consumed on first successful enrollment; a persisted container value is ignored on later restarts |
|
||||
| `OMNIVOICE_WORKER_ENDPOINT` | Control plane to dial when no code is being redeemed; normally remembered from the code |
|
||||
|
||||
`OMNIVOICE_ENGINE_IDLE_UNLOAD_SECONDS` and `OMNIVOICE_IDLE_SWEEP_SECONDS` exist
|
||||
@@ -333,10 +381,10 @@ scripts/verify-remote-worker.sh \
|
||||
`WORKER_ID`, `WORKER_SSH_TARGET`, `WORKER_START_COMMAND`, `VOICESTUDIO_API`,
|
||||
and `WORKER_CONTROL_PORT` are equivalent environment variables. Pass
|
||||
`--worker-start-command` (or its environment equivalent) when the worker does
|
||||
not start with `OMNIVOICE_WORKER_MODE=1 omnivoice`; it is printed only in the
|
||||
manual worker-loss procedure. The worker id is optional only when exactly one
|
||||
worker is connected. The script requires an SSH target so it can verify the
|
||||
worker's OS and NVIDIA GPU before accepting any result.
|
||||
not use the headless command documented above; it is printed only in the manual
|
||||
worker-loss procedure. The worker id is optional only when exactly one worker
|
||||
is connected. The script requires an SSH target so it can verify the worker's
|
||||
OS and NVIDIA GPU before accepting any result.
|
||||
|
||||
The check never deletes model caches or user data. It selects an engine the
|
||||
worker itself reports as absent for the missing-model check. Operations that
|
||||
|
||||
@@ -2,6 +2,23 @@
|
||||
|
||||
## TL;DR
|
||||
|
||||
> **Persistence amendment (2026-08-24, #1636):** the unified in-memory model and
|
||||
> Zustand actions below are unchanged, but the localStorage-only storage design
|
||||
> is superseded. Schema v9 stores unbounded `storyTracks`, `cast`,
|
||||
> `storyProjects`, `script`, `meta`, `lexicon`, and `voiceCast` in the local
|
||||
> IndexedDB database `omnivoice.longform`; `omnivoice.app` keeps only bounded
|
||||
> preferences and ids. Upgrading a v8 envelope commits its long-form payload to
|
||||
> IndexedDB before trimming localStorage. If compacting localStorage hits quota,
|
||||
> revision markers choose the newest full fallback without discarding a newer
|
||||
> edit. Startup retries IndexedDB reads three times; a persistent failure gates
|
||||
> the main UI behind localized Retry and explicit-clear actions. Explicit clear
|
||||
> first persists a local tombstone, so permanently blocked IndexedDB does not
|
||||
> trap startup or later resurrect inaccessible projects. Malformed records fail
|
||||
> closed into that gate, and retries reopen invalid connections. Lifecycle
|
||||
> commits force-flush their compact local envelope; preference-only reset keeps
|
||||
> revisioned and legacy full fallbacks plus pending-clear tombstones. Content
|
||||
> reset also clears bounded references to deleted files.
|
||||
|
||||
Today there are two long-form text-to-speech editors — **Stories** (`frontend/src/components/StoriesEditor.jsx`, multi-voice cast + per-line tracks) and **Audiobook** (`frontend/src/pages/AudiobookTab.jsx`, single raw-text script + book metadata) — and they share *nothing* at the data layer. Stories has a persisted project model in `storiesSlice` (cast/tracks/projects in localStorage via the root zustand persist); Audiobook has **zero persistence** (title/author/narrator/genre/cover/lexicon/format/loudness/text all live in component `useState` at `AudiobookTab.jsx:21-48` and evaporate on tab switch or reload).
|
||||
|
||||
This spec introduces **one project concept** — a `LongformProject` — that both editors bind to. It carries the shared book identity (`title/author/narrator/genre/year/description` + cover + lexicon) once, alongside the structured content (cast + tracks + raw script), plus a `mode` discriminator (`'stories' | 'audiobook'`). This (a) gives Audiobook the persistence it lacks today, (b) lets a single saved project be re-opened in either editor, turning **#24 (Story ⇄ Audiobook convert)** from a data-shuttling problem into a `mode` toggle, and (c) consolidates the two divergent persistence stories into one slice with a tested localStorage migration that is backward-compatible with existing `storyProjects`.
|
||||
@@ -30,7 +47,7 @@ This spec introduces **one project concept** — a `LongformProject` — that bo
|
||||
### Non-goals (explicitly deferred)
|
||||
- **#24's convert UI / mode-toggle button** — not built here. We ship the store + the `convertMode` action; the toggle button and the "open this project in the other editor" UX is #24.
|
||||
- **Backend changes.** `/longform/render`, `/audiobook` + `/audiobook/{plan,preview,cover,import}`, `/stories/encode` endpoints are untouched. This is a frontend store + binding change only. The SSE event vocabulary (pinned below in API/data shapes → SSE) is **read-only contract** here — neither editor's event handling changes.
|
||||
- **No alembic / `omnivoice_data/` change.** LongformProject is browser localStorage state (zustand persist via `store/index.ts:55-132`), not backend DB. The "backward-compatible project data → schema change goes through alembic" constraint applies to the **Python/`omnivoice_data/` SQLite path only**; the *localStorage analog* of that same constraint (no data loss, no manual migration) is satisfied here by the **versioned lazy `migrate` fn** (zustand persist's equivalent of an alembic upgrade — see Constraints → Backward-compatible data). Confirmed: Audiobook metadata/cover are client-side only today (the only server contact is `audiobookUploadCover` → `POST /audiobook/cover`, which returns `{ path: string }`, not stored project state).
|
||||
- **No alembic / `omnivoice_data/` change.** LongformProject is browser-owned state: Zustand supplies the in-memory/versioned shape, while schema v9 splits bounded fields into localStorage and unbounded documents into IndexedDB. The alembic rule applies only to the Python `omnivoice_data/` SQLite path, which is untouched. Browser upgrades are ordered and versioned in the storage adapter. Audiobook cover references remain client-side; `POST /audiobook/cover` returns a local backend path rather than storing project state.
|
||||
- **No change to the rendered output** of either editor — same spans, same SSE handling (`utils/sseParse.js`).
|
||||
- Merging the actual editor *components* into one — they stay two components; only the store unifies. (Component merge, if ever, is downstream of #22 shared `<VoiceSelector>`.)
|
||||
- Cover-file persistence as bytes. localStorage can't hold the blob; we persist only a re-uploadable reference + filename (see API/data shapes).
|
||||
@@ -79,29 +96,29 @@ This contract is the single most important completeness invariant in the task an
|
||||
### Why one slice, not a sibling slice for Audiobook
|
||||
The two editors share the Projects list. `ProjectsPage` (`frontend/src/pages/Projects.jsx`) already takes a `storyProjects` prop (default `[]` at `:79`) and an `onOpenStory` callback (`:82`), iterates `storyProjects` to build story cards (`:146-160`), and is rendered from `App.jsx:1046-1056` with `storyProjects={storyProjects}` (`:1051`) and `onOpenStory={(id) => { loadStoryProject(id); setMode('stories'); }}` (`:1054`). A unified store means the existing Projects UI lists *both* kinds of projects with no new plumbing — Audiobook projects simply appear there for free.
|
||||
|
||||
> **Disambiguation (verified in code):** `Projects.jsx` *also* maintains a separate `longformJobs` local state (`:102`, rendered at `:198-211`) — these are completed **backend** audiobook render jobs fetched from the server, of `type: 'audiobooks'`, NOT client-side projects. Do **not** conflate `longformJobs` with the new client `longformProjects`. The renamed projects array (`storyProjects` → `longformProjects`) flows in via the existing `storyProjects` prop (rename the prop or keep the prop name + pass the new array — see Integration); `longformJobs` is untouched.
|
||||
> **Disambiguation (verified in code):** `Projects.jsx` *also* maintains a separate `longformJobs` local state (`:102`, rendered at `:198-211`) — these are completed **backend** audiobook render jobs fetched from the server, of `type: 'audiobooks'`, NOT client-side projects. Do **not** conflate `longformJobs` with the client `storyProjects` collection. The project record type is unified as `LongformProject`, while the collection keeps its compatibility key and flows through the existing `storyProjects` prop; `longformJobs` is untouched.
|
||||
|
||||
A second slice would require a second projects list, two `currentProjectId`s, and bespoke routing — the opposite of unification.
|
||||
|
||||
### Migration strategy (zustand persist `version` 4 → 5)
|
||||
The root store's persist config (`store/index.ts:55-132`) bumps `version: 4` (`:115`) → `version: 5`, and the `migrate` function (`:120-130`) gains a `version < 5` branch that:
|
||||
1. Renames nothing destructively — old keys `storyProjects`, `storyTracks`, `cast`, `currentProjectId` (persisted at `index.ts:109-113`) are **read** and mapped forward.
|
||||
2. For each old `StoryProject` (`{id,name,tracks,cast,updatedAt}`), produces a `LongformProject` by spreading it over defaults: `mode: 'stories'`, `meta: {}`, `lexicon: {}`, `outputFormat: 'm4b'`, `loudness: 'off'`, `coverRef: null`, `defaultVoice: null`, `script: ''`. The spread `...sp` goes **last** so original `id/name/cast/tracks/updatedAt` always win over defaults.
|
||||
2. For each old `StoryProject` (`{id,name,tracks,cast,updatedAt}`), produces a `LongformProject`, then validates nested fields so `null` or a wrong type cannot override usable defaults. Non-empty string ids/names and array tracks/cast survive; invalid values become a generated id, `Untitled`, and empty arrays.
|
||||
3. Maps the loose working state: `storyTracks` and `cast` pass through unchanged (keep `storyTracks`/`cast` as the persisted/working key names to minimize component churn — see Working-state naming); seeds the new working metadata fields (`mode`, `meta`, `script`, `lexicon`, `coverRef`, `outputFormat`, `loudness`, `defaultVoice`) to defaults.
|
||||
4. Returns the upgraded partial. The migrate fn must **never throw** — `index.ts:121` already returns `{}` for non-object input, and `:122-128` passes old shapes through (`version < 4` branch). **Upgrade > crash** is this codebase's stated migration philosophy (the comment at `index.ts:116-119`, verified verbatim: *"Drop old persisted shapes rather than crashing the app… Upgrade > crash."*); we keep it: anything missing falls through to slice init defaults.
|
||||
|
||||
**Migration must defend against malformed persisted data (enumerated):**
|
||||
- `p.storyProjects` is **not an array** (corrupted blob, hand-edited localStorage, partial write) → `Array.isArray` guard, treat as `[]`; do not throw.
|
||||
- An individual project entry is **not an object** (e.g. `null`, a string) → skip it (filter to objects) rather than spreading a non-object; never produce a `longformProjects` entry that lacks `id`/`name`.
|
||||
- A project entry is **missing `tracks`/`cast`/`id`/`name`/`updatedAt`** → defaults fill them (`tracks: []`, `cast: []`, generated `id`, `name: 'Untitled'`, `updatedAt: 0`). A project with no `id` is still openable (we synthesize one) rather than silently dropped — but if `id` synthesis would collide, last-write-wins on save is acceptable (matches existing `saveProject` upsert at `:95-98`).
|
||||
- An individual project entry is **not an object** (e.g. `null`, a string) → skip it (filter to objects) rather than spreading a non-object; never produce a `storyProjects` entry that lacks `id`/`name`.
|
||||
- A project entry has missing, null, or wrong-typed `tracks`/`cast`/`id`/`name`/`updatedAt` → post-spread normalization fills them (`tracks: []`, `cast: []`, generated `id`, `name: 'Untitled'`, `updatedAt: 0`). A project with no usable `id` is still openable rather than silently dropped.
|
||||
- `currentProjectId` points to a project that **no longer exists** post-migration (e.g. was the malformed one we skipped) → leave `currentProjectId` as-is; the load-time guard in `loadProject` (`if (!p) return`, `:103`) already no-ops harmlessly, and `currentProject` (`StoriesEditor.jsx:217`) resolves to `null` → blank working state, no crash. Optionally null it during migration; either is safe.
|
||||
- `version` is `> 5` (user downgraded the app, then re-upgraded) → the function returns `persisted` unchanged via the final passthrough (`:129`); new fields already present, no double-migration.
|
||||
- `version` is `< 4` → existing `version < 4` branch (`:122-128`) runs **first**, then falls through to the `version < 5` branch in the same call (ordering: handle `< 4` passthrough, then `< 5` upgrade). Confirm a v2/v3 blob (no `storyProjects` key at all) yields `longformProjects: []`, not a throw.
|
||||
- `version` is `< 4` → existing `version < 4` branch (`:122-128`) runs **first**, then falls through to the `version < 5` branch in the same call (ordering: handle `< 4` passthrough, then `< 5` upgrade). Confirm a v2/v3 blob (no `storyProjects` key at all) yields `storyProjects: []`, not a throw.
|
||||
|
||||
Backward-compat guarantee: a user on v4 with three saved Stories projects reloads → sees the same three projects, same names, same cast/tracks, `mode: 'stories'` — opens identically in Stories. Nothing prompts them. New Audiobook projects they save afterward coexist in the same list.
|
||||
|
||||
### Working-state naming (minimize component churn)
|
||||
`StoriesEditor.jsx:115-127` binds **13 store selectors** (`storyTracks` `:115`, `setStoryTracks` `:116`, `cast` `:117`, `setCast` `:118`, `upsertCastMember` `:119`, `removeCastMember` `:120`, `setCharacterVoice` `:121`, `storyProjects` `:122`, `currentProjectId` `:123`, `saveProject` `:124`, `loadProject` `:125`, `newProject` `:126`, `deleteProject` `:127`), and `setStoryTracks` is re-wrapped at `:130-133` as `setTracks`. To keep these call sites stable, the **working content field names stay** (`storyTracks`, `cast`, `setStoryTracks`, `setCast`, `upsertCastMember`, etc., unchanged signatures from `storiesSlice.ts:71-83`). We **add** new working fields/actions for the shared metadata that Audiobook binds to. The *project record* type is renamed (`StoryProject` → `LongformProject`) and the projects array is renamed (`storyProjects` → `longformProjects`) with a deprecated alias to avoid breaking `ProjectsPage`/`App.jsx`/`StoriesEditor.jsx` in the same PR (see Integration points). This keeps the diff reviewable: Stories' line-editing code is unchanged; only project save/load gains metadata, and Audiobook gains store binding.
|
||||
`StoriesEditor.jsx:115-127` binds **13 store selectors** (`storyTracks` `:115`, `setStoryTracks` `:116`, `cast` `:117`, `setCast` `:118`, `upsertCastMember` `:119`, `removeCastMember` `:120`, `setCharacterVoice` `:121`, `storyProjects` `:122`, `currentProjectId` `:123`, `saveProject` `:124`, `loadProject` `:125`, `newProject` `:126`, `deleteProject` `:127`), and `setStoryTracks` is re-wrapped at `:130-133` as `setTracks`. To keep these call sites stable, the **working content field names stay** (`storyTracks`, `cast`, `setStoryTracks`, `setCast`, `upsertCastMember`, etc., unchanged signatures from `storiesSlice.ts:71-83`). We **add** new working fields/actions for the shared metadata that Audiobook binds to. The *project record* type is renamed (`StoryProject` → `LongformProject`), while the projects array retains the `storyProjects` compatibility key used by `ProjectsPage`/`App.jsx`/`StoriesEditor.jsx` (see Integration points). This keeps the diff reviewable: Stories' line-editing code is unchanged; only project save/load gains metadata, and Audiobook gains store binding.
|
||||
|
||||
> **Note on `storyProjects` consumers:** a repo-wide grep confirms exactly 6 files reference `storyProjects`: `store/storiesSlice.ts`, `store/storiesSlice.test.ts`, `store/index.ts`, `components/StoriesEditor.jsx`, `pages/Projects.jsx`, `App.jsx`. All must be accounted for in the rename + alias bridge.
|
||||
|
||||
@@ -135,8 +152,8 @@ This section enumerates every state the feature must handle and the exact behavi
|
||||
- **D1. Corrupted / non-object persisted blob** → `migrate` returns `{}` (`index.ts:121` guard); store boots to all slice defaults. No crash. *Test #6 garbage case.*
|
||||
- **D2. `storyProjects` non-array** → treated as `[]` (Array.isArray guard). *See Migration → malformed defenses.*
|
||||
- **D3. Malformed individual project entries** → object-filter + default-fill; never drop silently except non-objects. *See Migration.*
|
||||
- **D4. localStorage write fails (quota exceeded / private-mode).** zustand persist's `setItem` throwing is swallowed by the middleware (it logs, does not crash the app). Working state still functions in-memory for the session; it just won't survive reload. We add no new failure surface here, but the spec **acknowledges** large books (50-chapter scripts as a single `script` string + tracks) push localStorage size; quota is realistically multi-MB so a single book is fine, but a user with dozens of saved big projects could hit it. Mitigation note: no eviction policy in v1 (Stories already persists projects unbounded); flag for v0.4 if reported.
|
||||
- **D5. Two app windows / tabs open simultaneously** (Tauri can have multiple webviews) both writing the persisted store → last-write-wins, standard localStorage behavior; zustand persist does not cross-tab-sync by default. No regression vs. today (Stories already has this). Out of scope to fix; note it. (Identical on macOS/Windows/Linux — not a platform divergence; see Constraints → Cross-platform parity.)
|
||||
- **D4. IndexedDB write fails (quota/policy/private mode).** The adapter writes a revisioned full fallback to localStorage and synchronously verifies the coalesced flush. If that fallback also cannot be written, the orderly flush rejects instead of falsely acknowledging durability; intentional recovery reload/exit actions log the failure and continue. Preference reset preserves even an unrevisioned legacy full envelope because it may be the last copy left after adding the revision itself exceeded quota.
|
||||
- **D5. Two browser main tabs write the same revision.** Each tab can start from revision N and independently produce N+1. A full local fallback with a revision equal to IndexedDB wins on hydration (`>=`), because that local write necessarily followed the durable record it is recovering from. Desktop's widget remains read-only and cannot create this collision.
|
||||
- **D6. Migration runs but `cast` working field is missing on a v4 blob that only had `storyProjects`** (user never touched the editor) → working `cast` falls through to slice init `DEFAULT_CAST` (`:68`). Confirmed safe.
|
||||
|
||||
### E. Routing / open-from-Projects edge cases
|
||||
@@ -181,18 +198,18 @@ This section enumerates every state the feature must handle and the exact behavi
|
||||
| `frontend/src/store/index.ts:36-37` | Update import: `import type { StoriesSlice } from './storiesSlice'` / `import { createStoriesSlice } from './storiesSlice'` → `import type { LongformSlice } from './longformSlice'` / `import { createLongformSlice, genProjectId, SLICE_DEFAULTS } from './longformSlice'` (keep alias export from the slice so other importers don't break; import `genProjectId`+`SLICE_DEFAULTS` for the migrate fn). |
|
||||
| `frontend/src/store/index.ts:45` | `AppStore` type: `… & StoriesSlice & …` → `… & LongformSlice & …`. |
|
||||
| `frontend/src/store/index.ts:63` | `...createStoriesSlice(set, get, api)` → `...createLongformSlice(set, get, api)`. |
|
||||
| `frontend/src/store/index.ts:107-113` | `partialize`: today persists `storyTracks` (stripped at `:109-110`), `cast` (`:111`), `storyProjects` (`:112`), `currentProjectId` (`:113`). Rename `storyProjects` → `longformProjects`; keep `storyTracks`/`cast`/`currentProjectId`; **add** loose working `meta`/`script`/`lexicon`/`coverRef`/`outputFormat`/`loudness`/`defaultVoice`/`mode` so an unsaved Audiobook session survives reload — matching how `storyTracks` is already persisted loose (`:109`). Keep the existing transient strip on `storyTracks` (`:109-110`). **Do NOT persist any transient render flag** (no `generating`/`output`/`progress`/`exporting`/`exportPct` — they live in component `useState`, not the slice; confirm none sneak in) so no ghost-spinner on reload (F3). Exact partialize block diff pinned in API/data shapes → partialize. |
|
||||
| `frontend/src/store/index.ts:107-113` | `partialize`: today persists `storyTracks` (stripped at `:109-110`), `cast` (`:111`), `storyProjects` (`:112`), `currentProjectId` (`:113`). Keep `storyProjects`/`storyTracks`/`cast`/`currentProjectId`; **add** loose working `meta`/`script`/`lexicon`/`coverRef`/`outputFormat`/`loudness`/`defaultVoice`/`mode` so an unsaved Audiobook session survives reload — matching how `storyTracks` is already persisted loose (`:109`). Keep the existing transient strip on `storyTracks` (`:109-110`). **Do NOT persist any transient render flag** (no `generating`/`output`/`progress`/`exporting`/`exportPct` — they live in component `useState`, not the slice; confirm none sneak in) so no ghost-spinner on reload (F3). Exact partialize block diff pinned in API/data shapes → partialize. |
|
||||
| `frontend/src/store/index.ts:115` | `version: 4` → `version: 5`. |
|
||||
| `frontend/src/store/index.ts:120-130` | Add a `version < 5` branch in `migrate` (after the existing `version < 4` branch at `:122-128`); see Design → Migration **and** the malformed-data defenses (D1–D6). Keep the non-object guard at `:121` and the final passthrough at `:129`. Ensure `version < 4` blobs flow through the `< 4` branch *then* the `< 5` branch in the same call. Full migrate fn body pinned in API/data shapes → Migration. **This `migrate` fn IS the localStorage analog of an alembic upgrade** — Test #3 is the data-integrity gate. |
|
||||
| `frontend/src/pages/AudiobookTab.jsx:21-48` | Replace local `useState` for `text` (`:21`), `defaultVoice` (`:22`), `format` (`:34`), `loudness` (`:35`), `meta` (`:36-38`) with `useAppStore` bindings (selectors pinned in API/data shapes → component binding). **Lexicon stays partially local (B3):** the editable *rows* array (with in-progress blanks) is component-local `useState`, hydrated from store `lexicon` on mount, flushed to store via `setLexicon(lexDict())` on change. Cover stays partially local: `coverFile`/`coverPreview` blob (`:39-40`) is transient (matches the `URL.revokeObjectURL` cleanup at `:57-64`); the persisted *reference* is the new `coverRef` (filename + server path). `clearCover` (`:57-61`) must also `setCoverRef(null)` (C5). The `lexDict()` helper (`:43-45`) and `setMetaField`/`setLexRow`/`addLexRow`/`removeLexRow` (`:46-49`) adapt to read/write the store (meta via `setProjectMeta` merge, lexicon via local-rows→`setLexicon`). The `format`/`loudness` `onChange` handlers (`:231`, `:239`) route through `setOutputPrefs`. **All new/changed user-facing strings via `t()` only** — no hardcoded labels; this file is **not** on the CJK-test allowlist, so any non-English literal here would fail CI (Constraints → Localization). |
|
||||
| `frontend/src/pages/AudiobookTab.jsx:113-175` (`onCreate`) | Read `meta`/`lexicon`/`format`/`loudness`/`text`/`defaultVoice` from store instead of local state. The `audiobookGenerate` body construction (`:130-138`) is **unchanged in shape** — only the *source* of each field moves to the store; keep the wire transforms in place (whitespace-filter J1, `loudness→null` J2). **Cover logic (C2):** `cover_path = coverFile ? (await audiobookUploadCover(coverFile)).path : (coverRef?.serverPath ?? null)` — and on a fresh upload, `setCoverRef({ filename: coverFile.name, serverPath: cover_path })`. Keep the upload-failure path (C4) intact via the existing `try/catch` (`:170-172`). All requests stay local (`/audiobook/*`); no new external endpoint (Constraints → Local-first). SSE event handling (`:151-167`) unchanged — same event vocabulary (API/data shapes → SSE). |
|
||||
| `frontend/src/pages/AudiobookTab.jsx:80-95, :97-111` (`onImport`, `onPreviewChapter`) | These set/read `text`/`defaultVoice`/`lex` (`:87-88`, `:101-104`); update to read/write the store-bound equivalents. `onImport` sets store `script` (was `setText(r.text)`, `:88`) and clears `plan` (transient, stays local). `onPreviewChapter` reads store `text`/`defaultVoice` + computed lexicon dict; the `audiobookPreviewChapter` body shape (`audiobook.ts:40`: `{text, chapter_index, default_voice, lexicon}`) is unchanged. |
|
||||
| `frontend/src/components/StoriesEditor.jsx:115-127` | Store-binding selectors. `storyProjects` (`:122`) → select `longformProjects` (via alias or renamed selector). The other 12 selectors (`:115-121`, `:123-127`) stay by name. `currentProject` computed at `:217` and the projects panel render follow the rename. (This file IS on the CJK allowlist — line-editing code may keep its existing functional CJK; no *new* hardcoded UI CJK should be added regardless.) |
|
||||
| `frontend/src/components/StoriesEditor.jsx:115-127` | Store-binding selectors. Keep `storyProjects` (`:122`) as the compatibility collection key. The other 12 selectors (`:115-121`, `:123-127`) stay by name. `currentProject` computed at `:217` and the projects panel render use the same collection. (This file IS on the CJK allowlist — line-editing code may keep its existing functional CJK; no *new* hardcoded UI CJK should be added regardless.) |
|
||||
| `frontend/src/components/StoriesEditor.jsx:135-140` (`_trackId` reseed) | The mount-only reseed effect (dep array `[]` at `:140`) doesn't re-run on cross-session `loadProject` (latent collision bug now reachable). Change dep array to `[currentProjectId]` or reseed in the add path. See Working-state naming note. |
|
||||
| `frontend/src/components/StoriesEditor.jsx:360-402` (`generateAll`) | **Now can attach metadata**: currently `:368-371` passes only `{ chapters, format }` to `longformRender`. Extend to pass `metadata`/`loudness`/`cover_path` from the unified project — `LongformRenderBody` (`audiobook.ts:100-108`) already accepts all three. Apply the **same empty-filter** as `onCreate` so a metadata-less Stories project still sends `metadata: null` (J3 — no wire regression). Map `loudness: 'off' → null`, reuse `coverRef.serverPath` if present. Exact extended body pinned in API/data shapes → Stories export. This is the immediate payoff of unification — Stories full export stops dropping book metadata. (Note: chapter-level `speed` is **not** a `LongformRenderBody.chapters[]` field; `speed` rides at span level only — `LongformRenderBody.chapters[].spans[].speed?` at `audiobook.ts:101`, set by `storyToSpans.js:40,47`.) SSE event handling (`:383-391`) unchanged. |
|
||||
| `frontend/src/App.jsx:251-252` | `const storyProjects = useAppStore(s => s.storyProjects)` (`:251`) → select `longformProjects`; `const loadStoryProject = useAppStore(s => s.loadProject)` (`:252`) — `loadProject` now also sets working `mode` from the record. |
|
||||
| `frontend/src/App.jsx:251-252` | Keep `const storyProjects = useAppStore(s => s.storyProjects)` (`:251`); `const loadStoryProject = useAppStore(s => s.loadProject)` (`:252`) now also sets working `mode` from the record. |
|
||||
| `frontend/src/App.jsx:1046-1056` (`ProjectsPage` render) | `storyProjects={storyProjects}` (`:1051`) passes the renamed array (keep the `storyProjects` *prop name* on `Projects.jsx` OR rename it there too — pick one and apply consistently). Make `onOpenStory` (`:1054`) **mode-aware AND load-result-aware (E5)** — exact body pinned in API/data shapes → Routing. `setMode` is the store action selected at `App.jsx:110`. Routing behaves identically on macOS/Windows/Linux (no platform branch). |
|
||||
| `frontend/src/pages/Projects.jsx:79, :146-160, :228` | If renaming the prop: `storyProjects = []` (`:79`) → `longformProjects = []`; the iteration at `:146-160` and the `useMemo` dep at `:228` follow. Story cards (`type:'stories'`, `:150`) can optionally branch icon/label on `sp.mode` (e.g. `BookMarked` for audiobook vs `BookOpen` for stories, mirroring the existing `:157` icon). **Subtitle for audiobook-mode projects** (`:153-154` currently builds `[story_lines, story_voices].filter(Boolean).join(' · ')`) should show a sensible audiobook summary (e.g. char count from `script.length` or a chapter hint) or fall back gracefully when `tracks`/`cast` are empty — don't render an empty subtitle for an audiobook. Any new card label/subtitle text via `t()` (Constraints → Localization). Leave `longformJobs` (`:102, :198-211`) untouched. |
|
||||
| `frontend/src/pages/Projects.jsx:79, :146-160, :228` | Keep the `storyProjects = []` prop (`:79`), its iteration at `:146-160`, and its `useMemo` dependency at `:228`. Story cards (`type:'stories'`, `:150`) can optionally branch icon/label on `sp.mode` (e.g. `BookMarked` for audiobook vs `BookOpen` for stories, mirroring the existing `:157` icon). **Subtitle for audiobook-mode projects** (`:153-154` currently builds `[story_lines, story_voices].filter(Boolean).join(' · ')`) should show a sensible audiobook summary (e.g. char count from `script.length` or a chapter hint) or fall back gracefully when `tracks`/`cast` are empty — don't render an empty subtitle for an audiobook. Any new card label/subtitle text via `t()` (Constraints → Localization). Leave `longformJobs` (`:102, :198-211`) untouched. |
|
||||
| `frontend/src/store/storiesSlice.test.ts` | Rename to `longformSlice.test.ts`; keep all existing assertions green (they're the backward-compat contract — `:16-114`); add new-field tests. The test harness (`:4-10`) calls `createStoriesSlice` directly (`:8`) and reads `storyTracks`/`cast`/`storyProjects`/`currentProjectId` — keep those names working via aliases. (Test files `*.test.ts` are **excluded from CodeQL** by `security.yml:102-105` and from the CJK source scan by the test-fixture exemption in `tests/test_no_hardcoded_cjk.py:90-94` — so the ported suite carries no SAST/CJK burden.) |
|
||||
| `frontend/src/i18n/locales/*.json` (**21 files**: ar, de, en, es, fr, hi, id, it, ja, ko, nl, pl, pt, ru, sv, th, tr, uk, vi, zh-CN, zh-TW — count verified: `ls locales/ \| wc -l` = 21) | No *new* user-facing strings strictly required for 31a/31b core — reuse existing keys: Audiobook keys live under `audiobook.*` (en.json `:112` block); Stories keys under `stories.*` (en.json `:26` block: `untitled:54`, `projectSaved:58`); Projects card keys under `projects.*` (`untitled_story`, `story_lines`, `story_voices` — used at `Projects.jsx:152-154`). **New keys (each × 21 locales) if 31b adds the affordances:** an Audiobook "Save project" button label + "Saved" toast (`audiobook.save_project`, `audiobook.project_saved`), an `audiobook.untitled` default name (B1), and optionally a "cover set: {filename}" label for the reloaded-cover state (C1). Each must go through `t()` and into **all 21** locales in the **same PR** (Docs-sync + Localization hard rules); a key present in `en.json` but missing in the other 20 is a CI/UX failure, not a follow-up. |
|
||||
|
||||
@@ -253,7 +270,7 @@ export interface LongformProject {
|
||||
}
|
||||
```
|
||||
|
||||
Concrete persisted-record example (a saved Audiobook-mode project, as it lives inside `omnivoice.app` localStorage → `state.longformProjects[0]`):
|
||||
Concrete project-record example (a saved Audiobook-mode project, as it lives in the `omnivoice.longform` IndexedDB payload → `storyProjects[0]`):
|
||||
```json
|
||||
{
|
||||
"id": "p_4f9ab2c1",
|
||||
@@ -306,7 +323,7 @@ export interface LongformSlice {
|
||||
defaultVoice: string | null;
|
||||
mode: LongformMode; // mode of the working project
|
||||
// --- projects ---
|
||||
longformProjects: LongformProject[]; // was storyProjects (storiesSlice.ts:39,69)
|
||||
storyProjects: LongformProject[]; // compatibility key retained (storiesSlice.ts:39,69)
|
||||
currentProjectId: string | null; // unchanged (storiesSlice.ts:40,70)
|
||||
|
||||
// --- existing actions (UNCHANGED signatures, storiesSlice.ts:71-83) ---
|
||||
@@ -357,17 +374,17 @@ saveProject: (name) => set((s) => {
|
||||
defaultVoice: s.defaultVoice,
|
||||
updatedAt: ts,
|
||||
};
|
||||
const exists = s.longformProjects.some((p) => p.id === id);
|
||||
const exists = s.storyProjects.some((p) => p.id === id);
|
||||
return {
|
||||
longformProjects: exists
|
||||
? s.longformProjects.map((p) => (p.id === id ? proj : p))
|
||||
: [...s.longformProjects, proj],
|
||||
storyProjects: exists
|
||||
? s.storyProjects.map((p) => (p.id === id ? proj : p))
|
||||
: [...s.storyProjects, proj],
|
||||
currentProjectId: id,
|
||||
};
|
||||
}),
|
||||
|
||||
loadProject: (id) => {
|
||||
const p = get().longformProjects.find((x) => x.id === id);
|
||||
const p = get().storyProjects.find((x) => x.id === id);
|
||||
if (!p) return; // E5: no-op when id missing
|
||||
set({
|
||||
storyTracks: (p.tracks || []).map((t) => ({ ...t })),
|
||||
@@ -504,39 +521,41 @@ type LongformSSEEvent =
|
||||
export type StoryProject = LongformProject; // alias for storiesSlice.ts:28
|
||||
export const createStoriesSlice = createLongformSlice; // alias for storiesSlice.ts:66
|
||||
export type StoriesSlice = LongformSlice; // alias for storiesSlice.ts:36
|
||||
// NOTE: the working-state field `storyProjects` was renamed to `longformProjects`.
|
||||
// Either (a) keep a `storyProjects` getter alias in the slice, or (b) update all 6
|
||||
// consumer files in the same PR. The test harness (storiesSlice.test.ts) reads
|
||||
// `.storyProjects` directly (test :19,:67,:79,:103), so (a) is the lowest-churn
|
||||
// bridge for one PR. (Getter alias: in createLongformSlice return, add
|
||||
// `get storyProjects() { return get().longformProjects; }` is NOT possible in a
|
||||
// plain object literal across set/get — instead provide a derived selector OR
|
||||
// keep `storyProjects` as a duplicated key kept in sync; simplest one-PR path is
|
||||
// (b): rename in all 6 files + port the test to `.longformProjects`.)
|
||||
// The record type was unified, but the persisted collection key remains
|
||||
// `storyProjects` for compatibility. Schema v9 moves that collection to the
|
||||
// `omnivoice.longform` IndexedDB payload without renaming it.
|
||||
```
|
||||
|
||||
### `partialize` block (exact diff vs `index.ts:107-113`)
|
||||
### Pre-adapter Zustand projection vs durable storage
|
||||
|
||||
`partialize` still projects the complete in-memory long-form state so Zustand's
|
||||
migration and slice API remain unchanged:
|
||||
|
||||
```ts
|
||||
// Stories/Longform Editor — persist the project; strip transient runtime fields
|
||||
// (generating, audioUrl) so a dead blob: URL / stuck spinner never rehydrates.
|
||||
storyTracks: s.storyTracks.map(({ id, character, text, profileId, emotion, speed }) =>
|
||||
({ id, character, text, profileId, emotion, speed })), // UNCHANGED (:109-110)
|
||||
cast: s.cast, // UNCHANGED (:111)
|
||||
longformProjects: s.longformProjects, // RENAMED from storyProjects (:112)
|
||||
currentProjectId: s.currentProjectId, // UNCHANGED (:113)
|
||||
// NEW loose working fields (unsaved Audiobook session survives reload, F3-safe):
|
||||
script: s.script,
|
||||
meta: s.meta,
|
||||
lexicon: s.lexicon,
|
||||
coverRef: s.coverRef,
|
||||
outputFormat: s.outputFormat,
|
||||
loudness: s.loudness,
|
||||
defaultVoice: s.defaultVoice,
|
||||
mode: s.mode,
|
||||
// DO NOT add: generating/output/progress/exporting/exportPct — they are component
|
||||
// useState (AudiobookTab.jsx:25-30, StoriesEditor.jsx:147-148), not slice fields.
|
||||
{
|
||||
storyTracks: s.storyTracks,
|
||||
cast: s.cast,
|
||||
storyProjects: s.storyProjects, // compatibility key retained
|
||||
script: s.script,
|
||||
meta: s.meta,
|
||||
lexicon: s.lexicon,
|
||||
voiceCast: s.voiceCast,
|
||||
currentProjectId: s.currentProjectId,
|
||||
coverRef: s.coverRef,
|
||||
outputFormat: s.outputFormat,
|
||||
loudness: s.loudness,
|
||||
defaultVoice: s.defaultVoice,
|
||||
projectMode: s.projectMode,
|
||||
}
|
||||
```
|
||||
|
||||
The schema-v9 storage adapter then splits that projection. IndexedDB
|
||||
`omnivoice.longform/documents/workspace` owns the unbounded payload
|
||||
(`storyTracks`, `cast`, `storyProjects`, `script`, `meta`, `lexicon`,
|
||||
`voiceCast`). localStorage `omnivoice.app` receives only bounded fields such as
|
||||
`currentProjectId`, `coverRef`, output preferences, app preferences, and ids.
|
||||
Transient track fields are stripped while constructing the IndexedDB record.
|
||||
|
||||
### Migration (v4 → v5), FULL `migrate` fn body (replaces `index.ts:120-130`)
|
||||
```ts
|
||||
migrate: (persisted, version) => {
|
||||
@@ -548,18 +567,19 @@ migrate: (persisted, version) => {
|
||||
}
|
||||
if (version < 5) {
|
||||
const rawProjects = Array.isArray(p.storyProjects) ? p.storyProjects : []; // (D2)
|
||||
p.longformProjects = rawProjects
|
||||
p.storyProjects = rawProjects
|
||||
.filter((sp: any) => sp && typeof sp === 'object') // (D3) drop non-objects
|
||||
.map((sp: any) => ({
|
||||
// defaults FIRST…
|
||||
id: genProjectId(), name: 'Untitled', mode: 'stories',
|
||||
cast: [], tracks: [], script: '', meta: {}, lexicon: {},
|
||||
coverRef: null, outputFormat: 'm4b', loudness: 'off',
|
||||
defaultVoice: null, updatedAt: 0,
|
||||
// …then real fields win (spread LAST): id/name/cast/tracks/updatedAt
|
||||
...sp,
|
||||
// Validate AFTER spread: malformed values cannot erase usable defaults.
|
||||
id: typeof sp.id === 'string' && sp.id.trim() ? sp.id : genProjectId(),
|
||||
name: typeof sp.name === 'string' && sp.name.trim() ? sp.name : 'Untitled',
|
||||
mode: sp.mode === 'audiobook' ? 'audiobook' : 'stories',
|
||||
cast: Array.isArray(sp.cast) ? sp.cast : [],
|
||||
tracks: Array.isArray(sp.tracks) ? sp.tracks : [],
|
||||
updatedAt: typeof sp.updatedAt === 'number' && Number.isFinite(sp.updatedAt)
|
||||
? sp.updatedAt : 0,
|
||||
}));
|
||||
delete p.storyProjects;
|
||||
// Loose working fields seed to defaults; storyTracks/cast pass through (D6).
|
||||
p.mode = 'stories';
|
||||
// script/meta/lexicon/coverRef/outputFormat/loudness/defaultVoice are absent →
|
||||
@@ -576,34 +596,42 @@ migrate: (persisted, version) => {
|
||||
### `onOpenStory` routing (App.jsx:1054 — exact body, E1–E5)
|
||||
```ts
|
||||
onOpenStory={(id) => {
|
||||
const rec = useAppStore.getState().longformProjects.find((x) => x.id === id);
|
||||
const rec = useAppStore.getState().storyProjects.find((x) => x.id === id);
|
||||
if (!rec) return; // E5: id no longer resolves → stay on Projects
|
||||
loadStoryProject(id); // = loadProject(id); sets working mode from rec
|
||||
setMode(rec.mode === 'audiobook' ? 'audiobook' : 'stories'); // E1/E2/E3 default-safe
|
||||
}}
|
||||
```
|
||||
|
||||
### Persisted-blob shape (whole `omnivoice.app` localStorage value)
|
||||
### Persisted storage shapes (schema v9)
|
||||
```jsonc
|
||||
// localStorage["omnivoice.app"] = { state: {...}, version: 5 }
|
||||
// v5 `state` (relevant slice keys only — other slices' keys unchanged):
|
||||
// localStorage["omnivoice.app"] — bounded compact envelope
|
||||
{
|
||||
"state": {
|
||||
"storyTracks": [ /* StoryTrack[] working */ ],
|
||||
"cast": [ { "id": "narrator", "name": "Narrator", "color": "#fabd2f", "profileId": null } ],
|
||||
"longformProjects": [ /* LongformProject[] — see record example above */ ],
|
||||
"currentProjectId": "p_4f9ab2c1",
|
||||
"script": "",
|
||||
"meta": {},
|
||||
"lexicon": {},
|
||||
"coverRef": null,
|
||||
"outputFormat": "m4b",
|
||||
"loudness": "off",
|
||||
"defaultVoice": null,
|
||||
"mode": "stories"
|
||||
// … plus all other slices' persisted keys (translateQuality, etc.)
|
||||
"projectMode": "stories"
|
||||
// … plus bounded fields from other slices; no manuscript/project arrays
|
||||
},
|
||||
"version": 5
|
||||
"version": 9
|
||||
}
|
||||
|
||||
// IndexedDB omnivoice.longform / documents / workspace — durable payload
|
||||
{
|
||||
"schema": 1,
|
||||
"revision": 12,
|
||||
"payload": {
|
||||
"storyTracks": [ /* sanitized StoryTrack[] working */ ],
|
||||
"cast": [ /* CastMember[] */ ],
|
||||
"storyProjects": [ /* LongformProject[]; compatibility key retained */ ],
|
||||
"script": "",
|
||||
"meta": {},
|
||||
"lexicon": {},
|
||||
"voiceCast": {}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -611,21 +639,21 @@ onOpenStory={(id) => {
|
||||
|
||||
All frontend tests run via `bunx vitest run` (package.json script `"test": "vitest run"` at `frontend/package.json:14`; `vitest ^4.1.5` at `:75`). Per MEMORY: the local loop must include `bunx vitest run`. pytest is irrelevant here — no backend change — **except** the project-wide `tests/test_no_hardcoded_cjk.py` gate (runs in CI), which scans git-tracked `.jsx`/`.ts` source: the renamed `longformSlice.ts` and modified `AudiobookTab.jsx` are scanned (not allowlisted), so the local loop should also include a quick CJK self-check on changed files before push (see Constraints → Localization).
|
||||
|
||||
1. **`longformSlice.test.ts` — backward-compat (must stay green):** the entire existing `storiesSlice.test.ts` suite (`:16-114` — two `describe` blocks: `storiesSlice` `:16-60`, `storiesSlice — projects` `:62-114`) ported verbatim against the renamed slice. The harness (`:4-10`) and `track()` helper (`:12-14`) carry over (rename `createStoriesSlice` import → `createLongformSlice`, and `.storyProjects` reads → `.longformProjects` if not aliased). Proves cast/track/project behavior is unchanged. Specifically the transient-strip test (`:107-113`) and the in-place-update test (`:73-81`) are load-bearing.
|
||||
1. **`longformSlice.test.ts` — backward-compat (must stay green):** the entire existing `storiesSlice.test.ts` suite (`:16-114` — two `describe` blocks: `storiesSlice` `:16-60`, `storiesSlice — projects` `:62-114`) ported verbatim against the renamed slice. The harness (`:4-10`) and `track()` helper (`:12-14`) carry over (rename `createStoriesSlice` import → `createLongformSlice`, while `.storyProjects` reads remain unchanged). Proves cast/track/project behavior is unchanged. Specifically the transient-strip test (`:107-113`) and the in-place-update test (`:73-81`) are load-bearing.
|
||||
2. **New project-shape tests:**
|
||||
- `saveProject` snapshots `meta`, `lexicon`, `script`, `coverRef`, `outputFormat`, `loudness`, `defaultVoice`, `mode` into the record (extends the existing snapshot test at `:63-71`). Assert the exact saved-record shape matches `LongformProject` (every field present).
|
||||
- Saving twice with `currentProjectId` updates metadata in place, no duplicate (H1) — `longformProjects.length === 1` after two saves with mutated `meta`.
|
||||
- Saving twice with `currentProjectId` updates metadata in place, no duplicate (H1) — `storyProjects.length === 1` after two saves with mutated `meta`.
|
||||
- `loadProject` restores all of the above into working state *and* sets working `mode` from the record (extends `:83-94`).
|
||||
- `newProject('audiobook')` seeds `mode==='audiobook'` + blank working surface; `newProject()` defaults to `'stories'` and clears tracks/cast (preserve the no-arg behavior of `:106`).
|
||||
- transient track-field stripping still holds (port of `:107-113`; `snapshotTracks` at `storiesSlice.ts:62-64`).
|
||||
- `convertMode('audiobook')` flips working `mode` without mutating `cast`/`tracks`/`script` (G1); `convertMode` to the current mode is a no-op (G2); `convertMode('bogus' as any)` is ignored (G3).
|
||||
- `setProjectMeta({title})` merges (does not clear `author`) (I1); `setOutputPrefs({loudness})` merges (`outputFormat`/`defaultVoice` untouched) (I2); `setOutputPrefs({defaultVoice: null})` overwrites to `null` (the `!== undefined` rule); `setLexicon(dict)` replaces (I3).
|
||||
3. **Migration test (the load-bearing one — the localStorage "alembic upgrade" gate):** construct a v4 persisted blob `{ storyProjects: [{id:'x',name:'A',tracks:[...],cast:[...],updatedAt:1}], storyTracks, cast, currentProjectId:'x' }`, run the `migrate` fn with `version=4`, assert: `longformProjects[0]` has the original `id/name/cast/tracks/updatedAt` plus `mode:'stories'` + metadata defaults (`script:''`, `meta:{}`, `lexicon:{}`, `coverRef:null`, `outputFormat:'m4b'`, `loudness:'off'`, `defaultVoice:null`); `storyProjects` is gone; no throw. **Malformed-input cases (D1–D6, each asserts no-throw):**
|
||||
3. **Migration test (the load-bearing one — the localStorage "alembic upgrade" gate):** construct a v4 persisted blob `{ storyProjects: [{id:'x',name:'A',tracks:[...],cast:[...],updatedAt:1}], storyTracks, cast, currentProjectId:'x' }`, run the `migrate` fn with `version=4`, assert: `storyProjects[0]` has the original `id/name/cast/tracks/updatedAt` plus `mode:'stories'` + metadata defaults (`script:''`, `meta:{}`, `lexicon:{}`, `coverRef:null`, `outputFormat:'m4b'`, `loudness:'off'`, `defaultVoice:null`); the compatibility key remains; no throw. **Malformed-input cases (D1–D6, each asserts no-throw):**
|
||||
- non-object `persisted` (string / null / number) → `{}`.
|
||||
- `storyProjects` not an array (an object / string) → `longformProjects: []`.
|
||||
- `storyProjects` not an array (an object / string) → `storyProjects: []`.
|
||||
- `storyProjects` array containing `null` / a string / `{}` → non-objects dropped; the `{}` entry gets a synthesized `id` (truthy `p_…`) + `name:'Untitled'` + default content.
|
||||
- a v2/v3 blob with **no** `storyProjects` key → `longformProjects: []`.
|
||||
- `version: 6` blob → returned unchanged (no re-migration; `longformProjects` left as-is).
|
||||
- a v2/v3 blob with **no** `storyProjects` key → `storyProjects: []`.
|
||||
- `version: 6` blob → returned unchanged (no re-migration; `storyProjects` left as-is).
|
||||
- `currentProjectId` pointing to a now-missing project → no throw; a later `loadProject(id)` no-ops.
|
||||
- Note: `migrate` is defined inline in the persist config (`index.ts:120-130`) — extract it to a named export (e.g. `migrateAppStore`) or test via a re-created store to exercise it.
|
||||
4. **AudiobookTab persistence test (component or store-level):** set `meta.title`, switch mode away and back (simulate by re-reading store) → title persists. Assert `audiobookGenerate` is still called with the **exact body shape** of `AudiobookGenerateBody` (`text`/`default_voice`/`format`/`loudness`/`cover_path`/`metadata`/`lexicon`) — only the *source* of each field changed. Assert the whitespace-only-meta field is dropped from the wire body (J1), `loudness:'off'→null` (J2), and a CJK `meta.title` round-trips byte-for-byte through the store (J4).
|
||||
@@ -640,11 +668,11 @@ All frontend tests run via `bunx vitest run` (package.json script `"test": "vite
|
||||
|
||||
Each relevant hard rule, and exactly how this task satisfies it:
|
||||
|
||||
- **Backward-compatible project data (no manual migration; DB→alembic / localStorage→versioned `migrate`).** Existing `storyProjects` in localStorage (persisted at `index.ts:112`, persist `version: 4`) must keep working with zero user action. **No `omnivoice_data/` SQLite / alembic touched** — this state is browser localStorage, so the alembic clause does not literally apply; its *intent* (versioned, tested, lossless upgrade with no manual step) is satisfied by zustand-persist's `version`-bump + lazy `migrate` fn — the localStorage analog of an alembic upgrade. Bump `version: 4 → 5` (`index.ts:115`), add a `version < 5` branch (`:120-130`, full body pinned in API/data shapes → Migration) that maps every old `StoryProject` forward (spread `...sp` **last** so original fields always win), seeds defaults, and **never throws** (matches the existing non-object guard `:121` and the "Upgrade > crash" philosophy `:116-119`). Malformed/corrupt blobs degrade to defaults, never a white screen (D1–D6). The alias bridge keeps the 6 `storyProjects` consumers compiling. The migrate is the data-integrity gate — Test #3 enumerates the realistic v4 blob + all malformed cases.
|
||||
- **Backward-compatible project data (no manual migration; DB→alembic / browser storage→versioned migration).** Existing v4–v8 `storyProjects` envelopes keep working with zero user action. The v5 shape migration validates nested project fields after applying defaults; schema v9 then commits the complete long-form payload to IndexedDB before compacting localStorage. Neither browser schema touches `omnivoice_data/` SQLite, so no alembic revision applies. Malformed data fails to defaults or the recovery gate, never a white screen.
|
||||
|
||||
- **Cross-platform parity (default behavior identical on macOS/Windows/Linux; platform-only features behind opt-in).** This task is **pure frontend store/UI logic in the Tauri webview** — JS/TS only, **zero platform branches, zero OS/shell/path APIs, no new Tauri permissions**. localStorage, zustand persist, React controlled inputs, and `setMode` routing behave **identically** on all three platforms. No default-feature divergence → no P0 platform risk, and nothing here is platform-only so nothing needs an opt-in gate. The two pre-existing cross-platform-identical behaviors we *inherit* (multi-window last-write-wins D5; localStorage quota D4) are the same on every OS and are not regressions. The Audiobook cover `serverPath` is a *local* backend path returned by `/audiobook/cover` and validated by `_safe_cover_path` (`audiobook.py:451`) — same code path on all platforms.
|
||||
- **Cross-platform parity (default behavior identical on macOS/Windows/Linux; platform-only features behind opt-in).** The split localStorage/IndexedDB adapter is standard browser behavior on every supported shell. Tauri's widget is explicitly read-only; browser main tabs use revision arbitration for collisions. The Windows-only WebView repair clears cache-only directories and preserves both storage databases.
|
||||
|
||||
- **Local-first guarantee (no cloud, no accounts, no API keys, no third-party telemetry; app fully functional offline).** Every new field lives in **browser localStorage**; the only network contact is the **existing local backend** (`/audiobook/cover`, `/audiobook`, `/longform/render`) — **no new endpoint, no external host, no telemetry, no credential**. Nothing in this task phones home; persistence and migration work fully offline. The cover `serverPath` references the user's own local VoiceStudio backend, never a remote store.
|
||||
- **Local-first guarantee (no cloud, no accounts, no API keys, no third-party telemetry; app fully functional offline).** Bounded fields live in localStorage and manuscript/project payloads live in local IndexedDB. Neither layer makes a network call. Existing cover/render requests still target the configured VoiceStudio backend.
|
||||
|
||||
- **CodeQL — `js/redos` / polynomial-ReDoS on user-input regex (and `py/*`).** Verified: `security.yml:94-96` runs CodeQL `security-and-quality` on **both `python` and `javascript-typescript`** (`:74`), with `*.test.{ts,tsx,js,jsx}` excluded from analysis (`:102-105`). This is a **frontend-only** task, so the relevant query is the **JS/TS ReDoS** one, not `py/*`. **This task introduces no new regex over user-pasted text** — the only new "validation" is the two-value `rec.mode === 'audiobook' ? … : …` literal-equality check (E3) and the `convertMode` membership guard (G3), and the migrate fn uses only `typeof`/`Array.isArray`/spread (no regex). Existing parsers (`parseScript.js:14-39` builds `attributionName` regexes from a `TAG_VERBS` alternation and runs them on user-pasted script via `autoCast`) are **untouched** here. ⚠️ **Hand-off flag for #24:** #24 *will* call `parseScript`/`importStory` on the `script` string to populate `tracks`. If #24 adds or modifies any regex that consumes the user `script`, it must run the ReDoS discipline from MEMORY (`codeql-redos-regex`): no overlapping `\s*`/`.+`, exclude both delimiters in `[^x]*`, atomic groups OK on py≥3.11 — and re-audit the existing `parseScript` alternations under the same lens. #31 carries the ReDoS *zero-delta*; #24 inherits the *duty*.
|
||||
|
||||
@@ -669,14 +697,14 @@ Each relevant hard rule, and exactly how this task satisfies it:
|
||||
|---|---|---|
|
||||
| Migration drops/corrupts existing saved Stories projects (the localStorage data-integrity / "no manual migration" constraint) | High | Dedicated migration unit test with a realistic v4 blob **plus** malformed-input cases (non-object, non-array `storyProjects`, null/garbage entries, missing keys, version >5, dangling `currentProjectId` — D1–D6); spread `...sp` *last* so original fields always win over defaults; `migrate` returns upgraded partial, never throws (matches the existing guard at `index.ts:121` + philosophy at `:116-119`). Full fn body pinned in API/data shapes → Migration. This `migrate` is the localStorage analog of an alembic upgrade — Test #3 is the gate. |
|
||||
| Stale metadata leaks across project loads (the new working fields aren't reset) | High | The working-state reset contract: `loadProject`/`newProject` (re)set the **entire** working surface with default-fill (`SLICE_DEFAULTS`) every time (exact action bodies pinned); dedicated stale-carryover test (Test #7, edge cases A1–A4). This is the subtlest regression the task introduces. |
|
||||
| Renaming `storyProjects`→`longformProjects` breaks the 6 consumers | Medium | One-PR `storyProjects` alias OR update all 6 files (`storiesSlice.ts`, `storiesSlice.test.ts`, `index.ts`, `StoriesEditor.jsx`, `Projects.jsx`, `App.jsx`) in the same PR; the test harness reads `.storyProjects` directly (`test:19,:67,:79,:103`) so either alias or port the test; grep `storyProjects` across `frontend/src` before merge. |
|
||||
| Changing the `storyProjects` compatibility key breaks the 6 consumers | Medium | Retain `storyProjects` across the slice, persistence adapter, `StoriesEditor.jsx`, `Projects.jsx`, and `App.jsx`; the test harness reads `.storyProjects` directly (`test:19,:67,:79,:103`), so its unchanged assertions pin the contract. |
|
||||
| Hardcoded CJK / missing-locale string sneaks into the renamed/changed source and trips CI | Medium | `tests/test_no_hardcoded_cjk.py` scans the (non-allowlisted) `longformSlice.ts` + `AudiobookTab.jsx`; route every new label through `t()`, add new keys to **all 21** locales in-PR, run the gate locally (Test #10). User-typed CJK metadata is exempt runtime data (J4). |
|
||||
| Conflating client `longformProjects` with `Projects.jsx` server-side `longformJobs` | Medium | These are distinct: `longformJobs` (`Projects.jsx:102,198-211`) are completed backend render jobs; `longformProjects` are client localStorage projects flowing through the `storyProjects` prop (`:79,146-160`). Do not touch `longformJobs`. |
|
||||
| Conflating client `storyProjects` with `Projects.jsx` server-side `longformJobs` | Medium | These are distinct: `longformJobs` (`Projects.jsx:102,198-211`) are completed backend render jobs; `storyProjects` are client projects flowing through the same-named prop (`:79,146-160`) and stored durably in IndexedDB. Do not touch `longformJobs`. |
|
||||
| Cover image can't round-trip through localStorage; stale `serverPath` | Medium | Persist only `CoverRef` (filename + `serverPath` from `audiobookUploadCover` `{path}`, `audiobook.ts:85-90`), not bytes. On reload show the filename; re-render reuses `serverPath` as `cover_path` (C2). Stale/GC'd/invalid `serverPath` (C3) surfaces a generic backend error SSE event (`audiobook.py:474`) via the existing channel, never crashes; re-picking always works. `clearCover` nulls `coverRef` (C5). Blob preview stays transient by design (matches `AudiobookTab.jsx:57-64`). serverPath is a *local* path — no cloud (Local-first). |
|
||||
| `_trackId` reseed effect (`StoriesEditor.jsx:135-140`) doesn't re-run on cross-session `loadProject` → new-line id collision | Medium | Latent today (dep array `[]`, single-mount), made reachable by Projects-list load. Change dep array to `[currentProjectId]` or reseed in the add path; called out in Integration. |
|
||||
| #24 imports the existing `parseScript` regexes onto the unified `script` without a ReDoS pass | Medium | #31 carries **zero** new user-input regex (Constraints → CodeQL). Explicitly hand the ReDoS-review duty to #24/#27 in Dependencies + a code comment on the `convertMode` seam so a future PR doesn't wire `parseScript` over user `script` blind to the `js/redos` gate. |
|
||||
| In-flight render abandoned on tab switch / reload (no resume) | Low | Explicitly NOT resumed in v1 (F1–F3); inputs survive via store, transient render flags (component `useState`, never in `partialize`) never persist, so no ghost spinner. Backend-job resume is `longformJobs` territory, out of scope. |
|
||||
| Persisting loose working metadata bloats localStorage | Low | Same pattern already used for `storyTracks` (`index.ts:109-110`). Strings + small dicts; a single big book is well under quota. No eviction policy in v1 (Stories already unbounded); D4 flags dozens-of-big-projects for v0.4 if reported. |
|
||||
| Long manuscripts or many projects exhaust localStorage quota | High | Schema v9 moves every unbounded long-form field to IndexedDB and leaves only bounded preferences/ids in `omnivoice.app`. Migration commits IndexedDB before trimming the v8 fallback, and quota-failure tests prove the durable record still restores. |
|
||||
| Controlled-input `undefined → string` warning on v4-migrated projects | Low | Every metadata `<input value>` binds to a guaranteed default (A4) — read `meta.title ?? ''`, never bare; slice init + load-time default-fill ensure `meta.*` are always strings, `lexicon`/`coverRef` always object/null. |
|
||||
| Scope creep into #24 | Medium | Hard line: 31c (mode toggle/convert UI + content transforms `scriptToTracks`/`tracksToScript`) is explicitly out; ship 31a/31b only. The "defer until #24" caveat is honored by shipping just the store + Audiobook binding whose shape is already determined by the two existing editors. |
|
||||
| Two content models (`tracks` vs `script`) drift in a single project | Low (v1) | v1 keeps one authoritative per `mode`; only the matching editor binds working state. Cross-editing is #24's concern with explicit convert. `convertMode` flips the flag only (G1). |
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
# Local speech platform
|
||||
|
||||
VoiceStudio is both a desktop dictation app and a headless local speech
|
||||
service. The desktop remains one app: its bundled Rust control sidecar owns
|
||||
microphone activation, focused-target capture, clipboard safety, and native
|
||||
insertion; the Python backend keeps ASR models warm and exposes the audio data
|
||||
plane.
|
||||
|
||||
This split lets an integration choose how much it owns:
|
||||
|
||||
```text
|
||||
Herdr / terminal / desktop app ── start, stop, toggle ──> Rust control :3902
|
||||
│
|
||||
├─ captures target
|
||||
├─ opens VoiceStudio mic
|
||||
└─ inserts final text
|
||||
|
||||
VS Code / custom GUI / remote mic ── PCM or WebM ───────> WS/HTTP :3900
|
||||
│
|
||||
└─ partial/final text
|
||||
reserve target / insert final ─> Rust control :3902
|
||||
|
||||
Claude Code / Codex / Pi / agents ── MCP HTTP/stdio ───> MCP :3900
|
||||
```
|
||||
|
||||
The Rust sidecar is part of the VoiceStudio process, not a second application.
|
||||
It starts with the desktop app and binds only to `127.0.0.1`.
|
||||
|
||||
## Discover capabilities
|
||||
|
||||
Desktop/native discovery:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:3902/.well-known/voicestudio-speech
|
||||
```
|
||||
|
||||
Engine/data-plane discovery:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:3900/.well-known/voicestudio-speech
|
||||
```
|
||||
|
||||
Both return `voicestudio.speech.v1`. The desktop document includes absolute
|
||||
control, batch, streaming, output-session, and MCP endpoints. The backend
|
||||
document uses relative URLs so it also works behind Tailscale or a reverse
|
||||
proxy; it advertises native control only when launched by the desktop app.
|
||||
|
||||
## Use VoiceStudio capture from any app
|
||||
|
||||
These calls use VoiceStudio's existing microphone, model selection, pill,
|
||||
refinement, and session-bound insertion. The app under the cursor remains the
|
||||
destination.
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:3902/v1/dictation/start
|
||||
curl -X POST http://127.0.0.1:3902/v1/dictation/stop
|
||||
curl -X POST http://127.0.0.1:3902/v1/dictation/toggle
|
||||
```
|
||||
|
||||
JSON-RPC clients use the same actions:
|
||||
|
||||
```json
|
||||
{"jsonrpc":"2.0","id":1,"method":"dictation.toggle"}
|
||||
```
|
||||
|
||||
Send that object to `POST http://127.0.0.1:3902/rpc`. The installed
|
||||
VoiceStudio executable also accepts `--dictate-start`, `--dictate-stop`, and
|
||||
`--dictate-toggle`; the single-instance bridge forwards them to the running
|
||||
app without opening the Studio window.
|
||||
|
||||
The dependency-free Python bridge is convenient for hooks and TUIs:
|
||||
|
||||
```bash
|
||||
python -m backend.speech_client status
|
||||
python -m backend.speech_client toggle
|
||||
python -m backend.speech_client transcribe recording.wav
|
||||
python -m backend.speech_client transcribe recording.wav --insert
|
||||
```
|
||||
|
||||
`--insert` captures the focused destination before transcription starts and
|
||||
uses the same clipboard-preserving native delivery as the global shortcut.
|
||||
|
||||
## Bring your own capture interface
|
||||
|
||||
An editor extension or GUI can own the microphone and consume live text.
|
||||
Connect to:
|
||||
|
||||
```text
|
||||
ws://127.0.0.1:3900/v1/audio/transcriptions/stream
|
||||
```
|
||||
|
||||
Send binary WebM/Opus frames by default. For raw signed 16-bit mono PCM, use
|
||||
`?pcm=1&sr=16000`. Finish without closing the socket by sending:
|
||||
|
||||
```json
|
||||
{"type":"input_audio.end"}
|
||||
```
|
||||
|
||||
Every response carries `protocol` and `session_id`:
|
||||
|
||||
```json
|
||||
{"type":"session.started","protocol":"voicestudio.speech.v1","session_id":"..."}
|
||||
{"type":"partial","text":"hello wor...","session_id":"..."}
|
||||
{"type":"final","final_kind":"summary","text":"Hello world.","session_id":"..."}
|
||||
```
|
||||
|
||||
Streaming Sherpa models can also emit `final_kind: "utterance"` before the
|
||||
authoritative whole-session `summary`. Existing `/ws/transcribe` clients keep
|
||||
their unchanged legacy frames and `EOF` control.
|
||||
|
||||
To reuse native insertion with a custom capture client:
|
||||
|
||||
1. `POST /v1/output/sessions` on port 3902 before opening the microphone.
|
||||
2. Stream audio and receive the final text on port 3900.
|
||||
3. `POST /v1/output/sessions/{id}/insert` with `{"text":"..."}`.
|
||||
4. If capture is cancelled, `DELETE /v1/output/sessions/{id}`.
|
||||
|
||||
Only one output session can own a focused destination at a time. Stale IDs are
|
||||
rejected instead of inserting into a newer target.
|
||||
|
||||
## Batch and agent protocols
|
||||
|
||||
| Transport | Endpoint | Use |
|
||||
|---|---|---|
|
||||
| OpenAI-compatible HTTP | `POST :3900/v1/audio/transcriptions` | Files, scripts, existing SDKs |
|
||||
| WebSocket | `:3900/v1/audio/transcriptions/stream` | Partial and final live text |
|
||||
| MCP Streamable HTTP | `POST :3900/mcp` | Modern agent clients |
|
||||
| MCP stdio | `python -m backend.mcp_shim` | Claude Code, Codex, and stdio-only clients |
|
||||
| JSON-RPC | `POST :3902/rpc` | Native dictation control |
|
||||
| Native CLI | VoiceStudio `--dictate-*` flags | Hooks and plugin actions |
|
||||
|
||||
## Integration map
|
||||
|
||||
| Interface | Recommended connection |
|
||||
|---|---|
|
||||
| Any desktop text field | Existing global shortcut or Rust `dictation.toggle` |
|
||||
| Herdr | Merge [the example command bindings](../examples/speech-platform/herdr-config.toml) into Herdr's config; detached commands call the Rust API while the pane stays focused |
|
||||
| Pi, Claude Code, Codex, Antigravity CLI | Dictate into the focused prompt through Rust; add MCP when the agent also needs file transcription or speech tools |
|
||||
| VS Code | Call Rust HTTP from the extension host for app-wide dictation, or stream editor-owned mic audio over the versioned WebSocket |
|
||||
| TUI or shell script | `python -m backend.speech_client` or HTTP/JSON-RPC |
|
||||
| Browser/WebView UI | Stream audio to the Python data plane; browser pages cannot silently call native control |
|
||||
| Remote microphone + local/remote GPU | Capture at the client edge and use the authenticated WebSocket/OpenAI endpoint |
|
||||
|
||||
Loopback clients need no credential. Remote native WebSocket clients can send
|
||||
the configured bearer key. Browser clients should exchange that key for a
|
||||
short-lived session, mint a path-bound ticket at `/api/auth/ws-ticket`, and
|
||||
connect with `?ws_ticket=...`; see [API authentication](api-auth.md).
|
||||
Keep remote endpoints restricted to a trusted network; an API key authenticates
|
||||
a client but does not provide network isolation. Beyond a fully trusted LAN,
|
||||
use HTTPS/WSS and never send bearer credentials or ticket exchanges over
|
||||
plaintext HTTP/WebSocket.
|
||||
|
||||
## Security and privacy
|
||||
|
||||
- The native control sidecar binds only to IPv4 loopback and rejects untrusted
|
||||
browser `Origin` headers, blocking ordinary websites from turning on the mic.
|
||||
- Native control never accepts audio and is never exposed through Network
|
||||
Sharing. Remote ASR stays on the existing API-key boundary.
|
||||
- Microphones stay at the interface edge. A remote GPU backend never assumes
|
||||
it owns the user's input device.
|
||||
- No protocol adds a required network call, account, analytics event, or cloud
|
||||
provider.
|
||||
|
||||
## Research basis
|
||||
|
||||
The design survey covered five pages of GitHub's
|
||||
[`speech-to-text` topic](https://github.com/topics/speech-to-text):
|
||||
[1](https://github.com/topics/speech-to-text?page=1),
|
||||
[2](https://github.com/topics/speech-to-text?page=2),
|
||||
[3](https://github.com/topics/speech-to-text?page=3),
|
||||
[4](https://github.com/topics/speech-to-text?page=4), and
|
||||
[5](https://github.com/topics/speech-to-text?page=5).
|
||||
|
||||
The platform keeps the strongest reusable ideas without copying their UI
|
||||
boundaries:
|
||||
|
||||
| Source | Adopted idea |
|
||||
|---|---|
|
||||
| [Handy](https://github.com/cjpais/Handy) | Cross-platform offline dictation, external toggle control, VAD-oriented capture |
|
||||
| [WhisperLiveKit](https://github.com/QuentinFuxa/WhisperLiveKit) | Live local transcription and compatibility-oriented serving |
|
||||
| [RealtimeSTT](https://github.com/KoljaB/RealtimeSTT) | Low-latency partials, endpointing, and warm recognizers |
|
||||
| [sherpa-onnx](https://github.com/k2-fsa/sherpa-onnx) | Portable CPU streaming models and WebSocket-friendly audio framing |
|
||||
| [FunASR](https://github.com/modelscope/FunASR) | OpenAI-compatible and MCP-facing serving |
|
||||
| [Vexa](https://github.com/Vexa-ai/vexa) | WebSocket transcripts plus agent access |
|
||||
| [Voquill](https://github.com/voquill/voquill) | Provider independence, refinement, and personal-vocabulary direction |
|
||||
| [Muesli](https://github.com/Muesli-HQ/muesli) | Machine-readable CLI contracts and session-safe automation |
|
||||
| [Herdr](https://github.com/motionharvest/herdr) | One local control surface behind CLI, socket, hooks, and plugin integrations |
|
||||
|
||||
The differentiator is the connection layer: one bundled app offers native
|
||||
capture/output control and a protocol-neutral ASR service, so every interface
|
||||
does not rebuild model loading, desktop permissions, and insertion safety.
|
||||
@@ -2,13 +2,31 @@
|
||||
|
||||
| Field | Decision |
|
||||
| --- | --- |
|
||||
| Status | Implemented in draft PR #1541; CI and review pending |
|
||||
| Status | Historical coalescer plan; storage and hydration ownership superseded by PR #1636 |
|
||||
| Target | One focused frontend PR |
|
||||
| Priority | P1 responsiveness and data-safety hardening |
|
||||
| Risk | Medium: persistence timing changes, persisted formats do not |
|
||||
| Dependencies | None |
|
||||
| Rollback | Revert the PR; the existing keys and schemas remain readable |
|
||||
|
||||
> **Successor amendment (2026-08-24, PR #1636):** the body below is retained as
|
||||
> the historical PR #1541 coalescer design. Its uses of “current”, `{version: 7}`,
|
||||
> and synchronous hydration describe that 2026-08-13 baseline, not today's
|
||||
> runtime. PR #1636 retains the coalescer for bounded browser state, moves
|
||||
> unbounded long-form data to IndexedDB, and makes store hydration explicitly
|
||||
> asynchronous after window ownership is known.
|
||||
|
||||
### Persisted-version ownership after PR #1636
|
||||
|
||||
| Version/era | Owner | Data and migration contract |
|
||||
| --- | --- | --- |
|
||||
| Historical Zustand v4 → v5 | `frontend/src/store/index.ts` → `migrateAppStore` | The retained `version < 5` migration normalizes legacy v4 project records into the unified long-form shape. This is an old-envelope migration, not the current storage version. |
|
||||
| Historical coalescer baseline v7 | PR #1541 and this plan | The complete persisted Zustand projection lived in synchronous `localStorage` as `{ state, version: 7 }`. References to v7 and synchronous `getItem` below are historical acceptance criteria. |
|
||||
| Current Zustand envelope v9 | `frontend/src/store/index.ts` plus `frontend/src/utils/longformPersistence.ts` | `omnivoice.app` remains the Zustand key, but its normal localStorage copy is a bounded v9 envelope. During upgrade, the split adapter commits any legacy full long-form payload before compacting that envelope. |
|
||||
| Current IndexedDB schema 1 | `frontend/src/utils/indexedDbLongformStore.ts` | `omnivoice.longform` owns the unbounded workspace payload and its writer revision. Bootstrap uses `skipHydration`, resolves main/widget ownership, and awaits the async split-store rehydrate before rendering. |
|
||||
|
||||
The Zustand envelope version (`9`) and IndexedDB database schema (`1`) are independent counters with separate owners; “schema v9” must not be used as a name for the IndexedDB format.
|
||||
|
||||
## Executive decision
|
||||
|
||||
The first optimization PR should remove synchronous JSON serialization and `localStorage` writes from high-frequency interaction paths. It should preserve the existing `omnivoice.app` and `omni_ui` contracts, coalesce each burst to the latest value, flush within a bounded window, and prevent deferred writes from undoing Factory Reset.
|
||||
@@ -52,7 +70,7 @@ Representative serialized sizes were approximately 156 KB for `omnivoice.app` an
|
||||
|
||||
## Goal
|
||||
|
||||
For a rapid sequence of edits, perform no JSON serialization or physical storage write in the originating interaction task and persist only the newest value after the burst, while retaining synchronous hydration and the current recovery formats.
|
||||
At the PR #1541 baseline, perform no JSON serialization or physical storage write in the originating interaction task and persist only the newest value after the burst, while retaining that baseline's synchronous hydration and recovery formats. The successor matrix above records the later hydration and format changes.
|
||||
|
||||
## Scope
|
||||
|
||||
@@ -77,19 +95,19 @@ For a rapid sequence of edits, perform no JSON serialization or physical storage
|
||||
- New dependencies, user-visible strings, locale files, or an app version bump.
|
||||
- Hardware-sensitive timing assertions in CI.
|
||||
|
||||
## Compatibility and safety invariants
|
||||
## Historical PR #1541 compatibility and safety invariants
|
||||
|
||||
The implementation must preserve all of the following:
|
||||
|
||||
| Contract | Required invariant |
|
||||
| --- | --- |
|
||||
| Zustand key | `omnivoice.app` |
|
||||
| Zustand envelope | `{ state, version: 7 }`, serialized with normal `JSON.stringify` semantics |
|
||||
| Zustand envelope at the #1541 baseline | `{ state, version: 7 }`, serialized with normal `JSON.stringify` semantics |
|
||||
| Zustand projection | Existing `partialize` fields and transient-field stripping remain semantically unchanged |
|
||||
| Zustand migration | Existing v1-v7 migration behavior remains unchanged |
|
||||
| Zustand migration at the #1541 baseline | Existing v1-v7 migration behavior remains unchanged |
|
||||
| Legacy recovery key | `omni_ui` |
|
||||
| Legacy recovery shape | Exact current field names, omission behavior, and `sanitizeOmniUi` restore path |
|
||||
| Hydration | Synchronous; no loading gate or async race is introduced |
|
||||
| Hydration at the #1541 baseline | Synchronous; no loading gate or async race is introduced |
|
||||
| Durability | When serialization/storage succeeds and the browser runs timers, a dirty key is attempted within 1,000 ms of its first unflushed change |
|
||||
| Lifecycle | `pagehide` and hidden-document events attempt pending values; both events together cause at most one physical write per unchanged generation |
|
||||
| Reset | A removed preference key cannot be recreated by old or newly queued work before the reset reload |
|
||||
@@ -97,7 +115,7 @@ The implementation must preserve all of the following:
|
||||
| Privacy | Logs may contain a key and error name, never persisted user content |
|
||||
| Platform parity | Same default behavior on macOS, Windows, Linux, browser, and Docker |
|
||||
|
||||
Direct consumers such as `utils/donationMoments.js`, E2E state seeding, long-form recovery, and the preference-key registry must continue to parse the existing envelope without changes. The donation opt-out's primary `omnivoice.donate.optOut` flag remains an immediate, separate write; add a compatibility assertion that its immediate behavior and the flushed legacy-envelope fallback both remain valid.
|
||||
For PR #1541, direct consumers such as `utils/donationMoments.js`, E2E state seeding, long-form recovery, and the preference-key registry had to continue parsing the then-existing envelope without changes. The donation opt-out's primary `omnivoice.donate.optOut` flag remained an immediate, separate write; its immediate behavior and flushed legacy-envelope fallback required compatibility coverage.
|
||||
|
||||
Concurrent browser/Docker tabs are explicitly not promoted to a coordinated multi-writer system in this PR. They retain unsupported last-physical-writer-wins behavior. The PR description must state that boundary; adding cross-tab revisions or `BroadcastChannel` arbitration would be a separate data-consistency design.
|
||||
|
||||
@@ -171,6 +189,7 @@ After `detectIsWidget()` resolves, `bootstrapApp()` should configure the role an
|
||||
- Flush on `visibilitychange` only when `document.visibilityState === 'hidden'`.
|
||||
- Do not add `beforeunload`; it is unnecessary and can interfere with back/forward caching.
|
||||
- Lifecycle flush uses the same generation/cancellation checks as timer flushes. If hidden visibility and `pagehide` both fire, the second invocation observes a clean generation and performs no second serialization/write.
|
||||
- Treat page lifecycle events as best-effort only for asynchronous IndexedDB work. On desktop, prevent the first native exit request, ask the main webview to await the long-form commit and then drain the compact local envelope, and acknowledge exit afterward; a three-second native timeout must still close or relaunch if the webview cannot respond. Intentional frontend reload/relaunch actions await the same ordered helper before navigating.
|
||||
|
||||
### 5. Zustand integration
|
||||
|
||||
@@ -465,7 +484,7 @@ The PR is ready for review only when all are true:
|
||||
|
||||
## Rollback plan
|
||||
|
||||
No data rollback or migration is required. Reverting the adapter wiring restores immediate writes, and both old and new builds read the same `omnivoice.app` v7 envelope and `omni_ui` object. If a release-only issue appears, revert the PR rather than introducing a second persistence mode or format.
|
||||
This was PR #1541's rollback plan: no data migration was then required because both builds read the same v7 envelope. It is not a current rollback instruction after PR #1636; reverting the split v9/IndexedDB storage requires its migration and downgrade guarantees rather than assuming a v7-only localStorage layout.
|
||||
|
||||
## Follow-up queue
|
||||
|
||||
@@ -474,6 +493,6 @@ These are intentionally not part of the first PR:
|
||||
1. **Incremental dub scheduling.** Add a 300 ms debounce, pass `AbortController.signal` through `apiPost`, use a monotonic request revision, cancel outside Dub, and prove one request per burst plus stale-response rejection.
|
||||
2. **Transactional dub undo.** Profile `pushUndo`, which currently stringifies the complete segment array per edit and retains up to 50 snapshots. If material, group edits by segment/field and focus or idle boundary while preserving one-step undo behavior.
|
||||
3. **Workspace isolation.** Profile React commits after persistence remediation; then extract one workspace at a time, moving heavy hooks/imports behind lazy boundaries. Source length and selector count alone are not success metrics.
|
||||
4. **Document storage migration.** Consider IndexedDB or a worker only if representative post-PR flushes remain over budget. That work requires an independent migration, downgrade, reset, quota, and async-hydration design.
|
||||
4. **Document storage migration — completed by successor PR #1636.** The Zustand envelope is now v9, while IndexedDB schema 1 stores unbounded long-form documents with migration, downgrade, reset, quota, and async-hydration coverage. No second migration is pending from this historical plan.
|
||||
|
||||
Each follow-up must begin from a fresh trace. None should be pulled into this PR merely because it is nearby.
|
||||
|
||||
@@ -66,7 +66,7 @@ StoryProject = { id, name, tracks: StoryTrack[], cast: CastMember[], updatedAt }
|
||||
|
||||
- Effective voice for a track = `track.profileId ?? cast[track.character].profileId ?? null` (→ /generate default).
|
||||
- Effective emotion/speed = per-line override, else cast/global default.
|
||||
- Persistence: zustand `persist` → `localStorage` (key `omnivoice.app`, via `partialize`). **No DB / no alembic** — satisfies the backward-compatible-project-data constraint trivially. Transient fields (`generating`, `audioUrl`) are stripped on persist.
|
||||
- Persistence: zustand `persist` splits unbounded project/manuscript data into the local IndexedDB database `omnivoice.longform`; `omnivoice.app` localStorage retains bounded preferences and ids. **No backend DB / no alembic.** Transient fields (`generating`, `audioUrl`) are stripped at the durable commit boundary. Existing v8 localStorage data migrates only after IndexedDB commits successfully (#1636).
|
||||
|
||||
## 5. Backend touchpoints
|
||||
|
||||
@@ -99,7 +99,7 @@ Each phase is independently shippable to `main` and leaves the editor more usefu
|
||||
- **Per-character stems:** one WAV per cast voice (mute others), zipped.
|
||||
- **Chapter markers:** explicit "Chapter" break lines → a sidecar cue sheet (and embedded if M4B).
|
||||
- **MP3/M4B export:** optional backend `/stories/encode` via ffmpeg.
|
||||
- **Named projects:** multiple saved stories (rename/duplicate/delete), still localStorage.
|
||||
- **Named projects:** multiple saved stories (rename/duplicate/delete), persisted in IndexedDB.
|
||||
- **Regenerate-one + export panel** with format/quality choices.
|
||||
|
||||
## 7. File structure
|
||||
@@ -127,8 +127,8 @@ Each phase is independently shippable to `main` and leaves the editor more usefu
|
||||
|
||||
## 10. Constraints honored
|
||||
|
||||
- **No DB / alembic:** localStorage persistence only.
|
||||
- **Default features work on every platform:** Web Audio, localStorage, native DnD, `/generate`, `spawn_subprocess` (cross-platform-hardened) — no platform divergence; no opt-in needed.
|
||||
- **No backend DB / alembic:** browser-owned projects use IndexedDB; bounded preferences/ids remain in localStorage.
|
||||
- **Default features work on every platform:** Web Audio, IndexedDB, localStorage, native DnD, `/generate`, `spawn_subprocess` (cross-platform-hardened) — no platform divergence; no opt-in needed.
|
||||
- **PIN/LAN-share safe:** all synth goes through `apiFetch` (API base + PIN).
|
||||
- **Localization hard rule:** all user-facing strings via i18n; any CJK only in `i18n/locales/zh-CN.json`.
|
||||
- **No new runtime deps** for Phases 1–3 (Web Audio + hand-rolled WAV + native DnD). Phase 2 `.epub` and Phase 4 MP3 may add one small dep each, evaluated at that phase.
|
||||
|
||||
+3
-2
@@ -1,12 +1,14 @@
|
||||
# VoiceStudio Examples
|
||||
|
||||
This directory contains scripts and configs for training, fine-tuning, and evaluating VoiceStudio.
|
||||
This directory contains integration examples plus scripts and configs for
|
||||
training, fine-tuning, and evaluating VoiceStudio.
|
||||
|
||||
| Use Case | Script | Description |
|
||||
|---|---|---|
|
||||
| Training from scratch | [run_emilia.sh](run_emilia.sh) | Full pipeline on the Emilia dataset (data check, tokenization, training) |
|
||||
| Fine-tuning | [run_finetune.sh](run_finetune.sh) | Fine-tune from a pretrained checkpoint using your own JSONL data |
|
||||
| Evaluation | [run_eval.sh](run_eval.sh) | Evaluate WER, speaker similarity, and UTMOS on standard test sets |
|
||||
| Herdr dictation | [speech-platform/herdr-config.toml](speech-platform/herdr-config.toml) | Trigger the bundled Rust dictation sidecar from detached Herdr command bindings |
|
||||
|
||||
---
|
||||
|
||||
@@ -115,4 +117,3 @@ bash examples/run_eval.sh
|
||||
```
|
||||
|
||||
> See [docs/evaluation.md](../docs/evaluation.md) for metrics details, test set preparation, and running individual metrics.
|
||||
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
# Merge these detached command bindings into ~/.config/herdr/config.toml.
|
||||
# Herdr leaves the current pane focused while VoiceStudio captures its native
|
||||
# insertion target, so the transcript returns to the pane that started it.
|
||||
|
||||
[[keys.command]]
|
||||
key = "prefix+alt+d"
|
||||
type = "shell"
|
||||
command = "curl -fsS -X POST http://127.0.0.1:3902/v1/dictation/toggle"
|
||||
description = "toggle VoiceStudio dictation"
|
||||
|
||||
[[keys.command]]
|
||||
key = "prefix+alt+s"
|
||||
type = "shell"
|
||||
command = "curl -fsS -X POST http://127.0.0.1:3902/v1/dictation/start"
|
||||
description = "start VoiceStudio dictation"
|
||||
|
||||
[[keys.command]]
|
||||
key = "prefix+alt+x"
|
||||
type = "shell"
|
||||
command = "curl -fsS -X POST http://127.0.0.1:3902/v1/dictation/stop"
|
||||
description = "stop VoiceStudio dictation"
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "omnivoice-studio",
|
||||
"version": "0.5.0",
|
||||
"version": "0.5.1",
|
||||
"private": true,
|
||||
"license": "AGPL-3.0-only",
|
||||
"type": "module",
|
||||
|
||||
Generated
+1
-1
@@ -2964,7 +2964,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "omnivoice-studio"
|
||||
version = "0.5.0"
|
||||
version = "0.5.1"
|
||||
dependencies = [
|
||||
"arboard",
|
||||
"dirs-next",
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
# launcher's pkill matches `omnivoice-studio` and must never match a user's
|
||||
# installed app. Renaming it would collapse that distinction.
|
||||
name = "omnivoice-studio"
|
||||
version = "0.5.0"
|
||||
version = "0.5.1"
|
||||
description = "VoiceStudio – AI voice cloning & dubbing desktop app"
|
||||
authors = ["Debpalash"]
|
||||
license = "AGPL-3.0-only"
|
||||
@@ -90,7 +90,7 @@ windows-core = "0.61"
|
||||
# `HWND` type — no second copy of the crate enters the dependency graph.
|
||||
# Win32_System_Registry: check_microphone reads the CapabilityAccessManager
|
||||
# ConsentStore mic toggle (RegGetValueW) for the permissions UX.
|
||||
windows = { version = "0.61", features = ["Win32_Foundation", "Win32_UI_WindowsAndMessaging", "Win32_System_Registry", "Win32_System_Threading"] }
|
||||
windows = { version = "0.61", features = ["Win32_Foundation", "Win32_UI_WindowsAndMessaging", "Win32_System_Registry", "Win32_System_Threading", "Win32_System_JobObjects", "Win32_System_Diagnostics_ToolHelp", "Win32_Security"] }
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
@@ -5,7 +5,7 @@ use std::io::BufRead;
|
||||
use std::io::BufReader;
|
||||
use std::net::{TcpStream, ToSocketAddrs};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -199,18 +199,17 @@ fn raw_http_get(url: &str, timeout: Duration) -> Result<String, String> {
|
||||
/// Keep in sync with `_EXIT_PORT_IN_USE` there.
|
||||
pub const EXIT_PORT_IN_USE: i32 = 78;
|
||||
|
||||
/// Kill whoever holds `port`, then confirm it actually came free.
|
||||
/// Confirm whether `port` is free. An unowned listener is never killed by
|
||||
/// numeric PID: even a successful HTTP identity probe cannot make a reusable
|
||||
/// PID/process-group identifier safe to signal.
|
||||
///
|
||||
/// #1223: every caller used to kill-then-sleep-then-spawn unconditionally, so
|
||||
/// a holder we cannot kill — a different user's process, a `taskkill` blocked
|
||||
/// by policy, a socket sitting in TIME_WAIT that the Windows `netstat`
|
||||
/// LISTENING filter can't even see — was indistinguishable from success. The
|
||||
/// backend then died on the bind with a raw errno and the user got "Backend
|
||||
/// died (exit code 1)".
|
||||
/// a holder we do not own was indistinguishable from success. The backend then
|
||||
/// died on the bind with a raw errno and the user got "Backend died (exit code
|
||||
/// 1)".
|
||||
///
|
||||
/// Returns true when the port is free afterwards. Polls rather than sleeping a
|
||||
/// flat interval: the common case (our own orphan) frees in well under 500ms,
|
||||
/// and the uncommon case deserves longer than one guess.
|
||||
/// Returns true when the port is free afterwards. Polling accommodates the
|
||||
/// short close handoff after a contained backend has just been drained.
|
||||
pub fn free_port_or_report(port: u16) -> bool {
|
||||
kill_orphan_on_port(port);
|
||||
for _ in 0..20 {
|
||||
@@ -220,69 +219,23 @@ pub fn free_port_or_report(port: u16) -> bool {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
log::error!(
|
||||
"Port {} is still held after attempting to kill its owner — the \
|
||||
backend cannot bind it. Another application (or a process owned by a \
|
||||
different user) is using the port.",
|
||||
"Port {} is still held by an unowned listener — the backend cannot \
|
||||
bind it. Quit the other VoiceStudio instance or application and try \
|
||||
again.",
|
||||
port
|
||||
);
|
||||
false
|
||||
}
|
||||
|
||||
/// Kill whatever process owns the port.
|
||||
#[cfg(unix)]
|
||||
/// An HTTP response can justify attaching to a healthy same-version backend,
|
||||
/// but never grants process ownership. Deliberately refuse orphan cleanup:
|
||||
/// signalling a PID discovered through lsof/netstat has an unavoidable reuse
|
||||
/// race, and a matching foreign service must never be terminated.
|
||||
pub fn kill_orphan_on_port(port: u16) {
|
||||
if let Ok(out) = Command::new("lsof")
|
||||
.args(["-ti", &format!(":{}", port)])
|
||||
.output()
|
||||
{
|
||||
if out.status.success() {
|
||||
let pids = String::from_utf8_lossy(&out.stdout);
|
||||
for pid in pids.split_whitespace() {
|
||||
if let Ok(pid_n) = pid.parse::<i32>() {
|
||||
log::warn!("Killing orphan process {} on port {}", pid_n, port);
|
||||
unsafe {
|
||||
libc::kill(pid_n, libc::SIGKILL);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
pub fn kill_orphan_on_port(port: u16) {
|
||||
// `netstat -ano` lists listening sockets with their owning PID.
|
||||
// Parse the output to find the process listening on exactly `port`.
|
||||
// no_window: this orphan-kill probe runs on every launch; without it a
|
||||
// netstat console window flashes each time the app starts.
|
||||
let out = match crate::tools::no_window(Command::new("netstat").args(["-ano", "-p", "TCP"])).output() {
|
||||
Ok(o) => o,
|
||||
Err(_) => return,
|
||||
};
|
||||
let stdout = String::from_utf8_lossy(&out.stdout);
|
||||
// Match the local address ending in ":PORT" exactly to avoid false
|
||||
// positives (e.g. :3900 must not match port 39000).
|
||||
let port_suffix = format!(":{}", port);
|
||||
for line in stdout.lines() {
|
||||
if !line.to_uppercase().contains("LISTENING") {
|
||||
continue;
|
||||
}
|
||||
// Local address is the second whitespace-delimited field.
|
||||
// Format: " TCP 0.0.0.0:3900 0.0.0.0:0 LISTENING 1234"
|
||||
let local_addr = line.split_whitespace().nth(1).unwrap_or("");
|
||||
if !local_addr.ends_with(&port_suffix) {
|
||||
continue;
|
||||
}
|
||||
let parts: Vec<&str> = line.split_whitespace().collect();
|
||||
if let Some(pid_str) = parts.last() {
|
||||
if let Ok(pid) = pid_str.parse::<u32>() {
|
||||
log::warn!("Killing orphan process {} on port {} (Windows)", pid, port);
|
||||
let _ = crate::tools::no_window(
|
||||
Command::new("taskkill").args(["/PID", &pid.to_string(), "/F"]),
|
||||
)
|
||||
.output();
|
||||
}
|
||||
}
|
||||
if port_in_use(port) {
|
||||
log::warn!(
|
||||
"Refusing to signal the unowned listener on port {port}; only desktop-contained backends are terminable"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -550,7 +503,10 @@ fn backend_cmd_override() -> Option<Vec<String>> {
|
||||
parse_backend_cmd_override(&std::env::var("OMNIVOICE_BACKEND_CMD").ok()?)
|
||||
}
|
||||
|
||||
pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Option<&Arc<Mutex<BootstrapStage>>>) -> Option<Child> {
|
||||
pub(crate) fn spawn_backend<R: tauri::Runtime>(
|
||||
app: &tauri::AppHandle<R>,
|
||||
progress: Option<&Arc<Mutex<BootstrapStage>>>,
|
||||
) -> Option<crate::tools::ContainedChild> {
|
||||
let log_path = backend_log_path();
|
||||
let err_path = log_path.with_file_name("backend_err.log");
|
||||
log::info!(
|
||||
@@ -601,7 +557,12 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
);
|
||||
}
|
||||
|
||||
let mut env: Vec<(String, String)> = vec![("PYTHONUNBUFFERED".into(), "1".into())];
|
||||
let mut env: Vec<(String, String)> = vec![
|
||||
("PYTHONUNBUFFERED".into(), "1".into()),
|
||||
// Backend-managed engines/installers must inherit the desktop-owned
|
||||
// process group/Job rather than escaping into a new session.
|
||||
("OMNIVOICE_DESKTOP_CONTAINED".into(), "1".into()),
|
||||
];
|
||||
// Pin the child's OMNIVOICE_PORT to the value Rust resolved so Python's
|
||||
// network_share.backend_port() always agrees with the uvicorn --port we
|
||||
// pass below — otherwise a user-set OMNIVOICE_PORT would change the
|
||||
@@ -673,18 +634,6 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
for (k, v) in &env {
|
||||
cmd.env(k, v);
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
use std::os::windows::process::CommandExt;
|
||||
// CREATE_NO_WINDOW (0x08000000) | CREATE_NEW_PROCESS_GROUP (0x00000200).
|
||||
// The backend used to inherit the app's console context, so OS console
|
||||
// CLOSE/LOGOFF events could reach it and MKL's Fortran runtime aborted
|
||||
// the process (`forrtl: error (200)`, exit 2 / 0xC000013A — #1153
|
||||
// class). No console + own process group = no console events, ever.
|
||||
// stdout/stderr are piped above, so nothing is lost. Same flag the
|
||||
// nvidia-smi probe already uses (setup.rs).
|
||||
cmd.creation_flags(0x0800_0000 | 0x0000_0200);
|
||||
}
|
||||
match cmd_override {
|
||||
Some(ref argv) => {
|
||||
cmd.args(&argv[1..]);
|
||||
@@ -703,16 +652,13 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
]);
|
||||
}
|
||||
}
|
||||
let mut child = match cmd
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
{
|
||||
cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
|
||||
let mut contained = match crate::tools::spawn_process_tree(&mut cmd) {
|
||||
Ok(c) => {
|
||||
log::info!(
|
||||
"Backend started via venv python {} (pid {})",
|
||||
python.display(),
|
||||
c.id()
|
||||
c.child.id()
|
||||
);
|
||||
c
|
||||
}
|
||||
@@ -736,7 +682,7 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(stdout_pipe) = child.stdout.take() {
|
||||
if let Some(stdout_pipe) = contained.child.stdout.take() {
|
||||
let app_clone = app.clone();
|
||||
let mut out_file = stdout_file;
|
||||
std::thread::spawn(move || {
|
||||
@@ -752,7 +698,7 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(stderr_pipe) = child.stderr.take() {
|
||||
if let Some(stderr_pipe) = contained.child.stderr.take() {
|
||||
let app_clone = app.clone();
|
||||
// Tracked (not detached): the next spawn joins this handle so this
|
||||
// run's buffered tail flushes before the next run's offset is taken.
|
||||
@@ -773,7 +719,7 @@ pub fn spawn_backend<R: tauri::Runtime>(app: &tauri::AppHandle<R>, progress: Opt
|
||||
}
|
||||
}
|
||||
|
||||
Some(child)
|
||||
Some(contained)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+737
-195
File diff suppressed because it is too large
Load Diff
@@ -1256,13 +1256,15 @@ pub fn reveal_host_path(app: tauri::AppHandle, path: String) -> Result<(), Strin
|
||||
// The splash's recovery panel (Windows-only affordance, error-state only)
|
||||
// calls `clear_webview_cache_and_relaunch` to fix it in one click.
|
||||
//
|
||||
// Deleting EBWebView from inside a running app fails — the WebView2 browser
|
||||
// Deleting caches from inside a running app fails — the WebView2 browser
|
||||
// processes hold locks on the profile — so this is a two-step dance:
|
||||
// 1. the command writes a marker file next to the cache and relaunches;
|
||||
// 1. the command writes a marker file, then requests a relaunch through the
|
||||
// bounded frontend persistence handshake;
|
||||
// 2. the fresh process calls `clear_webview_cache_if_marked()` at the very
|
||||
// top of `run()`, before any webview exists, and deletes the cache
|
||||
// there — retrying briefly while the old instance's WebView2 children
|
||||
// finish exiting.
|
||||
// top of `run()`, before any webview exists, and deletes cache-only
|
||||
// subdirectories there — retrying briefly while the old instance's
|
||||
// WebView2 children finish exiting. Local Storage and IndexedDB are never
|
||||
// touched: they contain the user's settings and long-form projects.
|
||||
//
|
||||
// Everything below compiles on every platform (runtime `cfg!` guards, not
|
||||
// `#[cfg]`) so a macOS/Linux `cargo check` validates the whole path; the
|
||||
@@ -1270,21 +1272,34 @@ pub fn reveal_host_path(app: tauri::AppHandle, path: String) -> Result<(), Strin
|
||||
// elsewhere.
|
||||
|
||||
const CLEAR_WEBVIEW_MARKER: &str = ".clear-webview-cache";
|
||||
const WEBVIEW_CACHE_DIR: &str = "EBWebView";
|
||||
/// Retry budget for step 2: `app.restart()` spawns the new process before the
|
||||
/// old one has fully exited, so its WebView2 children may still hold locks on
|
||||
/// the profile — 20 × 500 ms rides out that handoff.
|
||||
const WEBVIEW_PROFILE_DIR: &str = "EBWebView";
|
||||
const WEBVIEW_CACHE_RELATIVE_DIRS: &[&str] = &[
|
||||
"Default/Cache",
|
||||
"Default/Code Cache",
|
||||
"Default/GPUCache",
|
||||
"Default/DawnCache",
|
||||
"Default/Service Worker/CacheStorage",
|
||||
"Default/Service Worker/ScriptCache",
|
||||
"GPUCache",
|
||||
"DawnCache",
|
||||
"ShaderCache",
|
||||
"GrShaderCache",
|
||||
"GraphiteDawnCache",
|
||||
];
|
||||
/// Retry budget for step 2: the requested restart spawns the new process
|
||||
/// before the old one has fully exited, so its WebView2 children may still
|
||||
/// hold locks on the profile — 20 × 500 ms rides out that handoff.
|
||||
const CLEAR_WEBVIEW_ATTEMPTS: u32 = 20;
|
||||
const CLEAR_WEBVIEW_RETRY_DELAY: Duration = Duration::from_millis(500);
|
||||
|
||||
/// (marker file, cache dir) under the pre-app local data dir. Mirrors
|
||||
/// (marker file, WebView profile dir) under the pre-app local data dir. Mirrors
|
||||
/// `config::config_path_pre_app()` — `%LOCALAPPDATA%\<identifier>` on
|
||||
/// Windows — because step 2 runs before an `AppHandle` exists.
|
||||
fn webview_cache_paths() -> Option<(PathBuf, PathBuf)> {
|
||||
let base = dirs_next::data_local_dir()?.join(crate::config::BUNDLE_IDENTIFIER);
|
||||
Some((
|
||||
base.join(CLEAR_WEBVIEW_MARKER),
|
||||
base.join(WEBVIEW_CACHE_DIR),
|
||||
base.join(WEBVIEW_PROFILE_DIR),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -1293,7 +1308,7 @@ pub fn clear_webview_cache_and_relaunch(app: tauri::AppHandle) -> Result<(), Str
|
||||
if !cfg!(target_os = "windows") {
|
||||
return Err("WebView cache repair is only available on Windows (WebView2)".into());
|
||||
}
|
||||
let (marker, cache) = webview_cache_paths()
|
||||
let (marker, profile) = webview_cache_paths()
|
||||
.ok_or_else(|| "could not resolve the local app data directory".to_string())?;
|
||||
if let Some(parent) = marker.parent() {
|
||||
let _ = fs::create_dir_all(parent);
|
||||
@@ -1304,10 +1319,10 @@ pub fn clear_webview_cache_and_relaunch(app: tauri::AppHandle) -> Result<(), Str
|
||||
)
|
||||
.map_err(|e| format!("write {}: {e}", marker.display()))?;
|
||||
log::warn!(
|
||||
"WebView cache repair requested (#879) — relaunching to clear {}",
|
||||
cache.display()
|
||||
"WebView cache repair requested (#879) — relaunching to clear caches under {}",
|
||||
profile.display()
|
||||
);
|
||||
app.restart()
|
||||
crate::persistence_exit::request_restart(&app)
|
||||
}
|
||||
|
||||
/// Startup half of the repair: if the previous run left the marker, delete
|
||||
@@ -1318,12 +1333,12 @@ pub fn clear_webview_cache_if_marked() {
|
||||
if !cfg!(target_os = "windows") {
|
||||
return;
|
||||
}
|
||||
let Some((marker, cache)) = webview_cache_paths() else {
|
||||
let Some((marker, profile)) = webview_cache_paths() else {
|
||||
return;
|
||||
};
|
||||
clear_webview_cache_at(
|
||||
&marker,
|
||||
&cache,
|
||||
&profile,
|
||||
CLEAR_WEBVIEW_ATTEMPTS,
|
||||
CLEAR_WEBVIEW_RETRY_DELAY,
|
||||
);
|
||||
@@ -1336,41 +1351,56 @@ pub fn clear_webview_cache_if_marked() {
|
||||
/// - no marker → nothing is touched;
|
||||
/// - the marker is consumed FIRST, unconditionally — one-shot, so a failing
|
||||
/// repair can never loop across launches;
|
||||
/// - Local Storage and IndexedDB survive every repair;
|
||||
/// - a missing cache dir is success; a locked one is retried, then given up
|
||||
/// on with an error log — startup is never bricked over a failed repair.
|
||||
fn clear_webview_cache_at(marker: &Path, cache: &Path, attempts: u32, retry_delay: Duration) {
|
||||
fn clear_webview_cache_at(marker: &Path, profile: &Path, attempts: u32, retry_delay: Duration) {
|
||||
if !marker.exists() {
|
||||
return;
|
||||
}
|
||||
let _ = fs::remove_file(marker);
|
||||
if !cache.exists() {
|
||||
if !profile.exists() {
|
||||
return;
|
||||
}
|
||||
// `app.restart()` spawns the new process before the old one has fully
|
||||
// exited, so its WebView2 children may still hold locks — retry briefly.
|
||||
|
||||
let mut cleared_any = false;
|
||||
// The requested restart may start the new process while old WebView2 cache
|
||||
// handles are closing. Retry all remaining cache-only directories as one
|
||||
// bounded batch; never remove the profile root or persistent storage dirs.
|
||||
for attempt in 1..=attempts {
|
||||
match fs::remove_dir_all(cache) {
|
||||
Ok(()) => {
|
||||
log::warn!(
|
||||
"cleared WebView2 profile cache at {} (attempt {attempt}) — issue #879 repair",
|
||||
cache.display()
|
||||
);
|
||||
return;
|
||||
let mut retry_needed = false;
|
||||
for relative in WEBVIEW_CACHE_RELATIVE_DIRS {
|
||||
let cache = profile.join(relative);
|
||||
if !cache.exists() {
|
||||
continue;
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return,
|
||||
Err(e) if attempt < attempts => {
|
||||
log::debug!("WebView2 cache still locked ({e}) — retrying");
|
||||
std::thread::sleep(retry_delay);
|
||||
}
|
||||
Err(e) => {
|
||||
// Never brick startup over a failed repair: WebView2 rebuilds
|
||||
// whatever subset survived, and the user can retry.
|
||||
log::error!(
|
||||
"could not fully clear WebView2 cache at {}: {e} — continuing startup",
|
||||
cache.display()
|
||||
);
|
||||
match fs::remove_dir_all(&cache) {
|
||||
Ok(()) => cleared_any = true,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
|
||||
Err(e) if attempt < attempts => {
|
||||
retry_needed = true;
|
||||
log::debug!("WebView2 cache still locked at {} ({e})", cache.display());
|
||||
}
|
||||
Err(e) => {
|
||||
// Never brick startup over a failed repair: WebView2 rebuilds
|
||||
// whatever subset was cleared, and the user can retry.
|
||||
log::error!(
|
||||
"could not clear WebView2 cache at {}: {e} — continuing startup",
|
||||
cache.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !retry_needed {
|
||||
if cleared_any {
|
||||
log::warn!(
|
||||
"cleared WebView2 caches under {} (attempt {attempt}) — issue #879 repair",
|
||||
profile.display()
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
std::thread::sleep(retry_delay);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1389,24 +1419,58 @@ mod webview_cache_repair_tests {
|
||||
/// under a fresh temp dir.
|
||||
fn seed(dir: &Path) -> (PathBuf, PathBuf) {
|
||||
let marker = dir.join(super::CLEAR_WEBVIEW_MARKER);
|
||||
let cache = dir.join(super::WEBVIEW_CACHE_DIR);
|
||||
let cache = dir.join(super::WEBVIEW_PROFILE_DIR);
|
||||
fs::write(&marker, b"test").unwrap();
|
||||
fs::create_dir_all(cache.join("Default/Cache")).unwrap();
|
||||
fs::write(cache.join("Default/Cache/data_0"), b"x").unwrap();
|
||||
fs::create_dir_all(cache.join("Default/Code Cache/js")).unwrap();
|
||||
fs::write(cache.join("Default/Code Cache/js/data_0"), b"x").unwrap();
|
||||
fs::create_dir_all(cache.join("Default/Local Storage/leveldb")).unwrap();
|
||||
fs::write(
|
||||
cache.join("Default/Local Storage/leveldb/data"),
|
||||
b"settings",
|
||||
)
|
||||
.unwrap();
|
||||
fs::create_dir_all(cache.join("Default/IndexedDB/omnivoice.longform.leveldb")).unwrap();
|
||||
fs::write(
|
||||
cache.join("Default/IndexedDB/omnivoice.longform.leveldb/data"),
|
||||
b"projects",
|
||||
)
|
||||
.unwrap();
|
||||
(marker, cache)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn marker_present_clears_cache_and_consumes_marker_once() {
|
||||
fn marker_present_clears_only_caches_and_consumes_marker_once() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (marker, cache) = seed(dir.path());
|
||||
clear_webview_cache_at(&marker, &cache, FEW, NO_WAIT);
|
||||
assert!(!cache.exists(), "cache dir must be removed");
|
||||
assert!(
|
||||
!cache.join("Default/Cache").exists(),
|
||||
"HTTP cache must be removed"
|
||||
);
|
||||
assert!(
|
||||
!cache.join("Default/Code Cache").exists(),
|
||||
"compiled-code cache must be removed"
|
||||
);
|
||||
assert!(
|
||||
cache.join("Default/Local Storage/leveldb/data").exists(),
|
||||
"settings in Local Storage must survive"
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.join("Default/IndexedDB/omnivoice.longform.leveldb/data")
|
||||
.exists(),
|
||||
"projects in IndexedDB must survive"
|
||||
);
|
||||
assert!(!marker.exists(), "marker must be consumed");
|
||||
// One-shot: with the marker gone, a rebuilt cache is left alone.
|
||||
fs::create_dir_all(cache.join("Default")).unwrap();
|
||||
fs::create_dir_all(cache.join("Default/Cache")).unwrap();
|
||||
clear_webview_cache_at(&marker, &cache, FEW, NO_WAIT);
|
||||
assert!(cache.exists(), "second call without a marker is a no-op");
|
||||
assert!(
|
||||
cache.join("Default/Cache").exists(),
|
||||
"second call without a marker is a no-op"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1425,7 +1489,7 @@ mod webview_cache_repair_tests {
|
||||
fn missing_cache_dir_still_consumes_marker_and_returns() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let marker = dir.path().join(super::CLEAR_WEBVIEW_MARKER);
|
||||
let cache = dir.path().join(super::WEBVIEW_CACHE_DIR);
|
||||
let cache = dir.path().join(super::WEBVIEW_PROFILE_DIR);
|
||||
fs::write(&marker, b"test").unwrap();
|
||||
clear_webview_cache_at(&marker, &cache, FEW, NO_WAIT);
|
||||
assert!(
|
||||
@@ -1445,15 +1509,19 @@ mod webview_cache_repair_tests {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let (marker, cache) = seed(dir.path());
|
||||
// Deny writes on the cache dir so its entries can't be unlinked.
|
||||
fs::set_permissions(&cache, fs::Permissions::from_mode(0o555)).unwrap();
|
||||
let default_dir = cache.join("Default");
|
||||
fs::set_permissions(&default_dir, fs::Permissions::from_mode(0o555)).unwrap();
|
||||
clear_webview_cache_at(&marker, &cache, FEW, NO_WAIT);
|
||||
assert!(
|
||||
!marker.exists(),
|
||||
"one-shot: marker consumed even on failure"
|
||||
);
|
||||
assert!(cache.exists(), "a locked cache survives the failed repair");
|
||||
assert!(
|
||||
cache.join("Default/Cache").exists(),
|
||||
"a locked cache survives the failed repair"
|
||||
);
|
||||
// Restore permissions so TempDir can clean up.
|
||||
fs::set_permissions(&cache, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
fs::set_permissions(&default_dir, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -35,6 +35,12 @@ use serde::{Deserialize, Serialize};
|
||||
/// How many crash markers to retain (newest first).
|
||||
pub const MAX_MARKERS: usize = 3;
|
||||
|
||||
/// Windows debugger notification used when an external debugger terminates a
|
||||
/// process. It is not a Python/native crash code and carries no backend fault
|
||||
/// evidence; the supervisor may restart the child, but must not turn this into
|
||||
/// a crash report or spend the crash-loop budget (#1658/#1663).
|
||||
pub const WINDOWS_DBG_TERMINATE_PROCESS: i32 = 0x4001_0004;
|
||||
|
||||
// ── Exit-status decomposition ──────────────────────────────────────────────
|
||||
|
||||
/// Structured view of how the backend child ended: the numeric exit code (or
|
||||
@@ -74,6 +80,11 @@ impl BackendExit {
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether an observed child exit represents crash evidence worth persisting.
|
||||
pub fn should_record_backend_crash(exit: &BackendExit) -> bool {
|
||||
exit.code != Some(WINDOWS_DBG_TERMINATE_PROCESS)
|
||||
}
|
||||
|
||||
// ── Marker model ───────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
@@ -461,6 +472,23 @@ mod tests {
|
||||
assert_eq!(unknown.label(), "try_wait error: gone");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windows_debugger_termination_is_not_a_backend_crash() {
|
||||
let debugger_stop = BackendExit {
|
||||
code: Some(WINDOWS_DBG_TERMINATE_PROCESS),
|
||||
signal: None,
|
||||
description: "exit code 1073807364".into(),
|
||||
};
|
||||
assert!(!should_record_backend_crash(&debugger_stop));
|
||||
|
||||
let access_violation = BackendExit {
|
||||
code: Some(-1073741819),
|
||||
signal: None,
|
||||
description: "exit code -1073741819".into(),
|
||||
};
|
||||
assert!(should_record_backend_crash(&access_violation));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn backend_exit_decomposes_real_exit_statuses() {
|
||||
|
||||
+92
-118
@@ -15,8 +15,10 @@ pub mod config;
|
||||
pub mod crash;
|
||||
pub mod dictation_output;
|
||||
pub mod dictation_shortcut;
|
||||
pub mod persistence_exit;
|
||||
pub mod reset;
|
||||
pub mod setup;
|
||||
pub mod speech_sidecar;
|
||||
pub mod tools;
|
||||
pub mod uninstall;
|
||||
pub mod updater_channel;
|
||||
@@ -27,7 +29,6 @@ use std::collections::VecDeque;
|
||||
use std::process::Child;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use tauri::menu::{MenuBuilder, MenuItemBuilder};
|
||||
use tauri::tray::TrayIconBuilder;
|
||||
@@ -51,14 +52,44 @@ pub fn backend_port() -> u16 {
|
||||
// ── Shared state types ────────────────────────────────────────────────────
|
||||
|
||||
pub struct BackendState {
|
||||
/// Serializes every desktop lifecycle owner from its first port/process
|
||||
/// probe through child tracking and readiness. Bootstrap, Retry, reset,
|
||||
/// setup/uninstall, shutdown, and the crash supervisor never overlap.
|
||||
pub lifecycle: Mutex<()>,
|
||||
pub process: Mutex<Option<Child>>,
|
||||
/// Stable OS containment for a desktop-spawned backend. Unix keeps the
|
||||
/// root unreaped until its inherited process group is drained; Windows
|
||||
/// owns a kill-on-close Job handle. Neither path signals a reusable PID.
|
||||
pub owned_tree: Mutex<Option<tools::OwnedProcessTree>>,
|
||||
/// A healthy same-version backend which predates this desktop launch.
|
||||
/// It is health-supervised but deliberately never killed by PID: there is
|
||||
/// no safe way to adopt ownership of an arbitrary external process tree.
|
||||
pub attached: AtomicBool,
|
||||
/// Consecutive deep-health failures for an unowned attachment. A single
|
||||
/// busy/slow response is not process death; the supervisor confirms a
|
||||
/// sustained outage before considering a safe replacement.
|
||||
pub attached_health: Mutex<AttachedHealthState>,
|
||||
/// When the tracked child was spawned — feeds the crash marker's
|
||||
/// `uptime_s` (#941). Set alongside `process` in bootstrap.rs.
|
||||
pub spawned_at: Mutex<Option<std::time::Instant>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct AttachedHealthState {
|
||||
pub failures: u32,
|
||||
pub unhealthy_since: Option<std::time::Instant>,
|
||||
}
|
||||
|
||||
pub struct AppFlags {
|
||||
pub quitting: AtomicBool,
|
||||
/// A destructive uninstall is stopping the backend but is not itself a
|
||||
/// terminal app exit until the purge succeeds. Keeping this separate from
|
||||
/// `quitting` lets CloseRequested keep the main window alive and lets a
|
||||
/// failed purge recover without overwriting a concurrent real exit.
|
||||
pub uninstalling: AtomicBool,
|
||||
/// Generation which owns `uninstalling`. A stale join/panic finalizer may
|
||||
/// only release its own claim, never a newer uninstall attempt.
|
||||
pub uninstall_owner: std::sync::atomic::AtomicU64,
|
||||
/// Whether dictation is currently recording. The tray's Start/Stop item
|
||||
/// used to infer this from `widget.is_visible()`, which stopped meaning
|
||||
/// anything once the widget became a permanently hidden host. The frontend
|
||||
@@ -452,6 +483,22 @@ mod pill_noactivate_tests {
|
||||
|
||||
// ── Tauri entry ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Production `ExitRequested` teardown, exposed so the real-child lifecycle
|
||||
/// harness exercises the exact shutdown path used by the desktop event loop.
|
||||
#[doc(hidden)]
|
||||
pub fn shutdown_backend_for_exit<R: tauri::Runtime>(app_handle: &tauri::AppHandle<R>) {
|
||||
// Raise the quitting flag FIRST: exits that don't pass through the tray
|
||||
// Quit item (macOS ⌘Q, OS session end) would otherwise let a death watcher
|
||||
// observe our own termination and record a false crash marker (#941).
|
||||
app_handle
|
||||
.state::<AppFlags>()
|
||||
.quitting
|
||||
.store(true, Ordering::SeqCst);
|
||||
if let Err(error) = bootstrap::with_backend_stopped(app_handle, || {}) {
|
||||
log::warn!("Could not fully stop the backend during app exit: {error}");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
// #879: if the previous run requested a WebView cache repair (splash
|
||||
@@ -484,7 +531,12 @@ pub fn run() {
|
||||
|
||||
let app = tauri::Builder::default()
|
||||
// Single-instance MUST be registered first.
|
||||
.plugin(tauri_plugin_single_instance::init(move |app, _argv, _cwd| {
|
||||
.plugin(tauri_plugin_single_instance::init(move |app, argv, _cwd| {
|
||||
if let Some(action) = speech_sidecar::cli_dictation_action(&argv) {
|
||||
log::info!("Second-instance dictation control: {action:?}");
|
||||
speech_sidecar::dispatch_action(app, action);
|
||||
return;
|
||||
}
|
||||
log::info!("Second instance attempted — focusing existing window");
|
||||
// Always the studio window, never the widget. In pill mode this
|
||||
// used to target "widget" and show() it — which is precisely the
|
||||
@@ -553,6 +605,7 @@ pub fn run() {
|
||||
commands::open_input_monitoring_settings,
|
||||
commands::set_tray_recording,
|
||||
commands::quit_app,
|
||||
persistence_exit::confirm_persistence_flush,
|
||||
commands::save_text_file,
|
||||
commands::reveal_host_path,
|
||||
commands::get_dictation_shortcut,
|
||||
@@ -670,10 +723,26 @@ pub fn run() {
|
||||
|
||||
app.manage(AppFlags {
|
||||
quitting: AtomicBool::new(false),
|
||||
uninstalling: AtomicBool::new(false),
|
||||
uninstall_owner: std::sync::atomic::AtomicU64::new(0),
|
||||
dictating: AtomicBool::new(false),
|
||||
capture: Mutex::new(CaptureDispatchState::default()),
|
||||
output: dictation_output::DictationOutput::default(),
|
||||
});
|
||||
match speech_sidecar::start(app.handle().clone()) {
|
||||
Ok(sidecar) => {
|
||||
log::info!("Speech control API ready on port {}", sidecar.port);
|
||||
app.manage(sidecar);
|
||||
}
|
||||
Err(error) => {
|
||||
log::warn!("Speech control API unavailable: {error}");
|
||||
}
|
||||
}
|
||||
let initial_args: Vec<String> = std::env::args().collect();
|
||||
if let Some(action) = speech_sidecar::cli_dictation_action(&initial_args) {
|
||||
speech_sidecar::dispatch_action(app.handle(), action);
|
||||
}
|
||||
app.manage(persistence_exit::PersistenceExitState::default());
|
||||
app.manage(TrayHandle {
|
||||
tray: Mutex::new(None),
|
||||
dictate: Mutex::new(None),
|
||||
@@ -819,13 +888,10 @@ pub fn run() {
|
||||
let mut cfg = crate::config::load_config(app);
|
||||
cfg.launch_as_widget = false;
|
||||
crate::config::save_config(app, &cfg);
|
||||
if let Ok(exe) = std::env::current_exe() {
|
||||
let _ = std::process::Command::new(exe).spawn();
|
||||
if let Err(error) = persistence_exit::request_spawned_relaunch(app, vec![])
|
||||
{
|
||||
log::error!("Could not switch to studio mode: {error}");
|
||||
}
|
||||
app.state::<AppFlags>()
|
||||
.quitting
|
||||
.store(true, Ordering::SeqCst);
|
||||
app.exit(0);
|
||||
}
|
||||
"switch_to_pill" => {
|
||||
// Mirror of "open_studio" but the other direction:
|
||||
@@ -834,15 +900,12 @@ pub fn run() {
|
||||
let mut cfg = crate::config::load_config(app);
|
||||
cfg.launch_as_widget = true;
|
||||
crate::config::save_config(app, &cfg);
|
||||
if let Ok(exe) = std::env::current_exe() {
|
||||
let _ = std::process::Command::new(exe)
|
||||
.arg("--pill")
|
||||
.spawn();
|
||||
if let Err(error) = persistence_exit::request_spawned_relaunch(
|
||||
app,
|
||||
vec!["--pill".into()],
|
||||
) {
|
||||
log::error!("Could not switch to pill mode: {error}");
|
||||
}
|
||||
app.state::<AppFlags>()
|
||||
.quitting
|
||||
.store(true, Ordering::SeqCst);
|
||||
app.exit(0);
|
||||
}
|
||||
"dictate" => {
|
||||
// Toggle start/stop. This used to ask the widget
|
||||
@@ -955,7 +1018,11 @@ pub fn run() {
|
||||
let stage_handle = bootstrap_state.stage.clone();
|
||||
app.manage(bootstrap_state);
|
||||
app.manage(BackendState {
|
||||
lifecycle: Mutex::new(()),
|
||||
process: Mutex::new(None),
|
||||
owned_tree: Mutex::new(None),
|
||||
attached: AtomicBool::new(false),
|
||||
attached_health: Mutex::new(AttachedHealthState::default()),
|
||||
spawned_at: Mutex::new(None),
|
||||
});
|
||||
|
||||
@@ -975,68 +1042,11 @@ pub fn run() {
|
||||
set_stage(&stage_handle, BootstrapStage::AwaitingSetup);
|
||||
return;
|
||||
}
|
||||
match backend::running_backend_version(backend_port()) {
|
||||
Some(v) if backend::same_app_version(&v) => {
|
||||
if backend::backend_deep_healthy(backend_port()) {
|
||||
log::info!(
|
||||
"Port {} already serving VoiceStudio backend v{} — attaching",
|
||||
backend_port(), v
|
||||
);
|
||||
set_stage(&stage_handle, BootstrapStage::Ready);
|
||||
return;
|
||||
}
|
||||
// Same version but a DB-touching probe fails: a backend whose
|
||||
// install was wiped/corrupted while it kept running. Attaching
|
||||
// would look alive and 500 on everything — replace it.
|
||||
log::warn!(
|
||||
"Port {} serves VoiceStudio v{} but failed the deep health probe — replacing it",
|
||||
backend_port(), v
|
||||
);
|
||||
backend::kill_orphan_on_port(backend_port());
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
}
|
||||
Some(v) => {
|
||||
// Healthy-but-stale backend from a previous version —
|
||||
// the post-update orphan that made new installs run
|
||||
// old backend code. Replace it (see backend.rs
|
||||
// same_app_version for the full story).
|
||||
log::warn!(
|
||||
"Port {} serves a stale VoiceStudio backend (v{} != app v{}) — replacing it",
|
||||
backend_port(),
|
||||
if v.is_empty() { "<unknown>" } else { v.as_str() },
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
);
|
||||
backend::kill_orphan_on_port(backend_port());
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
if backend::port_in_use(backend_port()) {
|
||||
log::warn!(
|
||||
"Port {} in use — taking ownership (killing whatever's there)",
|
||||
backend_port()
|
||||
);
|
||||
backend::kill_orphan_on_port(backend_port());
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
}
|
||||
// First-run gate: never auto-install. With nothing on disk to
|
||||
// attach to, park on the setup screen and wait for the user to
|
||||
// confirm an install plan — `complete_setup` restarts the
|
||||
// bootstrap from there. Existing pre-setup-screen installs
|
||||
// (venv present) are migrated here — the bootstrap thread is
|
||||
// the only place that write happens — then pass straight
|
||||
// through the read-only is_first_run check.
|
||||
setup::migrate_existing_install_if_needed(&app_handle);
|
||||
if setup::is_first_run(&app_handle) {
|
||||
log::info!("First run — awaiting setup screen confirmation before installing");
|
||||
set_stage(&stage_handle, BootstrapStage::AwaitingSetup);
|
||||
return;
|
||||
}
|
||||
// Spawn + health-poll loop shared with the Retry button —
|
||||
// includes the #314 broken-venv self-heal (quarantine the
|
||||
// venv and rebuild once when the backend exits with
|
||||
// "No pyvenv.cfg file" / code 106).
|
||||
bootstrap::spawn_backend_and_wait(&app_handle, &stage_handle);
|
||||
// Probe/attach, the first-run gate, spawn, child tracking, and
|
||||
// readiness are one serialized lifecycle operation. A Retry
|
||||
// arriving during launch waits and attaches instead of
|
||||
// creating a second backend on the same port (#1635).
|
||||
bootstrap::spawn_initial_backend_and_wait(&app_handle, &stage_handle);
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
@@ -1065,47 +1075,11 @@ pub fn run() {
|
||||
.expect("error while building tauri application");
|
||||
|
||||
app.run(|app_handle, event| {
|
||||
if let tauri::RunEvent::ExitRequested { .. } = event {
|
||||
// Raise the quitting flag FIRST: exits that don't pass through the
|
||||
// tray Quit item (macOS ⌘Q, OS session end) would otherwise let a
|
||||
// death watcher observe our own SIGTERM below and record a false
|
||||
// "backend crashed" marker (#941).
|
||||
app_handle
|
||||
.state::<AppFlags>()
|
||||
.quitting
|
||||
.store(true, Ordering::SeqCst);
|
||||
if let Ok(mut lock) = app_handle.state::<BackendState>().process.lock() {
|
||||
if let Some(ref mut child) = *lock {
|
||||
let pid = child.id();
|
||||
log::info!("Shutting down backend (pid {})", pid);
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
unsafe {
|
||||
libc::kill(pid as i32, libc::SIGTERM);
|
||||
}
|
||||
let start = std::time::Instant::now();
|
||||
loop {
|
||||
match child.try_wait() {
|
||||
Ok(Some(_)) => break,
|
||||
Ok(None) if start.elapsed() < Duration::from_secs(2) => {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
_ => {
|
||||
log::warn!("Backend didn't exit in 2 s — SIGKILL");
|
||||
let _ = child.kill();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
let _ = child.kill();
|
||||
}
|
||||
let _ = child.wait();
|
||||
}
|
||||
if let tauri::RunEvent::ExitRequested { code, api, .. } = event {
|
||||
if !persistence_exit::handle_exit_requested(app_handle, code, &api) {
|
||||
return;
|
||||
}
|
||||
shutdown_backend_for_exit(app_handle);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
//! Orderly desktop exit/relaunch handshake for async browser persistence.
|
||||
//!
|
||||
//! `pagehide` cannot keep an IndexedDB transaction alive. For ordinary exits
|
||||
//! we therefore prevent the first `ExitRequested`, ask the main webview to
|
||||
//! flush, and exit only after its acknowledgement. A deadline prevents a dead
|
||||
//! or not-yet-mounted webview from trapping the native process indefinitely.
|
||||
|
||||
use std::process::Command;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use tauri::{Emitter, Manager};
|
||||
|
||||
use crate::AppFlags;
|
||||
|
||||
pub const PERSISTENCE_FLUSH_EVENT: &str = "persistence://flush-requested";
|
||||
const EXIT_FLUSH_TIMEOUT: Duration = Duration::from_secs(3);
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
enum ExitAction {
|
||||
Exit,
|
||||
Restart,
|
||||
Spawn(Vec<String>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct ExitPlan {
|
||||
code: i32,
|
||||
action: ExitAction,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum ExitRequestDecision {
|
||||
BeginFlush,
|
||||
WaitForFlush,
|
||||
Allow,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
enum ExitPhase {
|
||||
Idle,
|
||||
Waiting,
|
||||
Ready,
|
||||
}
|
||||
|
||||
struct ExitStateInner {
|
||||
phase: ExitPhase,
|
||||
code: i32,
|
||||
action: ExitAction,
|
||||
}
|
||||
|
||||
impl Default for ExitStateInner {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
phase: ExitPhase::Idle,
|
||||
code: 0,
|
||||
action: ExitAction::Exit,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct PersistenceExitState {
|
||||
inner: Mutex<ExitStateInner>,
|
||||
}
|
||||
|
||||
impl PersistenceExitState {
|
||||
fn request(&self, code: Option<i32>) -> ExitRequestDecision {
|
||||
let Ok(mut inner) = self.inner.lock() else {
|
||||
// A poisoned coordination lock must never make VoiceStudio
|
||||
// impossible to close.
|
||||
return ExitRequestDecision::Allow;
|
||||
};
|
||||
match inner.phase {
|
||||
ExitPhase::Idle => {
|
||||
inner.phase = ExitPhase::Waiting;
|
||||
inner.code = code.unwrap_or(0);
|
||||
ExitRequestDecision::BeginFlush
|
||||
}
|
||||
ExitPhase::Waiting => ExitRequestDecision::WaitForFlush,
|
||||
ExitPhase::Ready => ExitRequestDecision::Allow,
|
||||
}
|
||||
}
|
||||
|
||||
fn queue_action(&self, action: ExitAction) -> Result<(), String> {
|
||||
let mut inner = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| "persistence exit state lock poisoned".to_string())?;
|
||||
if inner.phase != ExitPhase::Idle {
|
||||
return Err("an application exit is already in progress".into());
|
||||
}
|
||||
inner.action = action;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn complete(&self) -> Option<ExitPlan> {
|
||||
let Ok(mut inner) = self.inner.lock() else {
|
||||
log::warn!("Persistence exit state lock poisoned");
|
||||
return None;
|
||||
};
|
||||
if inner.phase != ExitPhase::Waiting {
|
||||
return None;
|
||||
}
|
||||
inner.phase = ExitPhase::Ready;
|
||||
Some(ExitPlan {
|
||||
code: inner.code,
|
||||
action: inner.action.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn perform_exit_plan(app: &tauri::AppHandle, plan: ExitPlan, source: &str) {
|
||||
log::info!("Persistence exit handshake completed via {source}");
|
||||
match plan.action {
|
||||
ExitAction::Exit => app.exit(plan.code),
|
||||
// `restart()` may bypass RunEvent delivery when invoked by a command
|
||||
// on the main thread. `request_restart()` reliably returns through the
|
||||
// ExitRequested callback so the tracked backend is shut down first.
|
||||
ExitAction::Restart => app.request_restart(),
|
||||
ExitAction::Spawn(args) => {
|
||||
match std::env::current_exe()
|
||||
.and_then(|executable| Command::new(executable).args(args).spawn())
|
||||
{
|
||||
Ok(_) => {}
|
||||
Err(error) => log::error!("Could not launch the requested app mode: {error}"),
|
||||
}
|
||||
app.exit(plan.code);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn complete_pending_exit(app: &tauri::AppHandle, source: &str) -> bool {
|
||||
let Some(plan) = app.state::<PersistenceExitState>().complete() else {
|
||||
return false;
|
||||
};
|
||||
perform_exit_plan(app, plan, source);
|
||||
true
|
||||
}
|
||||
|
||||
fn schedule_timeout(app: tauri::AppHandle) {
|
||||
// The frontend materializes its latest full localStorage fallback before
|
||||
// awaiting IndexedDB. A blocked transaction may outlive this deadline,
|
||||
// but the timeout cannot overtake the only recoverable project copy.
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(EXIT_FLUSH_TIMEOUT);
|
||||
if complete_pending_exit(&app, "native timeout") {
|
||||
log::warn!(
|
||||
"Persistence flush acknowledgement did not arrive within {} ms; exiting anyway",
|
||||
EXIT_FLUSH_TIMEOUT.as_millis()
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Return true only when this exit request may proceed to native teardown.
|
||||
pub(crate) fn handle_exit_requested(
|
||||
app: &tauri::AppHandle,
|
||||
code: Option<i32>,
|
||||
api: &tauri::ExitRequestApi,
|
||||
) -> bool {
|
||||
// Tauri explicitly ignores prevent_exit() for its restart exit code. All
|
||||
// intentional frontend restarts flush before calling relaunch, while the
|
||||
// native cache-repair restart is queued through this module below.
|
||||
if code == Some(tauri::RESTART_EXIT_CODE) {
|
||||
return true;
|
||||
}
|
||||
|
||||
match app.state::<PersistenceExitState>().request(code) {
|
||||
ExitRequestDecision::Allow => true,
|
||||
ExitRequestDecision::WaitForFlush => {
|
||||
api.prevent_exit();
|
||||
false
|
||||
}
|
||||
ExitRequestDecision::BeginFlush => {
|
||||
api.prevent_exit();
|
||||
if let Err(error) = app.emit_to("main", PERSISTENCE_FLUSH_EVENT, ()) {
|
||||
log::warn!("Could not request a frontend persistence flush: {error}");
|
||||
}
|
||||
schedule_timeout(app.clone());
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn request_action(app: &tauri::AppHandle, action: ExitAction) -> Result<(), String> {
|
||||
app.state::<PersistenceExitState>().queue_action(action)?;
|
||||
// Stop backend supervision while the bounded frontend flush is pending.
|
||||
app.state::<AppFlags>()
|
||||
.quitting
|
||||
.store(true, Ordering::SeqCst);
|
||||
app.exit(0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn request_restart(app: &tauri::AppHandle) -> Result<(), String> {
|
||||
request_action(app, ExitAction::Restart)
|
||||
}
|
||||
|
||||
pub fn request_spawned_relaunch(app: &tauri::AppHandle, args: Vec<String>) -> Result<(), String> {
|
||||
request_action(app, ExitAction::Spawn(args))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn confirm_persistence_flush(app: tauri::AppHandle) -> bool {
|
||||
complete_pending_exit(&app, "frontend acknowledgement")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
ExitAction, ExitPhase, ExitRequestDecision, PersistenceExitState, EXIT_FLUSH_TIMEOUT,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn first_exit_waits_for_flush_and_only_acknowledgement_allows_the_next() {
|
||||
let state = PersistenceExitState::default();
|
||||
|
||||
assert_eq!(state.request(Some(17)), ExitRequestDecision::BeginFlush);
|
||||
assert_eq!(
|
||||
state.request(Some(99)),
|
||||
ExitRequestDecision::WaitForFlush,
|
||||
"a repeated quit must not replace the original exit code"
|
||||
);
|
||||
let plan = state.complete().expect("waiting exit should complete");
|
||||
assert_eq!(plan.code, 17);
|
||||
assert_eq!(plan.action, ExitAction::Exit);
|
||||
assert_eq!(state.request(None), ExitRequestDecision::Allow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_and_timeout_can_complete_the_same_exit_only_once() {
|
||||
let state = PersistenceExitState::default();
|
||||
assert_eq!(state.request(None), ExitRequestDecision::BeginFlush);
|
||||
|
||||
assert!(state.complete().is_some());
|
||||
assert!(state.complete().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn poisoned_lock_never_synthesizes_an_unrequested_exit() {
|
||||
let state = PersistenceExitState::default();
|
||||
let _ = std::panic::catch_unwind(|| {
|
||||
let _guard = state.inner.lock().unwrap();
|
||||
panic!("poison test lock");
|
||||
});
|
||||
|
||||
assert!(state.complete().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restart_and_mode_switch_are_deferred_until_flush_completion() {
|
||||
let restart = PersistenceExitState::default();
|
||||
restart.queue_action(ExitAction::Restart).unwrap();
|
||||
assert_eq!(restart.request(None), ExitRequestDecision::BeginFlush);
|
||||
assert_eq!(restart.complete().unwrap().action, ExitAction::Restart);
|
||||
|
||||
let mode_switch = PersistenceExitState::default();
|
||||
mode_switch
|
||||
.queue_action(ExitAction::Spawn(vec!["--pill".into()]))
|
||||
.unwrap();
|
||||
assert_eq!(mode_switch.request(None), ExitRequestDecision::BeginFlush);
|
||||
assert_eq!(
|
||||
mode_switch.complete().unwrap().action,
|
||||
ExitAction::Spawn(vec!["--pill".into()])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn actions_cannot_change_once_exit_has_started() {
|
||||
let state = PersistenceExitState::default();
|
||||
assert_eq!(state.request(None), ExitRequestDecision::BeginFlush);
|
||||
assert!(state.queue_action(ExitAction::Restart).is_err());
|
||||
assert_eq!(state.inner.lock().unwrap().phase, ExitPhase::Waiting);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_fallback_is_short_and_bounded() {
|
||||
assert!(EXIT_FLUSH_TIMEOUT >= std::time::Duration::from_millis(500));
|
||||
assert!(EXIT_FLUSH_TIMEOUT <= std::time::Duration::from_secs(5));
|
||||
}
|
||||
}
|
||||
@@ -40,7 +40,7 @@ use serde::Serialize;
|
||||
use tauri::Manager;
|
||||
|
||||
use crate::bootstrap::BootstrapState;
|
||||
use crate::{backend_port, AppFlags};
|
||||
use crate::AppFlags;
|
||||
|
||||
/// Every scope the UI can offer. Two of them (`ui_prefs`, `history`) own no
|
||||
/// files — they are listed here so the frontend has one registry to render, but
|
||||
@@ -374,6 +374,14 @@ pub fn purge_scopes(roots: &Roots, wanted: &[String], home: Option<&Path>) -> Re
|
||||
report
|
||||
}
|
||||
|
||||
async fn run_retained_reset<T: Send + 'static>(
|
||||
worker: impl FnOnce() -> Result<T, String> + Send + 'static,
|
||||
) -> Result<T, String> {
|
||||
tauri::async_runtime::spawn_blocking(worker)
|
||||
.await
|
||||
.map_err(|error| format!("reset failed: {error}"))?
|
||||
}
|
||||
|
||||
/// Delete the selected scopes, then bring the backend back.
|
||||
///
|
||||
/// Unknown or frontend-only scope names are ignored rather than erroring: the
|
||||
@@ -391,45 +399,85 @@ pub async fn reset_purge(app: tauri::AppHandle, scopes: Vec<String>) -> Result<R
|
||||
return Ok(report);
|
||||
}
|
||||
|
||||
// Stop the backend first. `set_backend_kill_intended` tells the #941/#567
|
||||
// supervisor this death is deliberate, so it neither writes a crash marker
|
||||
// nor races us by respawning a backend into the directories we are deleting.
|
||||
// Note we do NOT set `flags.quitting` — that is the uninstall path, and it
|
||||
// would stop us from starting the backend again at the end.
|
||||
crate::bootstrap::set_backend_kill_intended(true);
|
||||
crate::backend::kill_orphan_on_port(backend_port());
|
||||
|
||||
let purge_app = app.clone();
|
||||
let mut report = tauri::async_runtime::spawn_blocking(move || {
|
||||
// Give the process a moment to actually exit and drop its file handles;
|
||||
// on Windows a mapped weights file stays locked until it does.
|
||||
std::thread::sleep(std::time::Duration::from_millis(600));
|
||||
run_retained_reset(move || {
|
||||
// The lifecycle guard spans stop + purge. Bootstrap, Retry, and the
|
||||
// supervisor cannot spawn into directories while they are deleted.
|
||||
let mut report = crate::bootstrap::with_backend_stopped(&purge_app, || {
|
||||
// Give Windows mapped weights and their drainer threads a moment
|
||||
// to release file handles after the tracked child has exited.
|
||||
std::thread::sleep(std::time::Duration::from_millis(600));
|
||||
|
||||
let roots = roots_for(&purge_app);
|
||||
let home = dirs_next::home_dir();
|
||||
purge_scopes(&roots, &wanted, home.as_deref())
|
||||
let roots = roots_for(&purge_app);
|
||||
let home = dirs_next::home_dir();
|
||||
purge_scopes(&roots, &wanted, home.as_deref())
|
||||
})?;
|
||||
|
||||
// Finalization belongs to this retained worker, not the invoking IPC
|
||||
// future. A reload/closed webview can drop the await without stranding
|
||||
// the stopped backend.
|
||||
let flags = purge_app.state::<AppFlags>();
|
||||
if !flags.quitting.load(std::sync::atomic::Ordering::SeqCst)
|
||||
&& !flags.uninstalling.load(std::sync::atomic::Ordering::SeqCst)
|
||||
{
|
||||
let state = purge_app.state::<BootstrapState>();
|
||||
crate::bootstrap::respawn_backend(
|
||||
purge_app.clone(),
|
||||
state.stage.clone(),
|
||||
state.logs.clone(),
|
||||
);
|
||||
report.restarted = true;
|
||||
} else {
|
||||
crate::bootstrap::set_backend_kill_intended(false);
|
||||
}
|
||||
Ok(report)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| format!("reset failed: {e}"))?;
|
||||
|
||||
// Back up. The fresh backend re-runs ensure_dirs() and alembic, so a deleted
|
||||
// database returns empty instead of missing. If the app is on its way out
|
||||
// anyway, don't fight the shutdown.
|
||||
let flags = app.state::<AppFlags>();
|
||||
if !flags.quitting.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
let state = app.state::<BootstrapState>();
|
||||
crate::bootstrap::respawn_backend(app.clone(), state.stage.clone(), state.logs.clone());
|
||||
report.restarted = true;
|
||||
} else {
|
||||
crate::bootstrap::set_backend_kill_intended(false);
|
||||
}
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn aborted_reset_future_does_not_cancel_retained_finalization() {
|
||||
let entered = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let release = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let finalized = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let entered2 = entered.clone();
|
||||
let release2 = release.clone();
|
||||
let finalized2 = finalized.clone();
|
||||
|
||||
let task = tauri::async_runtime::spawn(run_retained_reset(move || {
|
||||
entered2.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
while !release2.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
std::thread::sleep(std::time::Duration::from_millis(5));
|
||||
}
|
||||
finalized2.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(())
|
||||
}));
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
while !entered.load(std::sync::atomic::Ordering::SeqCst)
|
||||
&& std::time::Instant::now() < deadline
|
||||
{
|
||||
std::thread::sleep(std::time::Duration::from_millis(5));
|
||||
}
|
||||
assert!(entered.load(std::sync::atomic::Ordering::SeqCst));
|
||||
task.abort();
|
||||
release.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
|
||||
while !finalized.load(std::sync::atomic::Ordering::SeqCst)
|
||||
&& std::time::Instant::now() < deadline
|
||||
{
|
||||
std::thread::sleep(std::time::Duration::from_millis(5));
|
||||
}
|
||||
assert!(
|
||||
finalized.load(std::sync::atomic::Ordering::SeqCst),
|
||||
"dropping the IPC await must not cancel backend restart/finalization"
|
||||
);
|
||||
}
|
||||
|
||||
fn roots(tmp: &Path) -> Roots {
|
||||
Roots {
|
||||
data: tmp.join("OmniVoice"),
|
||||
|
||||
+168
-39
@@ -230,6 +230,59 @@ fn write_portable_pointer(base: &Path) -> Option<&'static str> {
|
||||
Some(flavour)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FileSnapshot {
|
||||
path: PathBuf,
|
||||
contents: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
fn snapshot_file(path: PathBuf) -> Result<FileSnapshot, String> {
|
||||
let contents = match fs::read(&path) {
|
||||
Ok(contents) => Some(contents),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
|
||||
Err(error) => return Err(format!("Could not snapshot {}: {error}", path.display())),
|
||||
};
|
||||
Ok(FileSnapshot { path, contents })
|
||||
}
|
||||
|
||||
fn snapshot_portable_location() -> Result<Vec<FileSnapshot>, String> {
|
||||
[portable_pointer_path(), config::config_path_for_machine()]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.map(snapshot_file)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn restore_file(snapshot: FileSnapshot) -> Result<(), String> {
|
||||
match snapshot.contents {
|
||||
Some(contents) => {
|
||||
if let Some(parent) = snapshot.path.parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.map_err(|error| format!("Could not restore {}: {error}", parent.display()))?;
|
||||
}
|
||||
fs::write(&snapshot.path, contents)
|
||||
.map_err(|error| format!("Could not restore {}: {error}", snapshot.path.display()))
|
||||
}
|
||||
None => match fs::remove_file(&snapshot.path) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
|
||||
Err(error) => Err(format!("Could not restore {}: {error}", snapshot.path.display())),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn restore_portable_location(snapshots: Vec<FileSnapshot>) -> Result<(), String> {
|
||||
let errors: Vec<String> = snapshots
|
||||
.into_iter()
|
||||
.filter_map(|snapshot| restore_file(snapshot).err())
|
||||
.collect();
|
||||
if errors.is_empty() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(errors.join("; "))
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist where a relocated portable folder lives, so the next launch finds
|
||||
/// it. Prefers the pointer file (portability preserved); falls back to the
|
||||
/// per-user config when the app folder is read-only.
|
||||
@@ -882,10 +935,27 @@ pub fn check_install_target(path: String) -> TargetCheck {
|
||||
/// parked) bootstrap. Any `Err` keeps the app in `AwaitingSetup` with the
|
||||
/// message surfaced on the setup screen — nothing was installed.
|
||||
#[tauri::command]
|
||||
pub fn complete_setup(
|
||||
pub async fn complete_setup(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, BootstrapState>,
|
||||
plan: InstallPlan,
|
||||
) -> Result<(), String> {
|
||||
let owned_state = BootstrapState {
|
||||
stage: state.stage.clone(),
|
||||
logs: state.logs.clone(),
|
||||
};
|
||||
tauri::async_runtime::spawn_blocking(move || complete_setup_blocking(app, owned_state, plan))
|
||||
.await
|
||||
.map_err(|error| {
|
||||
log::error!("Setup task failed to join: {error}");
|
||||
"setup_task_failed".to_string()
|
||||
})?
|
||||
}
|
||||
|
||||
fn complete_setup_blocking(
|
||||
app: tauri::AppHandle,
|
||||
state: BootstrapState,
|
||||
plan: InstallPlan,
|
||||
) -> Result<(), String> {
|
||||
if !matches!(plan.install_mode.as_str(), "installed" | "portable") {
|
||||
return Err(format!("Unknown install mode: {}", plan.install_mode));
|
||||
@@ -942,34 +1012,12 @@ pub fn complete_setup(
|
||||
}
|
||||
cfg.locale = plan.locale.clone().filter(|l| !l.is_empty());
|
||||
|
||||
if plan.install_mode == "portable" {
|
||||
// Create the portable folder and seed config.json INSIDE it first, so
|
||||
// `config_path` resolves portable from here on and the whole install
|
||||
// (env + data + config) travels as one folder.
|
||||
let base = planned_portable_base(&plan).ok_or("Portable anchor disappeared")?;
|
||||
fs::create_dir_all(&base).map_err(|e| format!("Could not create {}: {e}", base.display()))?;
|
||||
// Record WHERE it is before seeding it — `portable_base()` has to
|
||||
// resolve to this folder on the next launch, and the config inside it
|
||||
// cannot say so (nothing would know where to look).
|
||||
if base != portable_anchor().map(|a| a.join(PORTABLE_DIR_NAME)).unwrap_or_default() {
|
||||
let how = record_portable_dir(&app, &base)?;
|
||||
log::info!("Portable folder relocated — recorded via {how}");
|
||||
} else {
|
||||
// Back to the default: drop any earlier relocation so a stale
|
||||
// pointer can't outrank it.
|
||||
clear_portable_dir(&app);
|
||||
}
|
||||
config::save_config_at(&base.join("config.json"), &cfg)?;
|
||||
} else {
|
||||
if plan.install_mode != "portable" {
|
||||
for (dir, _) in &targets {
|
||||
fs::create_dir_all(dir).map_err(|e| format!("Could not create {}: {e}", dir.display()))?;
|
||||
fs::create_dir_all(dir)
|
||||
.map_err(|e| format!("Could not create {}: {e}", dir.display()))?;
|
||||
}
|
||||
}
|
||||
// The plan must actually persist before bootstrap starts — a swallowed
|
||||
// write error here would bootstrap into a stale layout from disk while
|
||||
// the UI reports success.
|
||||
let cfg_path = config::config_path(&app).ok_or("Could not resolve the config file path")?;
|
||||
config::save_config_at(&cfg_path, &cfg)?;
|
||||
|
||||
// Custom paths are home-relative PII — log default-vs-custom flags, not
|
||||
// the raw locations.
|
||||
@@ -982,22 +1030,80 @@ pub fn complete_setup(
|
||||
custom(&cfg.models_dir),
|
||||
);
|
||||
|
||||
// `--setup` re-entry: a backend from the previous configuration may
|
||||
// still be serving. retry_bootstrap would attach to it and the new
|
||||
// env/mirror/layout settings would never apply — tear it down so the
|
||||
// restart spawns with the just-saved plan. (No-op on a true first run:
|
||||
// nothing is listening yet.)
|
||||
if crate::backend::port_in_use(crate::backend_port()) {
|
||||
log::info!(
|
||||
"Backend still running on port {} — restarting it so the new setup applies",
|
||||
crate::backend_port()
|
||||
);
|
||||
crate::backend::kill_orphan_on_port(crate::backend_port());
|
||||
std::thread::sleep(std::time::Duration::from_millis(500));
|
||||
// `--setup` re-entry: stop the previous backend before committing the new
|
||||
// durable layout. A failed stop must leave the old config/pointer active;
|
||||
// otherwise its automatic recovery would launch against a half-applied
|
||||
// setup plan.
|
||||
let persisted = crate::bootstrap::with_backend_stopped(&app, || -> Result<(), String> {
|
||||
let previous_location = if plan.install_mode == "portable" {
|
||||
Some(snapshot_portable_location()?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let result = (|| -> Result<(), String> {
|
||||
if plan.install_mode == "portable" {
|
||||
// Create the portable folder and seed config.json INSIDE it first,
|
||||
// so `config_path` resolves portable from here on and the whole
|
||||
// install (env + data + config) travels as one folder.
|
||||
let base = planned_portable_base(&plan).ok_or("Portable anchor disappeared")?;
|
||||
fs::create_dir_all(&base)
|
||||
.map_err(|e| format!("Could not create {}: {e}", base.display()))?;
|
||||
// Record WHERE it is before seeding it — `portable_base()` has to
|
||||
// resolve to this folder on the next launch, and the config inside
|
||||
// it cannot say so (nothing would know where to look).
|
||||
if base
|
||||
!= portable_anchor()
|
||||
.map(|a| a.join(PORTABLE_DIR_NAME))
|
||||
.unwrap_or_default()
|
||||
{
|
||||
let how = record_portable_dir(&app, &base)?;
|
||||
log::info!("Portable folder relocated — recorded via {how}");
|
||||
} else {
|
||||
// Back to the default: drop any earlier relocation so a stale
|
||||
// pointer can't outrank it.
|
||||
clear_portable_dir(&app);
|
||||
}
|
||||
config::save_config_at(&base.join("config.json"), &cfg)?;
|
||||
}
|
||||
// The plan must actually persist before bootstrap starts — a swallowed
|
||||
// write error here would bootstrap into a stale layout from disk while
|
||||
// the UI reports success.
|
||||
let cfg_path =
|
||||
config::config_path(&app).ok_or("Could not resolve the config file path")?;
|
||||
config::save_config_at(&cfg_path, &cfg)
|
||||
})();
|
||||
match (result, previous_location) {
|
||||
(Err(error), Some(snapshot)) => match restore_portable_location(snapshot) {
|
||||
Ok(()) => Err(error),
|
||||
Err(rollback_error) => Err(format!(
|
||||
"{error}; could not restore the previous portable location: {rollback_error}"
|
||||
)),
|
||||
},
|
||||
(result, _) => result,
|
||||
}
|
||||
});
|
||||
match persisted {
|
||||
Err(error) => {
|
||||
// with_backend_stopped already re-arms the old backend for every
|
||||
// non-terminal caller.
|
||||
log::warn!("Setup could not stop the previous backend: {error}");
|
||||
return Err("backend_stop_failed".into());
|
||||
}
|
||||
Ok(Err(error)) => {
|
||||
// The stop succeeded but persistence did not; restore service on
|
||||
// the surviving layout instead of leaving the settings UI down.
|
||||
crate::bootstrap::respawn_backend(
|
||||
app,
|
||||
state.stage.clone(),
|
||||
state.logs.clone(),
|
||||
);
|
||||
return Err(error);
|
||||
}
|
||||
Ok(Ok(())) => {}
|
||||
}
|
||||
|
||||
set_stage(&state.stage, BootstrapStage::Checking);
|
||||
crate::bootstrap::retry_bootstrap(app, state);
|
||||
crate::bootstrap::respawn_backend(app, state.stage, state.logs);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1036,6 +1142,29 @@ mod tests {
|
||||
assert_eq!(flavour, "pointer-absolute");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_snapshot_restores_present_and_absent_files() {
|
||||
let root = std::env::temp_dir().join(format!(
|
||||
"ov-portable-snapshot-{}-{:?}",
|
||||
std::process::id(),
|
||||
std::thread::current().id()
|
||||
));
|
||||
fs::create_dir_all(&root).unwrap();
|
||||
let present = root.join("present");
|
||||
let absent = root.join("absent");
|
||||
fs::write(&present, b"old").unwrap();
|
||||
let present_snapshot = snapshot_file(present.clone()).unwrap();
|
||||
let absent_snapshot = snapshot_file(absent.clone()).unwrap();
|
||||
|
||||
fs::write(&present, b"new").unwrap();
|
||||
fs::write(&absent, b"new").unwrap();
|
||||
restore_portable_location(vec![present_snapshot, absent_snapshot]).unwrap();
|
||||
|
||||
assert_eq!(fs::read(&present).unwrap(), b"old");
|
||||
assert!(!absent.exists());
|
||||
fs::remove_dir_all(root).unwrap();
|
||||
}
|
||||
|
||||
/// Portable-folder relocation (#1403 follow-up). One test fn, not several:
|
||||
/// it mutates `APPIMAGE`, which is process-global, and Rust tests run in
|
||||
/// parallel — same rationale as the uv-env test below.
|
||||
|
||||
@@ -0,0 +1,557 @@
|
||||
//! Headless loopback control plane for VoiceStudio dictation.
|
||||
//!
|
||||
//! The Python backend owns ASR data-plane protocols. This small Rust server
|
||||
//! owns desktop authority: start/stop capture and the session-bound native
|
||||
//! insertion target. Native integrations can therefore control the bundled
|
||||
//! dictation service without embedding a WebView or depending on Tauri IPC.
|
||||
|
||||
use std::io::{Read, Write};
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use tauri::Manager;
|
||||
|
||||
use crate::dictation_output::CaptureOrigin;
|
||||
use crate::{backend_port, dispatch_dictation_capture, AppFlags};
|
||||
|
||||
const DEFAULT_SIDECAR_PORT: u16 = 3902;
|
||||
const MAX_BODY_BYTES: usize = 1024 * 1024;
|
||||
const MAX_REQUEST_BYTES: usize = MAX_BODY_BYTES + 16 * 1024;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum DictationAction {
|
||||
Start,
|
||||
Stop,
|
||||
Toggle,
|
||||
}
|
||||
|
||||
impl DictationAction {
|
||||
fn wire_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Start => "start",
|
||||
Self::Stop => "stop",
|
||||
Self::Toggle => "toggle",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SpeechSidecarState {
|
||||
pub port: u16,
|
||||
stop: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Drop for SpeechSidecarState {
|
||||
fn drop(&mut self) {
|
||||
self.stop.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
struct HttpRequest {
|
||||
method: String,
|
||||
path: String,
|
||||
origin: Option<String>,
|
||||
body: Vec<u8>,
|
||||
}
|
||||
|
||||
struct HttpResponse {
|
||||
status: u16,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
impl HttpResponse {
|
||||
fn ok(body: Value) -> Self {
|
||||
Self { status: 200, body }
|
||||
}
|
||||
|
||||
fn error(status: u16, code: &str, message: &str) -> Self {
|
||||
Self {
|
||||
status,
|
||||
body: json!({"error": {"code": code, "message": message}}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sidecar_port() -> u16 {
|
||||
std::env::var("VOICESTUDIO_SPEECH_PORT")
|
||||
.ok()
|
||||
.and_then(|value| value.parse::<u16>().ok())
|
||||
.filter(|port| *port > 0)
|
||||
.unwrap_or(DEFAULT_SIDECAR_PORT)
|
||||
}
|
||||
|
||||
pub fn cli_dictation_action(args: &[String]) -> Option<DictationAction> {
|
||||
if args.iter().any(|arg| arg == "--dictate-start") {
|
||||
return Some(DictationAction::Start);
|
||||
}
|
||||
if args.iter().any(|arg| arg == "--dictate-stop") {
|
||||
return Some(DictationAction::Stop);
|
||||
}
|
||||
if args
|
||||
.iter()
|
||||
.any(|arg| arg == "--dictate" || arg == "--dictate-toggle")
|
||||
{
|
||||
return Some(DictationAction::Toggle);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn plan_action(
|
||||
requested: DictationAction,
|
||||
is_recording: bool,
|
||||
has_session: bool,
|
||||
) -> Option<DictationAction> {
|
||||
match requested {
|
||||
DictationAction::Start if !is_recording && !has_session => Some(DictationAction::Start),
|
||||
DictationAction::Stop if is_recording || has_session => Some(DictationAction::Stop),
|
||||
DictationAction::Toggle if is_recording || has_session => Some(DictationAction::Stop),
|
||||
DictationAction::Toggle => Some(DictationAction::Start),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dispatch_action(app: &tauri::AppHandle, action: DictationAction) -> Value {
|
||||
let flags = app.state::<AppFlags>();
|
||||
let before = flags.dictating.load(Ordering::SeqCst);
|
||||
// A start event creates its output session synchronously, before the
|
||||
// WebView can report `dictating=true`. Use both signals so two rapid
|
||||
// toggle calls mean start-then-stop instead of duplicate starts.
|
||||
let has_session = flags.output.current_session_id().is_some();
|
||||
let planned = plan_action(action, before, has_session);
|
||||
if let Some(planned) = planned {
|
||||
dispatch_dictation_capture(app, planned.wire_name());
|
||||
}
|
||||
let session_id = app.state::<AppFlags>().output.current_session_id();
|
||||
json!({
|
||||
"accepted": true,
|
||||
"action": action.wire_name(),
|
||||
"dispatched_action": planned.map(DictationAction::wire_name),
|
||||
"already_in_requested_state": planned.is_none(),
|
||||
"recording_before_request": before,
|
||||
"session_id": session_id,
|
||||
})
|
||||
}
|
||||
|
||||
fn status(app: &tauri::AppHandle, port: u16) -> Value {
|
||||
let flags = app.state::<AppFlags>();
|
||||
let engine_port = backend_port();
|
||||
json!({
|
||||
"schema": "voicestudio.speech-control-status",
|
||||
"protocol": "voicestudio.speech.v1",
|
||||
"service": "VoiceStudio",
|
||||
"service_version": env!("CARGO_PKG_VERSION"),
|
||||
"recording": flags.dictating.load(Ordering::SeqCst),
|
||||
"session_id": flags.output.current_session_id(),
|
||||
"control_url": format!("http://127.0.0.1:{port}"),
|
||||
"engine_url": format!("http://127.0.0.1:{engine_port}"),
|
||||
})
|
||||
}
|
||||
|
||||
fn capabilities(port: u16) -> Value {
|
||||
let engine_port = backend_port();
|
||||
json!({
|
||||
"schema": "voicestudio.speech-capabilities",
|
||||
"protocol": "voicestudio.speech.v1",
|
||||
"protocol_version": "1.0",
|
||||
"service": "VoiceStudio",
|
||||
"service_version": env!("CARGO_PKG_VERSION"),
|
||||
"local_first": true,
|
||||
"endpoints": {
|
||||
"status": format!("http://127.0.0.1:{port}/v1/status"),
|
||||
"dictation_start": format!("http://127.0.0.1:{port}/v1/dictation/start"),
|
||||
"dictation_stop": format!("http://127.0.0.1:{port}/v1/dictation/stop"),
|
||||
"dictation_toggle": format!("http://127.0.0.1:{port}/v1/dictation/toggle"),
|
||||
"output_sessions": format!("http://127.0.0.1:{port}/v1/output/sessions"),
|
||||
"json_rpc": format!("http://127.0.0.1:{port}/rpc"),
|
||||
"batch_transcription": format!("http://127.0.0.1:{engine_port}/v1/audio/transcriptions"),
|
||||
"streaming_transcription": format!("ws://127.0.0.1:{engine_port}/v1/audio/transcriptions/stream"),
|
||||
"mcp": format!("http://127.0.0.1:{engine_port}/mcp"),
|
||||
},
|
||||
"cli": {
|
||||
"start": "--dictate-start",
|
||||
"stop": "--dictate-stop",
|
||||
"toggle": "--dictate-toggle",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_output_session_path(path: &str) -> Option<(u64, bool)> {
|
||||
let suffix = path.strip_prefix("/v1/output/sessions/")?;
|
||||
if let Some(raw_id) = suffix.strip_suffix("/insert") {
|
||||
return raw_id.parse().ok().map(|id| (id, true));
|
||||
}
|
||||
suffix.parse().ok().map(|id| (id, false))
|
||||
}
|
||||
|
||||
fn begin_output_session(app: &tauri::AppHandle) -> HttpResponse {
|
||||
let flags = app.state::<AppFlags>();
|
||||
if flags.output.current_session_id().is_some() {
|
||||
return HttpResponse::error(
|
||||
409,
|
||||
"output_busy",
|
||||
"another dictation output session is active",
|
||||
);
|
||||
}
|
||||
let session_id = flags.output.begin_session(CaptureOrigin::Shortcut);
|
||||
HttpResponse::ok(json!({"session_id": session_id}))
|
||||
}
|
||||
|
||||
fn insert_output_session(app: &tauri::AppHandle, session_id: u64, body: &[u8]) -> HttpResponse {
|
||||
let request: Value = match serde_json::from_slice(body) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return HttpResponse::error(400, "invalid_json", "body must be JSON"),
|
||||
};
|
||||
let Some(text) = request.get("text").and_then(Value::as_str) else {
|
||||
return HttpResponse::error(400, "invalid_text", "body requires a string text field");
|
||||
};
|
||||
let output = app.state::<AppFlags>().output.clone();
|
||||
if output.current_session_id() != Some(session_id) {
|
||||
return HttpResponse::error(409, "stale_session", "output session is not active");
|
||||
}
|
||||
let result = output
|
||||
.activate_session(session_id)
|
||||
.and_then(|_| output.deliver(session_id, text));
|
||||
output.finish_session(session_id);
|
||||
match result {
|
||||
Ok(outcome) => HttpResponse::ok(json!({
|
||||
"session_id": session_id,
|
||||
"outcome": outcome,
|
||||
})),
|
||||
Err(error) => HttpResponse::error(500, "delivery_failed", &error),
|
||||
}
|
||||
}
|
||||
|
||||
fn cancel_output_session(app: &tauri::AppHandle, session_id: u64) -> HttpResponse {
|
||||
let output = app.state::<AppFlags>().output.clone();
|
||||
if output.current_session_id() != Some(session_id) {
|
||||
return HttpResponse::error(409, "stale_session", "output session is not active");
|
||||
}
|
||||
output.finish_session(session_id);
|
||||
HttpResponse::ok(json!({"session_id": session_id, "cancelled": true}))
|
||||
}
|
||||
|
||||
fn action_for_path(path: &str) -> Option<DictationAction> {
|
||||
match path {
|
||||
"/v1/dictation/start" => Some(DictationAction::Start),
|
||||
"/v1/dictation/stop" => Some(DictationAction::Stop),
|
||||
"/v1/dictation/toggle" => Some(DictationAction::Toggle),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn action_for_rpc_method(method: &str) -> Option<DictationAction> {
|
||||
match method {
|
||||
"dictation.start" => Some(DictationAction::Start),
|
||||
"dictation.stop" => Some(DictationAction::Stop),
|
||||
"dictation.toggle" => Some(DictationAction::Toggle),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_rpc(app: &tauri::AppHandle, body: &[u8]) -> HttpResponse {
|
||||
let request: Value = match serde_json::from_slice(body) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
return HttpResponse::ok(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": null,
|
||||
"error": {"code": -32700, "message": "Parse error"},
|
||||
}));
|
||||
}
|
||||
};
|
||||
let id = request.get("id").cloned().unwrap_or(Value::Null);
|
||||
let Some(method) = request.get("method").and_then(Value::as_str) else {
|
||||
return HttpResponse::ok(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"error": {"code": -32600, "message": "Invalid Request"},
|
||||
}));
|
||||
};
|
||||
let Some(action) = action_for_rpc_method(method) else {
|
||||
return HttpResponse::ok(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"error": {"code": -32601, "message": "Method not found"},
|
||||
}));
|
||||
};
|
||||
HttpResponse::ok(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": dispatch_action(app, action),
|
||||
}))
|
||||
}
|
||||
|
||||
fn origin_allowed(origin: Option<&str>) -> bool {
|
||||
let Some(origin) = origin else {
|
||||
return true;
|
||||
};
|
||||
matches!(
|
||||
origin,
|
||||
"tauri://localhost"
|
||||
| "http://tauri.localhost"
|
||||
| "https://tauri.localhost"
|
||||
| "http://localhost:3901"
|
||||
| "http://127.0.0.1:3901"
|
||||
)
|
||||
}
|
||||
|
||||
fn route(app: &tauri::AppHandle, request: HttpRequest, port: u16) -> HttpResponse {
|
||||
if !origin_allowed(request.origin.as_deref()) {
|
||||
return HttpResponse::error(
|
||||
403,
|
||||
"origin_denied",
|
||||
"browser origins cannot control dictation",
|
||||
);
|
||||
}
|
||||
let path = request.path.split('?').next().unwrap_or(&request.path);
|
||||
match (request.method.as_str(), path) {
|
||||
("GET", "/health") | ("GET", "/v1/status") => HttpResponse::ok(status(app, port)),
|
||||
("GET", "/.well-known/voicestudio-speech") | ("GET", "/v1/capabilities") => {
|
||||
HttpResponse::ok(capabilities(port))
|
||||
}
|
||||
("POST", "/v1/output/sessions") => begin_output_session(app),
|
||||
("POST", "/rpc") => handle_rpc(app, &request.body),
|
||||
("POST", path) if parse_output_session_path(path).is_some() => {
|
||||
let (session_id, is_insert) = parse_output_session_path(path).expect("guarded above");
|
||||
if !is_insert {
|
||||
return HttpResponse::error(405, "method_not_allowed", "use DELETE");
|
||||
}
|
||||
insert_output_session(app, session_id, &request.body)
|
||||
}
|
||||
("DELETE", path) if parse_output_session_path(path).is_some() => {
|
||||
let (session_id, is_insert) = parse_output_session_path(path).expect("guarded above");
|
||||
if is_insert {
|
||||
return HttpResponse::error(405, "method_not_allowed", "use POST");
|
||||
}
|
||||
cancel_output_session(app, session_id)
|
||||
}
|
||||
("POST", path) => match action_for_path(path) {
|
||||
Some(action) => HttpResponse::ok(dispatch_action(app, action)),
|
||||
None => HttpResponse::error(404, "not_found", "unknown speech-control endpoint"),
|
||||
},
|
||||
(_, "/v1/dictation/start" | "/v1/dictation/stop" | "/v1/dictation/toggle") => {
|
||||
HttpResponse::error(405, "method_not_allowed", "use POST")
|
||||
}
|
||||
_ => HttpResponse::error(404, "not_found", "unknown speech-control endpoint"),
|
||||
}
|
||||
}
|
||||
|
||||
fn header_end(bytes: &[u8]) -> Option<usize> {
|
||||
bytes.windows(4).position(|window| window == b"\r\n\r\n")
|
||||
}
|
||||
|
||||
fn read_request(stream: &mut TcpStream) -> Result<HttpRequest, &'static str> {
|
||||
stream
|
||||
.set_read_timeout(Some(Duration::from_secs(2)))
|
||||
.map_err(|_| "read timeout setup failed")?;
|
||||
let mut bytes = Vec::new();
|
||||
let mut chunk = [0_u8; 2048];
|
||||
let mut expected_len = None;
|
||||
loop {
|
||||
let read = stream.read(&mut chunk).map_err(|_| "request read failed")?;
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
bytes.extend_from_slice(&chunk[..read]);
|
||||
if bytes.len() > MAX_REQUEST_BYTES {
|
||||
return Err("request too large");
|
||||
}
|
||||
if let Some(end) = header_end(&bytes) {
|
||||
if expected_len.is_none() {
|
||||
let headers = String::from_utf8_lossy(&bytes[..end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
if content_length > MAX_BODY_BYTES {
|
||||
return Err("request body too large");
|
||||
}
|
||||
expected_len = Some(end + 4 + content_length);
|
||||
}
|
||||
if bytes.len() >= expected_len.unwrap_or(end + 4) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
let end = header_end(&bytes).ok_or("incomplete request headers")?;
|
||||
let headers = String::from_utf8_lossy(&bytes[..end]);
|
||||
let mut lines = headers.lines();
|
||||
let mut request_line = lines
|
||||
.next()
|
||||
.ok_or("missing request line")?
|
||||
.split_whitespace();
|
||||
let method = request_line.next().ok_or("missing method")?.to_owned();
|
||||
let path = request_line.next().ok_or("missing path")?.to_owned();
|
||||
let origin = lines.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("origin")
|
||||
.then(|| value.trim().to_owned())
|
||||
});
|
||||
Ok(HttpRequest {
|
||||
method,
|
||||
path,
|
||||
origin,
|
||||
body: bytes[end + 4..].to_vec(),
|
||||
})
|
||||
}
|
||||
|
||||
fn write_response(stream: &mut TcpStream, response: HttpResponse) {
|
||||
let body = response.body.to_string();
|
||||
let reason = match response.status {
|
||||
200 => "OK",
|
||||
403 => "Forbidden",
|
||||
404 => "Not Found",
|
||||
405 => "Method Not Allowed",
|
||||
409 => "Conflict",
|
||||
413 => "Payload Too Large",
|
||||
500 => "Internal Server Error",
|
||||
_ => "Bad Request",
|
||||
};
|
||||
let head = format!(
|
||||
"HTTP/1.1 {} {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nCache-Control: no-store\r\nConnection: close\r\n\r\n",
|
||||
response.status,
|
||||
reason,
|
||||
body.len(),
|
||||
);
|
||||
let _ = stream.write_all(head.as_bytes());
|
||||
let _ = stream.write_all(body.as_bytes());
|
||||
let _ = stream.flush();
|
||||
}
|
||||
|
||||
fn handle_connection(app: &tauri::AppHandle, mut stream: TcpStream, port: u16) {
|
||||
match read_request(&mut stream) {
|
||||
Ok(request) => write_response(&mut stream, route(app, request, port)),
|
||||
Err(message) => {
|
||||
let status = if message.contains("too large") {
|
||||
413
|
||||
} else {
|
||||
400
|
||||
};
|
||||
write_response(
|
||||
&mut stream,
|
||||
HttpResponse::error(status, "invalid_request", message),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn start(app: tauri::AppHandle) -> Result<SpeechSidecarState, String> {
|
||||
let port = sidecar_port();
|
||||
let listener = TcpListener::bind(("127.0.0.1", port))
|
||||
.map_err(|error| format!("could not bind 127.0.0.1:{port}: {error}"))?;
|
||||
listener
|
||||
.set_nonblocking(true)
|
||||
.map_err(|error| format!("could not configure speech sidecar: {error}"))?;
|
||||
|
||||
// The child Python backend inherits this and advertises the native control
|
||||
// endpoint only when the desktop shell actually owns it.
|
||||
std::env::set_var("VOICESTUDIO_SPEECH_CONTROL_PORT", port.to_string());
|
||||
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let thread_stop = stop.clone();
|
||||
thread::Builder::new()
|
||||
.name("speech-control-sidecar".into())
|
||||
.spawn(move || {
|
||||
log::info!("Speech control sidecar listening on 127.0.0.1:{port}");
|
||||
while !thread_stop.load(Ordering::SeqCst) {
|
||||
match listener.accept() {
|
||||
Ok((stream, _)) => handle_connection(&app, stream, port),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
|
||||
thread::sleep(Duration::from_millis(30));
|
||||
}
|
||||
Err(error) => {
|
||||
log::warn!("Speech control sidecar accept failed: {error}");
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.map_err(|error| format!("could not start speech sidecar thread: {error}"))?;
|
||||
|
||||
Ok(SpeechSidecarState { port, stop })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
action_for_path, action_for_rpc_method, cli_dictation_action, origin_allowed,
|
||||
parse_output_session_path, plan_action, DictationAction,
|
||||
};
|
||||
|
||||
fn args(values: &[&str]) -> Vec<String> {
|
||||
values.iter().map(|value| (*value).to_owned()).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_flags_map_to_idempotent_actions() {
|
||||
assert_eq!(
|
||||
cli_dictation_action(&args(&["VoiceStudio", "--dictate-start"])),
|
||||
Some(DictationAction::Start)
|
||||
);
|
||||
assert_eq!(
|
||||
cli_dictation_action(&args(&["VoiceStudio", "--dictate-stop"])),
|
||||
Some(DictationAction::Stop)
|
||||
);
|
||||
assert_eq!(
|
||||
cli_dictation_action(&args(&["VoiceStudio", "--dictate-toggle"])),
|
||||
Some(DictationAction::Toggle)
|
||||
);
|
||||
assert_eq!(cli_dictation_action(&args(&["VoiceStudio"])), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_and_json_rpc_share_the_same_action_vocabulary() {
|
||||
assert_eq!(
|
||||
action_for_path("/v1/dictation/start"),
|
||||
Some(DictationAction::Start)
|
||||
);
|
||||
assert_eq!(
|
||||
action_for_rpc_method("dictation.stop"),
|
||||
Some(DictationAction::Stop)
|
||||
);
|
||||
assert_eq!(action_for_rpc_method("dictation.delete"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_start_makes_a_second_toggle_stop() {
|
||||
assert_eq!(
|
||||
plan_action(DictationAction::Toggle, false, false),
|
||||
Some(DictationAction::Start)
|
||||
);
|
||||
assert_eq!(
|
||||
plan_action(DictationAction::Toggle, false, true),
|
||||
Some(DictationAction::Stop)
|
||||
);
|
||||
assert_eq!(plan_action(DictationAction::Start, false, true), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn browser_origins_cannot_silently_trigger_the_microphone() {
|
||||
assert!(origin_allowed(None));
|
||||
assert!(origin_allowed(Some("http://tauri.localhost")));
|
||||
assert!(origin_allowed(Some("http://localhost:3901")));
|
||||
assert!(!origin_allowed(Some("https://example.com")));
|
||||
assert!(!origin_allowed(Some("http://localhost.evil.test")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_session_paths_are_strictly_typed() {
|
||||
assert_eq!(
|
||||
parse_output_session_path("/v1/output/sessions/42"),
|
||||
Some((42, false))
|
||||
);
|
||||
assert_eq!(
|
||||
parse_output_session_path("/v1/output/sessions/42/insert"),
|
||||
Some((42, true))
|
||||
);
|
||||
assert_eq!(parse_output_session_path("/v1/output/sessions/nope"), None);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user