Skip to content

[Test] Add unit tests for Triton load watch - #36106

Open
chg073 wants to merge 1 commit into
sgl-project:mainfrom
chg073:test-triton-load-watch
Open

[Test] Add unit tests for Triton load watch#36106
chg073 wants to merge 1 commit into
sgl-project:mainfrom
chg073:test-triton-load-watch

Conversation

@chg073

@chg073 chg073 commented Aug 23, 2026

Copy link
Copy Markdown

Motivation

Part of #20865.

triton_load_watch.py monitors Triton compilations and kernel device-loads that occur after serving starts, but it does not have focused unit-test coverage. These paths are important because late compilation can stall serving and late device-loading can fail when GPU memory headroom is low.

Modifications

Add CPU-only tests covering:

  • Idempotent hook installation and preservation of an existing compilation listener
  • Missing triton.knobs and unavailable hook APIs
  • Arming diagnostics through mark_serving_started()
  • Compilation warning gates for startup, cache hits, fast compilation, and the exact threshold boundary
  • Kernel loads before serving, low-memory warnings, and the memory threshold boundary
  • CUDA memory-query failures and crash-on-late-load behavior
  • Isolation and restoration of module-level state between tests

The test mirrors the source path under test/registered/unit/utils/, uses CustomTestCase, and is registered with register_cpu_ci(est_time=2, suite="base-a-test-cpu").

Accuracy Tests

Not applicable. This PR only adds unit tests and does not change model outputs or production behavior.

Speed Tests and Profiling

Not applicable. This PR does not affect inference performance. The targeted unit test completed in 0.08 seconds locally.

Local Validation

Command:

pytest test/registered/unit/utils/test_triton_load_watch.py -v

Output:

============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.1.1, pluggy-1.6.0
collected 9 items

test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_compilation_chains_existing_listener PASSED [ 11%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_compilation_warning_gates PASSED [ 22%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_install_is_noop_when_required_hooks_are_unavailable PASSED [ 33%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_install_is_noop_when_triton_knobs_cannot_be_imported PASSED [ 44%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_install_registers_once_and_chains_existing_listener PASSED [ 55%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_kernel_load_before_serving_does_not_query_cuda PASSED [ 66%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_kernel_load_query_failure_is_safe_unless_crash_mode_is_enabled PASSED [ 77%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_kernel_load_warns_only_below_memory_threshold PASSED [ 88%]
test/registered/unit/utils/test_triton_load_watch.py::TestTritonLoadWatch::test_mark_serving_started_arms_diagnostics PASSED [100%]

=============== 9 passed, 1 warning, 9 subtests passed in 0.08s ================

Additional validation:

python scripts/lint/check_registered_tests.py
# passed

pre-commit run --files test/registered/unit/utils/test_triton_load_watch.py
# all applicable hooks passed

Name                                           Stmts   Miss  Cover
------------------------------------------------------------------
python/sglang/srt/utils/triton_load_watch.py      53      0   100%

Checklist

  • Format your code according to the Format code with pre-commit.
  • Add unit tests according to the Run and add unit tests.
  • Documentation updates are not applicable to this test-only change.
  • Accuracy and speed benchmarks are not applicable to this test-only change.
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #32661869803
Latest PR Test (Extra): ❌ Run #32661869683
Latest PR Test (AMD ROCm 7.2): ❌ Run #32661870034

This file contains CPU-only unit tests for Triton load diagnostics, including installation checks, compilation listener chaining, and memory warning conditions.
@chg073

chg073 commented Aug 23, 2026

Copy link
Copy Markdown
Author

@cctry @ispobock Could you please review this PR and trigger CI by adding the run-ci label? The base-a-test-cpu job is currently skipped because the required label is missing. Thank you!

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.

1 participant