Improve agent container resource-mounting performance (aiodocker + runtime layer)
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 15h 13m
- Merged PRs (30d)
- 368
Description
## Main idea
The agent's container resource-mounting pipeline (CPU, memory, GPU) works correctly but has accumulated several latent performance issues and half-landed optimizations. This epic tracks the work to address them and to prepare the runtime layer for CDI / containerd / Kata adoption.
Investigation summary (evidence below): the allocation pipeline — `KernelResourceSpec` → `allocate()` → per-plugin `generate_docker_args()` → `docker.containers.create()` — is well-architected, but stats collection polls the Docker API per-container, Docker connections are opened per-operation, NUMA memory pinning is implemented-but-commented-out, scratch-dir size collection is hard-disabled, and GPU mounting is tied to Docker-specific `DeviceRequests` (blocks non-Docker runtimes).
## Scope
### P0 — localized, low-risk, measurable wins
- [ ] **Enable** `CpusetMems` when allocation is NUMA-local
`src/ai/backend/agent/docker/intrinsic.py:429–442` — `CpusetMems` is commented out despite `AffinityMap` already computing NUMA nodes. Gate it on `len(unique_numa_nodes) == 1`. Expected 2–3× improvement on remote-memory access for NUMA hosts.
- [ ] **Pool the Docker daemon connection**
`src/ai/backend/agent/docker/agent.py:826–835, 1240` — `async with closing_async(Docker())` is called 5–10× per kernel start (once per computer plugin + creation). Keep a long-lived client on the agent.
- [ ] **Stream container stats instead of per-cycle polling**
`src/ai/backend/agent/docker/intrinsic.py:134–169` — `stats(stream=False)` with a 2 s per-container timeout (`_CONTAINER_STAT_TIMEOUT`) means a single hung container delays the stat cycle. Move to a streamed reader per container.
- [ ] **Default to sysfs-first stats on Linux; Docker API only for network/IO and fallback**
`src/ai/backend/agent/docker/intrinsic.py:262–347, 625–900` — dual path already exists (`StatModes.CGROUP` vs `DOCKER`). Drops CPU/memory stat cost from ~100–500 ms to <10 ms.
### P1 — blocked or requires design
- [ ] **Re-enable scratch directory size collection** via inotify-driven incremental tracking
`src/ai/backend/agent/docker/intrinsic.py:631–648` currently hard-returns `0` — comment notes the blocking `rglob()` caused "indefinite accumulation of stat collector tasks."
- [ ] **Cache container inspect JSON** with short TTL + event-based invalidation
`src/ai/backend/agent/docker/resources.py:64`, `agent.py:1625` — recovery paths redo `container.show()` each time.
- [ ] **Adopt CDI for NVIDIA GPU mounting**
`src/ai/backend/accelerator/cuda_open/plugin.py:367–419` — replace Docker-specific `DeviceRequests` with CDI device names (`nvidia.com/gpu=N`). Unblocks containerd/CRI and Kata without per-runtime branching.
- [ ] **Fractional GPU support (MIG / MPS)**
`FractionAllocMap` already exists in `agent/resources.py` but is not wired into the CUDA plugin. Auto-detect A100/H100 and enable MIG; fall back to MPS on older SKUs. Needs a design note first.
### P2 — larger refactors
- [ ] **Require cgroup v2 for new installs**, drop dual v1/v2 code paths at `intrinsic.py:273–287, 660–722`.
- [ ] **Direct containerd/CRI path** (skip dockerd). Depends on CDI above + a gRPC CRI client.
- [ ] **Kata / gVisor runtime selection hook** — today runtime is implicitly `runc`; no selection point exists in `start_container()`.
## Out of scope
- K8s agent (`src/ai/backend/agent/kubernetes/`) — tracked separately.
- Storage/volume-mount performance — different subsystem.
## Alternative ideas
- Instead of pooling `Docker()` in the current code, migrate agent container ops to a thin internal client that wraps containerd directly and keep aiodocker only for legacy Docker hosts. Larger blast radius — propose as a BEP if P0/P1 are not enough.
## Anything else?
Runtime support matrix as-is:
|Runtime|Status|
|---|---|
|Docker (runc)|First-class|
|Kubernetes (containerd)|Parallel agent, partial parity|
|nvidia DeviceRequests (Docker ≥19.3)|Supported|
|nvidia-container-runtime (legacy)|Supported|
|Kata / gVisor / Podman / CRI direct|Not supported|
Notable half-landed work found during investigation:
- `CpusetMems` — commented out (`intrinsic.py:440`)
- `get_scratch_size` — returns `0` unconditionally (`intrinsic.py:635`)
- `FractionAllocMap` — defined but unused for GPUs
JIRA Issue: BA-5856
Contributor guide
Assessment
This issue has not been assessed yet.