Skip to content

fix(policy): condition shutdown ordering on fast reap - #726

Draft
maximilianigl wants to merge 1 commit into
nvidia-cosmos:mainfrom
maximilianigl:fix/conditional-policy-shutdown
Draft

fix(policy): condition shutdown ordering on fast reap#726
maximilianigl wants to merge 1 commit into
nvidia-cosmos:mainfrom
maximilianigl:fix/conditional-policy-shutdown

Conversation

@maximilianigl

Copy link
Copy Markdown
Contributor

Summary

  • unregister RL policy workers before payload-packer/NCCL teardown when COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS is disabled, allowing the controller to broadcast rollout STOP and break the shutdown cycle
  • preserve the teardown-before-unregister fast-reap guarantee from fix(policy): abort NCCL + tear down before unregister on RL policy shutdown #709 when the setting is enabled
  • keep unregister and distributed teardown idempotent
  • extend the existing shutdown-order tests for both modes and repeated shutdown

Why

With fast reap disabled, payload-packer teardown can wait for rollout NCCL peers before the final policy unregister lets the controller send rollout STOP. The policy waits for rollout while rollout waits for policy departure.

Moving unregister first unconditionally would regress #709: with fast reap enabled, unregister arms controller SIGTERM before local policy cleanup completes. The shutdown order therefore follows the existing shared parsed setting.

Only rank 0 sends the unregister request; other policy ranks retain the existing idempotent no-op behavior.

Test plan

  • python -m pytest -q tests/test_multirank_shutdown.py — 64 passed, 2 subtests passed
  • ruff check cosmos_rl/policy/worker/rl_worker.py tests/test_multirank_shutdown.py
  • ruff format --check cosmos_rl/policy/worker/rl_worker.py tests/test_multirank_shutdown.py
  • git diff --check

Not re-run on a distributed cluster.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant