kvcache-ai/AgentENV

Fork 0

mirror of https://github.com/kvcache-ai/AgentENV.git synced 2026-11-05 01:17:36 +00:00

[PR #414] fix: leave time for graceful shutdown when draining a node #416

New issue

Open

opened 2026-10-29 14:23:08 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:23:08 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/AgentENV/pull/414
Author: @dushulin
Created: 10/22/2026
Status: 🔄 Open

Base: main ← Head: fix/shutdown-drain-budget

📝 Commits (2)

75f021c fix(orchestrator): pause sandboxes concurrently during shutdown
5619497 fix(k8s): cap the preStop drain to leave time for graceful shutdown

📊 Changes

3 files changed (+90 additions, -36 deletions)

View changed files

📝 deploy/k8s/base/agentenv-daemonset.yaml (+10 -0)
📝 src/orchestrator/service.rs (+43 -36)
📝 src/orchestrator/tests.rs (+37 -0)

📄 Description

What

Make sure a draining node leaves the server enough time to finish its graceful shutdown:

deploy/k8s/base: the preStop drain loop stops waiting after 3000s, leaving 600s of the 3600s terminationGracePeriodSeconds to the server's SIGTERM shutdown.
run_shutdown_cleanup: preserve (pause / delete / wait for) sandboxes with bounded concurrency (8) instead of one at a time.

Why

Today a busy node is usually SIGKILLed in the middle of its graceful shutdown:

The preStop loop only exits when GET /sandboxes reports 0. Sandboxes with a timeout longer than the grace period keep it running for the whole 3600s. kubelet then grants a one-off 2s extension, sends SIGTERM, and SIGKILLs 2s later.
Even when SIGTERM arrives on time, run_shutdown_cleanup awaits each pause before starting the next, so the time needed grows linearly with the number of running sandboxes. UblkDeviceManager::global().shutdown_daemon() only runs after it returns.

SIGKILL takes out the whole container cgroup. Sandboxes that were not paused yet are never persisted, and uvm-ublk-daemon dies with its ublk devices still open by Firecracker. That skips the graceful daemon shutdown from #266 / #335, and on kernels without torvalds/linux@f40139fde5 it can crash the host. Details in #413.

Closes #413

Scope and non-goals

Included:

A deadline in the preStop loop of deploy/k8s/base/agentenv-daemonset.yaml.
Concurrent sandbox preservation in run_shutdown_cleanup, plus a unit test.

Not included:

No new configuration knob. The deadline lives next to terminationGracePeriodSeconds in the same manifest, and the concurrency is a constant.
No change to what shutdown does per sandbox (pause, delete template builders, wait for in-flight transitions), to the 3-pass retry, or to the order of orchestrator shutdown → ublk daemon shutdown.
No change to scheduler behavior. A terminating node is already excluded from Schedule.
The resume-vs-shutdown race (#372 / #400) is a separate problem.

Design and behavior changes

preStop

deadline = now + 3000, checked at the top of each iteration.
When it is reached, the hook logs preStop: drain deadline reached, leaving remaining sandboxes to graceful shutdown and exits. kubelet then sends SIGTERM, and the server pauses and persists what is left.
Behavior before the deadline is unchanged.
The deadline also bounds the loop when the API cannot be queried. Before, that case retried forever.

run_shutdown_cleanup

Each pass now runs the per-sandbox work through futures::stream::iter(..).buffer_unordered(8) and collects the failures.
Per-sandbox logic and error reporting are unchanged.
Concurrency is safe: pausing different sandboxes concurrently is already what happens when clients call the pause API concurrently, and per-sandbox state transitions are unchanged.
The bound is there so concurrent memory snapshot writes do not saturate the node disk.

Compatibility and operations

Public API or generated protocol: N/A, no API change.
Configuration or defaults: the reference manifest's preStop now gives up after 3000s. Deployments that copied the manifest and changed terminationGracePeriodSeconds should keep the deadline below it, with enough margin for their sandbox count.
Snapshot manifest, artifact layout, or storage format: N/A.
Upgrade and rollback: the preStop change applies to pods created from the new template. The old pods being replaced during the rollout still run the old hook and the old shutdown code. Rolling back restores the previous behavior.
Host requirements, permissions, ports, or dependencies: none. futures is already a dependency of the crate.

Validation

make fmt (ran cargo fmt --all -- --check)
make clippy (ran for the touched crate only, see below)
make test-unit
Relevant Rust integration tests
make -C services test (required when services/ changes)
Generated clients/server regenerated with the documented make target
Documentation updated
Benchmarks or performance comparison completed

Commands and results:

$ cargo fmt --all -- --check
(clean)

$ cargo clippy -p agentenv --all-targets --all-features -- -D warnings
Finished `dev` profile [unoptimized + debuginfo] target(s)

$ cargo test -p agentenv --lib orchestrator::
test result: ok. 151 passed; 0 failed; 0 ignored; 0 measured; 871 filtered out

New test: orchestrator::service::tests::shutdown_pauses_running_sandboxes_concurrently.

It runs 4 running sandboxes whose mock pause takes 10s, under paused tokio time, and asserts that shutdown() finishes in less than 2 pause durations and that all sandboxes end up Paused.
Against the previous sequential code it fails with 4 pauses should overlap, took 40s.
The existing shutdown tests pass unchanged, including the persister call-order test.

The preStop snippet was run under dash (the runtime image's /bin/sh) against a local fake API with the deadline lowered to a few seconds:

it exits with the deadline message once the deadline passes;
it keeps retrying while the API is unreachable, until the deadline.

Skipped checks and reasons:

make clippy / make test-unit cover more crates than this change touches; I ran clippy and the orchestrator unit tests for agentenv only.
Integration tests and a real cluster drain were not run in this repository's setup.

Risks and reviewer notes

The 3000s / 600s split is a judgment call. With concurrency 8, 600s covers roughly 8 × 600s / pause time sandboxes. I did not add a knob, since the manifest is the place to tune both numbers together.
Pausing 8 sandboxes at once means up to 8 memory snapshots being written and packaged at the same time during shutdown. That is the reason for the bound rather than unbounded concurrency.
Main files: src/orchestrator/service.rs (run_shutdown_cleanup) and deploy/k8s/base/agentenv-daemonset.yaml.
Open PR #400 also touches service.rs / tests.rs, but in the resume path. It should not conflict textually.

Checklist

The PR contains one coherent change and no unrelated formatting or refactoring.
New behavior is covered by tests, or I explained why testing is impractical.
Logs and examples contain no credentials, tokens, or private registry information.
I did not manually edit generated code without updating its source and regenerating it.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/AgentENV/pull/414 **Author:** [@dushulin](https://github.com/dushulin) **Created:** 10/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/shutdown-drain-budget` --- ### 📝 Commits (2) - [`75f021c`](https://github.com/kvcache-ai/AgentENV/commit/75f021cf66d038d0c882b1c6c72d3b9f63510680) fix(orchestrator): pause sandboxes concurrently during shutdown - [`5619497`](https://github.com/kvcache-ai/AgentENV/commit/561949784223f11b15e5b789a760c3963ce33180) fix(k8s): cap the preStop drain to leave time for graceful shutdown ### 📊 Changes **3 files changed** (+90 additions, -36 deletions) <details> <summary>View changed files</summary> 📝 `deploy/k8s/base/agentenv-daemonset.yaml` (+10 -0) 📝 `src/orchestrator/service.rs` (+43 -36) 📝 `src/orchestrator/tests.rs` (+37 -0) </details> ### 📄 Description ## What Make sure a draining node leaves the server enough time to finish its graceful shutdown: - `deploy/k8s/base`: the preStop drain loop stops waiting after 3000s, leaving 600s of the 3600s `terminationGracePeriodSeconds` to the server's SIGTERM shutdown. - `run_shutdown_cleanup`: preserve (pause / delete / wait for) sandboxes with bounded concurrency (8) instead of one at a time. ## Why Today a busy node is usually SIGKILLed in the middle of its graceful shutdown: 1. The preStop loop only exits when `GET /sandboxes` reports 0. Sandboxes with a timeout longer than the grace period keep it running for the whole 3600s. kubelet then grants a one-off 2s extension, sends SIGTERM, and SIGKILLs 2s later. 2. Even when SIGTERM arrives on time, `run_shutdown_cleanup` awaits each pause before starting the next, so the time needed grows linearly with the number of running sandboxes. `UblkDeviceManager::global().shutdown_daemon()` only runs after it returns. SIGKILL takes out the whole container cgroup. Sandboxes that were not paused yet are never persisted, and `uvm-ublk-daemon` dies with its ublk devices still open by Firecracker. That skips the graceful daemon shutdown from #266 / #335, and on kernels without torvalds/linux@f40139fde527 it can crash the host. Details in #413. ## Related issue Closes #413 ## Scope and non-goals Included: - A deadline in the preStop loop of `deploy/k8s/base/agentenv-daemonset.yaml`. - Concurrent sandbox preservation in `run_shutdown_cleanup`, plus a unit test. Not included: - No new configuration knob. The deadline lives next to `terminationGracePeriodSeconds` in the same manifest, and the concurrency is a constant. - No change to what shutdown does per sandbox (pause, delete template builders, wait for in-flight transitions), to the 3-pass retry, or to the order of orchestrator shutdown → ublk daemon shutdown. - No change to scheduler behavior. A terminating node is already excluded from `Schedule`. - The resume-vs-shutdown race (#372 / #400) is a separate problem. ## Design and behavior changes **preStop** - `deadline = now + 3000`, checked at the top of each iteration. - When it is reached, the hook logs `preStop: drain deadline reached, leaving remaining sandboxes to graceful shutdown` and exits. kubelet then sends SIGTERM, and the server pauses and persists what is left. - Behavior before the deadline is unchanged. - The deadline also bounds the loop when the API cannot be queried. Before, that case retried forever. **`run_shutdown_cleanup`** - Each pass now runs the per-sandbox work through `futures::stream::iter(..).buffer_unordered(8)` and collects the failures. - Per-sandbox logic and error reporting are unchanged. - Concurrency is safe: pausing different sandboxes concurrently is already what happens when clients call the pause API concurrently, and per-sandbox state transitions are unchanged. - The bound is there so concurrent memory snapshot writes do not saturate the node disk. ## Compatibility and operations - Public API or generated protocol: N/A, no API change. - Configuration or defaults: the reference manifest's preStop now gives up after 3000s. Deployments that copied the manifest and changed `terminationGracePeriodSeconds` should keep the deadline below it, with enough margin for their sandbox count. - Snapshot manifest, artifact layout, or storage format: N/A. - Upgrade and rollback: the preStop change applies to pods created from the new template. The old pods being replaced during the rollout still run the old hook and the old shutdown code. Rolling back restores the previous behavior. - Host requirements, permissions, ports, or dependencies: none. `futures` is already a dependency of the crate. ## Validation - [x] `make fmt` (ran `cargo fmt --all -- --check`) - [x] `make clippy` (ran for the touched crate only, see below) - [ ] `make test-unit` - [ ] Relevant Rust integration tests - [ ] `make -C services test` (required when `services/` changes) - [ ] Generated clients/server regenerated with the documented `make` target - [ ] Documentation updated - [ ] Benchmarks or performance comparison completed Commands and results: ```text $ cargo fmt --all -- --check (clean) $ cargo clippy -p agentenv --all-targets --all-features -- -D warnings Finished `dev` profile [unoptimized + debuginfo] target(s) $ cargo test -p agentenv --lib orchestrator:: test result: ok. 151 passed; 0 failed; 0 ignored; 0 measured; 871 filtered out ``` New test: `orchestrator::service::tests::shutdown_pauses_running_sandboxes_concurrently`. - It runs 4 running sandboxes whose mock pause takes 10s, under paused tokio time, and asserts that `shutdown()` finishes in less than 2 pause durations and that all sandboxes end up `Paused`. - Against the previous sequential code it fails with `4 pauses should overlap, took 40s`. - The existing shutdown tests pass unchanged, including the persister call-order test. The preStop snippet was run under `dash` (the runtime image's `/bin/sh`) against a local fake API with the deadline lowered to a few seconds: - it exits with the deadline message once the deadline passes; - it keeps retrying while the API is unreachable, until the deadline. Skipped checks and reasons: - `make clippy` / `make test-unit` cover more crates than this change touches; I ran clippy and the orchestrator unit tests for `agentenv` only. - Integration tests and a real cluster drain were not run in this repository's setup. ## Risks and reviewer notes - The 3000s / 600s split is a judgment call. With concurrency 8, 600s covers roughly `8 × 600s / pause time` sandboxes. I did not add a knob, since the manifest is the place to tune both numbers together. - Pausing 8 sandboxes at once means up to 8 memory snapshots being written and packaged at the same time during shutdown. That is the reason for the bound rather than unbounded concurrency. - Main files: `src/orchestrator/service.rs` (`run_shutdown_cleanup`) and `deploy/k8s/base/agentenv-daemonset.yaml`. - Open PR #400 also touches `service.rs` / `tests.rs`, but in the resume path. It should not conflict textually. ## Checklist - [x] The PR contains one coherent change and no unrelated formatting or refactoring. - [x] New behavior is covered by tests, or I explained why testing is impractical. - [x] Logs and examples contain no credentials, tokens, or private registry information. - [x] I did not manually edit generated code without updating its source and regenerating it. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 14:23:08 +00:00

bsakaguchi referenced this issue

2026-10-29 14:23:11 +00:00

[PR #417] feat(cli): add aenv compose up #418

bsakaguchi referenced this issue

2026-10-29 14:23:12 +00:00

[PR #416] feat(compose): add the guest Compose runtime image #420