AgentClientPool leaks zmq.Context on connect failure, melting down managers over time
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 15h 13m
- Merged PRs (30d)
- 368
Description
## Summary
`AgentClientPool._create_entry` constructs a `PeerInvoker` (which allocates one `zmq.asyncio.Context` plus a background IO thread via callosum) and then calls `client.connect()`. When `connect()` raises (e.g. agent unreachable, TCP refused, version mismatch handshake failure), the peer is never closed and no `_CachedEntry` is stored. The leaked context lives until the manager process exits, and the next `acquire()` for the same agent re-creates a fresh peer and leaks another one.
In production (dogbowl cluster, 26.4.4rc6), the manager accumulated ~3,000 ZMQ background threads per worker over ~7 days (~12,000 per host × 3 hosts ≈ 36,000 total). Load averages reached 100+ on 4-core hosts and `/metrics` scrapes timed out; previously diagnosed at 2026-05-22 and again 2026-05-29 with the same signature. The acute trigger this round: `sokovan.scheduler.terminator.terminator` periodically calling `check_running(kernel_id)` for 4 orphan RUNNING kernels left on an intentionally-offline agent (`i-haplo03`, TCP refused). Each cycle of the terminator's stale-presence check called `acquire(i-haplo03)`, which failed `connect()`, which leaked one IO thread, and because nothing was cached as unhealthy the next cycle repeated the same RPC and leak. Any unreachable agent path produces the same effect.
## Root Cause
`pool.py:_create_entry` (current behaviour):
```python
peer: PeerInvoker = self._create_peer(...) # allocates zmq.asyncio.Context + IO thread
client = AgentClient(peer, agent_id)
try:
await client.connect()
except Exception as e:
raise AgentConnectionUnavailable(agent_id, str(e)) from e
# peer is discarded without close() → zmq.Context never destroyed → IO thread leaks
```
`_get_or_create` _does_ have an unhealthy short-circuit (`if entry is not None and not entry.is_healthy: raise`), but the creation-failure path raises before the entry is stored, so the short-circuit never engages for agents that fail on the very first connect.
Secondary: `AgentClient.close()` swallows all teardown exceptions silently (`except Exception: pass`), making future regressions in the callosum exit chain unobservable.
## Steps to Reproduce
1. Register an agent whose RPC port (default 6001) refuses TCP connections.
1. From manager: `await agent_client_pool.acquire(agent_id)` (or let `sokovan.terminator` run a cycle for a RUNNING kernel owned by that agent).
1. Repeat: each acquire allocates a fresh `PeerInvoker` and leaks one `zmq.asyncio.Context` + IO thread.
Diagnostic: `ls /proc//task | xargs -I{} cat /proc//task/{}/comm | grep -c '^ZMQ'`.
## Expected Behavior
- `_create_entry` connect-failure path closes the partially-initialized peer so its `zmq.asyncio.Context` is destroyed (`destroy(linger=50)`), releasing the IO thread.
- The failure is cached as an unhealthy `_CachedEntry` so subsequent `acquire()` calls short-circuit via the existing `is_healthy=False` branch, instead of repeatedly re-creating peers.
- `_health_check_loop` continues to evict unhealthy entries after `recovery_timeout`, allowing a fresh retry once the agent comes back.
## Fix
PR: branch `fix/agent-client-pool-zmq-context-leak`.
1. `src/ai/backend/manager/clients/agent/pool.py`
- In `_create_entry`'s except-branch: `await client.close()` to destroy the zmq context, then return an unhealthy `_CachedEntry` instead of raising.
- In `_get_or_create`: after `_create_entry` returns an unhealthy entry, store it in `self._entries` and raise `AgentConnectionUnavailable` so the current caller is notified, while subsequent acquires now hit the cached unhealthy short-circuit.
1. `src/ai/backend/manager/clients/agent/client.py`
- `AgentClient.close()` upgrades the silent `except Exception: pass` to `log.debug(...)` so future regressions in callosum's exit chain are at least observable.
1. New unit tests under `tests/unit/manager/clients/agent/test_pool.py`:
- `test_closes_peer_and_caches_unhealthy_on_connect_failure`
- `test_subsequent_acquire_short_circuits_without_recreating_peer`
## Out of Scope (Follow-ups)
- Defensive cleanup of orphan RUNNING kernels when an agent transitions to TERMINATED. The current `mark_agent_exit` only updates the agent row; orphan kernels can remain RUNNING indefinitely and keep being probed by the terminator (this triggered the leak in production, but the design choice is intentional to support "agent restarts while containers keep running" — a separate design discussion).
- WebUI display policy for sessions owned by non-ALIVE agents.
## Logs / Evidence (dogbowl, 2026-05-29 incident)
- Per-host thread counts before rolling restart: 12,134 / 12,092 / 11,872 (4 workers × ~3000 nlwp each).
- Per-host thread counts after rolling restart: 142 / 158 / 153 (4 workers × ~35 nlwp each).
- Trigger agent: `i-haplo03` (TERMINATED 2026-05-20 10:30 UTC, TCP `10.81.0.23:6001` refused).
- Orphan kernels driving terminator probes: 4 RUNNING kernels in session `8Etn8wR2-session` owned by `demo-admin@lablup.com`, dated 2026-05-19.
JIRA Issue: BA-6239
Contributor guide
Assessment
This issue has not been assessed yet.