fix(api): keep /api/results listener open after benchmark completes (DYN-701)

[FrankD412]

Summary

• Bug: NVBug 6204962 / DYN-701 — clients polling /api/results never observe terminal status='complete' because the API server is torn down within ~150-200ms of results arriving. The window is shorter than any practical polling cadence, so complete is effectively unreachable.
• Fix: API service now owns its shutdown grace. New env var AIPERF_API_SERVER_POST_COMPLETE_GRACE (default 5.0s, range 0.0-60.0) controls how long _stop_api_server keeps uvicorn accepting connections before flipping should_exit = True. SystemController._stop_system_controller extends its pre-shutdown wait to max(0.5s, GRACE) when the API is enabled so the API child isn't SIGKILL'd mid-grace by the 2s per-process join timeout. Setting the env var to 0 recovers the old behavior.
• Attribution correction: the reporter's "AI suggestion" pointed at PR feat: YAML-native v2 config + adaptive sweep orchestrator with BO & search recipes #912 / commit 94a9102, but the immediate-stop pattern actually goes back to PR feat: GPU Telemetry #335 (GPU Telemetry), refactored by PR feat: Complete Timing System Overhaul and Sticky Credit Router #565 (Timing Overhaul). The /api/results endpoint itself was added later in PRs feat(api): add FastAPI service with plugin-based router system #717/feat(api): add progress, results, and workers routers with tests #721. This is a long-latent design mismatch surfaced by the API, not a regression.

Test plan

• Unit tests: env-var bounds (4 new), _stop_api_server grace ordering (2 new, asserts sleep happens before should_exit=True), SystemController wait extension (3 new — api-disabled keeps 0.5s, api-enabled stretches to grace, floor at 0.5s preserved).
• Full tests/unit/ suite passes (12,732 passing; one unrelated pre-existing mlflow_data_exporter flake reproducible on clean main).
• End-to-end verification against the NVBug's exact repro steps using tests/aiperf_mock_server:
  • GRACE=5s (default): running ... complete observed for ~5.1s, then connection refused. 31 successful complete responses across the grace window with 100ms polling.
  • GRACE=0 (env override = recovers buggy behavior): complete window collapses to ~169ms; only 1 observation caught. Confirms the env knob cleanly toggles back.
• Reviewer suggestion: smoke-test with --api-port set against a real LLM endpoint to confirm benchmark wall-clock now includes the expected ~5s shutdown tail.
• Cherry-pick to release/0.9.0 after merge.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added a configurable post-complete API listener grace period (default 5.0s, range 0–300s).

• Behavior

  • Shutdown sequencing now honors the configured grace period; post-broadcast delivery wait is floored to 0.5s when applicable and skipped in specified skip conditions.

• Documentation

  • Documented the environment variable, default, allowed range, and that 0 disables the grace window.

• Tests

  • Added unit tests covering grace behavior, skip conditions, validation, and controller shutdown timing.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

Try out this PR

Quick install:

pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@ca265f243ea67bcf15dcb41ae3f8ede76a0b2c99

Recommended with virtual environment (using uv):

uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@ca265f243ea67bcf15dcb41ae3f8ede76a0b2c99

---

Last updated for commit: ca265f2 • Browse code

[github-actions]

Fern Docs Preview: https://nvidia-preview-f36471d0-5d0f-4c2a-836e-6a85db17305a.docs.buildwithfern.com/aiperf/dev

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

Adds a configurable post-benchmark grace period (env AIPERF_API_SERVER_POST_COMPLETE_GRACE) that keeps the API listener open for up to the configured seconds so polling clients can observe final status; integrates config, API shutdown hold, controller coordination, tests, and documentation.

Changes

API Server Grace Window Feature

Layer / File(s)	Summary
Configuration & Validation src/aiperf/common/environment.py, tests/unit/common/test_environment.py	_APIServerSettings.POST_COMPLETE_GRACE is added with range validation (0.0–300.0 seconds, default 5.0). Tests cover default, environment-variable override, zero allowance, and out-of-range rejection.
API Server Grace Window Implementation src/aiperf/api/api_service.py, tests/unit/api/test_api_service.py	FastAPIService._stop_api_server reads POST_COMPLETE_GRACE, logs when >0, awaits asyncio.sleep(grace) to hold the listener open before shutdown. Tests verify non-zero grace wait, zero fast-path, and skipping when server task already completed.
System Controller Shutdown Coordination src/aiperf/controller/system_controller.py, tests/unit/controller/test_system_controller.py	SystemController adds _api_enabled and _is_api_service_alive() and computes delivery_grace during shutdown as max(0.5, POST_COMPLETE_GRACE) when API enabled and alive; otherwise uses 0.5s. Tests check disabled/enabled, minimum floor, and skip conditions.
User Documentation docs/environment-variables.md	Documented AIPERF_API_SERVER_POST_COMPLETE_GRACE with default, allowed range, and behavior (0 = immediate shutdown).

---

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

