Skip to content

[TRTLLM-12648][test] implement disagg cancellation load thread#15124

Merged
chienchunhung merged 1 commit into
NVIDIA:mainfrom
chienchunhung:codex/phase0-load-thread
Jun 9, 2026
Merged

[TRTLLM-12648][test] implement disagg cancellation load thread#15124
chienchunhung merged 1 commit into
NVIDIA:mainfrom
chienchunhung:codex/phase0-load-thread

Conversation

@chienchunhung

@chienchunhung chienchunhung commented Jun 8, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fifth thread body in the disaggregated cancellation stress-test suite - the marathon-style harness that gates regressions of the disaggregated KV-transfer cancellation / poison failure class. The skeleton + log scanner landed in the first harness change, followed by the metrics thread, injector thread, and canary thread; this change adds the load thread.

This PR ships:

  • load_thread body - replaces the no-op stub with a duration-bounded loop. The thread exits on stop_event / failed_event, runs until stress_config.duration_min elapses, and sets stop_event at normal end-of-marathon so canary, metrics, injector, and log-scanner threads wind down.
  • Steady / burst load scheduling - selects the active load shape from stress_config. Steady state uses base_concurrency and input_length; burst windows use bursts.interval_min, bursts.duration_s, bursts.concurrency, and bursts.input_length. Bursts begin after the first full interval so the marathon starts with a steady-state baseline.
  • run_cancel_stress_test wrapper - calls the existing disaggregated cancellation load generator one burst at a time via _run_cancel_stress_iteration(...). The import is lazy so unit tests can monkeypatch the method without importing the heavyweight disaggregated integration module.
  • Load result records - appends one _load_records entry per load-generator call with {timestamp, elapsed_s, mode, num_bursts, requests_per_burst, prompt_len_range, cancel_after_range, success, error, duration_s} and exposes those records from collect_results() as load_records.
  • Config parsing helpers - adds _parse_token_range, _parse_cancel_after_range, and _load_iteration_shape to validate load-shape knobs and fail fast on malformed load config, including malformed burst timing.
  • Lifecycle smoke update - keeps the current setup()-stub path fast: if no server endpoint is bound, load_thread logs a warning, sets stop_event, and exits, so the marathon pytest remains a lifecycle smoke until real cluster setup lands.

PR chain plan

Change Adds Marathon assertion it enables
skeleton + log scanner harness lifecycle + log_scanner_thread hard-zero log patterns
metrics thread metrics_thread KV-cache utilization growth bound
injector thread injector_thread SIGSTOP / SIGCONT / SIGKILL+respawn; injection_events count
canary thread canary_thread canary correctness plus error-rate / recovery-time inputs
this change load_thread sustained steady/burst cancellation load over the configured duration
next real setup_disagg_cluster invocation with save_log=True + endpoint binding marathon entry point becomes a real cluster-backed marathon
later marathon YAML/reference completion + test-list registration weekly stress CI enablement

This change still does not enable the marathon in CI. Registration in tests/integration/test_lists/qa/llm_function_stress.txt remains deferred.

Why one-burst wrapper calls

run_cancel_stress_test already owns the asyncio request-spam/cancel behavior used by existing disaggregated cancellation tests. Rather than reimplementing the HTTP client, load_thread calls it with num_bursts=1 in a loop and lets the harness decide which load shape is active for each iteration.

This keeps Step 5 small and reviewable:

  • the shared load generator remains the only implementation of client-side disconnect-during-prefill traffic;
  • the harness owns marathon scheduling, duration, stop/fail-fast coordination, and observability;
  • future load-generator improvements, such as honoring client_cancel_rate or output_length, can land separately without restructuring the harness.

Tests added

All deterministic, GPU-free, second-scale. Under tests/integration/defs/stress_test/disagg_cancel/.

  • test_load_thread.py - 10 unit tests across three layers:
    • Config helpers (5): token range defaults/validation, cancel-after range defaults/validation, steady-to-burst load-shape selection, and invalid burst interval/duration fail-fast.
    • Thread-body lifecycle (3): no-endpoint path signals stop_event, steady iterations record successful load_records, and externally set stop_event stops after the current runner call.
    • Burst/failure behavior (2): burst window uses burst concurrency/prompt range, and runner exceptions trip fail-fast while recording the failed load iteration.
  • test_disagg_cancel_stress.py - lifecycle smoke now asserts the load_records result bucket.

Run locally (no GPU, no cluster):

PYTHONPATH=tests/integration/defs:tests/integration/defs/disaggregated \
python3 -m pytest -c /dev/null -o addopts= \
  -p no:cacheprovider \
  --confcutdir=tests/integration/defs/stress_test \
  tests/integration/defs/stress_test/disagg_cancel/test_load_thread.py \
  tests/integration/defs/stress_test/disagg_cancel/test_disagg_cancel_stress.py -q

Full component suite:

