Pool the Docker daemon connection in the agent
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 15h 13m
- Merged PRs (30d)
- 368
Description
Parent epic: #11216
## Main idea
The agent currently opens and closes a fresh aiodocker connection for each container operation. `src/ai/backend/agent/docker/agent.py:1240` and the accelerator-config generation loop (`agent.py:826–835`) use `async with closing_async(Docker())`, which means a single kernel start opens on the order of 5–10 Docker connections (one for `generate_docker_args` per computer plugin, one for `generate_accelerator_configs` per plugin, one for `containers.create()`, one for `start()`, plus recovery/inspect paths). Each connection pays TCP/Unix-socket setup and HTTP parser initialization.
Replace the per-call pattern with a long-lived `Docker` client owned by the agent, passed into accelerator plugins instead of created inside them. Preserve the current `closing_async()` semantics only for the agent shutdown path.
Concrete places to change:
- `src/ai/backend/agent/docker/agent.py:826` (`apply_accelerator_allocation`)
- `src/ai/backend/agent/docker/agent.py` (`generate_accelerator_configs`)
- `src/ai/backend/agent/docker/agent.py:1240` (`start_container`)
- Computer plugin `generate_docker_args(docker, ...)` signature already accepts a `docker` arg — callers just need to stop creating throwaway clients.
## Alternative ideas
- Connection pool (max N idle) rather than a singleton — useful if concurrent kernel starts end up serialized on a single connection. Profile first before adding the complexity.
- Keep per-op connections but reduce the number of ops (batch `generate_docker_args` into one pass). Smaller win, does not help recovery paths.
## Anything else?
Needs a sanity check that aiodocker's `Docker` is safe to share across concurrent tasks (it is, as of current versions — connector is an `aiohttp.ClientSession`). Should pair with a small benchmark: kernel start latency before/after, measured at the agent.
JIRA Issue: BA-5858
Contributor guide
Assessment
This issue has not been assessed yet.