[Train] Add NCCL RAS health callback

[pseudo-rnd-thoughts]

Description

This PR adds an opt-in callback that detects hanging distributed training jobs by polling NCCL's RAS (Reliability/Availability/Serviceability) subsystem. NCCL ≥ 2.24 runs a monitoring thread inside every NCCL process that forms a peer mesh tracking per-rank health and collective op-counts. By querying the ncclras client and diffing collective-counts between successive polls, Ray Train can automatically detect and act on hangs that would otherwise stall a job until a torch process group timeout occurs.

As RAS gives a snapshot of the current NCCL state, to determine if the mismatch is a temporary issue or a permanent one, the callback polls RAS every N seconds. If a communicator is mismatched and makes no progress at all across multiple sequential polls, then this most likely indicates that a deadlock / wedge / hang is in progress for that communicator.

Its important to note that its possible for NCCL to continue (not hang) with incorrect data depending on the nccl operation, tensor size, nccl version, etc. In these cases, users will observe the collective-count continue to increase but often with a constant mismatch. In this version of the callback, we are not detecting those issues and focusing on purely NCCL collective mismatches which cause hangs.

The callback is disabled by default and enabled with RAY_TRAIN_ENABLE_NCCL_HANG_DETECTOR=1, as the detection heuristics still need validation against real-world hangs (spot preemptions, pipeline-parallel send/recv comms that legitimately sit skewed for long periods) before it can be considered for default-on.

Requirements

• RAS itself landed in 2.24, but the poller parses ncclras -f json, and the -f flag requires NCCL 2.28+. If the binary rejects -f, or isn't on PATH, the callback logs a warning once and disables itself for the rest of the run rather than retrying forever.
• ci/docker/base.gpu.Dockerfile pins and holds libnccl2/libnccl-dev at 2.28.9-1+cuda12.9 (previously whatever the CUDA base image shipped) so ncclras is present and JSON-capable in GPU CI. The build fails fast if the pin doesn't stick or the client binary is missing. This affects all GPU CI images, not just Train tests.
• py-spy (and SYS_PTRACE in the container) for native stack traces at hang time. When unavailable, the callback falls back to a Python-only traceback of every thread.

How it works

• NCCLRASCallback lives on the controller and polls ncclras on a worker (ras forms a mesh-network between ranks requiring that only one of the workers to be polled and receive the whole network's state). Workers are tried in turn until one returns a usable report, and the query runs on a background thread so it never blocks the controller's poll loop.
• A communicator's frozen-streak only advances when the RAS report shows a collective-count mismatch, every rank in that communicator is RUNNING, and no rank advanced any collective since the previous poll. Any progress on any op resets that communicator's streak; streaks are tracked per communicator so each one is confirmed on its own.
• A confirmed hang requires the streak to persist for RAY_TRAIN_NCCL_RAS_CONFIRM_DURATION_S (default 600s), converted to a consecutive-poll count using RAY_TRAIN_NCCL_RAS_MIN_POLL_INTERVAL_S (default 15s), 40 polls by default, to ensure that a mismatch isn't a snapshot issue.
• Before confirmation there is an escalation ladder: a WARNING when a communicator first crosses ~60s frozen, then a periodic WARNING every ~120s naming every still-frozen communicator, its stalled duration, and (in fail mode) the time remaining until NCCLHangError, alongside the human-readable ncclras -f text report. A communicator that resumes progress after being suspected logs an explicit recovery message.
• On a confirmed hang, the callback logs the ncclras -f text report and captures py-spy stack traces from all workers (uploaded to <experiment_path>/nccl_ras_hang_stack_traces/ as rank_<i>.log), then depending on the callback mode, in "fail" a NCCLHangError is raised and in "observe", a detailed log message is produced but no more. observe is the default; set RAY_TRAIN_NCCL_RAS_ACTION=fail to fail the run.
• NCCLHangError (a WorkerGroupError subclass, DeveloperAPI, exported as ray.train.NCCLHangError) is treated as non-retryable by the DefaultFailurePolicy regardless of the max_failures budget, since a desync hang is usually deterministic and a restart would just hang again.

Configuration

Env var	Default	Read on
NCCL_RAS_ADDR	localhost:28028	worker
RAY_TRAIN_ENABLE_NCCL_HANG_DETECTOR	0	driver
RAY_TRAIN_NCCL_RAS_ACTION	observe	driver
RAY_TRAIN_NCCL_RAS_MIN_POLL_INTERVAL_S	15	driver
RAY_TRAIN_NCCL_RAS_CONFIRM_DURATION_S	600	driver
RAY_TRAIN_NCCLRAS_PATH	ncclras (PATH lookup)	driver

The RAY_TRAIN_* knobs are read once when the callback is constructed on the driver validated eagerly (invalid values raise at construction); NCCL_RAS_ADDR is NCCL's variable and is read on the worker at query time.

Known limitations

• Symmetric in-collective hangs are not detected. RAS reports collective launch counts, so if every rank launches the same collective and the fabric wedges mid-op, all counts match and skew is zero.
• A rank that exited early is not detected. A communicator is only considered for detection when all of its ranks report RUNNING, so a FINALIZE/ABORT rank alongside frozen peers is currently skipped. This could be considered in this implementation but I believe the Ray health checker should catch this first.
• missing_ranks (unresponsive / considered-dead peers) is parsed around but not yet used as a signal.
• No metrics are emitted yet, and there is no culprit-rank attribution — NCCLHangError.worker_failures is empty; the RAS text report and stack traces are the only diagnostics.

Testing

• test_nccl_ras_callback.py — CPU unit coverage of RAS JSON parsing (including the malformed missing_ranks[] comma NCCL 2.28.9 emits), NCCL_RAS_ADDR parsing, frozen-vs- advancing classification, per-communicator streak independence, communicators appearing and disappearing between polls, poll throttling, confirm-duration → poll-count conversion, config validation, the escalation/observe messaging, and that a failed stack dump does not suppress the hang error.
• test_nccl_ras_hang_detection.py — GPU end-to-end scenarios (train_v2_gpu) inducing real NCCL desyncs in a TorchTrainer with the detector configured for fast confirmation. Requires ≥ 2 visible GPUs and ncclras on PATH; the multi-communicator subset test requires 4 GPUs and skips otherwise.
• test_failure_policy.py — added test_nccl_hang_error_is_non_retryable.

[gemini-code-assist]

Code Review

This pull request introduces a new NCCL RAS-based hang detector callback (NCCLRASCallback) for Ray Train v2. This callback monitors the NCCL Reliability/Availability/Serviceability (RAS) subsystem to detect deadlocks and unresponsive ranks by tracking collective op-counts across communicators. If a hang is confirmed over multiple consecutive polls, it can raise a non-retryable NCCLHangError and dump worker stack traces using py-spy for debugging. The PR also includes unit and end-to-end tests for this new callback. The review comments point out three important issues: a potential ValueError in parse_ras_addr when parsing bare hostnames or bracketed IPv6 addresses without a port, a potential TypeError in query_ras_on_workers when calling ray.cancel(None) if worker.execute_async fails, and a thread leak due to the ThreadPoolExecutor not being shut down when the worker group is stopped.

[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.

[JasonLi1909]

Thanks! Left some comments, main things:

• The current default observe mode does not collect stack traces. Users may still want to see the stack traces while also not wanting their job to terminate. It will also help us to with identifying false positives rates moving forward. Without the stack traces, we will not be able to tell if the detection was valid.
• This PR only detect hangs if their is a misalignment between collectives, but it's possible for the hang to occur inside the collectives while they're all aligned. We should investigate this case and if we can detect it without many false positives.
• Because the ncclras query is run on the train worker, it may contend with the health check that runs every ~2s. Wondering if we need to run this on the train worker or if it can just be done as a separate task.
• Can we provide an example in the pr description for when the hang would be detected, showing the nccl ras output?

[JasonLi1909]

nit: maybe rename to NCCL_RAS_MIN_POLL_INTERVAL_S because the 15.0 isn't the true poll interval but instead serves to throttle querying to no more than every 15s

[pseudo-rnd-thoughts]

Updated

[JasonLi1909]

Can we make this time based instead? I assume the 10 minute assumption is based off the 15s poll interval * 40 = 10 mins, but this is fragile as depends on the poll interval. Because we need to collect stack traces prior to the nccl watchdog error, the time guarantee here is useful to have. And maybe we can update the timer only when a report is collected and misalignment confirmed. Then the 10 mins can be the min amount of time the hang must have persisted.

[pseudo-rnd-thoughts]

Thats fair, I'll update that the env-vars use time in seconds for a confirmed hang allowing it to be modified more easily.

[JasonLi1909]

Why are we testing for these FAIL and MAYBE cases? I think we should keep the scope of the feature to just detecting the hang case for now as that is the highest value case. Other cases may be more nuanced and we can't say much about them. We can consider adding those in an extension

[pseudo-rnd-thoughts]

From testing, for the maybes, I found that depending on tensor size and how long you leave it to done, depends on if it would hang or not.

[pseudo-rnd-thoughts]

Thanks for the review

The current default observe mode does not collect stack traces. Users may still want to see the stack traces while also not wanting their job to terminate. It will also help us to with identifying false positives rates moving forward. Without the stack traces, we will not be able to tell if the detection was valid.

Yes, will fix

This PR only detect hangs if their is a misalignment between collectives, but it's possible for the hang to occur inside the collectives while they're all aligned. We should investigate this case and if we can detect it without many false positives.

With nccl-ras no, this is part of the reason that NVIDIA Resiliency Extension (NVRx) uses collective time taken instead as a proxy for hangs which might be more reliable. However, this requires users to update their implementation to hook into tracking this information. In comparison, ras runs asynchronously of the user without any setup requirements and works in most software based cases. The cases where a hang occurs with equal collective calls, I believe this primarily occurs for hardware based failures, therefore, the new DCGM view should help users more.

Because the ncclras query is run on the train worker, it may contend with the health check that runs every ~2s. Wondering if we need to run this on the train worker or if it can just be done as a separate task.

The query runs asynchronously in a different thread within the worker node so I'm aiming (and expecting) this callback to cause zero drag on other callbacks / health checks.

Can we provide an example in the pr description for when the hang would be detected, showing the nccl ras output?

Yes, will provide

[justinvyu]

Thanks!

[justinvyu]

can you include some snippets of the logs seen by the user from your e2e testing?

[justinvyu]

consdier putting this into a longrunning loop in a thread instead of submitting new tasks to the threadpool every once in a while

[justinvyu]

Could you clarify this comment? what is the "world communicator"?

[pseudo-rnd-thoughts]

When doing tensor parallelism (or other parallelism strategies) then you can have rank 0 and 1 using one set of tensors and rank 2 and 3 using another set of tensors. If you have a torch distributed process (i.e., communicator) that only exists between rank 0 and 1, and 2 and 3 then NCCL will only see 0 and 1, then 2 and 3 and not all the ranks together. This matters as RAS runs on top of NCCL and doesn't search out more ranks.

Therefore, the problem is that if you query RAS on rank 0, then it would only return the data for rank 0 and 1 (as it doesn't know that ranks 2 and 3 exist). In practice this shouldn't exist as parallelism strategies will have a world communicator between all ranks to share global information.
However, in the case where this isn't the case then users might run in this bug.

I hope that makes sense. I can rewrite the comment with more detail.

[justinvyu]

put this default in the constants.py

[pseudo-rnd-thoughts]

Reading RAS again, we don't need to pass address and nccl will just use its own default, therefore, I'm updated to let nccl use its own value unless the user overrides it

[pseudo-rnd-thoughts]

I've removed this and let ras use its internal default variable in case it changes in the future. Users need to intentionally override it for us to pass the variable to ras

[justinvyu]

can simplify by capturing the loop logic into a helper function so we don't have a continue on every branch of the try catch.

for worker in workers:
    value, reason = self._try_query_worker(worker, ras_format)
    if value is not None:
        return value
    last_failure_reason = reason

[justinvyu]

log that we weren't able to get these rank stack traces in time

[pseudo-rnd-thoughts]

Great idea, done

[pseudo-rnd-thoughts]

Great idea, done

[justinvyu]

hmm I would actually vote to remove this special case; nccl hangs could be caused by flaky hardware (nic) that would benefit from a restart. and plus we'll only retry up to max_failures times

[pseudo-rnd-thoughts]

I'll remove this and leave it up to a future PR

[justinvyu]

do we have worker specific stuff we could populate worker_failures with? ex: the rank's stack traces

[pseudo-rnd-thoughts]

Would this double the stack traces from both the controller and workers? Also worker_failures: Dict[int, Exception] and the stack trace isn't an exception object strictly

[justinvyu]

can you break up the logic in here into a few parts for easier readability and testing?

ex:

def evaluate_comm_mismatch(self, report):
    if self.prev_report is None:
        return

    op_diff = compute_report_op_diff(self.prev_report, report)

    # 1. Classify each communicator
    frozen = {}
    for comm_id in report.mismatched_comms:
        if comm_id not in op_diff:
            continue
        if all(d == 0 for deltas in op_diff[comm_id].values() for d in deltas.values()):
            frozen[comm_id] = self.comm_deadlock_count.get(comm_id, 0) + 1

    confirmed = [c for c, n in frozen.items() if n == self._confirm_poll_counts]
    recovered = [c for c in self.comm_deadlock_count if c not in frozen
                 and self.comm_deadlock_count[c] > self._suspicion_polls]

    # 2. Update state
    self.comm_deadlock_count = frozen

    # 3. Act
    for comm_id in recovered:
        self._log_recovery(comm_id, ...)
    if confirmed:
        self._handle_confirmed_hang(confirmed, report)
    elif frozen:
        self._log_suspicion_warnings(report)

[pseudo-rnd-thoughts]

I'll see

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 9c114bd. Configure here.}