[TRTLLM-14727][test] Create MX donor-receiver qualification test harness - #17222
Conversation
563863c to
8de5259
Compare
chienchunhung
changed the title
8de5259 to
68818df
Compare
|
/bot run --stage-list "DGX_H100-2_GPUs-PyTorch-ModelExpress-1" |
|
PR_Github #63881 [ run ] triggered by Bot. Commit: |
|
PR_Github #63881 [ run ] completed with state
|
68818df to
7bcea5a
Compare
7bcea5a to
9bf88cb
Compare
|
Rebased onto the latest The previous targeted H100 request was rejected before execution because this PR lacks |
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe pull request adds ModelExpress donor/receiver GPU qualification tests, a reusable worker, Kubernetes sidecar support, H100 CI stages, on-demand stage filtering, test configuration, and qualification documentation. ChangesModelExpress qualification
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant QualificationTest
participant MxE2EWorker
participant ModelExpress
participant Redis
QualificationTest->>MxE2EWorker: Start baseline, donor, and receiver
MxE2EWorker->>ModelExpress: Publish and consume model parameters
ModelExpress->>Redis: Coordinate synchronization
MxE2EWorker-->>QualificationTest: Write token IDs and transfer logs
QualificationTest->>QualificationTest: Compare outputs and validate transfers
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
tests/integration/defs/model_express/mx_e2e_worker.py (1)
100-105: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBound the donor wait loop.
The donor polls
stop_filewith no deadline. If the orchestrating test process dies before it writes the stop file, this process holds its GPUs and the loaded engine until the CI pod ends.Add a maximum wait so the donor exits on its own.
♻️ Proposed bounded wait
if args.role == "donor": assert args.ready_file is not None assert args.stop_file is not None args.ready_file.write_text("ready\n", encoding="utf-8") - while not args.stop_file.exists(): - time.sleep(0.2) + deadline = time.monotonic() + args.max_serve_seconds + while not args.stop_file.exists(): + if time.monotonic() >= deadline: + raise TimeoutError( + f"The stop file {args.stop_file} did not appear within " + f"{args.max_serve_seconds}s" + ) + time.sleep(0.2)Add the matching argument in
_parse_args:parser.add_argument("--max-serve-seconds", type=float, default=1800.0)🤖 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/model_express/mx_e2e_worker.py` around lines 100 - 105, Bound the donor polling loop in the role-handling flow by adding the proposed --max-serve-seconds argument in _parse_args with its default, then have the donor stop waiting when either stop_file appears or the configured deadline expires. Use a monotonic elapsed-time check around the existing time.sleep polling.jenkins/L0_Test.groovy (1)
3328-3390: 🩺 Stability & Availability | 🔵 TrivialConfirm native sidecar and ModelExpress image support.
Use Kubernetes 1.33 or later, or enable
SidecarContainerson an older supported cluster. Otherwise, the Redis init container can block pod startup.Confirm that
modelexpress-server:0.4.1provides/app/modelexpress-serverand accepts--port 8001,MX_METADATA_BACKEND=redis, andREDIS_URL=redis://127.0.0.1:6379. WithTRTLLM_MX_E2E_REQUIRED=1, an incompatible image fails the stage.🤖 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 `@jenkins/L0_Test.groovy` around lines 3328 - 3390, Update the test setup around serviceInitContainerConfig and serviceContainerConfig to explicitly require Kubernetes 1.33+ or enable the SidecarContainers feature for older supported clusters, preventing Redis from blocking pod startup. Verify the configured MODEL_EXPRESS_SERVER_IMAGE resolves to modelexpress-server:0.4.1 and supports /app/modelexpress-server with --port 8001, MX_METADATA_BACKEND=redis, and REDIS_URL=redis://127.0.0.1:6379; ensure TRTLLM_MX_E2E_REQUIRED=1 causes the stage to fail when these requirements are not met.tests/integration/defs/model_express/test_model_express.py (1)
372-399: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueAdd context to transfer-evidence assertions.
Include the failed marker, matched counts, and expected and observed rank sets in assertion messages so CI failures identify the failed condition.
Test coverage:
test_mx_donor_receivercovers TP1 and TP2 and is listed intests/integration/test_lists/test-db/l0_model_express.yml; verdict: sufficient.🤖 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/model_express/test_model_express.py` around lines 372 - 399, Update _assert_transfer_evidence to add assertion messages containing the rejected failure marker, matched parameter counts versus case.tp_size, and expected versus observed transferred rank sets. Keep the existing validation logic unchanged while ensuring each failure identifies the relevant condition and values.
🤖 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 `@docs/source/features/model-express.md`:
- Around line 104-109: Update the documentation around the
DGX_H100-2_GPUs-PyTorch-ModelExpress-1 command to state its gating: it runs when
a multi-GPU file changes, in post-merge pipelines, or via multi-GPU dispatch
with the ci: full pre-merge approved label. Remove or replace “recurring” so the
stage is not described as running in every pre-merge pipeline.
In `@tests/integration/defs/model_express/mx_e2e_worker.py`:
- Around line 62-67: Update the donor setup in the mx_config construction so it
bypasses ModelExpress source discovery rather than treating
server_query_timeout_s=0 as an immediate fallback; alternatively, revise the
comment to accurately describe the 5-second delay and add coverage for
zero-timeout behavior in test_mx_donor_receiver and its listed TP1/TP2
configurations.
In `@tests/integration/defs/model_express/test_model_express.py`:
- Around line 339-361: Replace the local process teardown helpers
`_signal_process_group` and `_stop_donor` with the canonical `popen()` and
`cleanup_process_tree()` helpers from `trt_test_alternative.py`. Apply this
consistently to the worker launch in `_run_worker` and the inline donor launch,
preserving the existing startup, stop-file, timeout, and return-code behavior
while ensuring all descendant process groups are cleaned up.
---
Nitpick comments:
In `@jenkins/L0_Test.groovy`:
- Around line 3328-3390: Update the test setup around serviceInitContainerConfig
and serviceContainerConfig to explicitly require Kubernetes 1.33+ or enable the
SidecarContainers feature for older supported clusters, preventing Redis from
blocking pod startup. Verify the configured MODEL_EXPRESS_SERVER_IMAGE resolves
to modelexpress-server:0.4.1 and supports /app/modelexpress-server with --port
8001, MX_METADATA_BACKEND=redis, and REDIS_URL=redis://127.0.0.1:6379; ensure
TRTLLM_MX_E2E_REQUIRED=1 causes the stage to fail when these requirements are
not met.
In `@tests/integration/defs/model_express/mx_e2e_worker.py`:
- Around line 100-105: Bound the donor polling loop in the role-handling flow by
adding the proposed --max-serve-seconds argument in _parse_args with its
default, then have the donor stop waiting when either stop_file appears or the
configured deadline expires. Use a monotonic elapsed-time check around the
existing time.sleep polling.
In `@tests/integration/defs/model_express/test_model_express.py`:
- Around line 372-399: Update _assert_transfer_evidence to add assertion
messages containing the rejected failure marker, matched parameter counts versus
case.tp_size, and expected versus observed transferred rank sets. Keep the
existing validation logic unchanged while ensuring each failure identifies the
relevant condition and values.
🪄 Autofix
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: Enterprise
Run ID: 93c35419-6654-430c-9bc6-b9c72b77f101
📒 Files selected for processing (5)
docs/source/features/model-express.mdjenkins/L0_Test.groovytests/integration/defs/model_express/mx_e2e_worker.pytests/integration/defs/model_express/test_model_express.pytests/integration/test_lists/test-db/l0_model_express.yml
|
/bot run --disable-fail-fast |
|
PR_Github #65133 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/model_express/test_model_express.py (1)
399-487: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the test to the QA list.
test_mx_donor_receiveris registered for TP=1 and TP=2 inl0_model_express.yml. No QA-list entries exist for these cases.The added test function has sufficient CI registration but insufficient manual-QA coverage. Add both cases to the appropriate QA list and run the GPU integration test with
LLM_MODELS_ROOTset.🤖 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/model_express/test_model_express.py` around lines 399 - 487, Add QA-list entries for both TP=1 and TP=2 registrations of test_mx_donor_receiver, matching the existing model-express QA-list format and appropriate cases. Then run the GPU integration test with LLM_MODELS_ROOT configured to validate both entries.Source: Path instructions
🤖 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/integration/defs/model_express/test_model_express.py`:
- Around line 381-387: Update the assertions around _MATCHED_PARAMS_PATTERN and
matched_params so complete parameter-match summaries are associated with
distinct tensor-parallel ranks. Capture each rank from the receiver logs and
assert every expected rank has exactly one valid complete match, rather than
relying on the total summary count; retain the existing positive matched/total
validation.
---
Outside diff comments:
In `@tests/integration/defs/model_express/test_model_express.py`:
- Around line 399-487: Add QA-list entries for both TP=1 and TP=2 registrations
of test_mx_donor_receiver, matching the existing model-express QA-list format
and appropriate cases. Then run the GPU integration test with LLM_MODELS_ROOT
configured to validate both entries.
🪄 Autofix
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: Enterprise
Run ID: c28a2555-e1c1-4f66-8386-2f0244e068bb
📒 Files selected for processing (3)
jenkins/L0_Test.groovytests/integration/defs/model_express/mx_e2e_worker.pytests/integration/defs/model_express/test_model_express.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/integration/defs/model_express/mx_e2e_worker.py
- jenkins/L0_Test.groovy
|
/bot run --disable-fail-fast |
|
PR_Github #65145 [ run ] triggered by Bot. Commit: |
|
PR_Github #65133 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #66342 [ run ] triggered by Bot. Commit: |
|
PR_Github #66342 [ run ] completed with state
|
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
fa7b0a1 to
8e4a7e7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #66391 [ run ] triggered by Bot. Commit: |
|
PR_Github #66391 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66560 [ run ] triggered by Bot. Commit: |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Revisit after the earlier round — all three prior asks are genuinely addressed in the current code:
- The three ModelExpress paths are in
getMultiGpuFileChanged()(jenkins/L0_MergeRequest.groovy:1097, jenkins/L0_MergeRequest.groovy:1130, jenkins/L0_MergeRequest.groovy:1160), so harness/profile/test-db changes now enable pre-merge multi-GPU stages. - The preflight parses the
MODELEXPRESS_VERSION=line withpackaging.version.Version, rejects older/prerelease/malformed versions against the 0.4.1 floor, and routessubprocess.TimeoutExpiredthrough_skip_or_fail(tests/integration/defs/model_express/test_model_express.py:213). - The two QA rows are gone from
llm_function_core.txt; coverage now lives only inl0_model_express.yml, where the stages provision the services and setTRTLLM_MX_E2E_REQUIRED=1.
I also re-traced the OnDemand exclusion after the new parallelJobsFiltered -= onDemandJobs backstop: the multi-GPU trigger, auto-trigger tags, and CBTS paths all exclude it, while --stage-list/--extra-stage still select from the full parallelJobs map — reachable exactly as the docs describe.
One description-hygiene point: the Scope section says the PR adds "one diagnostic-only loader change", but there are two loader changes and the second is behavioral. The zero-timeout early fallback (tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py:425) now returns to disk loading before prepare_post_transform_receiver can mutate the module graph, and it applies to any receiver configured with query_timeout_s == 0, not just this harness. The change looks correct and is unit-tested (test_zero_timeout_falls_back_before_receiver_preparation), and per-load state is reset at load_weights entry so the early return can't leak stale preload flags — please just update the description so the behavior delta is on record for anyone bisecting MX fallback changes later.
The two inline comments are optional cleanups, not blockers.
|
PR_Github #66560 [ run ] completed with state |
6 participants
Summary
What This Verifies
Scope
MX_TRANSFER_LOG_DIRis set, enable upstream ModelExpress INFO records so the harness can verify per-rank transfer evidence.query_timeout_s == 0, fall back to disk when no source is registered, before post-transform receiver preparation can mutate the module graph.[mx]extra.Extending Model Coverage
MxE2ECaserow with the checkpoint path, canonical-cache prefix, and TP size.tests/integration/test_lists/test-db/l0_model_express.yml.Follow-ups
serviceInitContainerConfigdeclaration and interpolation fromjenkins/L0_Test.groovy; it currently emits only a blank line and does not affect generated pod specifications (review thread).TRTLLM_MX_E2E_TIMEOUT_Sindocs/source/features/model-express.md; it defaults to 1200 seconds and bounds each worker and the donor-readiness wait (review thread).Related PRs
Validation
8e4a7e7bf8; all required stages succeeded.736f77f056: 2 passed, 0 failed, 0 skipped.DGX_H100-2_GPUs-PyTorch-ModelExpress-1passedmodel_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1].DGX_H100-4_GPUs-PyTorch-ModelExpress-OnDemand-1passedmodel_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2].pytest -q --confcutdir=tests/unittest/tools tests/unittest/tools/test_test_to_stage_mapping.py- 9 passed, 3 skipped.