---

Poem

🐇 Five seconds more I softly stay,
a listening ear for callers' say.
A patient pause before I close,
then gentle threads unbind their bows.
Hooray — the graceful hop, repose.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title directly and accurately summarizes the main change: implementing a grace period to keep the API results listener open after benchmark completion.
Docstring Coverage	✅ Passed	Docstring coverage is 86.21% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unit/api/test_api_service.py`:
- Around line 228-233: The type annotation for sleep_calls is wrong: change its
type from list[float] to list[tuple[float, bool]] (or list[tuple[float, bool]]
using typing.Tuple if preferred) so it matches how values are appended in
fake_sleep and later unpacked; update the declaration of sleep_calls near the
fake_sleep async function (which appends (seconds,
bool(mock_server.should_exit))) to use the tuple type.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3d8b0986-b62a-4dd5-89ad-a4457e0e6c74

📥 Commits

Reviewing files that changed from the base of the PR and between 8b56760 and 1737be7.

📒 Files selected for processing (7)

• docs/environment-variables.md
• src/aiperf/api/api_service.py
• src/aiperf/common/environment.py
• src/aiperf/controller/system_controller.py
• tests/unit/api/test_api_service.py
• tests/unit/common/test_environment.py
• tests/unit/controller/test_system_controller.py

✅ Files skipped from review due to trivial changes (1)

• docs/environment-variables.md

[codecov]

Codecov Report

❌ Patch coverage is 95.45455% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/aiperf/controller/system_controller.py	93.75%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[coderabbitai]

🧹 Nitpick comments (1)

tests/unit/controller/test_system_controller.py (1)

515-577: ⚡ Quick win

Rename new tests to match the required test naming schema.

The newly added test names don’t follow the required test_<function>_<scenario>_<expected> format.

Proposed renames

-async def test_extends_to_grace_when_api_enabled_and_alive(
+async def test_stop_system_controller_api_enabled_and_alive_uses_post_complete_grace(

-async def test_floors_at_05s_when_api_alive_with_small_grace(
+async def test_stop_system_controller_api_alive_with_small_grace_uses_05s_floor(

-async def test_skips_grace_when_api_enabled_but_never_registered(
+async def test_stop_system_controller_api_enabled_never_registered_uses_05s(

-async def test_skips_grace_when_api_failed(
+async def test_stop_system_controller_api_failed_uses_05s(

-async def test_skips_grace_when_api_stopped(
+async def test_stop_system_controller_api_stopped_uses_05s(

As per coding guidelines: tests/**/*.py: Name test functions as test_<function>_<scenario>_<expected>.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/controller/test_system_controller.py` around lines 515 - 577,
Rename the tests to follow the test_<function>_<scenario>_<expected> naming
convention: update test_extends_to_grace_when_api_enabled_and_alive to
test_drive_stop_api_enabled_registered_running_extends_to_post_complete_grace,
test_floors_at_05s_when_api_alive_with_small_grace to
test_drive_stop_api_alive_small_grace_floors_0_5s,
test_skips_grace_when_api_enabled_but_never_registered to
test_drive_stop_api_enabled_never_registered_skips_grace,
test_skips_grace_when_api_failed to
test_drive_stop_api_registered_failed_skips_grace, and
test_skips_grace_when_api_stopped to
test_drive_stop_api_registered_stopped_skips_grace so they match the required
test_<function>_<scenario>_<expected> schema; locate and rename each function by
its current name in the diff (e.g.,
test_extends_to_grace_when_api_enabled_and_alive,
test_floors_at_05s_when_api_alive_with_small_grace, etc.).

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unit/controller/test_system_controller.py`:
- Around line 515-577: Rename the tests to follow the
test_<function>_<scenario>_<expected> naming convention: update
test_extends_to_grace_when_api_enabled_and_alive to
test_drive_stop_api_enabled_registered_running_extends_to_post_complete_grace,
test_floors_at_05s_when_api_alive_with_small_grace to
test_drive_stop_api_alive_small_grace_floors_0_5s,
test_skips_grace_when_api_enabled_but_never_registered to
test_drive_stop_api_enabled_never_registered_skips_grace,
test_skips_grace_when_api_failed to
test_drive_stop_api_registered_failed_skips_grace, and
test_skips_grace_when_api_stopped to
test_drive_stop_api_registered_stopped_skips_grace so they match the required
test_<function>_<scenario>_<expected> schema; locate and rename each function by
its current name in the diff (e.g.,
test_extends_to_grace_when_api_enabled_and_alive,
test_floors_at_05s_when_api_alive_with_small_grace, etc.).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 56c6306d-7403-4dac-80f3-f70eeee178e0

📥 Commits

Reviewing files that changed from the base of the PR and between f3646a1 and 3e17497.

📒 Files selected for processing (2)

• src/aiperf/controller/system_controller.py
• tests/unit/controller/test_system_controller.py

[dynamo-ops]

Previous review comments have been addressed. Approving.