PR #13109: Add flag for termination on nccl error
PR #13109: Add flag for termination on nccl error

Imported from GitHub PR https://github.com/openxla/xla/pull/13109

This introduces a flag for termination on NCCL async error. With the flag on, XLA will terminate the process on NCCL error. With the flag off, the existing behavior should remain unchanged.

The patch is motivated by several problems:

- Without this patch, the heartbeat monitor only checks communicators that are currently not use by the running executable (because it obtains the communicators with TryAcquire). Since NCCL errors cause a hang in the running communicator, most failing communicators are locked, so their async errors just go undetected. As a result, XLA often hangs until Grpc timeout even in cases when ncclCommGetAsyncError would report an error.

- Ideally we would recover by aborting the faulty communicators, but that seems to be unreliable (aborts can cause hangs if NCCL currently hangs on a different communicator than the one being aborted). NCCL team is aware of this and working on a fix (https://github.com/NVIDIA/nccl/issues/1013). At the moment, there does not seem to be a reliable fast recovery mechanism short of process termination.

We propose to expose a flag for terminating the process on failure so that there is some way to detect and recover from a NCCL failure. Once the comm-abort works reliably, we will use that and propagate the error to the API user.

The patch is based on a PoC from pshamis@nvidia.com and vincentz@nvidia.com.
Copybara import of the project:

--
858aeacb2d689e4b03f4e3bcc0595223119143d5 by Jaroslav Sevcik <jsevcik@nvidia.com>:

Add flag for termination on nccl error

Merging this change closes #13109

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/13109 from jaro-sevcik:terminate-on-nccl-error 858aeacb2d689e4b03f4e3bcc0595223119143d5
