Improve async-safety of agents for indefinite blocking of kernel runners
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 17h 7m
- Merged PRs (30d)
- 358
Description
## Summary
- Motivation: [https://lablup.atlassian.net/browse/BA-6508](https://lablup.atlassian.net/browse/BA-6508)
## Steps to Reproduce
1. Create a new session with default settings (no GPU required).
1. Spawn many many processes (more than 5000, for example) within the `main1` container.
1. Observe degradation of the agent service on the host running that specific container.
1. Frequent timeouts of statistics collection and heartbeats
1. Unresponsive kernel management operations on the host
## Expected Behavior
- It is possible to get performance degradation of operations against a specific container when the continer goes rogue, but it should have no/minimal impacts to other containers and the agent itself.
- Even when OOM occurs, the agent and the kernel runner should be the last ones to get killed.
## Actual Behavior
- We could observe unresponsive agents or very slow system reactions.
## Directions to Fix
- [https://lablup.atlassian.net/browse/BA-6521](https://lablup.atlassian.net/browse/BA-6521)
- This issue covers:
- **kernel-runner**: Wrap `scan_proc_stats()` with `asyncio.run_in_executor()` to prevent it from blocking the event loop, which may cause indefinite hangs of the agent-to-kernel ping messages.
- We also need to update the lifecycle-long asyncio tasks spawning and termination in the base kernel to use `aiotools.cancel_and_wait()` pattern.
- **agent**: Ensure asynchronously concurrent collection of per-container process metrics. It should be written using `as_completed_safe()` to prevent HoL blocking by a specific slow container. Or, we should spawn a separate process-metric-collection timer loop for each individual container.
- It would be better to use `asyncio.Semaphore` to limit the maximum concurrency of stat collection when there are many containers.
- **agent**: Let’s improve the “poor-man’s RPC” between the agent and the kernel runner:
- Add an explicit timeout to ping invocation in the agent side.
- Add per-message ID tracking to avoid “off-by-one” problem when there is a late response that comes after the agent-side timeouts, as the current RPC channel between agent and kernel-runner assumes synchronous, continuous message streaming.
- The agent-side should have per-kernel message loops to make request-response pairs asynchronous.
- The agent-side client could have a dict to keep track of message IDs with their `asyncio.Future` instances.
- Let the kernel methods use this common client invoker.
- Create a `future` via `loop.create_future()`, add it to the tracker dict, and `await` it to return its result.
- When the agent’s kernel-runner message loop receives a response:
- Check the message ID availability, and silently drop and exit when it does not exist in the tracker dict.
- Check `future.done()`, and exit silently when it is.
- Call `future.set_result()`, and clear the dict entry.
- When the method invoker times out or gets cancelled:
- Check `future.done()`, and exit silently when it is.
- Call `future.cancel()`, and clear the dict entry. Raise/re-raise `asyncio.CancelledError`.
- Ensure timeout compsability by using `asyncio.timeout()` async-context manager.
- When the agent’s kernel-runner message loop terminates:
- Call `future.cancel()` to all dict entries unless they’re `done()`. Clear the entire dict.
- Use the serial number arithmetic as implemented in [https://github.com/lablup/callosum/blob/main/src/callosum/serial.py](https://github.com/lablup/callosum/blob/main/src/callosum/serial.py) (RFC 1982) to have unique integer IDs per message with minimal computation overheads.
- We could consider adopting Callosum to avoid reinvention of the wheel, but to ensure backward compatibility as described below, it would be better to extend the current implementation.
- **Backward compatibility:** To allow seamless upgrade of the agent when there are running kernels, we need to have the following fallback mechanisms:
- Add a distinguishing label when creating a new container. (e.g., `ai.backend.kernel.rpc.version=2`)
- Use the ID-based RPC against the containers having the explicit label only, while skipping ID checks for legacy containers.
- This will allow the agent to make concurrent requests against a single kernel.
- For example, ping would not be blocked by another slow agent-to-krunner request such as OOM notification (BA-6508 case).
- Currently, it is already possible to make concurrent requests to different kernels.
\*\* **kernel-runner**: Update the existing message loop to become a similar async message loop to make the handlers asynchronous.
\*\*\* Instead of using `asyncio.Future` instances, spawn `asyncio.create_task()` per request and keep track of the task instances.
\*\*\* Other implemnetations are mostly same! It just replaces the remote network I/O with local task completion.
\*\*\* Add an explicit message type to propagate cancellation in the kernel-runner side (e.g., abrupt container restart or termination) and update the agent-side loop to resolve it as future cancellation.
JIRA Issue: BA-6522
Contributor guide
Assessment
This issue has not been assessed yet.