Refactor

Signed-off-by: Shunkang <[email protected]>

Author: Shunkang

cpp/include/tensorrt_llm/common/tllmException.h

@@ -37,7 +37,7 @@
 #define NEW_TLLM_EXCEPTION(...)                                                                                        \
     tensorrt_llm::common::TllmException(__FILE__, __LINE__, tensorrt_llm::common::fmtstr(__VA_ARGS__).c_str())
 
-#define NEW_TLLM_REQUEST_SPECIFIC_EXCEPTION_WITH_ERROR_CODE(requestID, errorCode, ...)                                 \
+#define TLLM_REQUEST_EXCEPTION(requestID, errorCode, ...)                                                              \
     tensorrt_llm::common::RequestSpecificException(                                                                    \
         __FILE__, __LINE__, tensorrt_llm::common::fmtstr(__VA_ARGS__).c_str(), requestID, errorCode)
 
@@ -84,6 +84,7 @@ class TllmException : public std::runtime_error
 {
     throw TllmException(file, line, fmtstr("[TensorRT-LLM][ERROR] Assertion failed: %s", info.c_str()).c_str());
 }
+
 class RequestSpecificException : public std::runtime_error
 {
 public:

cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

@@ -194,8 +194,7 @@ class DataResponder::Impl
         catch (tensorrt_llm::common::RequestSpecificException const& e)
         {
             TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s ", e.what());
-            auto new_exception
-                = NEW_TLLM_REQUEST_SPECIFIC_EXCEPTION_WITH_ERROR_CODE(id, e.getErrorCode(), "%s", e.what());
+            auto new_exception = TLLM_REQUEST_EXCEPTION(id, e.getErrorCode(), "%s", e.what());
             resp.mPromise.set_exception(std::make_exception_ptr(new_exception));
         }
         catch (std::exception const& e)
@@ -506,10 +505,10 @@ class DataRequester::Impl
                 }
                 catch (tensorrt_llm::common::RequestSpecificException const& err)
                 {
-                    TLLM_LOG_ERROR("Exception in DataRequester request(): request id:%ld , request context id:%ld : %s",
+                    TLLM_LOG_ERROR("Exception in DataRequester request(): request id:%zu , request context id:%zu : %s",
                         requestAndPromise.mRequest->mRequestId,
                         requestAndPromise.mRequest->getContextPhaseParams().value().getReqId(), err.what());
-                    auto new_exception = NEW_TLLM_REQUEST_SPECIFIC_EXCEPTION_WITH_ERROR_CODE(
+                    auto new_exception = TLLM_REQUEST_EXCEPTION(
                         requestAndPromise.mRequest->mRequestId, err.getErrorCode(), "%s", err.what());
                     requestAndPromise.mPromise->set_exception(std::make_exception_ptr(new_exception));
                 }

cpp/tensorrt_llm/common/tllmException.cpp

@@ -17,6 +17,7 @@
 #include "tensorrt_llm/common/tllmException.h"
 #include "tensorrt_llm/common/stringUtils.h"
 
+#include <cinttypes>
 #include <cstdlib>
 #if !defined(_MSC_VER)
 #include <cxxabi.h>
@@ -108,7 +109,7 @@ std::string TllmException::demangle(char const* name)
 
 RequestSpecificException::RequestSpecificException(
     std::string const& file, std::size_t line, char const* msg, uint64_t requestID, RequestErrorCode errorCode)
-    : std::runtime_error{fmtstr("%s (Request ID: %lu, Error Code: %u) (%s:%zu)", msg, requestID,
+    : std::runtime_error{fmtstr("%s (Request ID: %" PRIu64 ", Error Code: %u) (%s:%zu)", msg, requestID,
         static_cast<uint32_t>(errorCode), file.c_str(), line)}
     , mRequestID{requestID}
     , mErrorCode{errorCode}

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -8,7 +8,7 @@
 import traceback
 import weakref
 from contextlib import contextmanager
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from typing import Iterable, List, Optional, Tuple, Union
 
 import torch
 
@@ -1552,8 +1552,10 @@ def _handle_errors(self,
                 client_id=request.py_client_id)
 
         if request_ids is not None:
-            for req_id in request_ids:
-                self.active_requests.remove(req_id)
+            req_id_set = set(request_ids)
+            for request in self.active_requests:
+                if request.py_request_id in req_id_set:
+                    self.active_requests.remove(request)
         else:
             self.active_requests.clear()
         self._enqueue_responses(error_responses)