Summary

The goal of this PR is to make it so that when users Ctrl C their ray train driver script, we gracefully terminate the run and mark train runs and train run attempts as ABORTED.

Implementation Details

Here are the changes to each component:

• driver: sigint handler that calls controller.abort in a blocking fashion
• controller/worker_group: mark train run and train run attempts as aborted and exit
• workers: no change - for now we rely on ray's object reference counting to clean up workers. This is safe because the workers cannot modify train run and train run attempt state if the controller has exited. We can consider gracefully terminating the workers in the future but we don't for now since there is some risk of hanging e.g. when we call destroy_process_group on an active group.

We decided to implement abort as an async method instead of a separate thread to avoid race conditions when setting state. See the diagrams below.

Testing

Here is my workspace

When I Ctrl C my Ray Train run I see these logs.

which results in this train run page.

Regular training runs still work - here is the overall ray train dashboard page the aborted run and a successful run.

Misc notes

• I moved controller._start from run to __init__ because the new abort method calls before_controller_abort which assumes that after_controller_start was called; for example, the state manager callback can only mark a run as aborted after it was already created.
• I made the driver sigint handler catch the ActorDiedError and sys.exit(0) to indicate that this is expected behavior. However, even after doing this, the driver exits with exit code 1. Let me know if this is ok.