Add request parameters API (#174)

* Add request parameters API

* Fix typos

* Fix up

* Fix name of InferenceRequestSet*

* Rename key->value

* Fix build variants

Author: Tabrizian

include/triton/core/tritonbackend.h

@@ -27,6 +27,7 @@
 
 #include <stddef.h>
 #include <stdint.h>
+
 #include "triton/core/tritonserver.h"
 
 #ifdef __cplusplus
@@ -93,7 +94,7 @@ struct TRITONBACKEND_Batcher;
 ///   }
 ///
 #define TRITONBACKEND_API_VERSION_MAJOR 1
-#define TRITONBACKEND_API_VERSION_MINOR 11
+#define TRITONBACKEND_API_VERSION_MINOR 12
 
 /// Get the TRITONBACKEND API version supported by Triton. This value
 /// can be compared against the TRITONBACKEND_API_VERSION_MAJOR and
@@ -406,6 +407,36 @@ TRITONBACKEND_RequestCorrelationIdString(
 TRITONBACKEND_DECLSPEC TRITONSERVER_Error* TRITONBACKEND_RequestFlags(
     TRITONBACKEND_Request* request, uint32_t* flags);
 
+/// Get the number of parameters specified in the inference request.
+///
+/// \param request The inference request.
+/// \param count Returns the number of parameters.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error* TRITONBACKEND_RequestParameterCount(
+    TRITONBACKEND_Request* request, uint32_t* count);
+
+/// Get a request parameters by index. The order of parameters in a given
+/// request is not necessarily consistent with other requests, even if
+/// the requests are in the same batch. As a result, you can not
+/// assume that an index obtained from one request will point to the
+/// same parameter in a different request.
+///
+/// The lifetime of the returned parameter object matches that of the
+/// request and so the parameter object should not be accessed after the
+/// request object is released.
+///
+/// \param request The inference request.
+/// \param index The index of the parameter. Must be 0 <= index <
+/// count, where count is the value returned by
+/// TRITONBACKEND_RequestParameterCount.
+/// \param key Returns the key of the parameter.
+/// \param type Returns the type of the parameter.
+/// \param vvalue Returns a pointer to the parameter value.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error* TRITONBACKEND_RequestParameter(
+    TRITONBACKEND_Request* request, const uint32_t index, const char** key,
+    TRITONSERVER_ParameterType* type, const void** vvalue);
+
 /// Get the number of input tensors specified in the request.
 ///
 /// \param request The inference request.
@@ -639,7 +670,7 @@ TRITONBACKEND_DECLSPEC TRITONSERVER_Error*
 TRITONBACKEND_ResponseSetIntParameter(
     TRITONBACKEND_Response* response, const char* name, const int64_t value);
 
-/// Set an boolean parameter in the response.
+/// Set a boolean parameter in the response.
 ///
 /// \param response The response.
 /// \param name The name of the parameter.

include/triton/core/tritonserver.h

@@ -91,7 +91,7 @@ struct TRITONSERVER_MetricFamily;
 ///   }
 ///
 #define TRITONSERVER_API_VERSION_MAJOR 1
-#define TRITONSERVER_API_VERSION_MINOR 20
+#define TRITONSERVER_API_VERSION_MINOR 21
 
 /// Get the TRITONBACKEND API version supported by the Triton shared
 /// library. This value can be compared against the
@@ -1283,6 +1283,36 @@ TRITONSERVER_InferenceRequestSetResponseCallback(
     TRITONSERVER_InferenceResponseCompleteFn_t response_fn,
     void* response_userp);
 
+/// Set a string parameter in the request.
+///
+/// \param request The request.
+/// \param key The name of the parameter.
+/// \param value The value of the parameter.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetStringParameter(
+    TRITONSERVER_InferenceRequest* request, const char* key, const char* value);
+
+/// Set an integer parameter in the request.
+///
+/// \param request The request.
+/// \param key The name of the parameter.
+/// \param value The value of the parameter.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetIntParameter(
+    TRITONSERVER_InferenceRequest* request, const char* key, const int64_t value);
+
+/// Set a boolean parameter in the request.
+///
+/// \param request The request.
+/// \param key The name of the parameter.
+/// \param value The value of the parameter.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetBoolParameter(
+    TRITONSERVER_InferenceRequest* request, const char* key, const bool value);
+
 /// TRITONSERVER_InferenceResponse
 ///
 /// Object representing an inference response. The inference response

src/backend_model.cc

@@ -27,6 +27,7 @@
 #include "backend_model.h"
 
 #include <vector>
+
 #include "backend_config.h"
 #include "backend_model_instance.h"
 #include "dynamic_batch_scheduler.h"
@@ -1236,6 +1237,43 @@ TRITONBACKEND_ResponseSend(
   return nullptr;  // success
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONBACKEND_RequestParameterCount(
+    TRITONBACKEND_Request* request, uint32_t* count)
+{
+  InferenceRequest* lrequest = reinterpret_cast<InferenceRequest*>(request);
+
+  const auto& parameters = lrequest->Parameters();
+  *count = parameters.size();
+
+  return nullptr;  // Success
+}
+
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error*
+TRITONBACKEND_RequestParameter(
+    TRITONBACKEND_Request* request, const uint32_t index, const char** key,
+    TRITONSERVER_ParameterType* type, const void** vvalue)
+{
+  InferenceRequest* lrequest = reinterpret_cast<InferenceRequest*>(request);
+
+  const auto& parameters = lrequest->Parameters();
+  if (index >= parameters.size()) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        ("out of bounds index " + std::to_string(index) +
+            std::string(": request has ") + std::to_string(parameters.size()) +
+            " parameters").c_str());
+  }
+
+  const InferenceParameter& param = parameters[index];
+
+  *key = param.Name().c_str();
+  *type = param.Type();
+  *vvalue = param.ValuePointer();
+
+  return nullptr;  // Success
+}
+
 ///
 /// TRITONBACKEND_Input
 ///

