[dingying-ai]

Summary

Supersedes #2782 with a simpler, zero-Python-dependency teardown mechanism for multi-slice GKE runs:

• PID-1 Bash Teardown (jobset.cpu.yaml / yaml_generator.py): When --teardown_gang_on_exit is passed alongside --gang_id, the orchestrator container's PID-1 bash wrapper runs kubectl delete jobsets.jobset.x-k8s.io -n <namespace> -l scheduling.x-k8s.io/gang-id=<gang-id> --field-selector=metadata.name!=<orch-name> --ignore-not-found=true --wait=false immediately after wait $PID, before exiting with $EXIT_CODE. Because it runs in PID 1 after wait $PID, it executes on both normal completion (EXIT_CODE=0) and unhandled Python exceptions/crashes (EXIT_CODE!=0), while leaving the orchestrator JobSet itself intact so Kubernetes records Completed vs Failed accurately.
• Increased Grace Period: Bumps terminationGracePeriodSeconds in jobset.cpu.yaml from 30 to 300 so _sigterm has sufficient time to let Python flush logs and run the kubectl delete call before SIGKILL.
• Standalone kubectl in Dockerfile: Installs the official standalone /usr/local/bin/kubectl binary in Dockerfile so orchestrator containers have kubectl in PATH out of the box.
• Automatic Enablement in K8s Launchers: Updates deepswe_dist/k8s_launcher.sh and math_gsm8k_dist/k8s_launcher.sh to pass ${GANG_ID:+--teardown_gang_on_exit} in start_orchestrator() whenever GANG_ID is set.
• In-Process Shutdown Timeout Guard (orchestrator.py): Wraps best-effort worker cluster.close() in a daemon thread with a 30s timeout so hung gRPC channels to dead/preempted workers cannot block orchestrator exit.

Test

1. Unit Tests:
   • pytest tests/experimental/distributed/deployment/yaml_generator_test.py tests/experimental/orchestrator/orchestrator_test.py (32 passed)
2. End-to-End GKE Verification (Qwen/Qwen3-0.6B, two 2x2x1 TPU v5p slices = 8 chips):
   • Normal completion (dy-pr2825-v2-ok, MAX_STEPS=2): Completed 2 GRPO steps; PID-1 wrapper deleted sibling *-roll and *-train JobSets in <1s and exited with EXIT_CODE=0 (*-orch JobSet marked Completed).
   • Unhandled error exit (dy-pr2825-v2-err): Triggered WeightSyncError on step 1; PID-1 wrapper caught non-zero Python exit, deleted sibling *-roll and *-train JobSets in 0.21s, and exited with EXIT_CODE=1 (*-orch JobSet marked Failed, zero orphaned worker pods/slices).

Orchestrator PID-1 Teardown Logs (Normal & Error Exit)

Scenario A (EXIT_CODE=0):

2026-10-12 23:19:39,077 - [Orchestrator] === GRPO Training Finished Successfully ===
2026-10-12 23:19:39,110 - [Orchestrator] close.started
2026-10-12 23:19:39,111 - [Orchestrator] close.complete
2026-10-12 23:19:39,112 - [Orchestrator] Clearing JAX backend caches.
Tearing down sibling JobSets for gang 'dy-pr2825-v2-ok' in namespace 'trellis'...
jobset.jobset.x-k8s.io "dy-pr2825-v2-ok-roll" deleted from trellis namespace
jobset.jobset.x-k8s.io "dy-pr2825-v2-ok-train" deleted from trellis namespace
Program End: Mon Oct 12 23:19:45 UTC 2026
EXIT_CODE=0

Scenario B (EXIT_CODE=1):

tunix.experimental.weight_sync.weight_sync_coordinator.WeightSyncError: round 0 (req_id wsync-v0-r0, uuid 1): bind/metadata/source-prepare failed before any destination was quiesced; no rollback needed; final state preparing
2026-10-12 23:20:55,349 - [Orchestrator] close.started
2026-10-12 23:20:55,349 - [Orchestrator] close.complete
2026-10-12 23:20:55,350 - [Orchestrator] Clearing JAX backend caches.
Tearing down sibling JobSets for gang 'dy-pr2825-v2-err' in namespace 'trellis'...
jobset.jobset.x-k8s.io "dy-pr2825-v2-err-roll" deleted from trellis namespace
jobset.jobset.x-k8s.io "dy-pr2825-v2-err-train" deleted from trellis namespace
Program End: Mon Oct 12 23:20:58 UTC 2026
EXIT_CODE=1

[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!