[train] Eager Data Resource Cleanup on Train Run Failures and Aborts

[JasonLi1909]

Following a worker failure or a user abort during a Train job, the execution of sharded datasets (provided through get_dataset_shard) is ungracefully shutdown. Consequently, any ongoing resource request made by a sharded dataset's SplitCoordinator to the AutoscalingRequester is not cancelled. This can result in resources being held for a preset timeout, leading to inefficient cluster utilization and slower train job turnarounds.

To address the issue, this PR:

• Implements an eager shutdown path to cleanup resource requests made to the AutoscalingRequester (depicted below)
• Adds new WorkerGroupCallback hooks(after_worker_group_abort and after_worker_group_shutdown) to DatasetsSetupCallback for the new shutdown path
• Implements tests for the new cleanup path

Note on new WorkerGroupCallback hooks:
The new WorkerGroupCallback hooks after_worker_group_abort and after_worker_group_shutdown were added to ensure that the StreamingExecutor shutdown logic ran prior to the shutdown of train workers. This helps to avoid any race conditions and additional complexity related to timing executor shutdown while train workers are still alive.

Diagram of the new cleanup path:

[gemini-code-assist]

Code Review

This PR correctly renames DatasetsSetupCallback to DatasetsCallback. I've added a few suggestions to improve consistency in related parts of the code, such as a variable name, a test function name, and a docstring, to fully align with this change.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[justinvyu]

Thanks! Should be good after this. Nice tests

[justinvyu]

let's also add a call to after_worker_group_shutdown/abort and check the liveness of the actors?

[JasonLi1909]

The callbacks don't kill the SplitCoordinator actors as of now, they only shutdown their data executors. Ref counting should take care of them. That said, I added two more tests for the after_worker_group_shutdown/abort hooks.

[justinvyu]

🚢

[cursor]

Bug: Shutdown refs overwritten across worker group restarts

When a worker group is restarted due to failure or rescaling, _shutdown_data_executors is called for each worker group instance, but it overwrites self._shutdown_refs instead of appending to it. This means shutdown refs from previous worker group attempts are lost and never awaited in before_controller_shutdown, potentially causing incomplete resource cleanup from earlier failed attempts.

 

[iamjustinhsu]

is it safe to shutdown multiple times?

[JasonLi1909]

Yes, it is. That said, we only ever call it once per SplitCoordinator.