Ipod/custom srun bash

[podkidyshev]

Summary

• Allowing running custom bash commands at the start of llm/sglang services

Test Plan

• Automated CI
• Manual runs by ver. team

Additional Notes

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: d871f542-90c2-4c81-95ba-247a2eb01abe

📥 Commits

Reviewing files that changed from the base of the PR and between 4bdc465 and e760249.

📒 Files selected for processing (10)

• src/cloudai/workloads/common/llm_serving.py
• src/cloudai/workloads/sglang/sglang.py
• src/cloudai/workloads/vllm/slurm_command_gen_strategy.py
• src/cloudai/workloads/vllm/vllm.py
• tests/ref_data/sglang-disagg-2nodes.sbatch
• tests/ref_data/sglang-disagg.sbatch
• tests/ref_data/vllm-disagg-2nodes.sbatch
• tests/ref_data/vllm-disagg.sbatch
• tests/workloads/sglang/test_command_gen_strategy_slurm.py
• tests/workloads/vllm/test_command_gen_strategy_slurm.py

---

📝 Walkthrough

Walkthrough

The PR adds optional custom bash injection to LLM serving workloads. It introduces a CustomBash type alias, regex-based validation, and wrapping logic that injects custom bash setup before commands using bash -c execution. Integration is added to SGLang and vLLM test definitions, with corresponding environment variable exports to support disaggregated mode infrastructure.

Changes

Custom Bash Injection for LLM Serving

Layer / File(s)	Summary
Custom bash type and validation src/cloudai/workloads/common/llm_serving.py	Introduces CustomBash type alias (str | dict[str, str]) and validate_custom_bash_patterns() function that compiles regex keys from dict forms and raises ValueError on invalid patterns.
Command wrapping implementation src/cloudai/workloads/common/llm_serving.py	Adds _custom_bash_for_command() method to select matching bash snippets from string or regex-keyed dict custom_bash, and updates _with_custom_bash() to wrap commands with bash -c "<custom>; exec <command>" when a match is found.
Aggregated and disaggregated command generation integration src/cloudai/workloads/common/llm_serving.py	Updates aggregated mode to wrap serve and benchmark commands through _with_custom_bash(). Updates disaggregated mode to wrap prefill, decode, router, and benchmark commands consistently.
SGLang workload custom_bash support src/cloudai/workloads/sglang/sglang.py	Adds imports for validate_custom_bash_patterns, introduces custom_bash: CustomBash | None field to SglangTestDefinition, and adds Pydantic @field_validator that normalizes the field via validation.
vLLM workload custom_bash support src/cloudai/workloads/vllm/vllm.py	Adds imports for validate_custom_bash_patterns, introduces custom_bash: CustomBash | None field to VllmTestDefinition, and adds Pydantic @field_validator that normalizes the field via validation.
Environment variable exports src/cloudai/workloads/common/llm_serving.py, src/cloudai/workloads/vllm/slurm_command_gen_strategy.py, tests/ref_data/sglang-disagg.sbatch, tests/ref_data/sglang-disagg-2nodes.sbatch, tests/ref_data/vllm-disagg.sbatch, tests/ref_data/vllm-disagg-2nodes.sbatch	Exports disaggregated node-role variables (PREFILL_NODE, DECODE_NODE) and networking variables (PORT_OFFSET, PREFILL_NIXL_PORT, DECODE_NIXL_PORT) so they propagate as environment variables to child processes in disaggregated setups.
Test coverage for custom_bash behavior tests/workloads/sglang/test_command_gen_strategy_slurm.py, tests/workloads/vllm/test_command_gen_strategy_slurm.py	Adds SGLang tests for aggregated and disaggregated custom_bash wrapping with string and regex-to-string mappings. Adds vLLM tests for both modes with regex-targeted rules. Updates vLLM test expected output to reflect environment variable exports.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 Custom bash now flows through SGLang and vLLM,
Like carrots injected before each command,
Regex patterns match what we need,
Environment exports let subprocesses feed,
A hop, skip, and exec to the finish line! 🎯

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects the main feature being added—custom bash command injection into srun execution for LLM/SGLang services.
Description check	✅ Passed	The description clearly relates to the changeset by summarizing the feature (custom bash at service startup) and listing the test plan.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch ipod/custom-srun-bash

---

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