From 28f1c9b87135c591b8f6b9920a05ae40a7bd3402 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Fri, 24 Jul 2026 09:20:28 +0000 Subject: [PATCH] [NVBUG-6482589][fix] Make CuError pickle-safe across process boundaries CuError could not survive crossing a process boundary (e.g. disaggregated serving). The pure-Python CuError.__init__ stores a formatted message string in .args, so the default Exception pickling reconstructs it by calling the constructor with that string. The constructor then feeds the string to cuGetErrorString, which fails with "invalid literal for int() with base 10". Fixes: - Python backend: add CuError.__reduce__ so the exception is reconstructed from the numeric CUDA error code instead of the formatted message string. - C++ backend: bind CuError in nanobind (it was never bound and fell back to RuntimeError, losing type identity and error_code). Add a custom exception translator that carries the numeric error_code onto the Python instance. The bound type uses the plain Exception.__init__, so BaseException.__reduce__ round-trips both the message (.args) and error_code (__dict__) with no custom __reduce__ needed. - Dispatcher: promote CuError to a first-class C++ port (_cpp.CuError) instead of the getattr(..., RuntimeError) fallback. Signed-off-by: Yao Yao --- .../batch_manager/kvCacheManagerV2.cpp | 39 +++++++++++++++++++ .../runtime/kv_cache_manager_v2/__init__.py | 2 +- .../kv_cache_manager_v2/_exceptions.py | 3 ++ 3 files changed, 43 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a5fcbdd22c2e..62faf997cf1b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -512,6 +512,43 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) static nb::object sLogicError = nb::exception(m, "LogicError"); static nb::object sResourceBusyError = nb::exception(m, "ResourceBusyError"); static nb::object sOutOfPagesError = nb::exception(m, "OutOfPagesError"); + static nb::object sCuError = nb::exception(m, "CuError"); + // Default attribute so the class mirrors the pure-Python CuError surface. + sCuError.attr("error_code") = nb::none(); + + // Translate kv::CuError so the Python instance carries the numeric CUDA + // error code (mirrors the pure-Python CuError.error_code). Registered after + // the nb::exception auto-translator so it is tried first. + nb::register_exception_translator( + [](std::exception_ptr const& p, void*) + { + try + { + if (p) + { + std::rethrow_exception(p); + } + } + catch (kv::CuError const& e) + { + nb::object inst = sCuError(nb::str(e.what())); + // Match Python's error_code type (cuda.bindings.driver.CUresult) + // when available; fall back to a plain int otherwise. + nb::object code; + try + { + nb::object cuResult = nb::module_::import_("cuda.bindings.driver").attr("CUresult"); + code = cuResult(static_cast(e.errorCode)); + } + catch (nb::python_error const&) + { + PyErr_Clear(); + code = nb::cast(static_cast(e.errorCode)); + } + inst.attr("error_code") = code; + PyErr_SetObject(sCuError.ptr(), inst.ptr()); + } + }); // Map kv::AssertionError to Python's builtin AssertionError so shared tests see // the same exception type as the pure-Python backend (which uses `assert`). @@ -522,7 +559,9 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) try { if (p) + { std::rethrow_exception(p); + } } catch (kv::AssertionError const& e) { diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 2d5b90468bdf..7b9322b1f63c 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -215,11 +215,11 @@ class _KVCacheManagerConfigFieldSpec: _cpp_introspection = getattr(_cpp, "_introspection", None) _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names) PlannedDropHandle = _cpp.PlannedDropHandle + CuError = _cpp.CuError # Symbols added on main that are not yet ported to the C++ backend. # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) - CuError = getattr(_cpp, "CuError", RuntimeError) OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py index 20c59e903020..7c037cb36d3d 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py @@ -51,6 +51,9 @@ def __init__(self, error_code: drv.CUresult) -> None: err_str = "" super().__init__(f"CUDA driver error: {error_code} ({err_str})") + def __reduce__(self) -> tuple[type["CuError"], tuple[drv.CUresult]]: + return (self.__class__, (self.error_code,)) + class ResourceBusyError(Exception): pass