src/infer_request.cc

@@ -28,6 +28,7 @@
 
 #include <algorithm>
 #include <deque>
+
 #include "model.h"
 #include "model_config_utils.h"
 #include "server.h"
@@ -130,6 +131,27 @@ InferenceRequest::SetPriority(uint32_t p)
   }
 }
 
+Status
+InferenceRequest::AddParameter(const char* name, const char* value)
+{
+  parameters_.emplace_back(name, value);
+  return Status::Success;
+}
+
+Status
+InferenceRequest::AddParameter(const char* name, const int64_t value)
+{
+  parameters_.emplace_back(name, value);
+  return Status::Success;
+}
+
+Status
+InferenceRequest::AddParameter(const char* name, const bool value)
+{
+  parameters_.emplace_back(name, value);
+  return Status::Success;
+}
+
 #ifdef TRITON_ENABLE_TRACING
 Status
 InferenceRequest::TraceInputTensors(

src/infer_request.h

@@ -339,6 +339,16 @@ class InferenceRequest {
       TRITONSERVER_InferenceTraceActivity activity, const std::string& msg);
 #endif  // TRITON_ENABLE_TRACING
 
+  // Add an parameter to the request.
+  Status AddParameter(const char* name, const char* value);
+  Status AddParameter(const char* name, const int64_t value);
+  Status AddParameter(const char* name, const bool value);
+  const std::deque<InferenceParameter>& Parameters() const
+  {
+    return parameters_;
+  }
+
+
   // The original inputs are the inputs added to the request before
   // the inference execution (that is before
   // TRITONSERVER_ServerInferAsync is called). Once execution has
@@ -736,6 +746,10 @@ class InferenceRequest {
   // Whether the stats of the request should be collected.
   bool collect_stats_;
 
+  // The parameters of the request. Use a deque so that there is no
+  // reallocation.
+  std::deque<InferenceParameter> parameters_;
+
 #ifdef TRITON_ENABLE_STATS
   uint64_t request_start_ns_;
   InferenceStatsAggregator* secondary_stats_aggregator_ = nullptr;

src/server.cc

@@ -34,7 +34,6 @@
 #include <memory>
 #include <utility>
 #include <vector>
-
 #include "backend_manager.h"
 #include "constants.h"
 #include "cuda_utils.h"
@@ -90,6 +89,7 @@ InferenceServer::InferenceServer()
   extensions_.push_back("system_shared_memory");
   extensions_.push_back("cuda_shared_memory");
   extensions_.push_back("binary_tensor_data");
+  extensions_.push_back("parameters");
 #ifdef TRITON_ENABLE_STATS
   extensions_.push_back("statistics");
 #endif  // TRITON_ENABLE_STATS

src/tritonserver.cc

@@ -26,6 +26,7 @@
 
 #include <string>
 #include <vector>
+
 #include "buffer_attributes.h"
 #include "cuda_utils.h"
 #include "infer_parameter.h"
@@ -1877,6 +1878,34 @@ TRITONSERVER_InferenceRequestSetResponseCallback(
   return nullptr;  // Success
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetStringParameter(
+    TRITONSERVER_InferenceRequest* request, const char* name, const char* value)
+{
+  tc::InferenceRequest* tr = reinterpret_cast<tc::InferenceRequest*>(request);
+  RETURN_IF_STATUS_ERROR(tr->AddParameter(name, value));
+  return nullptr;  // success
+}
+
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetIntParameter(
+    TRITONSERVER_InferenceRequest* request, const char* name,
+    const int64_t value)
+{
+  tc::InferenceRequest* tr = reinterpret_cast<tc::InferenceRequest*>(request);
+  RETURN_IF_STATUS_ERROR(tr->AddParameter(name, value));
+  return nullptr;  // success
+}
+
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceRequestSetBoolParameter(
+    TRITONSERVER_InferenceRequest* request, const char* name, const bool value)
+{
+  tc::InferenceRequest* tr = reinterpret_cast<tc::InferenceRequest*>(request);
+  RETURN_IF_STATUS_ERROR(tr->AddParameter(name, value));
+  return nullptr;  // success
+}
+
 //
 // TRITONSERVER_InferenceResponse
 //

src/tritonserver_stub.cc

@@ -683,6 +683,26 @@ TRITONBACKEND_RequestRelease()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceRequestSetBoolParameter()
+{
+}
+TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceRequestSetIntParameter()
+{
+}
+TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceRequestSetStringParameter()
+{
+}
+TRITONAPI_DECLSPEC void
+TRITONBACKEND_RequestParameter()
+{
+}
+TRITONAPI_DECLSPEC void
+TRITONBACKEND_RequestParameterCount()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONBACKEND_ResponseFactoryNew()
 {
 }