test: Test safe GetElementCount and GetByteSize APIs at common and core repos (#8689)

Author: yinggeh

Dockerfile.QA

@@ -1,4 +1,4 @@
-# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2018-2026, 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
@@ -139,7 +139,7 @@ RUN mkdir -p qa/common && \
     mkdir qa/L0_data_compression/models && \
     cp -r docs/examples/model_repository/simple qa/L0_data_compression/models && \
     cp bin/data_compressor_test qa/L0_data_compression/. && \
-    cp bin/backend_tensor_size_test qa/L0_input_validation/. && \
+    cp bin/tensor_size_test qa/L0_input_validation/. && \
     cp bin/metrics_api_test qa/L0_metrics/. && \
     cp bin/response_cache_test qa/L0_response_cache/. && \
     cp bin/request_cancellation_test qa/L0_request_cancellation/. && \

qa/L0_input_validation/test.sh

@@ -162,15 +162,15 @@ if [ $? -ne 0 ]; then
 fi
 set -e
 
-# backend_tensor_size_test
-TEST_LOG="./backend_tensor_size_test.log"
-TEST_EXEC=./backend_tensor_size_test
+# tensor_size_test
+TEST_LOG="./tensor_size_test.log"
+TEST_EXEC=./tensor_size_test
 
 set +e
 LD_LIBRARY_PATH=/opt/tritonserver/lib:$LD_LIBRARY_PATH $TEST_EXEC >> $TEST_LOG 2>&1
 if [ $? -ne 0 ]; then
     cat $TEST_LOG
-    echo -e "\n***\n*** backend_tensor_size_test FAILED\n***"
+    echo -e "\n***\n*** tensor_size_test FAILED\n***"
     RET=1
 fi
 set -e

src/test/CMakeLists.txt

@@ -1,4 +1,4 @@
-# Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-2026, 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
@@ -78,43 +78,53 @@ if(${TRITON_ENABLE_HTTP} OR ${TRITON_ENABLE_METRICS} OR
 endif()
 
 #
-# Unit test for Backend API
+# Unit test for Backend + Common + Core tensor size APIs (GetElementCount,
+# GetByteSize). Core template headers are used directly; the few core symbols
+# not available via libtritonserver.so (hidden by linker script) are defined
+# in the test source.
 #
-add_executable(
-  backend_tensor_size_test
-  backend_tensor_size_test.cc
-)
+if(TARGET proto-library)
+  add_executable(
+    tensor_size_test
+    tensor_size_test.cc
+  )
 
-set_target_properties(
-  backend_tensor_size_test
-  PROPERTIES
-    SKIP_BUILD_RPATH TRUE
-    BUILD_WITH_INSTALL_RPATH TRUE
-    INSTALL_RPATH_USE_LINK_PATH FALSE
-    INSTALL_RPATH ""
-)
+  set_target_properties(
+    tensor_size_test
+    PROPERTIES
+      SKIP_BUILD_RPATH TRUE
+      BUILD_WITH_INSTALL_RPATH TRUE
+      INSTALL_RPATH_USE_LINK_PATH FALSE
+      INSTALL_RPATH "$ORIGIN/../lib"
+  )
 
-target_include_directories(
-  backend_tensor_size_test
-  PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/..
-    ${GTEST_INCLUDE_DIRS}
-)
+  target_include_directories(
+    tensor_size_test
+    PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}/..
+      ${GTEST_INCLUDE_DIRS}
+      ${repo-core_SOURCE_DIR}/src
+      $<TARGET_PROPERTY:proto-library,INTERFACE_INCLUDE_DIRECTORIES>
+  )
 
-target_link_libraries(
-  backend_tensor_size_test
-  PRIVATE
-    triton-backend-utils    # from repo-backend
-    triton-core-serverapi   # from repo-core
-    triton-core-serverstub  # from repo-core
-    GTest::gtest
-    GTest::gtest_main
-)
+  target_link_libraries(
+    tensor_size_test
+    PRIVATE
+      triton-backend-utils
+      triton-core-serverapi
+      triton-core-serverstub
+      triton-common-model-config
+      proto-library
+      protobuf::libprotobuf
+      GTest::gtest
+      GTest::gtest_main
+  )
 
-install(
-  TARGETS backend_tensor_size_test
-  RUNTIME DESTINATION bin
-)
+  install(
+    TARGETS tensor_size_test
+    RUNTIME DESTINATION bin
+  )
+endif()
 
 add_subdirectory(repoagent/relocation_repoagent repoagent/relocation_repoagent)
 

src/test/backend_tensor_size_test.cc src/test/tensor_size_test.ccsrc/test/backend_tensor_size_test.cc renamed to src/test/tensor_size_test.cc

@@ -1,4 +1,4 @@
-// Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2025-2026, 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
@@ -35,10 +35,25 @@
 #include <string>
 #include <vector>
 
+#include "model_config_utils.h"
+#undef RETURN_IF_ERROR  // core (Status) vs backend (TRITONSERVER_Error*); avoid
+                        // redefinition
 #include "triton/backend/backend_common.h"
-#include "triton/core/tritonserver.h"
+#include "triton/common/model_config.h"
 
 namespace tb = triton::backend;
+namespace tc = triton::common;
+namespace tcore = triton::core;
+
+// Core's linker script hides C++ symbols from libtritonserver.so.
+// Provide the small set of definitions the header-only templates need.
+const tcore::Status tcore::Status::Success(tcore::Status::Code::SUCCESS);
+
+inference::DataType
+tcore::TritonToDataType(const TRITONSERVER_DataType dtype)
+{
+  return static_cast<inference::DataType>(dtype);
+}
 
 namespace {
 
@@ -76,28 +91,39 @@ TRITONSERVER_ErrorMessage(TRITONSERVER_Error* error)
 namespace {
 
 enum class ErrorCode {
-  kInvalidDim = -2,
-  kOverflow = -3,
+  kInvalidDim = tc::INVALID_SIZE,
+  kOverflow = tc::OVERFLOW_SIZE,
 };
 
+static const std::string kTensorName{"input0"};
+
 void
 assert_get_element_count_success(
     std::vector<int64_t>& shape, int64_t expected_cnt)
 {
   int64_t cnt;
   TRITONSERVER_Error* err;
 
-  // assert old APIs
+  // Backend (old APIs)
   ASSERT_EQ(expected_cnt, tb::GetElementCount(shape.data(), shape.size()));
   ASSERT_EQ(expected_cnt, tb::GetElementCount(shape));
 
-  // assert new APIs
+  // Backend (new APIs)
   err = tb::GetElementCount(shape.data(), shape.size(), &cnt);
   ASSERT_EQ(err, nullptr);
   ASSERT_EQ(cnt, expected_cnt);
   err = tb::GetElementCount(shape, &cnt);
   ASSERT_EQ(err, nullptr);
   ASSERT_EQ(cnt, expected_cnt);
+
+  // Common
+  ASSERT_EQ(tc::GetElementCount(shape), expected_cnt);
+
+  // Core
+  cnt = 0;
+  auto status = tcore::GetElementCount(shape, kTensorName, &cnt);
+  ASSERT_TRUE(status.IsOk()) << status.Message();
+  ASSERT_EQ(cnt, expected_cnt);
 }
 
 void
@@ -108,13 +134,13 @@ assert_get_element_count_error(
   int64_t cnt;
   TRITONSERVER_Error* err;
 
-  // assert old APIs
+  // Backend (old APIs)
   ASSERT_EQ(
       static_cast<int>(error_code),
       tb::GetElementCount(shape.data(), shape.size()));
   ASSERT_EQ(static_cast<int>(error_code), tb::GetElementCount(shape));
 
-  // assert new APIs
+  // Backend (new APIs)
   err = tb::GetElementCount(shape.data(), shape.size(), &cnt);
   ASSERT_NE(err, nullptr);
   ASSERT_EQ(TRITONSERVER_ERROR_INVALID_ARG, TRITONSERVER_ErrorCode(err));
@@ -123,23 +149,50 @@ assert_get_element_count_error(
   ASSERT_NE(err, nullptr);
   ASSERT_EQ(TRITONSERVER_ERROR_INVALID_ARG, TRITONSERVER_ErrorCode(err));
   ASSERT_STREQ(error_msg.c_str(), TRITONSERVER_ErrorMessage(err));
+
+  // Common
+  ASSERT_EQ(tc::GetElementCount(shape), static_cast<int64_t>(error_code));
+
+  // Core
+  cnt = 0;
+  auto status = tcore::GetElementCount(shape, kTensorName, &cnt);
+  ASSERT_FALSE(status.IsOk());
+  ASSERT_EQ(status.StatusCode(), triton::core::Status::Code::INVALID_ARG);
+  ASSERT_TRUE(
+      std::string(status.Message())
+          .find(
+              error_code == ErrorCode::kInvalidDim
+                  ? "invalid dimension"
+                  : "exceeds maximum size") != std::string::npos);
 }
 
 void
 assert_get_byte_size_success(
     TRITONSERVER_DataType dtype, std::vector<int64_t>& shape,
-    int64_t expected_size)
+    int64_t expected_size, bool test_core = true)
 {
   int64_t size;
   TRITONSERVER_Error* err;
 
-  // assert old API
+  // Backend (old API)
   ASSERT_EQ(expected_size, tb::GetByteSize(dtype, shape));
 
-  // assert new API
+  // Backend (new API)
   err = tb::GetByteSize(dtype, shape, &size);
   ASSERT_EQ(err, nullptr);
   ASSERT_EQ(expected_size, size);
+
+  // Common
+  inference::DataType core_dtype = tcore::TritonToDataType(dtype);
+  ASSERT_EQ(tc::GetByteSize(core_dtype, shape), expected_size);
+
+  // Core
+  if (test_core) {
+    size = 0;
+    auto status = tcore::GetByteSize(core_dtype, shape, kTensorName, &size);
+    ASSERT_TRUE(status.IsOk()) << status.Message();
+    ASSERT_EQ(size, expected_size);
+  }
 }
 
 void
@@ -150,14 +203,33 @@ assert_get_byte_size_error(
   int64_t size;
   TRITONSERVER_Error* err;
 
-  // assert old API
+  // Backend (old API)
   ASSERT_EQ(static_cast<int>(error_code), tb::GetByteSize(dtype, shape));
 
-  // assert new API
+  // Backend (new API)
   err = tb::GetByteSize(dtype, shape, &size);
   ASSERT_NE(err, nullptr);
   ASSERT_EQ(TRITONSERVER_ERROR_INVALID_ARG, TRITONSERVER_ErrorCode(err));
   ASSERT_EQ(error_msg, TRITONSERVER_ErrorMessage(err));
+
+  // Common
+  inference::DataType core_dtype = tcore::TritonToDataType(dtype);
+  ASSERT_EQ(
+      tc::GetByteSize(core_dtype, shape), error_code == ErrorCode::kInvalidDim
+                                              ? tc::INVALID_SIZE
+                                              : tc::OVERFLOW_SIZE);
+
+  // Core
+  size = 0;
+  auto status = tcore::GetByteSize(core_dtype, shape, kTensorName, &size);
+  ASSERT_FALSE(status.IsOk());
+  ASSERT_EQ(status.StatusCode(), triton::core::Status::Code::INVALID_ARG);
+  ASSERT_TRUE(
+      std::string(status.Message())
+          .find(
+              error_code == ErrorCode::kInvalidDim
+                  ? "invalid dimension"
+                  : "exceeds maximum size") != std::string::npos);
 }
 
 class GetElementCountTest : public ::testing::Test {
@@ -186,10 +258,10 @@ TEST_F(GetElementCountTest, GetElementCount)
   assert_get_element_count_success(shape, expected_cnt);
 }
 
-TEST_F(GetElementCountTest, GetElementCountNegative)
+TEST_F(GetElementCountTest, GetElementCountWildcard)
 {
   std::vector<int64_t> shape;
-  int64_t expected_cnt = -1;
+  int64_t expected_cnt = tc::WILDCARD_SIZE;
 
   // Test 1: -1 dim
   shape = {-1};
@@ -221,7 +293,6 @@ TEST_F(GetElementCountTest, GetElementCountZero)
   shape = {1, 8, 0};
   assert_get_element_count_success(shape, expected_cnt);
 
-  // Test 2: one 0 dim
   shape = {0, 1, 8};
   assert_get_element_count_success(shape, expected_cnt);
 
@@ -307,11 +378,11 @@ TEST_F(GetByteSizeTest, GetByteSize)
   assert_get_byte_size_success(dtype, shape, expected_size);
 }
 
-TEST_F(GetByteSizeTest, GetByteSizeNegative)
+TEST_F(GetByteSizeTest, GetByteSizeWildcard)
 {
   TRITONSERVER_DataType dtype;
   std::vector<int64_t> shape;
-  int64_t expected_size = -1;
+  int64_t expected_size = tc::WILDCARD_SIZE;
 
   // Test 1: invalid dtype
   dtype = TRITONSERVER_TYPE_INVALID;
@@ -321,7 +392,14 @@ TEST_F(GetByteSizeTest, GetByteSizeNegative)
   // Test 2: bytes dtype
   dtype = TRITONSERVER_TYPE_BYTES;
   shape = {8, 8};
-  assert_get_byte_size_success(dtype, shape, expected_size);
+  assert_get_byte_size_success(dtype, shape, expected_size, false);
+  // test core explicitly as it treats string dtype size as 4
+  int64_t size = 0;
+  inference::DataType core_dtype = tcore::TritonToDataType(dtype);
+  auto status = tcore::GetByteSize(core_dtype, shape, kTensorName, &size);
+  ASSERT_TRUE(status.IsOk()) << status.Message();
+  ASSERT_EQ(size, sizeof(int32_t) * 8 * 8);
+
 
   // Test 3: invalid shape and element count overflows
   dtype = TRITONSERVER_TYPE_INVALID;
@@ -348,7 +426,6 @@ TEST_F(GetByteSizeTest, GetByteSizeZero)
   shape = {1, 8, 0};
   assert_get_byte_size_success(dtype, shape, expected_cnt);
 
-  // Test 2: one 0 dim
   shape = {0, 1, 8};
   assert_get_byte_size_success(dtype, shape, expected_cnt);
 
@@ -403,6 +480,7 @@ TEST_F(GetByteSizeTest, GetByteSizeOverflow)
   error_msg = "unexpected integer overflow while calculating byte size.";
   assert_get_byte_size_error(dtype, shape, ErrorCode::kOverflow, error_msg);
 }
+
 }  // namespace
 
 int