[https://nvbugs/6190759][fix] set env on some gb300 cluster

[chuangz0]

Summary by CodeRabbit

• Tests
  • Enhanced test infrastructure to properly configure GPU-specific transport settings across disaggregated worker and server processes.
  • Updated test waivers to reflect improved test coverage for disaggregated serving scenarios.
  • Improved environment variable handling for multi-node test scenarios on specialized hardware.

Description

Test Coverage

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.

[chuangz0]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR updates disaggregated integration and unit tests to compute GPU SM-version-specific UCX transport settings and propagate them to worker and server processes. A new helper function and pytest fixture in the integration test suite derive UCX configuration based on GPU model, which is then passed to all process launchers. Similar changes apply to unit tests for GB300/Aarch64 systems, and test waivers are updated to reflect the changes.

Changes

UCX Transport Configuration for Disaggregated Tests

Layer / File(s)	Summary
UCX Configuration Infrastructure Setup tests/integration/defs/disaggregated/test_auto_scaling.py	Import get_sm_version from defs.conftest. Add get_ucx_tls() helper function that returns SM-version-specific UCX transport strings. Create worker_env pytest fixture that builds a per-worker environment dictionary by copying os.environ and setting UCX_TLS from the helper.
Integration Test Fixture Application tests/integration/defs/disaggregated/test_auto_scaling.py	Four integration test functions (test_service_discovery, test_minimal_instances, test_worker_restart, test_disagg_server_restart) now accept worker_env fixture parameter. All worker and server process launcher calls within these tests are updated to pass env=worker_env, including initial spawns and subsequent restarts.
Unit Test UCX Configuration tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py	Import platform and get_sm_version. Extend env() helper function to conditionally set UCX_TLS environment variable for systems matching get_sm_version() == 103 and platform.machine().lower() == "aarch64".
Test Waiver List Updates tests/integration/test_lists/waives.txt	Remove existing disaggregated test waivers including test_auto_scaling cases and test_disaggregated_conditional. Add new waivers for disaggregated/test_disaggregated.py tests: test_disaggregated_cache_aware_balance, test_disaggregated_cancel_large_context_requests, and test_disaggregated_chat_completion_tool_calls.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description consists entirely of the template with no substantive content filled in—the author did not complete the Description, Test Coverage, or PR Checklist sections beyond marking the final checkbox.	Complete the Description section explaining the issue and solution, the Test Coverage section listing relevant tests, and provide details about how the changes address the GB300 cluster environment setup issue.
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the fix associated with NVBugs ID 6190759 and concisely summarizes the main change: setting environment variables on GB300 cluster systems.
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]

Actionable comments posted: 1

🧹 Nitpick comments (3)

tests/integration/defs/disaggregated/test_auto_scaling.py (2)

19-39: ⚡ Quick win

Add explicit return type annotations on the new helper and fixture.

Please annotate both functions’ return types to match repo Python typing guidelines.

Proposed patch

-def get_ucx_tls():
+def get_ucx_tls() -> str:
@@
 `@pytest.fixture`
-def worker_env():
+def worker_env() -> dict[str, str]:

As per coding guidelines "Static type checking via mypy is opt-in by submodule and highly recommended; always annotate Python functions with return types (use None if no return); make return type explicit".

🤖 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/disaggregated/test_auto_scaling.py` around lines 19 -
39, The helper get_ucx_tls and the fixture worker_env lack explicit return type
annotations; add return types per repo guidelines: annotate get_ucx_tls() to
return str and worker_env() to return typing.Dict[str, str] (or Mapping[str,
str]) and add the necessary typing import if missing; update the function
signatures for get_ucx_tls and worker_env accordingly so static type checkers
like mypy can pick them up.

---

55-57: QA test-list updates look unnecessary for this PR shape.

These edits adjust process env wiring inside existing tests and do not introduce new pytest node-ids, so updates under tests/integration/test_lists/qa/ are likely optional.

As per coding guidelines "If the change adds or materially alters an integration test under tests/integration/defs/ ... call out whether an entry is needed under tests/integration/test_lists/qa/."

🤖 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/disaggregated/test_auto_scaling.py` around lines 55 -
57, The QA test-list changes are likely unnecessary because only environment
wiring inside existing integration tests (e.g., test_service_discovery) was
modified without adding or renaming tests; either revert the edits under
tests/integration/test_lists/qa/ so the QA lists remain unchanged, or if you did
in fact add/rename a test node in defs/disaggregated (e.g.,
test_service_discovery) then add the corresponding new pytest node-id to the
appropriate file in tests/integration/test_lists/qa/; make the change
consistent: no test node-id changes -> revert QA list edits, otherwise add the
exact node-id for the new/renamed test.

tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py (1)

72-84: QA list update is unnecessary for this unit-test-only scope.

This change is limited to unittest environment setup behavior, so updating tests/integration/test_lists/qa/* is unnecessary for this file.

As per coding guidelines: “If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional.”

🤖 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/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py` around lines
72 - 84, This change only affects unit-test environment setup (the code that
builds the env dict and conditionally sets UCX_TLS based on get_sm_version() and
platform.machine()), so revert or omit any updates to the QA lists and
explicitly state in the PR/commit message that QA list updates are unnecessary
for this unittest-only change; ensure the PR description includes a short note
like "QA list update not required — change only touches unit-test setup
(get_sm_version()/UCX_TLS)".

🤖 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/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py`:
- Line 1: Add the required NVIDIA Apache-2.0 copyright/license header at the
very top of the file so it appears before the first import (the line "import
os"); include the full NVIDIA header block with the correct year(s) and
Apache-2.0 text as per project guidelines, and update the year if this is a
modified file rather than a new one.

---

Nitpick comments:
In `@tests/integration/defs/disaggregated/test_auto_scaling.py`:
- Around line 19-39: The helper get_ucx_tls and the fixture worker_env lack
explicit return type annotations; add return types per repo guidelines: annotate
get_ucx_tls() to return str and worker_env() to return typing.Dict[str, str] (or
Mapping[str, str]) and add the necessary typing import if missing; update the
function signatures for get_ucx_tls and worker_env accordingly so static type
checkers like mypy can pick them up.
- Around line 55-57: The QA test-list changes are likely unnecessary because
only environment wiring inside existing integration tests (e.g.,
test_service_discovery) was modified without adding or renaming tests; either
revert the edits under tests/integration/test_lists/qa/ so the QA lists remain
unchanged, or if you did in fact add/rename a test node in defs/disaggregated
(e.g., test_service_discovery) then add the corresponding new pytest node-id to
the appropriate file in tests/integration/test_lists/qa/; make the change
consistent: no test node-id changes -> revert QA list edits, otherwise add the
exact node-id for the new/renamed test.

In `@tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py`:
- Around line 72-84: This change only affects unit-test environment setup (the
code that builds the env dict and conditionally sets UCX_TLS based on
get_sm_version() and platform.machine()), so revert or omit any updates to the
QA lists and explicitly state in the PR/commit message that QA list updates are
unnecessary for this unittest-only change; ensure the PR description includes a
short note like "QA list update not required — change only touches unit-test
setup (get_sm_version()/UCX_TLS)".

🪄 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: Enterprise

Run ID: f31632da-f426-4535-a93a-9d18032e97e9

📥 Commits

Reviewing files that changed from the base of the PR and between 8791347 and d3f7443.

📒 Files selected for processing (3)

• tests/integration/defs/disaggregated/test_auto_scaling.py
• tests/integration/test_lists/waives.txt
• tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

[tensorrt-cicd]

PR_Github #49927 [ run ] triggered by Bot. Commit: d3f7443 Link to invocation

[tensorrt-cicd]

PR_Github #49927 [ run ] completed with state SUCCESS. Commit: d3f7443
/LLM/main/L0_MergeRequest_PR pipeline #39504 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chuangz0]

/bot skip --comment " all tests have passed"

[tensorrt-cicd]

PR_Github #50137 [ skip ] triggered by Bot. Commit: b8b02be Link to invocation

[tensorrt-cicd]

PR_Github #50137 [ skip ] completed with state SUCCESS. Commit: b8b02be
Skipping testing for commit b8b02be

Link to invocation