fix: prevent process hang when vLLM engine dies during async generation

[dinhxuanvu]

Summary

• Fix process hang caused by raise RuntimeError() inside asyncio.create_task() being silently swallowed
• Use os._exit(1) instead of sys.exit(1) to force-terminate from async task context
• Release capacity slot on CancelledError to prevent assertion failures in validate_state_at_epoch_end
• Reset slot_acquired after on_rollout_accepted() to avoid stale state across loop iterations

Problem

_run_generate_for_a_group_loop runs inside asyncio.create_task(). When vLLM's EngineDeadError occurs during generation, slot_acquired is True, so the handler executes:

raise RuntimeError("Generation workers should only run into error when they finish running.")

This exception is silently swallowed by asyncio (task exceptions only surface if you await the task). The sys.exit(1) on the next line is never reached because raise exits the except block.

Even if sys.exit(1) were reached, it raises SystemExit which is also captured by the asyncio task machinery and does not terminate the process.

The training process hangs indefinitely with all GPUs idle.

Fix

1. Remove the raise RuntimeError() guards that prevented process termination
2. Use os._exit(1) which bypasses asyncio and forces immediate termination
3. On CancelledError, call on_rollout_rejected() to release the capacity slot
4. Reset slot_acquired = False after on_rollout_accepted() for correct state tracking

Impact

Without this fix, a single EngineDeadError causes all training GPUs to sit idle indefinitely. Observed in production: 128 H100 GPUs were idle for 5+ days due to this bug.

Test plan

• Verified the diff is minimal (6 insertions, 10 deletions) and only touches exception handling logic
• Existing tests pass (no behavioral change in the happy path)
• Integration test: kill a vLLM engine during generation and confirm the process exits instead of hanging

[gemini-code-assist]

Code Review

This pull request modifies the exception handling in the _run_generate_for_a_group_loop method of the fully asynchronous trainer, replacing runtime errors with error logging. Review feedback highlights two critical issues: first, the failure to release capacity slots when a worker is cancelled, which leads to state inconsistency and assertion failures; and second, the use of sys.exit(1) within an asyncio task, which does not terminate the process and causes the training loop to hang. Suggestions were provided to use os._exit(1) and to ensure proper slot management during cancellation.