Training termination for grpo/main with finite step limit

[DNXie]

Previously the GRPO loop ran indefinitely with both training and rollout tasks active.
This PR adds a cfg.training.steps limit so training stops after the specified number of steps, then cleanly terminates rollout tasks.

Now logs end with messages like (I tested with cfg.training.steps=2):

Reached training limit (2 steps). Exiting continuous_training loop.
Shutting down...
WandbBackend global_controller: Finished run
Health loop stopped gracefully.
Health loop stopped gracefully.
Health loop stopped gracefully.
... 
Shutting down provisioner..

[felipemello1]

can we get the logger from src/forge/util/logging.py instead?

[joecummings]

We don't really use the logger elsewhere in the main script do we?

[DNXie]

The entire main script is using print instead of logger right now.

[joecummings]

Can you put this at the top of the main() loop? And also make it required please. The default can still be null / None, but this way it's very visible.

[DNXie]

Done!

[felipemello1]

can you explain this change a bit? before we would do gather in the try, and they would be impacted by KeyboardInterrupt. Now the gather happens after KeyboardInterrupt. Would the user possibly have to run 'KeyboardInterrupt' twice?

[casteryh]

can you explain this change a bit? before we would do gather in the try, and they would be impacted by KeyboardInterrupt. Now the gather happens after KeyboardInterrupt. Would the user possibly have to run 'KeyboardInterrupt' twice?

I think after the first KeyboardInterrupt, all the tasks are canceled and now we just gather on the canceled tasks (which should be fast to resolve) for graceful shutdown.

[DNXie]

The previous gather(*rollout_tasks, training_task) call blocked on rollouts indefinitely, even after training_task completed (e.g., once max_steps was reached).
Since we now have a finite step limit that cleanly terminates training_task, we shouldn’t continue waiting on rollout tasks.

With this change, rollout termination is handled explicitly in the finally block.
This doesn’t theoretically change how KeyboardInterrupt is handled, since an interrupt caught in the except would still flow into finally.
However, we’ve already seen issues with KeyboardInterrupt handling not working properly (see #360 (2)); I’ll look into that separately in a follow-up PR.

[casteryh]

I think we should make `cfg.training.steps' a required argument. And explicitly setting to -1 means run until interrupted.

[DNXie]

Made max_steps

max_steps = cfg.trainer.training.steps or -1

and the condition

while max_steps < 0 or training_step < max_steps:

[allenwang28]

I think we can/should make this part more graceful. Proposal:

Continuous rollouts takes a shutdown event:

async def continuous_rollouts(shutdown_evnt: asyncio.Event):
    ...
    while not shutdown_event.is_set(): # no more while True
        ...

    print("Rollout loop got shutdown event, shutting down...")

then in our finally we can do a 2-phased shutdown:

finally:
    print("Shutting down...")
    shutdown_event.set()

    try:
        # give tasks a chance to exit gracefully
        await asyncio.wait_for(
            asyncio.gather(*rollout_tasks, return_exceptions=True),
            timeout=5
        )
    except asyncio.TimeoutError:
        print("Forcing cancellation...")
        for t in rollout_tasks:
            t.cancel()
        await asyncio.gather(*rollout_tasks, return_exceptions=True)

[DNXie]

Done

[allenwang28]

LGTM as long as it's still runnig correctly!