Revert "fix: Improve validation for system shared memory register (#8336) (#8340)"

This reverts commit 6c1e449.

Author: mc-nv

qa/L0_shared_memory/shared_memory_test.py

@@ -566,45 +566,16 @@ def test_register_reserved_names(self):
         """
         # This matches kTritonSharedMemoryRegionPrefix in the server code.
         reserved_prefix = "triton_python_backend_shm_region_"
-        shm_name = "my_test_shm_name"
-
-        # The shared memory key cannot start with the reserved prefix,
-        # regardless of leading slashes.
-        shm_keys_to_test = [
-            f"{reserved_prefix}_my_test_shm_key",
-            f"/{reserved_prefix}_my_test_shm_key",
-            f"///{reserved_prefix}_my_test_shm_key",
-        ]
 
-        for shm_key in shm_keys_to_test:
-            with self.subTest(shm_key=shm_key):
-                expected_msg = f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as the key contains the reserved prefix '{reserved_prefix}'"
-                with self.assertRaisesRegex(
-                    utils.InferenceServerException, expected_msg
-                ):
-                    self.triton_client.register_system_shared_memory(
-                        shm_name, shm_key, 10000
-                    )
-
-    def test_register_invalid_shm_key(self):
-        """
-        Test that registration fails if attempting to use an invalid name for the shm key.
-        """
+        # The shared memory key cannot start with the reserved prefix.
         shm_name = "my_test_shm_name"
-        shm_keys_to_test = [
-            "/",
-            "///",
-        ]
+        shm_key = f"{reserved_prefix}_my_test_shm_key"
 
-        for shm_key in shm_keys_to_test:
-            with self.subTest(shm_key=shm_key):
-                expected_msg = f"cannot register shared memory region '{shm_name}' - invalid shm key '{shm_key}'"
-                with self.assertRaisesRegex(
-                    utils.InferenceServerException, expected_msg
-                ):
-                    self.triton_client.register_system_shared_memory(
-                        shm_name, shm_key, 10000
-                    )
+        with self.assertRaisesRegex(
+            utils.InferenceServerException,
+            f"cannot register shared memory region '{shm_name}' with key '{shm_key}' as the key contains the reserved prefix '{reserved_prefix}'",
+        ) as e:
+            self.triton_client.register_system_shared_memory(shm_name, shm_key, 10000)
 
 
 def callback(user_data, result, error):

qa/L0_shared_memory/test.sh

@@ -70,7 +70,6 @@ for i in \
         test_infer_integer_overflow \
         test_register_out_of_bound \
         test_register_reserved_names \
-        test_register_invalid_shm_key \
         test_python_client_leak; do
     for client_type in http grpc; do
         SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}"

src/common.cc

@@ -163,39 +163,4 @@ DecodeBase64(
   return nullptr;
 }
 
-TRITONSERVER_Error*
-ValidateSharedMemoryKey(const std::string& name, const std::string& shm_key)
-{
-  std::string_view key_view(shm_key);
-
-  // Find the index of the first character that is not a slash
-  const std::size_t first_non_slash = key_view.find_first_not_of('/');
-
-  // If the entire key is slashes
-  if (first_non_slash == std::string_view::npos) {
-    return TRITONSERVER_ErrorNew(
-        TRITONSERVER_ERROR_INVALID_ARG,
-        std::string(
-            "cannot register shared memory region '" + name +
-            "' - invalid shm key '" + shm_key + "'")
-            .c_str());
-  }
-
-  // Check whether the substring starting at first_non_slash starts with the
-  // reserved prefix
-  if (key_view.substr(first_non_slash)
-          .rfind(kTritonSharedMemoryRegionPrefix, 0) == 0) {
-    return TRITONSERVER_ErrorNew(
-        TRITONSERVER_ERROR_INVALID_ARG,
-        std::string(
-            "cannot register shared memory region '" + name + "' with key '" +
-            shm_key + "' as the key contains the reserved prefix '" +
-            kTritonSharedMemoryRegionPrefix + "'")
-            .c_str());
-  }
-
-  // Valid shm key
-  return nullptr;
-}
-
 }}  // namespace triton::server

src/common.h

@@ -196,16 +196,6 @@ TRITONSERVER_Error* DecodeBase64(
     const char* input, size_t input_len, std::vector<char>& decoded_data,
     size_t& decoded_size, const std::string& name);
 
-
-/// Validate shared memory key
-///
-/// \param name The name of the memory block.
-/// \param shm_key The name of the posix shared memory object
-/// \return The error status.
-TRITONSERVER_Error* ValidateSharedMemoryKey(
-    const std::string& name, const std::string& shm_key);
-
-
 /// Joins container of strings into a single string delimited by
 /// 'delim'.
 ///

src/shared_memory_manager.cc

@@ -355,7 +355,15 @@ SharedMemoryManager::RegisterSystemSharedMemory(
     const size_t byte_size)
 {
   // Check if the shared memory key starts with the reserved prefix
-  RETURN_IF_ERR(ValidateSharedMemoryKey(name, shm_key));
+  if (shm_key.rfind(kTritonSharedMemoryRegionPrefix, 0) == 0) {
+    return TRITONSERVER_ErrorNew(
+        TRITONSERVER_ERROR_INVALID_ARG,
+        std::string(
+            "cannot register shared memory region '" + name + "' with key '" +
+            shm_key + "' as the key contains the reserved prefix '" +
+            kTritonSharedMemoryRegionPrefix + "'")
+            .c_str());
+  }
 
   std::lock_guard<std::mutex> lock(mu_);