[https://nvbugs/6482589][fix] Make CuError pickle-safe across process boundaries

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -512,6 +512,43 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
     static nb::object sLogicError = nb::exception<kv::LogicError>(m, "LogicError");
     static nb::object sResourceBusyError = nb::exception<kv::ResourceBusyError>(m, "ResourceBusyError");
     static nb::object sOutOfPagesError = nb::exception<kv::OutOfPagesError>(m, "OutOfPagesError");
+    static nb::object sCuError = nb::exception<kv::CuError>(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<kv::CuError> 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<int>(e.errorCode));
+                }
+                catch (nb::python_error const&)
+                {
+                    PyErr_Clear();
+                    code = nb::cast(static_cast<int>(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)
             {

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)

tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py

@@ -51,6 +51,9 @@ def __init__(self, error_code: drv.CUresult) -> None:
             err_str = "<Failed to get error string with cuGetErrorString>"
         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