From 7e8366476a8ad3a2c8737b75cd5464918c63cd88 Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Thu, 3 Sep 2026 16:42:38 +0530 Subject: [PATCH] feat(errors): make backend-classified GPU faults catchable The backend now classifies GPU and NVSwitch faults into the ExecutionError it hands the SDK: the code becomes one of the gpufault codes, the driver's own sentence is prepended to the message, and the fault travels as typed data on ExecutionError.gpu_fault. Until now all of that arrived as an opaque RuntimeUserError or RuntimeSystemError, so a task that lost a GPU could not tell that apart from any other failure. Add a GPUFaultError family and select it from those codes. One except clause catches every GPU fault, and the xid, severity, gpu_uuid, node and device attributes are there to branch on. The concrete errors keep the kind the backend chose, GPUFaultUserError for a fault the workload caused and GPUFaultSystemError for one that condemned the hardware, which is also what decides whose retry budget paid for it. Attributes come from the typed fault where the failure carries one and from the prepended sentence otherwise, so a failure from a backend that predates the typed field still names the Xid. Reading them never fails the conversion: an unreadable message costs the details, not the error. Co-Authored-By: Claude Fable 5 --- src/flyte/_internal/runtime/convert.py | 152 ++++++++++- src/flyte/errors.py | 107 ++++++++ .../runtime/test_convert_gpu_fault.py | 235 ++++++++++++++++++ 3 files changed, 493 insertions(+), 1 deletion(-) create mode 100644 tests/flyte/internal/runtime/test_convert_gpu_fault.py diff --git a/src/flyte/_internal/runtime/convert.py b/src/flyte/_internal/runtime/convert.py index 3bf0de759..efff933c1 100644 --- a/src/flyte/_internal/runtime/convert.py +++ b/src/flyte/_internal/runtime/convert.py @@ -5,6 +5,7 @@ import contextvars import hashlib import inspect +import re from dataclasses import dataclass from types import NoneType from typing import Any, Dict, List, Optional, Tuple, Union, cast, get_args @@ -473,6 +474,147 @@ async def convert_outputs_to_native(interface: NativeInterface, outputs: Outputs return tuple(kwargs[k] for k in interface.outputs.keys()) +# The sentence the backend prepends to the failure message when it attributes the failure to a GPU fault, for example +# "[gpu-health] [CRITICAL] Xid 79 (GPU has fallen off the bus) on GPU 3 GPU-1a2b." or +# "[gpu-health] [CRITICAL] SXid 22 on NVSwitch 0000:3b:00.0.". It is the human half of the GPU health daemon's event +# message, and it is what the fault attributes are read from when a failure arrives with no typed fault on it, for +# example from a backend older than the one that added the field. The trailing full stop is matched only where a space +# or the end of the message follows it, so a PCI bus id keeps its own dots. +_GPU_FAULT_SENTENCE = re.compile( + r"\[gpu-health\]\s+\[(?P[A-Za-z]+)\]\s+" + r"(?:Xid\s+(?P\d+)\s+\((?P[^)]*)\)|SXid\s+(?P\d+))" + r"(?P.*?)\.(?=\s|$)" +) + +# The tail of the sentence naming the device the fault happened on, in the four shapes the backend renders it in. An +# NVSwitch SXid names the switch by bus id, a GPU Xid names the GPU by index and UUID and degrades to whichever of the +# three it could resolve. +_GPU_FAULT_LOCATION = re.compile( + r"^\s+on\s+(?:NVSwitch\s+(?P\S+)" + r"|GPU\s+(?:at\s+PCI\s+(?P\S+)" + r"|(?P\d+)\s+(?P\S+)" + r"|(?P\D\S*)" + r"|(?P\d+)))$" +) + +# The keys of the k=v tail the GPU health daemon writes after the sentence. Failure classification does not carry the +# tail today, but a message that does have one carries the node and the process, which the sentence never does. +_GPU_FAULT_TAIL_KEYS = frozenset({"xid", "sxid", "severity", "gpu_uuid", "gpu_index", "pci", "node", "process"}) + + +def _int_or_none(value: str | None) -> int | None: + if value is None: + return None + try: + return int(value) + except ValueError: + return None + + +def _parse_gpu_fault_message(message: str) -> Dict[str, Any]: + """ + Read whatever the backend's GPU fault sentence says about the fault. Returns the fields it could recover, which is + the empty dict for a message with no sentence in it. It never raises: a message that does not parse only costs the + attributes, and the error itself still has to be converted. + """ + match = _GPU_FAULT_SENTENCE.search(message or "") + if match is None: + return {} + + is_sxid = match.group("sxid") is not None + fields: Dict[str, Any] = { + "fault_kind": "sxid" if is_sxid else "xid", + "fault_code": _int_or_none(match.group("sxid") if is_sxid else match.group("xid")), + "fault_name": match.group("name") or None, + "severity": (match.group("severity") or "").lower() or None, + } + + location = _GPU_FAULT_LOCATION.match(match.group("where") or "") + if location is not None: + fields["gpu_uuid"] = location.group("indexed_uuid") or location.group("uuid") + fields["gpu_index"] = _int_or_none(location.group("index") or location.group("bare_index")) + fields["pci_bus_id"] = location.group("pci") or location.group("switch_pci") + + fields.update(_parse_gpu_fault_tail(message[match.end() :])) + return {k: v for k, v in fields.items() if v is not None} + + +def _parse_gpu_fault_tail(rest: str) -> Dict[str, Any]: + """ + Read the k=v tail that follows the sentence when the whole event message was carried over, ignoring every token + that is not one of the daemon's own keys. + """ + tail: Dict[str, Any] = {} + for token in rest.split(): + key, sep, value = token.partition("=") + if not sep or key not in _GPU_FAULT_TAIL_KEYS or not value: + continue + if key in ("xid", "sxid"): + tail["fault_kind"] = key + tail["fault_code"] = _int_or_none(value) + elif key == "gpu_index": + tail["gpu_index"] = _int_or_none(value) + elif key == "pci": + tail["pci_bus_id"] = value + else: + tail[key] = value + return {k: v for k, v in tail.items() if v is not None} + + +_GPU_FAULT_KINDS = { + execution_pb2.GpuFault.KIND_XID: "xid", + execution_pb2.GpuFault.KIND_SXID: "sxid", +} + +_GPU_FAULT_SEVERITIES = { + execution_pb2.GpuFault.SEVERITY_USER: "user", + execution_pb2.GpuFault.SEVERITY_WARN: "warn", + execution_pb2.GpuFault.SEVERITY_CRITICAL: "critical", +} + + +def _typed_gpu_fault_fields(err: execution_pb2.ExecutionError) -> Dict[str, Any]: + """ + Read the typed fault the backend attaches to the failure on ExecutionError.gpu_fault. An unset or unspecified value + is dropped rather than guessed at, so the caller can fall back to the message sentence for it. + """ + if not err.HasField("gpu_fault"): + return {} + fault = err.gpu_fault + + fields: Dict[str, Any] = { + "fault_kind": _GPU_FAULT_KINDS.get(fault.kind), + "fault_code": fault.code or None, + "fault_name": fault.name or None, + "severity": _GPU_FAULT_SEVERITIES.get(fault.severity), + "gpu_uuid": fault.gpu_uuid or None, + "node": fault.node or None, + "pci_bus_id": fault.pci_bus_id or None, + "process": fault.process or None, + } + # gpu_index is optional in the IDL because index 0 is a real GPU and an unresolved index is not. + if fault.HasField("gpu_index"): + fields["gpu_index"] = fault.gpu_index + return {k: v for k, v in fields.items() if v is not None} + + +def _gpu_fault_fields(err: execution_pb2.ExecutionError) -> Dict[str, Any]: + """ + Everything the failure says about the GPU fault behind it, preferring the typed fault and falling back to the + sentence the backend prepended to the message. Attributes are best effort, so a parse that finds nothing still + yields a GPU fault error, only one with less on it. + """ + try: + fields = _typed_gpu_fault_fields(err) + if fields: + return fields + return _parse_gpu_fault_message(err.message) + except Exception: + # The attributes are a convenience. Losing them costs the caller a detail, failing the conversion would cost + # it the error itself. + return {} + + def convert_error_to_native( err: execution_pb2.ExecutionError | Exception | Error, ) -> Exception | None: @@ -490,7 +632,11 @@ def convert_error_to_native( case execution_pb2.ExecutionError.UNKNOWN: return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker) case execution_pb2.ExecutionError.USER: - if "OOM" in err.code.upper(): + if user_code in flyte.errors.GPU_FAULT_CODES: + return flyte.errors.GPUFaultUserError( + code=user_code, message=err.message, worker=err.worker, **_gpu_fault_fields(err) + ) + elif "OOM" in err.code.upper(): return flyte.errors.OOMError(code=user_code, message=err.message, worker=err.worker) elif "Interrupted" in err.code: return flyte.errors.TaskInterruptedError(code=user_code, message=err.message, worker=err.worker) @@ -508,6 +654,10 @@ def convert_error_to_native( return flyte.errors.ImagePullBackOffError(code=user_code, message=err.message, worker=err.worker) return flyte.errors.RuntimeUserError(code=user_code, message=err.message, worker=err.worker) case execution_pb2.ExecutionError.SYSTEM: + if user_code in flyte.errors.GPU_FAULT_CODES: + return flyte.errors.GPUFaultSystemError( + code=user_code, message=err.message, worker=err.worker, **_gpu_fault_fields(err) + ) return flyte.errors.RuntimeSystemError(code=user_code, message=err.message, worker=err.worker) return None diff --git a/src/flyte/errors.py b/src/flyte/errors.py index 6b34c500f..ef849d69e 100644 --- a/src/flyte/errors.py +++ b/src/flyte/errors.py @@ -100,6 +100,113 @@ class OOMError(RuntimeUserError): """ +GPU_FAULT_CODES: tuple[str, ...] = ( + "GpuXidError", + "GpuFallenOffBus", + "GpuEccUncorrectable", + "GpuRowRemapPending", + "GpuNvlinkError", + "GpuGspError", +) +""" +The error codes the backend puts on a failure it attributed to a GPU or NVSwitch fault. GpuXidError is the catch-all +for a fault with no more specific code, the rest name a class of trouble a user or an operator can act on. Any of them +converts to a GPU fault error in the SDK. +""" + + +class GPUFaultError(BaseRuntimeError): + """ + This error is raised when the backend attributed the task failure to a GPU or NVSwitch fault that the GPU health + daemon observed on the node, such as an Xid 31 (a GPU memory page fault) or an Xid 79 (the GPU fell off the bus). + + Catch this class to handle every GPU fault. It is the base of both concrete errors, GPUFaultUserError for a fault + the workload caused and GPUFaultSystemError for a hardware fault, so one except clause covers both, and the code, + severity and xid attributes are there to branch on afterwards. + + The two do not reach user code on the same terms. A user severity Xid (13, 31, 43, 45) is the workload's own + doing, it will fault again if it is replayed unchanged, so the backend charges it to the task's own retry budget + and this error surfaces as soon as that budget is spent. A critical hardware fault is not the workload's doing, so + the platform retries it without charging the user's budget and reschedules onto other hardware where it can, which + means user code sees a critical fault only after platform policy has given up on it. Neither one is a signal to + retry in place: a user fault has already exhausted its own retries by the time it is raised, and a critical fault + has already been retried elsewhere. + + The fault attributes are best effort. They are read from the typed fault the backend attaches to the failure, and + where there is none, from the sentence the backend prepends to the failure message, which does not carry every + attribute. Any of them can be None, so read them defensively. + """ + + def __init__( + self, + code: str, + kind: ErrorKind, + message: str, + worker: str | None = None, + *, + fault_kind: str | None = None, + fault_code: int | None = None, + fault_name: str | None = None, + severity: str | None = None, + gpu_uuid: str | None = None, + gpu_index: int | None = None, + node: str | None = None, + pci_bus_id: str | None = None, + process: str | None = None, + ): + # Named explicitly rather than through super(): the concrete errors below mix this class with RuntimeUserError + # and RuntimeSystemError, whose own initializers fix the kind and take one argument fewer. + BaseRuntimeError.__init__(self, code, kind, message, worker) + self.fault_kind = fault_kind + self.fault_code = fault_code + self.fault_name = fault_name + self.severity = severity + self.gpu_uuid = gpu_uuid + self.gpu_index = gpu_index + self.node = node + self.pci_bus_id = pci_bus_id + self.process = process + + @property + def xid(self) -> int | None: + """ + The NVIDIA Xid number of the fault, or None when the fault was an NVSwitch SXid or when the number could not + be determined. Xid and SXid numbers share a numbering space but not a meaning, so a number alone never + identifies a fault, read fault_kind together with fault_code to tell them apart. + """ + return self.fault_code if self.fault_kind == "xid" else None + + @property + def sxid(self) -> int | None: + """ + The NVSwitch SXid number of the fault, or None when the fault was a GPU Xid or when the number could not be + determined. + """ + return self.fault_code if self.fault_kind == "sxid" else None + + +class GPUFaultUserError(GPUFaultError, RuntimeUserError): + """ + This error is raised when the GPU fault the backend attributed the failure to was the workload's own doing, for + example an out-of-bounds access that the driver reported as an Xid 31. The GPU itself is fine once the process is + gone, so the failure was charged to the task's own retry budget. + """ + + def __init__(self, code: str, message: str, worker: str | None = None, **fault): + GPUFaultError.__init__(self, code, "user", message, worker, **fault) + + +class GPUFaultSystemError(GPUFaultError, RuntimeSystemError): + """ + This error is raised when the GPU fault the backend attributed the failure to condemned the device or the node, + for example an uncorrectable ECC error or a GPU that fell off the bus. The workload did not cause it, so the + platform retried the task on its own budget before this error reached user code. + """ + + def __init__(self, code: str, message: str, worker: str | None = None, **fault): + GPUFaultError.__init__(self, code, "system", message, worker, **fault) + + class TaskInterruptedError(RuntimeUserError): """ This error is raised when the underlying task execution is interrupted. diff --git a/tests/flyte/internal/runtime/test_convert_gpu_fault.py b/tests/flyte/internal/runtime/test_convert_gpu_fault.py new file mode 100644 index 000000000..83be9af95 --- /dev/null +++ b/tests/flyte/internal/runtime/test_convert_gpu_fault.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +import pytest +from flyteidl2.core import execution_pb2 + +import flyte.errors +from flyte._internal.runtime.convert import Error, convert_error_to_native + +XID_SENTENCE = "[gpu-health] [CRITICAL] Xid 79 (GPU has fallen off the bus) on GPU 3 GPU-1a2b-3c." +SXID_SENTENCE = "[gpu-health] [CRITICAL] SXid 22 on NVSwitch 0000:3b:00.0." +USER_SENTENCE = "[gpu-health] [USER] Xid 31 (GPU memory page fault) on GPU 0 GPU-abc." + + +def _err(code: str, kind, message: str = "", worker: str = "worker-0") -> execution_pb2.ExecutionError: + return execution_pb2.ExecutionError(code=code, kind=kind, message=message, worker=worker) + + +@pytest.mark.parametrize("code", list(flyte.errors.GPU_FAULT_CODES)) +def test_user_kind_gpu_code_selects_user_fault_error(code): + exc = convert_error_to_native(_err(code, execution_pb2.ExecutionError.USER, USER_SENTENCE)) + + assert isinstance(exc, flyte.errors.GPUFaultUserError) + assert isinstance(exc, flyte.errors.RuntimeUserError) + assert exc.code == code + assert exc.kind == "user" + assert exc.worker == "worker-0" + + +@pytest.mark.parametrize("code", list(flyte.errors.GPU_FAULT_CODES)) +def test_system_kind_gpu_code_selects_system_fault_error(code): + exc = convert_error_to_native(_err(code, execution_pb2.ExecutionError.SYSTEM, XID_SENTENCE)) + + assert isinstance(exc, flyte.errors.GPUFaultSystemError) + assert isinstance(exc, flyte.errors.RuntimeSystemError) + assert exc.code == code + assert exc.kind == "system" + + +def test_both_kinds_are_caught_by_the_one_base_class(): + for kind in (execution_pb2.ExecutionError.USER, execution_pb2.ExecutionError.SYSTEM): + exc = convert_error_to_native(_err("GpuXidError", kind, USER_SENTENCE)) + with pytest.raises(flyte.errors.GPUFaultError) as raised: + raise exc + assert raised.value is exc + + +def test_server_injected_code_still_selects_the_gpu_error(): + exc = convert_error_to_native( + _err("RetriesExhaustedError|GpuEccUncorrectable", execution_pb2.ExecutionError.SYSTEM, XID_SENTENCE) + ) + + assert isinstance(exc, flyte.errors.GPUFaultError) + assert exc.code == "GpuEccUncorrectable" + + +def test_error_wrapper_is_unwrapped_like_any_other_failure(): + exc = convert_error_to_native(Error(err=_err("GpuFallenOffBus", execution_pb2.ExecutionError.SYSTEM, XID_SENTENCE))) + + assert isinstance(exc, flyte.errors.GPUFaultSystemError) + + +def test_fields_from_the_xid_sentence(): + exc = convert_error_to_native( + _err("GpuFallenOffBus", execution_pb2.ExecutionError.SYSTEM, XID_SENTENCE + " Pod terminated.") + ) + + assert exc.fault_kind == "xid" + assert exc.fault_code == 79 + assert exc.xid == 79 + assert exc.sxid is None + assert exc.fault_name == "GPU has fallen off the bus" + assert exc.severity == "critical" + assert exc.gpu_index == 3 + assert exc.gpu_uuid == "GPU-1a2b-3c" + assert str(exc) == XID_SENTENCE + " Pod terminated." + + +def test_fields_from_the_sxid_sentence(): + exc = convert_error_to_native(_err("GpuNvlinkError", execution_pb2.ExecutionError.SYSTEM, SXID_SENTENCE)) + + assert exc.fault_kind == "sxid" + assert exc.fault_code == 22 + assert exc.sxid == 22 + assert exc.xid is None + assert exc.severity == "critical" + assert exc.pci_bus_id == "0000:3b:00.0" + assert exc.gpu_uuid is None + + +@pytest.mark.parametrize( + "sentence, gpu_index, gpu_uuid, pci_bus_id", + [ + ("[gpu-health] [USER] Xid 31 (GPU memory page fault) on GPU 0 GPU-abc.", 0, "GPU-abc", None), + ("[gpu-health] [USER] Xid 31 (GPU memory page fault) on GPU GPU-abc.", None, "GPU-abc", None), + ("[gpu-health] [USER] Xid 31 (GPU memory page fault) on GPU 2.", 2, None, None), + ("[gpu-health] [USER] Xid 31 (GPU memory page fault) on GPU at PCI 0000:3b:00.0.", None, None, "0000:3b:00.0"), + ("[gpu-health] [USER] Xid 31 (GPU memory page fault).", None, None, None), + ], +) +def test_device_is_read_from_every_shape_the_sentence_takes(sentence, gpu_index, gpu_uuid, pci_bus_id): + exc = convert_error_to_native(_err("GpuXidError", execution_pb2.ExecutionError.USER, sentence)) + + assert exc.xid == 31 + assert exc.severity == "user" + assert exc.gpu_index == gpu_index + assert exc.gpu_uuid == gpu_uuid + assert exc.pci_bus_id == pci_bus_id + + +def test_machine_readable_tail_is_read_when_the_message_carries_one(): + message = ( + "(combined from similar events): [gpu-health] [USER] Xid 13 (Graphics Engine Exception) on GPU 1 GPU-x." + " xid=13 severity=user gpu_uuid=GPU-x gpu_index=1 pci=0000:3b:00.0 node=ip-10-0-0-1 pid=42 process=python3" + ) + exc = convert_error_to_native(_err("GpuXidError", execution_pb2.ExecutionError.USER, message)) + + assert exc.xid == 13 + assert exc.node == "ip-10-0-0-1" + assert exc.process == "python3" + assert exc.pci_bus_id == "0000:3b:00.0" + + +@pytest.mark.parametrize( + "message", + [ + "", + "container exited with code 137", + "[gpu-health] [CRITICAL] Xid but no number at all.", + ], +) +def test_a_message_without_a_readable_sentence_still_converts(message): + exc = convert_error_to_native(_err("GpuXidError", execution_pb2.ExecutionError.USER, message)) + + assert isinstance(exc, flyte.errors.GPUFaultUserError) + assert exc.xid is None + assert exc.severity is None + assert exc.gpu_uuid is None + assert exc.node is None + + +@pytest.mark.parametrize( + "code, kind, expected", + [ + ("OOMKilled", execution_pb2.ExecutionError.USER, flyte.errors.OOMError), + ("Interrupted", execution_pb2.ExecutionError.USER, flyte.errors.TaskInterruptedError), + ("SomeOtherError", execution_pb2.ExecutionError.USER, flyte.errors.RuntimeUserError), + ("SomeOtherError", execution_pb2.ExecutionError.SYSTEM, flyte.errors.RuntimeSystemError), + ("SomeOtherError", execution_pb2.ExecutionError.UNKNOWN, flyte.errors.RuntimeUnknownError), + ("GpuXidError", execution_pb2.ExecutionError.UNKNOWN, flyte.errors.RuntimeUnknownError), + ], +) +def test_codes_that_are_not_gpu_faults_are_converted_as_before(code, kind, expected): + exc = convert_error_to_native(_err(code, kind, XID_SENTENCE)) + + assert type(exc) is expected + assert not isinstance(exc, flyte.errors.GPUFaultError) + + +# --------------------------------------------------------------------------------------------------------------- +# The typed fault the backend attaches to the failure, which is what the attributes are read from whenever it is +# there. The sentence is only the fallback for a failure that arrives without one. +# --------------------------------------------------------------------------------------------------------------- + + +def test_typed_fault_fills_every_field(): + err = _err("GpuFallenOffBus", execution_pb2.ExecutionError.SYSTEM, XID_SENTENCE) + err.gpu_fault.CopyFrom( + execution_pb2.GpuFault( + kind=execution_pb2.GpuFault.KIND_XID, + code=79, + name="GPU has fallen off the bus", + severity=execution_pb2.GpuFault.SEVERITY_CRITICAL, + gpu_uuid="GPU-typed", + gpu_index=0, + pci_bus_id="0000:3b:00.0", + node="ip-10-0-0-7", + pid=42, + process="train.py", + ) + ) + + exc = convert_error_to_native(err) + + assert isinstance(exc, flyte.errors.GPUFaultSystemError) + assert exc.fault_kind == "xid" + assert exc.xid == 79 + assert exc.fault_name == "GPU has fallen off the bus" + assert exc.severity == "critical" + assert exc.gpu_uuid == "GPU-typed" + # Read from the typed fault, not from the sentence on the same failure, which says GPU 3. + assert exc.gpu_index == 0 + assert exc.pci_bus_id == "0000:3b:00.0" + assert exc.node == "ip-10-0-0-7" + assert exc.process == "train.py" + + +def test_typed_sxid_fault_is_not_reported_as_an_xid(): + err = _err("GpuNvlinkError", execution_pb2.ExecutionError.SYSTEM, SXID_SENTENCE) + err.gpu_fault.CopyFrom( + execution_pb2.GpuFault( + kind=execution_pb2.GpuFault.KIND_SXID, + code=22, + severity=execution_pb2.GpuFault.SEVERITY_CRITICAL, + pci_bus_id="0000:3b:00.0", + ) + ) + + exc = convert_error_to_native(err) + + assert exc.fault_kind == "sxid" + assert exc.sxid == 22 + assert exc.xid is None + # An unresolved GPU index is absent rather than zero, which is a GPU of its own. + assert exc.gpu_index is None + + +def test_no_typed_fault_falls_back_to_the_sentence(): + exc = convert_error_to_native(_err("GpuXidError", execution_pb2.ExecutionError.USER, USER_SENTENCE)) + + assert isinstance(exc, flyte.errors.GPUFaultUserError) + assert exc.xid == 31 + assert exc.gpu_uuid == "GPU-abc" + assert exc.node is None + + +def test_typed_fault_with_nothing_filled_in_reads_as_unknown(): + err = _err("GpuXidError", execution_pb2.ExecutionError.USER, "container exited with code 137") + err.gpu_fault.CopyFrom(execution_pb2.GpuFault()) + + exc = convert_error_to_native(err) + + assert isinstance(exc, flyte.errors.GPUFaultUserError) + assert exc.fault_kind is None + assert exc.xid is None + assert exc.severity is None