Skip to content

Commit 66f5e1e

Browse files
kthuikthui
authored
Enable Python execute() to return Triton error code (triton-inference-server#292)
* Add error code to pb error * Return error code on pb error * Add to error code param to Python * Move ErrorCode into TritonError * Expose ErrorCode internal in TritonError * Unify PbError constructors
1 parent 6f369ef commit 66f5e1e

5 files changed

Lines changed: 106 additions & 19 deletions

File tree

‎src/infer_response.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -243,8 +243,8 @@ InferResponse::Send(
243243
});
244244

245245
if (HasError()) {
246-
*response_error = TRITONSERVER_ErrorNew(
247-
TRITONSERVER_ERROR_INTERNAL, Error()->Message().c_str());
246+
*response_error =
247+
TRITONSERVER_ErrorNew(Error()->Code(), Error()->Message().c_str());
248248
return;
249249
}
250250

‎src/pb_error.cc‎

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -27,6 +27,13 @@
2727
#include "pb_error.h"
2828

2929
namespace triton { namespace backend { namespace python {
30+
31+
TRITONSERVER_Error_Code
32+
PbError::Code()
33+
{
34+
return code_;
35+
}
36+
3037
const std::string&
3138
PbError::Message()
3239
{
@@ -43,22 +50,36 @@ void
4350
PbError::SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool)
4451
{
4552
message_shm_ = PbString::Create(shm_pool, message_);
46-
shm_handle_ = message_shm_->ShmHandle();
53+
error_shm_ = shm_pool->Construct<PbErrorShm>();
54+
error_shm_.data_->code = code_;
55+
error_shm_.data_->message_shm_handle = message_shm_->ShmHandle();
56+
shm_handle_ = error_shm_.handle_;
4757
}
4858

4959
std::shared_ptr<PbError>
5060
PbError::LoadFromSharedMemory(
5161
std::unique_ptr<SharedMemoryManager>& shm_pool,
5262
bi::managed_external_buffer::handle_t shm_handle)
5363
{
54-
std::unique_ptr<PbString> message_shm =
55-
PbString::LoadFromSharedMemory(shm_pool, shm_handle);
56-
return std::shared_ptr<PbError>(new PbError(message_shm));
64+
AllocatedSharedMemory<PbErrorShm> error_shm =
65+
shm_pool->Load<PbErrorShm>(shm_handle);
66+
std::unique_ptr<PbString> message_shm = PbString::LoadFromSharedMemory(
67+
shm_pool, error_shm.data_->message_shm_handle);
68+
69+
TRITONSERVER_Error_Code code = error_shm.data_->code;
70+
std::string message = message_shm->String();
71+
72+
return std::shared_ptr<PbError>(new PbError(
73+
std::move(message_shm), std::move(error_shm), code, std::move(message)));
5774
}
5875

59-
PbError::PbError(std::unique_ptr<PbString>& message_shm)
76+
PbError::PbError(
77+
std::shared_ptr<PbString>&& message_shm,
78+
AllocatedSharedMemory<PbErrorShm>&& error_shm, TRITONSERVER_Error_Code code,
79+
std::string&& message)
80+
: message_shm_(std::move(message_shm)), error_shm_(std::move(error_shm)),
81+
code_(code), message_(std::move(message))
6082
{
61-
message_shm_ = std::move(message_shm);
62-
message_ = message_shm_->String();
6383
}
84+
6485
}}} // namespace triton::backend::python

‎src/pb_error.h‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -32,21 +32,44 @@
3232
#include "pb_utils.h"
3333

3434
namespace triton { namespace backend { namespace python {
35+
36+
struct PbErrorShm {
37+
TRITONSERVER_Error_Code code;
38+
bi::managed_external_buffer::handle_t message_shm_handle;
39+
};
40+
3541
class PbError {
3642
public:
37-
PbError(const std::string& message) : message_(message) {}
43+
PbError(
44+
const std::string& message,
45+
TRITONSERVER_Error_Code code = TRITONSERVER_ERROR_INTERNAL)
46+
: code_(code), message_(message)
47+
{
48+
}
49+
DISALLOW_COPY_AND_ASSIGN(PbError);
50+
51+
TRITONSERVER_Error_Code Code();
3852
const std::string& Message();
53+
3954
void SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool);
4055
bi::managed_external_buffer::handle_t ShmHandle();
56+
4157
static std::shared_ptr<PbError> LoadFromSharedMemory(
4258
std::unique_ptr<SharedMemoryManager>& shm_pool,
4359
bi::managed_external_buffer::handle_t handle);
44-
DISALLOW_COPY_AND_ASSIGN(PbError);
4560

4661
private:
47-
PbError(std::unique_ptr<PbString>& pb_error);
48-
std::string message_;
62+
PbError(
63+
std::shared_ptr<PbString>&& message_shm,
64+
AllocatedSharedMemory<PbErrorShm>&& error_shm,
65+
TRITONSERVER_Error_Code code, std::string&& message);
66+
4967
std::shared_ptr<PbString> message_shm_;
68+
AllocatedSharedMemory<PbErrorShm> error_shm_;
5069
bi::managed_external_buffer::handle_t shm_handle_;
70+
71+
TRITONSERVER_Error_Code code_;
72+
std::string message_;
5173
};
74+
5275
}}}; // namespace triton::backend::python

