Skip to content

Commit 8917c86

Browse files
pskiran1pskiran1
authored
Add Correlation Id string support for BLS (triton-inference-server#344)
* Add correlation id string support for BLS
1 parent 0413e46 commit 8917c86

File tree

9 files changed

+316
-70
lines changed

9 files changed

+316
-70
lines changed

‎CMakeLists.txt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# Copyright 2020-2024, 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
@@ -149,6 +149,8 @@ configure_file(src/libtriton_python.ldscript libtriton_python.ldscript COPYONLY)
149149

150150
set(
151151
COMMON_SRCS
152+
src/correlation_id.cc
153+
src/correlation_id.h
152154
src/infer_response.cc
153155
src/infer_response.h
154156
src/infer_request.cc

‎README.md‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
<!--
2-
# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -1125,6 +1125,8 @@ class TritonPythonModel:
11251125
# `pb_utils.InferenceRequest` supports request_id, correlation_id,
11261126
# model version, timeout and preferred_memory in addition to the
11271127
# arguments described above.
1128+
# Note: Starting from the 24.03 release, the `correlation_id` parameter
1129+
# supports both string and unsigned integer values.
11281130
# These arguments are optional. An example containing all the arguments:
11291131
# inference_request = pb_utils.InferenceRequest(model_name='model_name',
11301132
# requested_output_names=['REQUESTED_OUTPUT_1', 'REQUESTED_OUTPUT_2'],
@@ -1262,11 +1264,13 @@ class TritonPythonModel:
12621264
# `pb_utils.InferenceRequest` supports request_id, correlation_id,
12631265
# model version, timeout and preferred_memory in addition to the
12641266
# arguments described above.
1267+
# Note: Starting from the 24.03 release, the `correlation_id` parameter
1268+
# supports both string and unsigned integer values.
12651269
# These arguments are optional. An example containing all the arguments:
12661270
# inference_request = pb_utils.InferenceRequest(model_name='model_name',
12671271
# requested_output_names=['REQUESTED_OUTPUT_1', 'REQUESTED_OUTPUT_2'],
12681272
# inputs=[<list of pb_utils.Tensor objects>],
1269-
# request_id="1", correlation_id=4, model_version=1, flags=0, timeout=5,
1273+
# request_id="1", correlation_id="ex-4", model_version=1, flags=0, timeout=5,
12701274
# preferred_memory=pb_utils.PreferredMemory(
12711275
# pb_utils.TRITONSERVER_MEMORY_GPU, # or pb_utils.TRITONSERVER_MEMORY_CPU
12721276
# 0))

‎src/correlation_id.cc‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
// Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
//
3+
// Redistribution and use in source and binary forms, with or without
4+
// modification, are permitted provided that the following conditions
5+
// are met:
6+
// * Redistributions of source code must retain the above copyright
7+
// notice, this list of conditions and the following disclaimer.
8+
// * Redistributions in binary form must reproduce the above copyright
9+
// notice, this list of conditions and the following disclaimer in the
10+
// documentation and/or other materials provided with the distribution.
11+
// * Neither the name of NVIDIA CORPORATION nor the names of its
12+
// contributors may be used to endorse or promote products derived
13+
// from this software without specific prior written permission.
14+
//
15+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
#include "correlation_id.h"
28+
29+
namespace triton { namespace backend { namespace python {
30+
31+
CorrelationId::CorrelationId()
32+
: id_string_(""), id_uint_(0), id_type_(CorrelationIdDataType::UINT64)
33+
{
34+
}
35+
36+
CorrelationId::CorrelationId(const std::string& id_string)
37+
: id_string_(id_string), id_uint_(0),
38+
id_type_(CorrelationIdDataType::STRING)
39+
{
40+
}
41+
42+
CorrelationId::CorrelationId(uint64_t id_uint)
43+
: id_string_(""), id_uint_(id_uint), id_type_(CorrelationIdDataType::UINT64)
44+
{
45+
}
46+
47+
CorrelationId::CorrelationId(const CorrelationId& rhs)
48+
{
49+
id_uint_ = rhs.id_uint_;
50+
id_type_ = rhs.id_type_;
51+
id_string_ = rhs.id_string_;
52+
}
53+
54+
CorrelationId::CorrelationId(std::unique_ptr<CorrelationId>& correlation_id_shm)
55+
{
56+
id_uint_ = correlation_id_shm->id_uint_;
57+
id_type_ = correlation_id_shm->id_type_;
58+
id_string_ = correlation_id_shm->id_string_;
59+
}
60+
61+
CorrelationId&
62+
CorrelationId::operator=(const CorrelationId& rhs)
63+
{
64+
id_uint_ = rhs.id_uint_;
65+
id_type_ = rhs.id_type_;
66+
id_string_ = rhs.id_string_;
67+
return *this;
68+
}
69+
70+
void
71+
CorrelationId::SaveToSharedMemory(
72+
std::unique_ptr<SharedMemoryManager>& shm_pool)
73+
{
74+
AllocatedSharedMemory<CorrelationIdShm> correlation_id_shm =
75+
shm_pool->Construct<CorrelationIdShm>();
76+
correlation_id_shm_ptr_ = correlation_id_shm.data_.get();
77+
78+
std::unique_ptr<PbString> id_string_shm =
79+
PbString::Create(shm_pool, id_string_);
80+
81+
correlation_id_shm_ptr_->id_uint = id_uint_;
82+
correlation_id_shm_ptr_->id_string_shm_handle = id_string_shm->ShmHandle();
83+
correlation_id_shm_ptr_->id_type = id_type_;
84+
85+
// Save the references to shared memory.
86+
correlation_id_shm_ = std::move(correlation_id_shm);
87+
id_string_shm_ = std::move(id_string_shm);
88+
shm_handle_ = correlation_id_shm_.handle_;
89+
}
90+
91+
std::unique_ptr<CorrelationId>
92+
CorrelationId::LoadFromSharedMemory(
93+
std::unique_ptr<SharedMemoryManager>& shm_pool,
94+
bi::managed_external_buffer::handle_t handle)
95+
{
96+
AllocatedSharedMemory<CorrelationIdShm> correlation_id_shm =
97+
shm_pool->Load<CorrelationIdShm>(handle);
98+
CorrelationIdShm* correlation_id_shm_ptr = correlation_id_shm.data_.get();
99+
100+
std::unique_ptr<PbString> id_string_shm = PbString::LoadFromSharedMemory(
101+
shm_pool, correlation_id_shm_ptr->id_string_shm_handle);
102+
103+
return std::unique_ptr<CorrelationId>(
104+
new CorrelationId(correlation_id_shm, id_string_shm));
105+
}
106+
107+
CorrelationId::CorrelationId(
108+
AllocatedSharedMemory<CorrelationIdShm>& correlation_id_shm,
109+
std::unique_ptr<PbString>& id_string_shm)
110+
: correlation_id_shm_(std::move(correlation_id_shm)),
111+
id_string_shm_(std::move(id_string_shm))
112+
{
113+
correlation_id_shm_ptr_ = correlation_id_shm_.data_.get();
114+
shm_handle_ = correlation_id_shm_.handle_;
115+
id_string_ = id_string_shm_->String();
116+
id_uint_ = correlation_id_shm_ptr_->id_uint;
117+
id_type_ = correlation_id_shm_ptr_->id_type;
118+
}
119+
120+
}}}; // namespace triton::backend::python

‎src/correlation_id.h‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
// Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
//
3+
// Redistribution and use in source and binary forms, with or without
4+
// modification, are permitted provided that the following conditions
5+
// are met:
6+
// * Redistributions of source code must retain the above copyright
7+
// notice, this list of conditions and the following disclaimer.
8+
// * Redistributions in binary form must reproduce the above copyright
9+
// notice, this list of conditions and the following disclaimer in the
10+
// documentation and/or other materials provided with the distribution.
11+
// * Neither the name of NVIDIA CORPORATION nor the names of its
12+
// contributors may be used to endorse or promote products derived
13+
// from this software without specific prior written permission.
14+
//
15+
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
#pragma once
28+
29+
#include <string>
30+
31+
#include "pb_string.h"
32+
#include "pb_utils.h"
33+
34+
namespace triton { namespace backend { namespace python {
35+
36+
enum class CorrelationIdDataType { UINT64, STRING };
37+
38+
struct CorrelationIdShm {
39+
bi::managed_external_buffer::handle_t id_string_shm_handle;
40+
uint64_t id_uint;
41+
CorrelationIdDataType id_type;
42+
};
43+
44+
class CorrelationId {
45+
public:
46+
CorrelationId();
47+
CorrelationId(const std::string& id_string);
48+
CorrelationId(uint64_t id_uint);
49+
CorrelationId(const CorrelationId& rhs);
50+
CorrelationId(std::unique_ptr<CorrelationId>& correlation_id_shm);
51+
CorrelationId& operator=(const CorrelationId& rhs);
52+
53+
/// Save CorrelationId object to shared memory.
54+
/// \param shm_pool Shared memory pool to save the CorrelationId object.
55+
void SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool);
56+
57+
/// Create a CorrelationId object from shared memory.
58+
/// \param shm_pool Shared memory pool
59+
/// \param handle Shared memory handle of the CorrelationId.
60+
/// \return Returns the CorrelationId in the specified handle
61+
/// location.
62+
static std::unique_ptr<CorrelationId> LoadFromSharedMemory(
63+
std::unique_ptr<SharedMemoryManager>& shm_pool,
64+
bi::managed_external_buffer::handle_t handle);
65+
66+
// Function that help determine exact type of Correlation Id
67+
CorrelationIdDataType Type() const { return id_type_; }
68+
69+
// Get the value of the CorrelationId based on the type
70+
const std::string& StringValue() const { return id_string_; }
71+
uint64_t UnsignedIntValue() const { return id_uint_; }
72+
73+
bi::managed_external_buffer::handle_t ShmHandle() { return shm_handle_; }
74+
75+
private:
76+
// The private constructor for creating a CorrelationId object from shared
77+
// memory.
78+
CorrelationId(
79+
AllocatedSharedMemory<CorrelationIdShm>& correlation_id_shm,
80+
std::unique_ptr<PbString>& id_string_shm);
81+
82+
std::string id_string_;
83+
uint64_t id_uint_;
84+
CorrelationIdDataType id_type_;
85+
86+
// Shared Memory Data Structures
87+
AllocatedSharedMemory<CorrelationIdShm> correlation_id_shm_;
88+
CorrelationIdShm* correlation_id_shm_ptr_;
89+
bi::managed_external_buffer::handle_t shm_handle_;
90+
std::unique_ptr<PbString> id_string_shm_;
91+
};
92+
93+
}}}; // namespace triton::backend::python

0 commit comments

Comments
 (0)