fix: Fix L0_request_cancellation and L0_backend_output_detail byte si…

…ze checks (#363)
triton-inference-server · Jun 1, 2024 · c26629f · c26629f
1 parent bbcd781
commit c26629f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/src/test/backend_output_detail_test.cc b/src/test/backend_output_detail_test.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -183,24 +183,28 @@ class BackendOutputDetailTest : public ::testing::Test {
 
     std::vector<int64_t> input0_shape({16});
     std::vector<int64_t> input1_shape({16});
+
+    const auto input0_byte_size = sizeof(input0_data_[0]) * input0_data_.size();
+    const auto input1_byte_size = sizeof(input1_data_[0]) * input1_data_.size();
+
     FAIL_TEST_IF_ERR(
         TRITONSERVER_InferenceRequestAddInput(
-            irequest_, "INPUT0", TRITONSERVER_TYPE_FP32, &input0_shape[0],
+            irequest_, "INPUT0", TRITONSERVER_TYPE_FP32, input0_shape.data(),
             input0_shape.size()),
         "setting input0 for the request");
     FAIL_TEST_IF_ERR(
         TRITONSERVER_InferenceRequestAppendInputData(
-            irequest_, "INPUT0", &input0_data_[0], input0_data_.size(),
+            irequest_, "INPUT0", input0_data_.data(), input0_byte_size,
             TRITONSERVER_MEMORY_CPU, 0),
         "assigning INPUT data");
     FAIL_TEST_IF_ERR(
         TRITONSERVER_InferenceRequestAddInput(
-            irequest_, "INPUT1", TRITONSERVER_TYPE_FP32, &input1_shape[0],
+            irequest_, "INPUT1", TRITONSERVER_TYPE_FP32, input1_shape.data(),
             input1_shape.size()),
         "setting input1 for the request");
     FAIL_TEST_IF_ERR(
         TRITONSERVER_InferenceRequestAppendInputData(
-            irequest_, "INPUT1", &input1_data_[0], input1_data_.size(),
+            irequest_, "INPUT1", input1_data_.data(), input1_byte_size,
             TRITONSERVER_MEMORY_CPU, 0),
         "assigning INPUT1 data");
   }

diff --git a/src/test/request_cancellation_test.cc b/src/test/request_cancellation_test.cc
@@ -1,4 +1,4 @@
-// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -155,10 +155,12 @@ class RequestCancellationTest : public ::testing::Test {
 
     std::vector<int64_t> input0_shape({1, 1000});
     FAIL_TEST_IF_ERR(TRITONSERVER_InferenceRequestAddInput(
-        irequest_, "INPUT0", TRITONSERVER_TYPE_INT32, &input0_shape[0],
+        irequest_, "INPUT0", TRITONSERVER_TYPE_INT32, input0_shape.data(),
         input0_shape.size()));
+
+    const auto input0_byte_size = sizeof(input0_data_[0]) * input0_data_.size();
     FAIL_TEST_IF_ERR(TRITONSERVER_InferenceRequestAppendInputData(
-        irequest_, "INPUT0", &input0_data_[0], input0_data_.size(),
+        irequest_, "INPUT0", input0_data_.data(), input0_byte_size,
         TRITONSERVER_MEMORY_CPU, 0));
   }
 
@@ -175,7 +177,7 @@ class RequestCancellationTest : public ::testing::Test {
 };
 
 TRITONSERVER_Server* RequestCancellationTest::server_ = nullptr;
-std::vector<int32_t> RequestCancellationTest::input0_data_(16, 1);
+std::vector<int32_t> RequestCancellationTest::input0_data_(1000, 1);
 
 TEST_F(RequestCancellationTest, Cancellation)
 {