feat: Support custom repo_dir for SLURM script #6546
Conversation
📝 WalkthroughWalkthroughAdds a repo path argument to SLURM submit/launcher scripts, optionally installs TRT-LLM from that repo, adjusts log directory handling, updates benchmark runner to accept and iterate over a concurrency list with per-concurrency logs, and adds a shutdown routine. Argument propagation updated across submit scripts and benchmark invocation paths. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant submit.sh
participant SLURM as disaggr_torch.slurm
participant Container
participant TRTRepo as TRT-LLM Repo
participant Benchmark as run_benchmark.sh
User->>submit.sh: Define args (incl. repo_dir)
submit.sh->>SLURM: sbatch with args (..., trtllm_repo)
SLURM->>Container: Start job, set up log dirs
alt trtllm_repo provided
SLURM->>TRTRepo: pip install -e . (log to install.log)
TRTRepo-->>SLURM: Installed
SLURM->>SLURM: Derive TRT_LLM_GIT_COMMIT
end
SLURM->>Benchmark: Invoke with log_path and concurrency_list
loop for each concurrency
Benchmark->>Container: Run benchmark for concurrency N
Benchmark->>Container: Collect logs into concurrency_N/
end
Benchmark->>Container: Shutdown server/worker/trtllm-serve
Container-->>SLURM: Return status
SLURM-->>User: Logs in full_logdir
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~18 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
examples/disaggregated/slurm/run_benchmark.sh (1)
93-93: Quote$HOMEwhen exportingPATH.If
HOMEever contains spaces (rare, but possible on custom systems) the un-quoted expansion will break the export.-export PATH=${HOME}/.local/bin:${PATH} +export PATH="${HOME}/.local/bin:${PATH}"examples/disaggregated/slurm/submit.sh (2)
10-10: Provide a sane default or validation forrepo_dir.The placeholder
<repo_dir>is introduced but never validated later in the script. If the user forgets to replace it, the literal string<repo_dir>will be forwarded to the SLURM job and evaluated as a real path, leading to apushdfailure downstream.Add a quick guard:
if [[ "$repo_dir" == "<repo_dir>" ]]; then echo "ERROR: Please set repo_dir or leave it empty"; exit 1 fi
35-35: Quote array element to retain spaces in path.
$repo_dirshould be wrapped in double quotes when added to theargsarray to avoid word-splitting for paths containing spaces.- $repo_dir + "$repo_dir"examples/wide_ep/slurm_scripts/submit.sh (2)
12-12: Same validation concern as disaggregated submit script.Add a guard or a default for
repo_dirto prevent accidental propagation of the literal placeholder to SLURM.(See previous comment for a minimal snippet.)
43-44: Quote$repo_dirwhen pushing into the argument list.- $repo_dir + "$repo_dir"examples/disaggregated/slurm/start_worker.sh (1)
22-22: Quote$HOMEwhen exportingPATH.-export PATH=${HOME}/.local/bin:${PATH} +export PATH="${HOME}/.local/bin:${PATH}"examples/disaggregated/slurm/start_server.sh (1)
26-26: Quote$HOMEwhen exportingPATH.-export PATH=${HOME}/.local/bin:${PATH} +export PATH="${HOME}/.local/bin:${PATH}"
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
examples/disaggregated/slurm/disaggr_torch.slurm(2 hunks)examples/disaggregated/slurm/run_benchmark.sh(1 hunks)examples/disaggregated/slurm/start_server.sh(1 hunks)examples/disaggregated/slurm/start_worker.sh(1 hunks)examples/disaggregated/slurm/submit.sh(2 hunks)examples/wide_ep/slurm_scripts/submit.sh(2 hunks)
🧰 Additional context used
🧠 Learnings (1)
📓 Common learnings
Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
🪛 Shellcheck (0.10.0)
examples/disaggregated/slurm/start_worker.sh
[warning] 12-12: Use 'pushd ... || exit' or 'pushd ... || return' in case pushd fails.
(SC2164)
[warning] 20-20: Use 'popd ... || exit' or 'popd ... || return' in case popd fails.
(SC2164)
examples/disaggregated/slurm/start_server.sh
[warning] 15-15: Use 'pushd ... || exit' or 'pushd ... || return' in case pushd fails.
(SC2164)
[warning] 24-24: Use 'popd ... || exit' or 'popd ... || return' in case popd fails.
(SC2164)
🔇 Additional comments (4)
examples/disaggregated/slurm/disaggr_torch.slurm (4)
41-41: LGTM: Parameter definition follows established pattern.The
repo_dirparameter is correctly defined as the 23rd positional argument, following the existing pattern of parameter assignments in the script.
112-112: LGTM: Consistent parameter passing to start_worker.sh.The
repo_dirparameter is correctly appended as the 6th argument tostart_worker.sh, maintaining consistency with the other script invocations.
119-119: LGTM: Consistent parameter passing to start_server.sh.The
repo_dirparameter is correctly appended as the 3rd argument tostart_server.sh, following the established pattern.
125-125: LGTM: Consistent parameter passing to run_benchmark.sh.The
repo_dirparameter is correctly appended as the 9th argument torun_benchmark.sh, maintaining consistency across all script invocations.
xxi-nv
force-pushed
the
user/kaiyu/support_custom_repo
branch
from
4c91cb8 to
cfca53d
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🔭 Outside diff range comments (3)
examples/disaggregated/slurm/disaggr_torch.slurm (1)
76-82: Fix inconsistent directory tokens: dep/tep vs tp for gen_tp_sizeThe logdir suffix uses dep in one branch and tep in another. This looks like a typo and makes log paths inconsistent. Use a single token, e.g., tp.
Apply:
- full_logdir=${logdir}/ctx${num_ctx_servers}_gen${num_gen_servers}_dep${gen_tp_size}_batch${gen_batch_size}_eplb${eplb_num_slots}_mtp${mtp_size} + full_logdir=${logdir}/ctx${num_ctx_servers}_gen${num_gen_servers}_tp${gen_tp_size}_batch${gen_batch_size}_eplb${eplb_num_slots}_mtp${mtp_size} @@ - full_logdir=${logdir}/ctx${num_ctx_servers}_gen${num_gen_servers}_tep${gen_tp_size}_batch${gen_batch_size}_eplb${eplb_num_slots}_mtp${mtp_size} + full_logdir=${logdir}/ctx${num_ctx_servers}_gen${num_gen_servers}_tp${gen_tp_size}_batch${gen_batch_size}_eplb${eplb_num_slots}_mtp${mtp_size}Also applies to: 91-92
examples/disaggregated/slurm/run_benchmark.sh (2)
29-29: Avoid set -u crash if TRT_LLM_GIT_COMMIT is unsetWith set -u, referencing an unset env var will abort the script. Default to “unknown” when not exported.
-echo "TRT_LLM_GIT_COMMIT: ${TRT_LLM_GIT_COMMIT}" +echo "TRT_LLM_GIT_COMMIT: ${TRT_LLM_GIT_COMMIT:-unknown}"
66-81: Health check only verifies connectivity, not HTTP 200Current loop exits on any HTTP response. Check for 200 explicitly and keep the timeout behavior.
-while ! curl -s -o /dev/null -w "%{http_code}" http://${hostname}:${port}/health; do +while :; do + status="$(curl -s -o /dev/null -w "%{http_code}" "http://${hostname}:${port}/health" || echo 000)" + if [ "${status}" = "200" ]; then + break + fi current_time=$(date +%s) elapsed=$((current_time - start_time)) if [ $elapsed -ge $timeout ]; then - echo "Error: Server is not healthy after ${timeout} seconds" + echo "Error: Server is not healthy (last status: ${status}) after ${timeout} seconds" exit 1 fi if [ $((elapsed % 30)) -eq 0 ]; then - echo "Waiting for server to be healthy... (${elapsed}s elapsed)" + echo "Waiting for server to be healthy... (${elapsed}s elapsed, last status: ${status})" fi sleep 10 done
🧹 Nitpick comments (2)
examples/disaggregated/slurm/run_benchmark.sh (2)
59-64: Quote paths and patterns in wget/grep to avoid word splitting and globbingMinor but defensive. This keeps logs gathering stable.
-wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json -O ${shared_gpt_path} +wget "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" -O "${shared_gpt_path}" @@ - grep -a "'num_ctx_requests': 0, 'num_ctx_tokens': 0" ${worker_log_path} > ${output_folder}/gen_only.txt || true - grep -a "'num_generation_tokens': 0" ${worker_log_path} > ${output_folder}/ctx_only.txt || true + grep -a "'num_ctx_requests': 0, 'num_ctx_tokens': 0" "${worker_log_path}" > "${output_folder}/gen_only.txt" || true + grep -a "'num_generation_tokens': 0" "${worker_log_path}" > "${output_folder}/ctx_only.txt" || trueAlso applies to: 84-89
24-27: Optional: make SLURM_PROCID check robust outside SlurmIf someone runs this locally (no Slurm), set -u will error. Default to “0”.
-if [[ ${SLURM_PROCID} != "0" ]]; then +if [[ ${SLURM_PROCID:-0} != "0" ]]; then
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
examples/disaggregated/slurm/disaggr_torch.slurm(4 hunks)examples/disaggregated/slurm/run_benchmark.sh(2 hunks)examples/wide_ep/slurm_scripts/submit.sh(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/wide_ep/slurm_scripts/submit.sh
🧰 Additional context used
🪛 Shellcheck (0.10.0)
examples/disaggregated/slurm/run_benchmark.sh
[warning] 111-111: Quote this to prevent word splitting.
(SC2046)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (4)
examples/disaggregated/slurm/disaggr_torch.slurm (2)
75-82: LGTM on logdir semantics and commit echoing
- Removing the trailing slash from logdir and passing the directory (not dir/) downstream is clean and reduces path-join pitfalls.
- Printing TRT_LLM_GIT_COMMIT improves traceability.
Also applies to: 93-99
41-42: Validate positional args count and confirm submit scripts pass the newrepo_dirargumentTo prevent silent mis‐alignment when adding the 23rd parameter, add an argument‐count check in
disaggr_torch.slurmand verify that both submit scripts include that extra argument before invoking it.• In
examples/disaggregated/slurm/disaggr_torch.slurm, after the SBATCH header (or before the firstgetopts/shift), insert:# Expect 23 positional arguments (including the new repo_dir) if [ "$#" -lt 23 ]; then echo "Error: disaggr_torch.slurm expects 23 arguments, got: $#" exit 1 fi• Ensure the submit‐helper scripts pass the new
repo_diras the 23rd element in theirargsarrays. For example:grep -R "args\s*=" -n \ examples/disaggregated/slurm/submit.sh \ examples/wide_ep/slurm_scripts/submit.shConfirm each
args=( … )definition includes the<repo_dir>entry at the end.examples/disaggregated/slurm/run_benchmark.sh (2)
19-21: Param rename to concurrency_list looks goodInput contract is updated and reflected in usage. Ensure all call sites pass a space- or comma-separated list.
I can scan for invocations to confirm they pass lists instead of a single integer if you’d like.
118-123: Graceful shutdown looks goodClear and comprehensive cleanup with a verification step for lingering processes. Keep as-is.
kaiyux
force-pushed
the
user/kaiyu/support_custom_repo
branch
from
cfca53d to
a720c47
Compare
|
/bot skip --comment "the scripts are not protected by CI pipeline" |
Signed-off-by: Kaiyu Xie <[email protected]>
Signed-off-by: Kaiyu Xie <[email protected]>
Signed-off-by: Kaiyu Xie <[email protected]>
Signed-off-by: xxi <[email protected]> modified: examples/disaggregated/slurm/disaggr_torch.slurm modified: examples/disaggregated/slurm/run_benchmark.sh modified: examples/disaggregated/slurm/start_server.sh modified: examples/disaggregated/slurm/start_worker.sh modified: examples/wide_ep/slurm_scripts/submit.sh
kaiyux
force-pushed
the
user/kaiyu/support_custom_repo
branch
from
a720c47 to
d9f2b9c
Compare
|
/bot skip --comment "the scripts are not protected by CI pipeline" |
|
PR_Github #15043 [ skip ] triggered by Bot |
|
PR_Github #15043 [ skip ] completed with state |
Summary by CodeRabbit
New Features
Refactor
Description
Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.