Add request input to profile export data (#549)

* Add request_inputs field in profile export file

* Add request input data to profile export

* Add test cases

* Remove comment

* Add ticket and remove comments

* Add tests

* Remove first 4 bytes

* Add warning

* handle empty input

* Avoid access to buffer when empty

Author: nv-hwoo

src/c++/library/common.cc

@@ -182,6 +182,20 @@ InferInput::AppendFromString(const std::vector<std::string>& input)
   return AppendRaw(reinterpret_cast<const uint8_t*>(&sbuf[0]), sbuf.size());
 }
 
+Error
+InferInput::RawData(const uint8_t** buf, size_t* byte_size)
+{
+  if (bufs_.size()) {
+    // TMA-1775 - handle multi-batch case
+    *buf = bufs_[0];
+    *byte_size = buf_byte_sizes_[0];
+  } else {
+    *buf = nullptr;
+    *byte_size = 0;
+  }
+  return Error::Success;
+}
+
 Error
 InferInput::ByteSize(size_t* byte_size) const
 {

src/c++/library/common.h

@@ -334,6 +334,15 @@ class InferInput {
   /// \return Error object indicating success or failure.
   Error AppendFromString(const std::vector<std::string>& input);
 
+  /// Get access to the buffer holding raw input. Note the buffer is owned by
+  /// InferInput instance. Users can copy out the data if required to extend
+  /// the lifetime.
+  /// \param buf Returns the pointer to the start of the buffer.
+  /// \param byte_size Returns the size of buffer in bytes.
+  /// \return Error object indicating success or failure of the
+  /// request.
+  Error RawData(const uint8_t** buf, size_t* byte_size);
+
   /// Gets the size of data added into this input in bytes.
   /// \param byte_size The size of data added in bytes.
   /// \return Error object indicating success or failure.

src/c++/perf_analyzer/client_backend/client_backend.cc

@@ -505,6 +505,14 @@ InferInput::SetSharedMemory(
       pa::GENERIC_ERROR);
 }
 
+Error
+InferInput::RawData(const uint8_t** buf, size_t* byte_size)
+{
+  return Error(
+      "client backend of kind " + BackendKindToString(kind_) +
+          " does not support RawData() for InferInput",
+      pa::GENERIC_ERROR);
+}
 
 InferInput::InferInput(
     const BackendKind kind, const std::string& name,

src/c++/perf_analyzer/client_backend/client_backend.h

@@ -558,6 +558,15 @@ class InferInput {
   virtual Error SetSharedMemory(
       const std::string& name, size_t byte_size, size_t offset = 0);
 
+  /// Get access to the buffer holding raw input. Note the buffer is owned by
+  /// InferInput instance. Users can copy out the data if required to extend
+  /// the lifetime.
+  /// \param buf Returns the pointer to the start of the buffer.
+  /// \param byte_size Returns the size of buffer in bytes.
+  /// \return Error object indicating success or failure of the
+  /// request.
+  virtual Error RawData(const uint8_t** buf, size_t* byte_size);
+
  protected:
   InferInput(
       const BackendKind kind, const std::string& name,

src/c++/perf_analyzer/client_backend/openai/openai_infer_input.cc

@@ -71,6 +71,15 @@ OpenAiInferInput::AppendRaw(const uint8_t* input, size_t input_byte_size)
   return Error::Success;
 }
 
+Error
+OpenAiInferInput::RawData(const uint8_t** buf, size_t* byte_size)
+{
+  // TMA-1775 - handle multi-batch case
+  *buf = bufs_[0];
+  *byte_size = buf_byte_sizes_[0];
+  return Error::Success;
+}
+
 Error
 OpenAiInferInput::PrepareForRequest()
 {

src/c++/perf_analyzer/client_backend/openai/openai_infer_input.h

@@ -51,6 +51,8 @@ class OpenAiInferInput : public InferInput {
   Error Reset() override;
   /// See InferInput::AppendRaw()
   Error AppendRaw(const uint8_t* input, size_t input_byte_size) override;
+  /// See InferInput::RawData()
+  Error RawData(const uint8_t** buf, size_t* byte_size) override;
   /// Prepare the input to be in the form expected by an OpenAI client,
   /// must call before accessing the data.
   Error PrepareForRequest();

src/c++/perf_analyzer/client_backend/triton/triton_client_backend.cc

@@ -756,6 +756,13 @@ TritonInferInput::SetSharedMemory(
   return Error::Success;
 }
 
+Error
+TritonInferInput::RawData(const uint8_t** buf, size_t* byte_size)
+{
+  RETURN_IF_TRITON_ERROR(input_->RawData(buf, byte_size));
+  return Error::Success;
+}
+
 TritonInferInput::TritonInferInput(
     const std::string& name, const std::string& datatype)
     : InferInput(BackendKind::TRITON, name, datatype)

src/c++/perf_analyzer/client_backend/triton/triton_client_backend.h

@@ -283,6 +283,8 @@ class TritonInferInput : public InferInput {
   /// See InferInput::SetSharedMemory()
   Error SetSharedMemory(
       const std::string& name, size_t byte_size, size_t offset = 0) override;
+  /// See InferInput::RawData()
+  Error RawData(const uint8_t** buf, size_t* byte_size) override;
 
  private:
   explicit TritonInferInput(

src/c++/perf_analyzer/infer_context.cc

@@ -112,6 +112,10 @@ InferContext::SendRequest(
   }
 
   thread_stat_->num_sent_requests_++;
+
+  // Parse the request inputs to save in the profile export file
+  RequestRecord::RequestInput request_inputs{GetInputs()};
+
   if (async_) {
     uint64_t unique_request_id{(thread_id_ << 48) | ((request_id << 16) >> 16)};
     infer_data_.options_->request_id_ = std::to_string(unique_request_id);
@@ -120,6 +124,7 @@ InferContext::SendRequest(
       auto it = async_req_map_
                     .emplace(infer_data_.options_->request_id_, RequestRecord())
                     .first;
+      it->second.request_inputs_ = {request_inputs};
       it->second.start_time_ = std::chrono::system_clock::now();
       it->second.sequence_end_ = infer_data_.options_->sequence_end_;
       it->second.delayed_ = delayed;
@@ -149,10 +154,10 @@ InferContext::SendRequest(
         &results, *(infer_data_.options_), infer_data_.valid_inputs_,
         infer_data_.outputs_);
     thread_stat_->idle_timer.Stop();
-    RequestRecord::ResponseOutput response_output{};
+    RequestRecord::ResponseOutput response_outputs{};
     if (results != nullptr) {
       if (thread_stat_->status_.IsOk()) {
-        response_output = GetOutput(*results);
+        response_outputs = GetOutputs(*results);
         thread_stat_->status_ = ValidateOutputs(results);
       }
       delete results;
@@ -169,8 +174,9 @@ InferContext::SendRequest(
       std::lock_guard<std::mutex> lock(thread_stat_->mu_);
       auto total = end_time_sync - start_time_sync;
       thread_stat_->request_records_.emplace_back(RequestRecord(
-          start_time_sync, std::move(end_time_syncs), {response_output},
-          infer_data_.options_->sequence_end_, delayed, sequence_id, false));
+          start_time_sync, std::move(end_time_syncs), {request_inputs},
+          {response_outputs}, infer_data_.options_->sequence_end_, delayed,
+          sequence_id, false));
       thread_stat_->status_ =
           infer_backend_->ClientInferStat(&(thread_stat_->contexts_stat_[id_]));
       if (!thread_stat_->status_.IsOk()) {
@@ -180,15 +186,36 @@ InferContext::SendRequest(
   }
 }
 
+const RequestRecord::RequestInput
+InferContext::GetInputs()
+{
+  RequestRecord::RequestInput input{};
+  for (const auto& request_input : infer_data_.valid_inputs_) {
+    const uint8_t* buf{nullptr};
+    size_t byte_size{0};
+    std::string data_type{request_input->Datatype()};
+    request_input->RawData(&buf, &byte_size);
+
+    // The first 4 bytes of BYTES data is a 32-bit integer to indicate the size
+    // of the rest of the data (which we already know based on byte_size). It
+    // should be ignored here, as it isn't part of the actual request
+    if (data_type == "BYTES" && byte_size >= 4) {
+      buf += 4;
+      byte_size -= 4;
+    }
+    input.emplace(request_input->Name(), RecordData(buf, byte_size, data_type));
+  }
+  return input;
+}
+
 const RequestRecord::ResponseOutput
-InferContext::GetOutput(const cb::InferResult& infer_result)
+InferContext::GetOutputs(const cb::InferResult& infer_result)
 {
   RequestRecord::ResponseOutput output{};
   for (const auto& requested_output : infer_data_.outputs_) {
     const uint8_t* buf{nullptr};
     size_t byte_size{0};
     infer_result.RawData(requested_output->Name(), &buf, &byte_size);
-
     // The first 4 bytes of BYTES data is a 32-bit integer to indicate the size
     // of the rest of the data (which we already know based on byte_size). It
     // should be ignored here, as it isn't part of the actual response
@@ -282,7 +309,7 @@ InferContext::AsyncCallbackFuncImpl(cb::InferResult* result)
         }
         it->second.response_timestamps_.push_back(
             std::chrono::system_clock::now());
-        it->second.response_outputs_.push_back(GetOutput(*result));
+        it->second.response_outputs_.push_back(GetOutputs(*result));
         num_responses_++;
         if (is_null_response == true) {
           it->second.has_null_last_response_ = true;
@@ -296,9 +323,9 @@ InferContext::AsyncCallbackFuncImpl(cb::InferResult* result)
           has_received_final_response_ = is_final_response;
           thread_stat_->request_records_.emplace_back(
               it->second.start_time_, it->second.response_timestamps_,
-              it->second.response_outputs_, it->second.sequence_end_,
-              it->second.delayed_, it->second.sequence_id_,
-              it->second.has_null_last_response_);
+              it->second.request_inputs_, it->second.response_outputs_,
+              it->second.sequence_end_, it->second.delayed_,
+              it->second.sequence_id_, it->second.has_null_last_response_);
           infer_backend_->ClientInferStat(&(thread_stat_->contexts_stat_[id_]));
           thread_stat_->cb_status_ = ValidateOutputs(result);
           async_req_map_.erase(request_id);

src/c++/perf_analyzer/infer_context.h

@@ -185,7 +185,9 @@ class InferContext {
   std::function<void(uint32_t)> async_callback_finalize_func_ = nullptr;
 
  private:
-  const RequestRecord::ResponseOutput GetOutput(
+  const RequestRecord::RequestInput GetInputs();
+
+  const RequestRecord::ResponseOutput GetOutputs(
       const cb::InferResult& infer_result);
 
   const uint32_t id_{0};