[s-noghabi]

What

Follow-up to #2845 (separate first-batch dispatch stagger). Tests only, no behavior change:

• cpu_distributed_rl_integration_test.py: two new cases, test_initial_dispatch_stagger_end_to_end_*, check the start-up stagger on the real dispatch loop.
• mlperf_submission_launcher_test.py: two new tests parse the rendered orchestrator command with run_deepswe_dist._parse_args, so the stagger values are checked from launcher env to parsed args.

End-to-end dispatch on the real loop (cpu_distributed_rl_integration_test.py::test_initial_dispatch_stagger_end_to_end_*):

• Runs ClusterOrchestrator + StandardRLProgram against real trainer and rollout worker processes, with 3 start-up batches and max_staleness=2. It records when each batch first clears the dispatch gate.
• The first training step is held past the last stagger deadline, as on a real run, where batch 0's first stage outlasts the stagger.
• first=2.0, stagger=0.5 → batches at 0 / 2.0 / 2.5 s. first=None, stagger=0.5 → 0 / 0.5 / 1.0 s. Each must land within 0.5 s after its due time, and training must finish all 3 steps.
• Mutation check: with the old k × stagger offset, the first case fails (batch 1 dispatched early, 0.70 s instead of ≥ 2.0 s).

Flags, render to parser (mlperf_submission_launcher_test.py):

• Renders the submission recipe (DRY_RUN=true), shell-splits the orchestrator command and parses everything after the entry point with run_deepswe_dist._parse_args.
• Default: initial_dispatch_stagger_s=75.0, initial_dispatch_first_stagger_s=175.0, max_staleness=2, rollout_priority_scheduling=True, sandbox_wait_initial=False.
• Overrides: 60 / 200 arrive as such. INITIAL_DISPATCH_FIRST_STAGGER_S=0 drops the flag, so it parses to None (= stagger).
• run_deepswe_dist.py passes both values straight into StandardRLProgram(initial_dispatch_stagger_s=..., initial_dispatch_first_stagger_s=...). The launcher's orchestrator command is the only path that sets them.

Note: in mlperf_gbs1024_submission.sh an empty INITIAL_DISPATCH_FIRST_STAGGER_S still gives 175 (${VAR:-175}). Set it to 0 to fall back to the general stagger.

Testing

On the TPU VM (JAX_PLATFORMS=cpu), on top of atwigg/mlperf @ 763367487:

• launcher tests + the two new integration cases: 17 pass.
• With the PR feat(orchestrator): separate first-batch dispatch stagger; mlperf_gbs1024_submission holds batch 1 175 s #2845 test set (rl_program_test.py, grpo_recipe_wiring_test.py, mlperf_submission_launcher_test.py, cpu_distributed_rl_integration_test.py), before rebasing onto Simplify orchestrator exit cleanup via PID-1 wrapper gang JobSet teardown #2825: 234 pass, 11 fail. All 11 fail the same way without this change:
  • PaddedBatchAssemblerTest::* (5);
  • test_m2, test_m4_* and test_m5 (6): get_weights raises 'NoneType' object has no attribute 'model' in the trainer worker.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!