Skip to content

fix(runner): bound MPS worker failures - #2336

Draft
fml09 wants to merge 2 commits into
cocoindex-io:mainfrom
fml09:g/2333-mps-runner-resilience
Draft

fix(runner): bound MPS worker failures#2336
fml09 wants to merge 2 commits into
cocoindex-io:mainfrom
fml09:g/2333-mps-runner-resilience

Conversation

@fml09

@fml09 fml09 commented Aug 3, 2026

Copy link
Copy Markdown

Summary

Follow-up to #2335.

Adds a private, Python 3.11-compatible process supervisor for built-in MPS SentenceTransformer work. The supervisor owns the worker process and IPC lifecycle, which makes timeout and cancellation enforceable, guarantees that a failed worker is reaped before its GPU permit is released, and bounds hard-crash recovery to one replay.

Generic coco.GPU execution deliberately remains on the existing ProcessPoolExecutor path.

Stack and review scope

This is the upper change in a two-PR logical stack:

  1. Merge fix(ops): bound SentenceTransformer MPS memory聽#2335, which introduces the native MPS memory-bound path.
  2. Rebase this branch onto the updated main and force-push it.
  3. Mark this PR ready and merge it after the resilience-only diff is reviewed.

GitHub native stacked pull requests do not support branches across a fork boundary, so this PR is opened against main as a draft. Until #2335 merges, the GitHub diff includes both commits. Review the resilience layer independently at:

fml09/cocoindex@g/2333-native-mps-memory-safety...g/2333-mps-runner-resilience

Problem

ProcessPoolExecutor is suitable for the existing generic subprocess mode, but on Python 3.11 it does not expose ownership primitives that can reliably terminate one running call and prove that its worker has exited. That leaves three failure modes unbounded for an MPS worker:

  • A hung call can retain unified memory and a CocoIndex GPU permit indefinitely.
  • Cancelling the awaiting task does not guarantee that the child process stops.
  • Repeated hard crashes can repeatedly restart the pool without a per-call recovery bound.

A timeout wrapped only around the future would bound the caller wait, but not the lifetime of the GPU work. The implementation therefore needs to own the process and IPC that carry the call.

Failure contract

For private MPS calls, this PR establishes the following contract:

  • A five-minute absolute wall-clock deadline is shared by the original attempt and any replay. It includes supervisor queueing, worker startup, lazy model loading, execution, and replay.
  • An unexpected hard worker crash is replayed at most once. A second crash fails the call.
  • Timeout, cancellation, and user-code exceptions are never replayed.
  • Timeout and cancellation retire the active worker with terminate, join, kill, join escalation before the GPU permit is released.
  • Idle shutdown uses a graceful STOP and STOPPED handshake before hard-stop fallback.
  • Picklable user exceptions preserve their original type and receive the remote traceback; unpicklable exceptions use a stable transport fallback.

The hard-crash budget is intentionally per call, not a cross-call circuit breaker. This bounds one request without introducing global health state or new public policy.

Implementation

  • Add private _SingleProcessSupervisor, backed by a spawned multiprocessing.Process, a duplex pipe, and a dedicated manager thread.
  • Tag messages with worker generation and call identifiers so stale responses cannot complete a newer request.
  • Serialize request outcome transitions under one lock to make close and cancellation races exactly-once.
  • Replace only _MPS_GPU subprocess execution with the owned supervisor.
  • Preserve the generic coco.GPU subprocess implementation and its environment-controlled behavior.
  • Add no public API, configuration knob, or dependency.

Operational trade-offs

  • Replacing a crashed or timed-out worker incurs a cold model reload on the next attempt.
  • A hard-crash replay has at-least-once execution semantics. The scoped built-in embedding operation is computation-only, so it has no external commit boundary.
  • A process using both generic subprocess GPU execution and built-in MPS execution may keep two persistent child processes, one for each intentionally separate path.
  • Setting COCOINDEX_RUN_GPU_IN_SUBPROCESS=0 opts out of the private supervisor and therefore also opts out of these timeout and hard-reap guarantees.
  • A slow first model download can consume the five-minute budget; the deadline is an absolute execution bound, not heartbeat-based hang detection.

Testing

  • Added 15 real spawned-process tests covering success, exception transport, unpicklable exceptions, timeout, cancellation, idle and active close, graceful shutdown, crash replay, repeated crash failure, stale messages, cross-event-loop use, and shutdown races.
  • Repeated the real-process supervisor suite five times to exercise timing-sensitive paths.
  • Added runner and SentenceTransformer tests proving only the private MPS route changes.
  • uvx prek run --all-files passes, including Ruff, mypy, maturin build, Rust tests and checks, the full Python test suite, and generated CLI documentation checks.

fml09 added 2 commits August 3, 2026 12:02
Route built-in MPS embedding work through an isolated persistent runner, install conservative allocator watermarks, and reclaim caches under pressure without changing generic GPU runner defaults.

Preserve split-and-retry OOM behavior and document the environment overrides.

Fixes cocoindex-io#2333
Own the private MPS worker process and IPC lifecycle so timeout and cancellation can terminate and reap active work on Python 3.11.

Replay an unexpected hard crash at most once within a shared five-minute deadline, while leaving generic coco.GPU subprocess execution unchanged.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant