fix(deepswe): forward TRAINER_WARMUP_COMPILE and RETURN_ROUTED_EXPERTS in k8s_launcher.sh

[andytwigg]

Summary

Restores forwarding of --trainer_warmup_compile and --trainer_warmup_routed_experts in start_orchestrator() in tunix/experimental/examples/deepswe_dist/k8s_launcher.sh, which was accidentally dropped during conflict resolution when rebasing #2812 onto atwigg/mlperf.

• When TRAINER_WARMUP_COMPILE=true, passes --trainer_warmup_compile to run_deepswe_dist.py.
• When both TRAINER_WARMUP_COMPILE=true and RETURN_ROUTED_EXPERTS=true, also passes --trainer_warmup_routed_experts.
• When TRAINER_WARMUP_COMPILE=false (default), rendered orchestrator command is unchanged.
• Adds a wiring test in tests/experimental/orchestrator/grpo_recipe_wiring_test.py.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request restores critical configuration flag forwarding in the deepswe orchestrator launcher. By ensuring these flags are correctly passed to the training process, the changes enable proper warmup compilation and routed expert handling, improving the reliability and performance of the distributed training setup.

Highlights

• Restoration of Configuration Forwarding: Re-implemented the forwarding of --trainer_warmup_compile and --trainer_warmup_routed_experts flags in the k8s_launcher.sh script, which were inadvertently removed during a previous rebase.
• Enhanced Test Coverage: Added a new test case in grpo_recipe_wiring_test.py to verify that the orchestrator correctly wires these warmup flags and that the argument parser handles them as expected.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds trainer warmup compile flags (--trainer_warmup_compile and --trainer_warmup_routed_experts) to the Kubernetes launcher script and includes a corresponding unit test to verify their wiring. The feedback recommends refactoring the single-line concatenation of these flags in k8s_launcher.sh into separate, independent lines to improve readability and maintainability.

[gemini-code-assist]

The single-line concatenation of --initial_dispatch_stagger_s, --trainer_warmup_compile, and --trainer_warmup_routed_experts with nested command substitutions and manual leading spaces is extremely long and difficult to read or maintain. Splitting these into separate lines with clear, independent conditions significantly improves readability and aligns with the repository's core philosophy of maintainability and readability.

Suggested change

	$([[ "${INITIAL_DISPATCH_STAGGER_S}" != "0" && "${INITIAL_DISPATCH_STAGGER_S}" != "0.0" && -n "${INITIAL_DISPATCH_STAGGER_S}" ]] && echo --initial_dispatch_stagger_s=${INITIAL_DISPATCH_STAGGER_S})$([[ "${TRAINER_WARMUP_COMPILE}" == "true" || "${TRAINER_WARMUP_COMPILE}" == "True" || "${TRAINER_WARMUP_COMPILE}" == "1" ]] && echo " --trainer_warmup_compile$([[ "${RETURN_ROUTED_EXPERTS}" == "true" || "${RETURN_ROUTED_EXPERTS}" == "True" || "${RETURN_ROUTED_EXPERTS}" == "1" ]] && echo " --trainer_warmup_routed_experts")") \
	$([[ "${INITIAL_DISPATCH_STAGGER_S}" != "0" && "${INITIAL_DISPATCH_STAGGER_S}" != "0.0" && -n "${INITIAL_DISPATCH_STAGGER_S}" ]] && echo --initial_dispatch_stagger_s=${INITIAL_DISPATCH_STAGGER_S}) \
	$([[ "${TRAINER_WARMUP_COMPILE}" == "true" || "${TRAINER_WARMUP_COMPILE}" == "True" || "${TRAINER_WARMUP_COMPILE}" == "1" ]] && echo --trainer_warmup_compile) \
	$([[ ( "${TRAINER_WARMUP_COMPILE}" == "true" || "${TRAINER_WARMUP_COMPILE}" == "True" || "${TRAINER_WARMUP_COMPILE}" == "1" ) && ( "${RETURN_ROUTED_EXPERTS}" == "true" || "${RETURN_ROUTED_EXPERTS}" == "True" || "${RETURN_ROUTED_EXPERTS}" == "1" ) ]] && echo --trainer_warmup_routed_experts) \

References

1. Code should be easy to understand, modify, and extend for all maintainers and users. ^(link)

[ultrons]

LGTM. The restored line is byte-identical to the one on the branch our TPU runs validated with (ultrons:sivaibhav/trainer-warmup-compile at d5c6b488, before the rebase), and the test pins both flags. Cause: when I rebased #2812 onto the #2815 merge the orchestrator-args line conflicted in every commit that touched it; my resolution kept the first (stagger-only) version of the line and dropped the later version that appended the warm-up flags. Sorry for the breakage. Note for anyone launching from atwigg/mlperf between #2812 and this fix: TRAINER_WARMUP_COMPILE=true was exported but not passed, so those runs compiled the trainer on the first live micro-batch, as before #2812.