PYTHONPATH=tests/integration/defs:tests/integration/defs/disaggregated \
python3 -m pytest -c /dev/null -o addopts= \
  -p no:cacheprovider \
  --confcutdir=tests/integration/defs/stress_test \
  tests/integration/defs/stress_test/disagg_cancel/ -q

Validation

  • Focused load/lifecycle tests: 12 passed.
  • Full colocated component suite: 89 passed.
  • git diff --check.
  • Pre-commit passed after ruff-format reformatted harness.py. The local commit skipped only waive list check and validate-test-lists because the hook runner's Python failed while parsing the repo script annotation str | None in scripts/check_test_list.py.

Note: the full component suite was run outside the filesystem sandbox because existing canary/metrics tests bind local loopback HTTP servers.

Out of scope (deferred to later changes in this chain)

  • Real setup_disagg_cluster invocation and worker/server endpoint binding. setup() is still a stub, so the lifecycle smoke does not launch a real cluster yet.
  • Extending run_cancel_stress_test to honor client_cancel_rate and output_length; Step 5 intentionally reuses the existing cancellation load generator as-is.
  • End-of-marathon gate computation from _load_records, _canary_records, _kv_utilization_samples, and _injection_events.
  • configs/stress_canary_prompts.json plus tools/generate_canary_references.py.
  • Test-list registration for weekly stress CI.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions).

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities.

  • CODEOWNERS updated if ownership changes.

  • Documentation updated as needed.

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

  • Tests

    • Added comprehensive test coverage for load-thread behavior in stress tests, including validation of cancellation load configurations, iteration shape transitions, and error handling scenarios.
    • Enhanced existing stress-test assertions to verify load-record collection.
  • Documentation

    • Updated stress-test README with improved coverage details and explicit step-by-step run instructions.

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
@chienchunhung
chienchunhung force-pushed the codex/phase0-load-thread branch from de5c27d to d3906ca Compare June 8, 2026 23:01
@chienchunhung
chienchunhung marked this pull request as ready for review June 8, 2026 23:04
@chienchunhung
chienchunhung requested a review from a team as a code owner June 8, 2026 23:04
@chienchunhung

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This PR implements a functional load thread for the disaggregated cancellation stress harness, converting a no-op stub into a duration-bounded loop that executes cancellation stress iterations in configurable steady or burst load patterns, records detailed per-iteration results, and handles error scenarios with fail-fast behavior.

Changes

Load Thread Feature

Layer / File(s) Summary
Config Type Foundation and Load Shape Helpers
tests/integration/defs/stress_test/disagg_cancel/harness.py
StressConfig.duration_min type changed from int to float for YAML coercion. New parsing helpers validate token-length ranges, cancel_after_range configuration, and compute steady-vs-burst iteration shape based on elapsed marathon time.
Harness Constructor and Load Records State
tests/integration/defs/stress_test/disagg_cancel/harness.py
Constructor extended with load_duration_s and load_iteration_pause_s parameters. New _load_records bucket introduced to accumulate per-iteration observations during load execution.
Load Thread Core Implementation
tests/integration/defs/stress_test/disagg_cancel/harness.py
Load thread replaced with duration-bounded loop that derives iteration shape, executes cancellation stress iterations, records success/error and timing per iteration, applies bounded pause between cycles, marks fail-fast on errors, and signals stop_event on completion.
Harness API and Test Integration
tests/integration/defs/stress_test/disagg_cancel/harness.py, tests/integration/defs/stress_test/disagg_cancel/test_disagg_cancel_stress.py
start() docstring updated to describe thread-exit behavior when harness not fully wired. collect_results() schema includes load_records. Test docstring and inline comments clarified; assertion extended to validate load_records field presence.
Load Thread Unit Test Module
tests/integration/defs/stress_test/disagg_cancel/test_load_thread.py
New test module with YAML fixture helpers and harness factory. Comprehensive coverage: parsing validation (token and cancel-after ranges), iteration shape switching (steady → burst), thread lifecycle (missing endpoint, normal completion), steady/burst result recording, stop-event observation, and exception fail-fast with recorded failure reason and error context.
Documentation and Configuration Updates
tests/integration/defs/stress_test/disagg_cancel/README.md, tests/integration/defs/stress_test/disagg_cancel/configs/marathon_cpp_v1_deepseek.yaml
README expanded with load-thread status, file layout, and restructured "how to run" section into numbered test-execution steps. Link formatting, CI registration wording, and lifecycle header refreshed. YAML config comment updated to reference "in-flight cancellation fix" instead of specific PR.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested reviewers

  • pcastonguay
  • StanleySun639
  • jieli-matrix
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 51.85% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the main change: implementing the load thread for the disaggregated cancellation stress-test harness, which is the primary focus of this PR.
Description check ✅ Passed The PR description is comprehensive and follows the template structure with Summary, Test Coverage, PR Checklist sections, and additional context about the PR chain plan, rationale, validation steps, and deferred items.
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.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
tests/integration/defs/stress_test/disagg_cancel/harness.py (1)

