fix(runner): bound MPS worker failures - #2336
Draft
fml09 wants to merge 2 commits into
Draft
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.GPUexecution deliberately remains on the existingProcessPoolExecutorpath.Stack and review scope
This is the upper change in a two-PR logical stack:
mainand force-push it.GitHub native stacked pull requests do not support branches across a fork boundary, so this PR is opened against
mainas 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
ProcessPoolExecutoris 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 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:
STOPandSTOPPEDhandshake before hard-stop 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
_SingleProcessSupervisor, backed by a spawnedmultiprocessing.Process, a duplex pipe, and a dedicated manager thread._MPS_GPUsubprocess execution with the owned supervisor.coco.GPUsubprocess implementation and its environment-controlled behavior.Operational trade-offs
COCOINDEX_RUN_GPU_IN_SUBPROCESS=0opts out of the private supervisor and therefore also opts out of these timeout and hard-reap guarantees.Testing
uvx prek run --all-filespasses, including Ruff, mypy, maturin build, Rust tests and checks, the full Python test suite, and generated CLI documentation checks.