‎src/pb_stub.cc‎

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1346,9 +1346,52 @@ Logger::BackendLoggingActive()
13461346

13471347
PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
13481348
{
1349-
py::class_<PbError, std::shared_ptr<PbError>>(module, "TritonError")
1350-
.def(py::init<std::string>())
1351-
.def("message", &PbError::Message);
1349+
py::class_<PbError, std::shared_ptr<PbError>> triton_error(
1350+
module, "TritonError");
1351+
py::enum_<TRITONSERVER_Error_Code>(triton_error, "__ErrorCode")
1352+
.value("UNKNOWN", TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN)
1353+
.value("INTERNAL", TRITONSERVER_Error_Code::TRITONSERVER_ERROR_INTERNAL)
1354+
.value("NOT_FOUND", TRITONSERVER_Error_Code::TRITONSERVER_ERROR_NOT_FOUND)
1355+
.value(
1356+
"INVALID_ARG",
1357+
TRITONSERVER_Error_Code::TRITONSERVER_ERROR_INVALID_ARG)
1358+
.value(
1359+
"UNAVAILABLE",
1360+
TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNAVAILABLE)
1361+
.value(
1362+
"UNSUPPORTED",
1363+
TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNSUPPORTED)
1364+
.value(
1365+
"ALREADY_EXISTS",
1366+
TRITONSERVER_Error_Code::TRITONSERVER_ERROR_ALREADY_EXISTS)
1367+
.export_values();
1368+
triton_error.def_property_readonly_static(
1369+
"UNKNOWN",
1370+
[](py::object /* self */) { return TRITONSERVER_ERROR_UNKNOWN; });
1371+
triton_error.def_property_readonly_static(
1372+
"INTERNAL",
1373+
[](py::object /* self */) { return TRITONSERVER_ERROR_INTERNAL; });
1374+
triton_error.def_property_readonly_static(
1375+
"NOT_FOUND",
1376+
[](py::object /* self */) { return TRITONSERVER_ERROR_NOT_FOUND; });
1377+
triton_error.def_property_readonly_static(
1378+
"INVALID_ARG",
1379+
[](py::object /* self */) { return TRITONSERVER_ERROR_INVALID_ARG; });
1380+
triton_error.def_property_readonly_static(
1381+
"UNAVAILABLE",
1382+
[](py::object /* self */) { return TRITONSERVER_ERROR_UNAVAILABLE; });
1383+
triton_error.def_property_readonly_static(
1384+
"UNSUPPORTED",
1385+
[](py::object /* self */) { return TRITONSERVER_ERROR_UNSUPPORTED; });
1386+
triton_error.def_property_readonly_static(
1387+
"ALREADY_EXISTS",
1388+
[](py::object /* self */) { return TRITONSERVER_ERROR_ALREADY_EXISTS; });
1389+
triton_error.def(
1390+
py::init<const std::string&, TRITONSERVER_Error_Code>(),
1391+
py::arg("message").none(false),
1392+
py::arg("code").none(false) = TRITONSERVER_ERROR_INTERNAL);
1393+
triton_error.def("code", &PbError::Code);
1394+
triton_error.def("message", &PbError::Message);
13521395

13531396
py::class_<PreferredMemory, std::shared_ptr<PreferredMemory>>(
13541397
module, "PreferredMemory")

‎src/python_be.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1456,7 +1456,7 @@ ModelInstanceState::ProcessRequests(
14561456
false /* open_cuda_handle */);
14571457
if (infer_response->HasError()) {
14581458
TRITONSERVER_Error* err = TRITONSERVER_ErrorNew(
1459-
TRITONSERVER_ERROR_INTERNAL,
1459+
infer_response->Error()->Code(),
14601460
infer_response->Error()->Message().c_str());
14611461

14621462
LOG_IF_ERROR(

0 commit comments

Comments
 (0)