1227-1230: 💤 Low value

Consider catching more specific exceptions from the load generator.

The broad Exception catch here is flagged by static analysis (BLE001). While recording and failing fast on any runner error is reasonable, catching Exception may mask unexpected issues like KeyboardInterrupt being caught (though it's a BaseException). If the underlying run_cancel_stress_test raises known exception types (e.g., ConnectionError, TimeoutError, RuntimeError), consider narrowing the catch clause.

That said, since this is test infrastructure and the error is both logged and triggers fail-fast, the current behavior is defensible for robustness.

🤖 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/integration/defs/stress_test/disagg_cancel/harness.py` around lines
1227 - 1230, Replace the broad "except Exception as exc:" in the load thread
block with a narrow catch for the known runner errors (e.g., except
(ConnectionError, TimeoutError, RuntimeError) as exc:) so you only record and
call self.mark_failed for expected failures from run_cancel_stress_test;
additionally add a separate except BaseException: raise to avoid swallowing
system-level interrupts (KeyboardInterrupt/SystemExit). Ensure you still set
record["error"] and call self.mark_failed(f"load_thread runner failed:
{record['error']}") in the narrowed exception handler.

Source: Coding guidelines

tests/integration/defs/stress_test/disagg_cancel/README.md (1)

81-100: 💤 Low value

Optional: Consider clarifying step numbering.

The grouping of "Steps 1-2" for log_scanner + metrics, followed by separate Steps 3, 4, 5, is understandable but slightly inconsistent. Consider either:

  • Numbering all five components separately (Steps 1 through 5), or
  • Using "Step 1" for the grouped command since it's executed as a single pytest invocation

The current form is functional and the commands are correct, so this is purely a readability suggestion.

🤖 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/integration/defs/stress_test/disagg_cancel/README.md` around lines 81 -
100, The step numbering is inconsistent: "Steps 1-2 — log scanner + metrics"
then separate Steps 3–5; update the README to make numbering clear by either
splitting the grouped command into two numbered steps (Step 1:
tests/integration/defs/stress_test/disagg_cancel/test_log_scanner.py and Step 2:
tests/integration/defs/stress_test/disagg_cancel/test_metrics_thread.py) or
relabel the grouped pytest invocation as a single Step (e.g., "Step 1 — log
scanner + metrics (run together)") and then increment following steps
accordingly (refer to the commands invoking test_log_scanner.py,
test_metrics_thread.py, test_injector.py, test_canary.py, and
test_load_thread.py when editing).
🤖 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/integration/defs/stress_test/disagg_cancel/harness.py`:
- Around line 1227-1230: Replace the broad "except Exception as exc:" in the
load thread block with a narrow catch for the known runner errors (e.g., except
(ConnectionError, TimeoutError, RuntimeError) as exc:) so you only record and
call self.mark_failed for expected failures from run_cancel_stress_test;
additionally add a separate except BaseException: raise to avoid swallowing
system-level interrupts (KeyboardInterrupt/SystemExit). Ensure you still set
record["error"] and call self.mark_failed(f"load_thread runner failed:
{record['error']}") in the narrowed exception handler.

In `@tests/integration/defs/stress_test/disagg_cancel/README.md`:
- Around line 81-100: The step numbering is inconsistent: "Steps 1-2 — log
scanner + metrics" then separate Steps 3–5; update the README to make numbering
clear by either splitting the grouped command into two numbered steps (Step 1:
tests/integration/defs/stress_test/disagg_cancel/test_log_scanner.py and Step 2:
tests/integration/defs/stress_test/disagg_cancel/test_metrics_thread.py) or
relabel the grouped pytest invocation as a single Step (e.g., "Step 1 — log
scanner + metrics (run together)") and then increment following steps
accordingly (refer to the commands invoking test_log_scanner.py,
test_metrics_thread.py, test_injector.py, test_canary.py, and
test_load_thread.py when editing).

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f8012131-bef8-4d15-849e-7ed2dffe484c

📥 Commits

Reviewing files that changed from the base of the PR and between 9827c21 and d3906ca.

📒 Files selected for processing (5)
  • tests/integration/defs/stress_test/disagg_cancel/README.md
  • tests/integration/defs/stress_test/disagg_cancel/configs/marathon_cpp_v1_deepseek.yaml
  • tests/integration/defs/stress_test/disagg_cancel/harness.py
  • tests/integration/defs/stress_test/disagg_cancel/test_disagg_cancel_stress.py
  • tests/integration/defs/stress_test/disagg_cancel/test_load_thread.py

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52874 [ run ] triggered by Bot. Commit: d3906ca Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #52874 [ run ] completed with state SUCCESS. Commit: d3906ca
/LLM/main/L0_MergeRequest_PR pipeline #42126 completed with status: 'SUCCESS'

CI Report

Link to invocation

@chienchunhung
chienchunhung merged commit 736dc22 into NVIDIA:main Jun 9, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants