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