Let workers drain and restart cleanly before leaked memory takes down the host - #163
Let workers drain and restart cleanly before leaked memory takes down the host#163mihow wants to merge 5 commits into
Conversation
The worker's resident host memory grows with the amount of work done and is not released while the process lives; on hosts without swap the end state is a kernel OOM kill. Until the retention itself is found and fixed (#144, #145), bound the damage by giving the worker a clean way to exit and be restarted fresh by its process supervisor (#147): - SIGUSR1 asks the worker to finish the batch in flight, post its results, and exit with status 0. Previously the only recycle lever was SIGTERM, which kills mid-batch and forces queue redelivery of the in-flight work. - _process_job accepts a should_stop callable checked at every batch boundary, so a drain during a long job stops between batches: completed batches' results are posted and the job's remaining tasks stay queued for the next worker to claim. - New AMI_WORKER_MAX_RSS_MB setting (default 0 = disabled): resident memory is sampled between jobs — when the working set is at its idle baseline, so a high reading indicates retained memory rather than a busy job — and a reading over the cap triggers the same drain path. - On multi-GPU hosts the parent process forwards SIGUSR1 to the spawned per-GPU worker instances, since process supervisors signal only the parent. Closes #147 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
|
Warning Review limit reached
Next review available in: 54 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughWorkers can drain through ChangesWorker recycling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Supervisor
participant WorkerLoop
participant ProcessJob
participant ResultPoster
Supervisor->>WorkerLoop: Send SIGUSR1
WorkerLoop->>ProcessJob: Request stop at batch boundary
ProcessJob->>ResultPoster: Flush completed results
WorkerLoop->>Supervisor: Exit after drain
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
trapdata/antenna/worker.py (1)
61-61: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd types to the signal callback signatures.
Add parameter and return types to
handle_signaland_forward_drain. These callbacks receive a signal number and a frame object.As per coding guidelines, “Use type hints in function signatures to document expected types without requiring extensive documentation.”
Also applies to: 149-153
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@trapdata/antenna/worker.py` at line 61, Add type annotations to the signal callback methods handle_signal and _forward_drain, annotating the signal number, frame object, and return type according to the project’s typing conventions while preserving their existing callback behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.env.example:
- Around line 19-24: Update the worker recycling documentation around
AMI_WORKER_MAX_RSS_MB to require successful-exit restarts: specify supervisord
autorestart=true and Docker restart policy always or unless-stopped, while
retaining the existing systemd guidance.
In `@trapdata/antenna/worker.py`:
- Around line 137-153: Update the startup flow around mp.spawn and
_forward_drain to register a parent SIGUSR1 handler before spawning, retain
signals received before readiness, and forward them only after every child
reports its _DrainRequest.handle_signal is installed. Add parameter and return
type hints to both signal callbacks, and add a multi-GPU startup test covering
signals during initialization.
In `@trapdata/settings.py`:
- Line 48: Update the worker_max_rss_mb setting validation and _check_rss_cap
behavior to reject negative values, reserving only 0 to disable the RSS cap.
Preserve positive values as active limits and use the existing settings
validation mechanism.
---
Nitpick comments:
In `@trapdata/antenna/worker.py`:
- Line 61: Add type annotations to the signal callback methods handle_signal and
_forward_drain, annotating the signal number, frame object, and return type
according to the project’s typing conventions while preserving their existing
callback behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a4098012-540e-4571-b6f3-00f5a7d3d1c4
📒 Files selected for processing (4)
.env.exampletrapdata/antenna/tests/test_worker_drain.pytrapdata/antenna/worker.pytrapdata/settings.py
… job A fresh production measurement changed the leak's shape: a worker that had completed exactly one job held 13-15 GB of resident memory while idle, whether the job was ~400 images or ~3,200 — retention that looks per-job and roughly fixed-size, not a per-hour ramp and not proportional to image count. (Alternative not yet excluded: both samples were the process's first job, so part of that is legitimately resident model weights.) If retention is per-job, a job count predicts growth better than a memory threshold, so: - New AMI_WORKER_MAX_JOBS setting (default 0 = disabled): after that many jobs the worker drains and exits cleanly for its supervisor to restart — a deterministic bound that works before the retention's size or source is known. A job that raised mid-processing counts too, since it may retain memory like a completed one; claims that yielded no work do not. - The worker logs "Resident memory after job N: X MiB" between jobs, where the working set is at its idle baseline. This is the instrument for the open question above: whether RSS keeps climbing after the second and third job, or plateaus after the first model load, is now answerable from ordinary worker logs. - _check_rss_cap is folded into _after_job_check alongside the new trigger; the RSS cap behaves as before. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
|
This doesn't seem necessary right now. I recommend we focus on the leak. Also the tests are failing. |
…e caps Worker instances are started with the spawn method, so each begins life with SIGUSR1 at its default action of terminating the process, and installs its own handler only once the worker loop runs. Forwarding the signal from the parent could land in that window and kill a child that was still starting up, and torch responds by terminating its siblings. The parent was also still at SIGUSR1's default action itself until after mp.spawn returned, so a drain arriving while its children were already running killed the parent and orphaned workers holding GPU memory. The parent now owns SIGUSR1 on behalf of the whole group and publishes drains through a shared spawn-context event that each instance reads at the safe points it already checks. That removes both windows rather than narrowing them: a drain requested before an instance is ready is simply observed on its first check. Worker instances also install their drain handler before slower startup work. Recycle caps now reject negative values, which the threshold checks would otherwise read as "disabled", quietly turning off a cap an operator meant to set. Only 0 disables a cap. .env.example now states that recycling requires a process manager configured to restart the worker after a successful exit, since a drain exits with status 0. Supervisord's autorestart=unexpected, systemd's Restart=on-failure and Docker's on-failure policy all leave a cleanly exited worker stopped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
trapdata/antenna/worker.py (1)
303-307: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winAn idle worker can delay the drain exit by up to
SLEEP_TIME_SECONDS.
time.sleepis restarted after a Python signal handler returns, so a SIGUSR1 that arrives during the idle sleep does not shorten it. The loop only re-checksdrain.requestedafter the full sleep. Issue#147requires a prompt exit when no batch is active.Sleep in short increments and break when a drain is requested.
♻️ Proposed change
if not any_jobs and not drain.requested: logger.info( f"[GPU {gpu_id}] No jobs found, sleeping for {SLEEP_TIME_SECONDS} seconds" ) - time.sleep(SLEEP_TIME_SECONDS) + slept = 0.0 + while slept < SLEEP_TIME_SECONDS and not drain.requested: + time.sleep(min(1.0, SLEEP_TIME_SECONDS - slept)) + slept += 1.0🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@trapdata/antenna/worker.py` around lines 303 - 307, Update the idle-worker sleep in the loop handling no jobs and no drain request to use short, repeated intervals instead of one full SLEEP_TIME_SECONDS call; re-check drain.requested between intervals and break immediately once a drain is requested, while preserving the existing idle logging and total sleep duration when no drain occurs.
🧹 Nitpick comments (2)
trapdata/antenna/worker.py (1)
223-225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the return type annotation to
_worker_loop.The parameters are typed, but the return type is missing. The coding guidelines require type hints in function signatures.
♻️ Proposed change
def _worker_loop( gpu_id: int, pipelines: list[str], drain_event: EventType | None = None -): +) -> None:As per coding guidelines: "Use type hints in function signatures to document expected types without requiring extensive documentation".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@trapdata/antenna/worker.py` around lines 223 - 225, Add the appropriate return type annotation to the _worker_loop function signature, preserving its existing parameters and behavior.Source: Coding guidelines
trapdata/antenna/tests/test_worker_drain.py (1)
34-36: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a return type to
_spawn_event.Declare the event return type. This makes the helper contract explicit for
_DrainRequesttests.As per coding guidelines, “Use type hints in function signatures to document expected types without requiring extensive documentation.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@trapdata/antenna/tests/test_worker_drain.py` around lines 34 - 36, Add an explicit return type annotation to the _spawn_event function, using the appropriate multiprocessing event type while preserving its existing spawned-context Event construction.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@trapdata/settings.py`:
- Around line 45-48: Update the RSS limit descriptions in trapdata/settings.py
lines 45-48 and .env.example lines 21-23 to state that the worker exits when
resident memory reaches or exceeds the configured cap, matching the inclusive
check in _after_job_check.
---
Outside diff comments:
In `@trapdata/antenna/worker.py`:
- Around line 303-307: Update the idle-worker sleep in the loop handling no jobs
and no drain request to use short, repeated intervals instead of one full
SLEEP_TIME_SECONDS call; re-check drain.requested between intervals and break
immediately once a drain is requested, while preserving the existing idle
logging and total sleep duration when no drain occurs.
---
Nitpick comments:
In `@trapdata/antenna/tests/test_worker_drain.py`:
- Around line 34-36: Add an explicit return type annotation to the _spawn_event
function, using the appropriate multiprocessing event type while preserving its
existing spawned-context Event construction.
In `@trapdata/antenna/worker.py`:
- Around line 223-225: Add the appropriate return type annotation to the
_worker_loop function signature, preserving its existing parameters and
behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 820f0b63-873f-46a8-bd5c-105773aa29bf
📒 Files selected for processing (4)
.env.exampletrapdata/antenna/tests/test_worker_drain.pytrapdata/antenna/worker.pytrapdata/settings.py
The between-jobs check drains at `rss_mb >= max_rss_mb`, so a worker sitting exactly at the cap recycles. The setting's description, the operator guidance and the warning it logs all said "exceeds", which reads as strictly greater. A test pins the boundary so the documented behaviour and the check cannot drift apart. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
1 participant
Summary
The ML worker's host memory grows with the work it has done and is not released while the process lives. Two rounds of measurement on two production GPU hosts (2026-08-11):
Caveat, stated plainly: the second round is two samples, and both processes had loaded models for the first time in that window — a long-lived worker is supposed to hold resident model weights, so some of that 13–15 GB may be legitimate. The distinguishing test is whether RSS climbs further after a second and third job, which has not been run yet. This PR builds the instrument for it (see change 3 below). These hosts run without swap, so if per-job retention is real, the terminal state on a 68 GB host is a hard kernel OOM after roughly four to five jobs per process — an outage, not a slowdown.
This is a separate defect from the GPU out-of-memory failures addressed in PR #162. That one is a fixed peak working set colliding with a co-tenant process on a shared card (the same 426 MiB allocation failing 504 times at a stable 13.49 GiB allocated — no drift). This one is unbounded growth of ordinary host memory that scales with jobs processed; idle GPU memory retention measured at the same moment was only 124–722 MiB per process, and the DataLoader subprocesses suspected of the retention (#144 Part B, #145) never touch CUDA. Different memory, different processes, different mechanism, different fix; conflating them cost days of triage earlier.
This PR does not fix the leak — it bounds the damage, implementing the drain-and-exit proposal from #147. A worker can be told (or can decide on its own, via opt-in caps) to finish the batch in flight, post its results, and exit cleanly with status 0, so its process supervisor restarts a fresh process. Any leak — this one or a future one — becomes routine process churn instead of a host outage.
Why this change rather than the alternatives on the table: #144 Part A (LRU model eviction) manages GPU memory of idle models, not this host-RAM growth, and is gated on the leak audit; the root-cause work in #145 first needs instrumentation to find the retaining call sites; PR #148 improves DataLoader hygiene and by its own measurements removes the catastrophic per-job blow-up, but residual growth across jobs remains (its RSS figures do not return to baseline). And if retention really is per-job, recycling after N jobs caps it deterministically without the leak ever being located. Bounded process lifetime is the smallest change that does that, and it stays useful as defense in depth after the underlying bugs are fixed. It complements #148 rather than overlapping it — #148 explicitly lists worker recycling as out of scope.
List of Changes
supervisorctl signal USR1 <program>) makes the worker finish the batch it is processing, post those results, and exit with status 0 for the supervisor to restart. The only lever before this was SIGTERM, which kills mid-batch and forces queue redelivery of the in-flight work.AMI_WORKER_MAX_JOBSdrains the worker after that many jobs — the deterministic bound, since job count is what the measurements say predicts growth;AMI_WORKER_MAX_RSS_MBdrains when resident memory sampled between jobs exceeds the cap — the backstop for whatever the job count does not predict. A job that raised mid-processing counts toward the job cap (it may retain memory like a completed one); job claims that yielded no work do not.Resident memory after job N: X MiBbetween jobs, when the working set is at its idle baseline. This makes the open question above — does RSS keep climbing after jobs 2 and 3, or plateau after the first model load? — answerable from ordinary production logs, with no host instrumentation._process_jobaccepts ashould_stopcallable checked at every batch boundary. On a stop, results for completed batches are still flushed, and the job's remaining tasks stay queued for the next worker (or the restarted one) to claim.trapdata/antenna/tests/test_worker_drain.py, 16 tests): the signal handler and drain state; both recycle caps in both directions (fire past their threshold, stay silent at the 0 default — a cap that fired while disabled would recycle every deployment that did not opt in); the per-job log line; the batch-boundary stop with and without a stop request; and the polling loop exiting after a signal.What draining does and does not guarantee
What still needs verification
AMI_WORKER_MAX_JOBSin the low single digits is the operative bound on a 68 GB host running multiple workers; the RSS cap needs an observation cycle before recommending a number. Both need the distinguishing test first.Relationship to other open work
Test and lint results
uv run pytest --import-mode=importlib): 54 passed, 1 skipped, 1 failed — the failure istrapdata/api/tests/test_models.py::TestSourceImageSchema::test_url, which live-fetches an image from Wikimedia and received an HTTP 400 from the external server; pre-existing environmental flake, unrelated to this branch.🤖 Generated with Claude Code
https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF
Summary by CodeRabbit
0disables that limit.