[train] Cleanup Zombie RayTrainWorker Actors

[JasonLi1909]

Background

Leaking Ray Train actors have been observed occupying GPU memory following Train run termination, causing training failures/OOMs in subsequent train runs. Despite the train actors being marked DEAD by Ray Core, we find that upon ssh-ing into nodes, that the actor processes are still alive and occupying valuable GPU memory:

Upon further investigation, we find that the actor process is hanging on Python garbage collection following termination via __ray_terminate__. __ray_terminate__ marks the actor as DEAD but does not guarantee the termination of the actor process. And following __ray_terminate__ when circular references are cleaned during Python interpreter teardown/garbage collection, the execution of object finalizers hang if the worker state is unstable- preventing the actor from ever exiting.

Previous User Workaround

We find that if garbage collection is performed earlier via gc.collect() while the worker state is more stable prior to shutdown (e.g. at the end of user train_func), there are no zombie actors- a mitigation leveraged by users prior to this stable fix.

Changes

This PR:

• Replaces __ray_terminate__ with ray.kill in Train run shutdown and abort paths to guarantee the termination of train actors

With the changes, the same train job as above no longer has Zombie train actors:

ray.kill preserves shutdown patience

Previously, __ray_terminate__ was used for "graceful" termination of actors, scheduling the termination task with a timeout of 5 seconds to allow for pending tasks to complete, following which ray.kill would be called for non-terminated actors. While this PR does not schedule a termination task, it allows the worker shutdown hooks a (default) timeout of 5 seconds- preserving the amount of time shutdown hooks and other pending tasks have to complete. The difference is that after those 5 seconds,ray.kill is called to terminate every actor.

Reliability of ray.kill vs __ray_terminate__ Shutdown

ray.kill performs a forceful out-of-band shutdown of an actor process, bypassing atexit handlers and python shutdown hooks- similar in nature to a SIGKILL- whereas __ray_terminate__ raises an exception, allowing python shutdown hooks to run and allowing the opportunity for Train actors to hang. Because Train actors do not rely on said hooks or atexit handlers for functionality, ray.kill is more fitting for our use case.

[gemini-code-assist]

Code Review

This pull request addresses a critical issue of leaking Ray Train actors that could lead to out-of-memory errors in subsequent training runs. The change from __ray_terminate__ to ray.kill is a solid approach to guarantee actor process termination. The refactored _shutdown_workers function is now cleaner and more reliable. It correctly allows a grace period for shutdown hooks to execute before ensuring all workers are killed. I have one suggestion to improve the robustness of the timeout handling.

[justinvyu]

Great investigation! 🚀 Can we add your minimal repro as a test? Ex: can do a pid liveness check assertion

Also, can you expand on the reason why ray_terminate fails vs. ray.kill()? Ex: letting a python process exit itself, vs. externally force killing the process. And can you also add some details about the early gc.collect() fixing things?