refactor: Remove C shared memory shim. Refine shared memory utils behavior. Add unit test for shared memory

src/python/library/tests/test_shared_memory.py

@@ -0,0 +1,179 @@
+# Copyright 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
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import unittest
+
+import numpy
+import tritonclient.utils as utils
+import tritonclient.utils.shared_memory as shm
+
+
+class SharedMemoryTest(unittest.TestCase):
+    """
+    Testing shared memory utilities
+    """
+
+    def setUp(self):
+        self.shm_handles = []
+
+    def tearDown(self):
+        for shm_handle in self.shm_handles:
+            shm.destroy_shared_memory_region(shm_handle)
+
+    def test_lifecycle(self):
+        cpu_tensor = numpy.ones([4, 4], dtype=numpy.float32)
+        byte_size = 64
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", byte_size)
+        )
+
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
+
+        # Set data from Numpy array
+        shm.set_shared_memory_region(self.shm_handles[0], [cpu_tensor])
+        shm_tensor = shm.get_contents_as_numpy(
+            self.shm_handles[0], numpy.float32, [4, 4]
+        )
+
+        self.assertTrue(numpy.allclose(cpu_tensor, shm_tensor))
+
+        shm.destroy_shared_memory_region(self.shm_handles.pop(0))
+
+    def test_invalid_create_shm(self):
+        # Raises error since tried to create invalid system shared memory region
+        try:
+            self.shm_handles.append(
+                shm.create_shared_memory_region("dummy_data", "/dummy_data", -1)
+            )
+        except Exception as ex:
+            self.assertTrue(str(ex) == "unable to initialize the size")
+
+    def test_set_region_offset(self):
+        large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
+        large_size = 64
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", large_size)
+        )
+        shm.set_shared_memory_region(self.shm_handles[0], [large_tensor])
+        small_tensor = numpy.zeros([2, 4], dtype=numpy.float32)
+        small_size = 32
+        shm.set_shared_memory_region(
+            self.shm_handles[0], [small_tensor], offset=large_size - small_size
+        )
+        shm_tensor = shm.get_contents_as_numpy(
+            self.shm_handles[0], numpy.float32, [2, 4], offset=large_size - small_size
+        )
+
+        self.assertTrue(numpy.allclose(small_tensor, shm_tensor))
+
+    def test_set_region_oversize(self):
+        large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
+        small_size = 32
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", small_size)
+        )
+        with self.assertRaises(shm.SharedMemoryException):
+            shm.set_shared_memory_region(self.shm_handles[0], [large_tensor])
+
+    def test_duplicate_key(self):
+        # [NOTE] change in behavior:
+        # previous: okay to create shared memory region of the same key with different size
+        #           and the behavior is not being study clearly.
+        # now: return the same handle if existed, warning will be print if size is different
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", 32)
+        )
+        with self.assertRaises(shm.SharedMemoryException):
+            self.shm_handles.append(
+                shm.create_shared_memory_region(
+                    "shm_name", "shm_key", 32, create_only=True
+                )
+            )
+
+        # Get handle to the same shared memory region but with larger size requested,
+        # check if actual size is checked
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", 64)
+        )
+
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
+
+        large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
+        with self.assertRaises(shm.SharedMemoryException):
+            shm.set_shared_memory_region(self.shm_handles[-1], [large_tensor])
+
+    def test_destroy_duplicate(self):
+        # [NOTE] change in behavior:
+        # previous: raise exception if underlying shared memory has been unlinked
+        # now: no exception as unlink only happen when last managed handle is destroyed
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 0)
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", 64)
+        )
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", 32)
+        )
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", 32)
+        )
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
+
+        shm.destroy_shared_memory_region(self.shm_handles.pop(0))
+        shm.destroy_shared_memory_region(self.shm_handles.pop(0))
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
+
+        shm.destroy_shared_memory_region(self.shm_handles.pop(0))
+        self.assertEqual(len(shm.mapped_shared_memory_regions()), 0)
+
+    def test_numpy_bytes(self):
+        int_tensor = numpy.arange(start=0, stop=16, dtype=numpy.int32)
+        bytes_tensor = numpy.array(
+            [str(x).encode("utf-8") for x in int_tensor.flatten()], dtype=object
+        )
+        bytes_tensor = bytes_tensor.reshape(int_tensor.shape)
+        bytes_tensor_serialized = utils.serialize_byte_tensor(bytes_tensor)
+        byte_size = utils.serialized_byte_size(bytes_tensor_serialized)
+
+        self.shm_handles.append(
+            shm.create_shared_memory_region("shm_name", "shm_key", byte_size)
+        )
+
+        # Set data from Numpy array
+        shm.set_shared_memory_region(self.shm_handles[0], [bytes_tensor_serialized])
+
+        shm_tensor = shm.get_contents_as_numpy(
+            self.shm_handles[0],
+            numpy.object_,
+            [
+                16,
+            ],
+        )
+
+        self.assertTrue(numpy.array_equal(bytes_tensor, shm_tensor))
+
+
+if __name__ == "__main__":
+    unittest.main()

src/python/library/tritonclient/utils/shared_memory/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2019-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-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
@@ -28,6 +28,7 @@
 
 import os
 import struct
+import warnings
 from ctypes import *
 
 import numpy as np
@@ -72,6 +73,7 @@ def from_param(cls, value):
 _cshm_shared_memory_region_destroy.argtypes = [c_void_p]
 
 mapped_shm_regions = []
+_key_mapping = {}
 
 
 def _raise_if_error(errno):
@@ -90,7 +92,18 @@ def _raise_error(msg):
     raise ex
 
 
-def create_shared_memory_region(triton_shm_name, shm_key, byte_size):
+class SharedMemoryRegion:
+    def __init__(
+        self,
+        triton_shm_name: str,
+        shm_key: str,
+    ) -> None:
+        self._triton_shm_name = triton_shm_name
+        self._shm_key = shm_key
+        self._c_handle = c_void_p()
+
+
+def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only=False):
     """Creates a system shared memory region with the specified name and size.
 
     Parameters
@@ -101,10 +114,15 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size):
         The unique key of the shared memory object.
     byte_size : int
         The size in bytes of the shared memory region to be created.
+    create_only : bool
+        Whether a shared memory region must be created. If False and
+        a shared memory region of the same name exists, a handle to that
+        shared memory region will be returned and user must be aware that
+        the shared memory size can be different from the size requested.
 
     Returns
     -------
-    shm_handle : c_void_p
+    shm_handle : SharedMemoryRegion
         The handle for the system shared memory region.
 
     Raises
@@ -113,15 +131,48 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size):
         If unable to create the shared memory region.
     """
 
-    shm_handle = c_void_p()
-    _raise_if_error(
-        c_int(
-            _cshm_shared_memory_region_create(
-                triton_shm_name, shm_key, byte_size, byref(shm_handle)
+    if create_only and shm_key in mapped_shm_regions:
+        raise SharedMemoryException(
+            "unable to create the shared memory region, already exists"
+        )
+
+    shm_handle = SharedMemoryRegion(triton_shm_name, shm_key)
+    # Has been created
+    if shm_key in _key_mapping:
+        shm_handle._c_handle = _key_mapping[shm_key][0]
+        _key_mapping[shm_key][1] += 1
+        # check on the size
+        shm_fd = c_int()
+        region_offset = c_uint64()
+        shm_byte_size = c_uint64()
+        shm_addr = c_char_p()
+        c_shm_key = c_char_p()
+        _raise_if_error(
+            c_int(
+                _cshm_get_shared_memory_handle_info(
+                    shm_handle._c_handle,
+                    byref(shm_addr),
+                    byref(c_shm_key),
+                    byref(shm_fd),
+                    byref(region_offset),
+                    byref(shm_byte_size),
+                )
             )
         )
-    )
-    mapped_shm_regions.append(shm_key)
+        if byte_size > shm_byte_size.value:
+            warnings.warn(
+                f"reusing shared memory region with key '{shm_key}', region size is {shm_byte_size.value} instead of requested {byte_size}"
+            )
+    else:
+        _raise_if_error(
+            c_int(
+                _cshm_shared_memory_region_create(
+                    triton_shm_name, shm_key, byte_size, byref(shm_handle._c_handle)
+                )
+            )
+        )
+        _key_mapping[shm_key] = [shm_handle._c_handle, 1]
+        mapped_shm_regions.append(shm_key)
 
     return shm_handle
 
@@ -131,7 +182,7 @@ def set_shared_memory_region(shm_handle, input_values, offset=0):
 
     Parameters
     ----------
-    shm_handle : c_void_p
+    shm_handle : SharedMemoryRegion
         The handle for the system shared memory region.
     input_values : list
         The list of numpy arrays to be copied into the shared memory region.
@@ -160,7 +211,7 @@ def set_shared_memory_region(shm_handle, input_values, offset=0):
             _raise_if_error(
                 c_int(
                     _cshm_shared_memory_region_set(
-                        shm_handle,
+                        shm_handle._c_handle,
                         c_uint64(offset_current),
                         c_uint64(byte_size),
                         cast(input_value, c_void_p),
@@ -172,7 +223,7 @@ def set_shared_memory_region(shm_handle, input_values, offset=0):
             _raise_if_error(
                 c_int(
                     _cshm_shared_memory_region_set(
-                        shm_handle,
+                        shm_handle._c_handle,
                         c_uint64(offset_current),
                         c_uint64(byte_size),
                         input_value.ctypes.data_as(c_void_p),
@@ -189,7 +240,7 @@ def get_contents_as_numpy(shm_handle, datatype, shape, offset=0):
 
     Parameters
     ----------
-    shm_handle : c_void_p
+    shm_handle : SharedMemoryRegion
         The handle for the system shared memory region.
     datatype : np.dtype
         The datatype of the array to be returned.
@@ -213,7 +264,7 @@ def get_contents_as_numpy(shm_handle, datatype, shape, offset=0):
     _raise_if_error(
         c_int(
             _cshm_get_shared_memory_handle_info(
-                shm_handle,
+                shm_handle._c_handle,
                 byref(shm_addr),
                 byref(shm_key),
                 byref(shm_fd),
@@ -276,39 +327,24 @@ def destroy_shared_memory_region(shm_handle):
 
     Parameters
     ----------
-    shm_handle : c_void_p
+    shm_handle : SharedMemoryRegion
         The handle for the system shared memory region.
 
     Raises
     ------
     SharedMemoryException
         If unable to unlink the shared memory region.
     """
-    shm_fd = c_int()
-    offset = c_uint64()
-    byte_size = c_uint64()
-    shm_addr = c_char_p()
-    shm_key = c_char_p()
-    _raise_if_error(
-        c_int(
-            _cshm_get_shared_memory_handle_info(
-                shm_handle,
-                byref(shm_addr),
-                byref(shm_key),
-                byref(shm_fd),
-                byref(offset),
-                byref(byte_size),
-            )
-        )
-    )
     # It is safer to remove the shared memory key from the list before
     # deleting the shared memory region because if the deletion should
     # fail, a re-attempt could result in a segfault. Secondarily, if we
     # fail to delete a region, we should not report it back to the user
     # as a valid memory region.
-    mapped_shm_regions.remove(shm_key.value.decode("utf-8"))
-    _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle)))
-    return
+    _key_mapping[shm_handle._shm_key][1] -= 1
+    if _key_mapping[shm_handle._shm_key][1] == 0:
+        mapped_shm_regions.remove(shm_handle._shm_key)
+        _key_mapping.pop(shm_handle._shm_key)
+        _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle._c_handle)))
 
 
 class SharedMemoryException(Exception):
@@ -328,6 +364,7 @@ def __init__(self, err):
             -4: "unable to read/mmap the shared memory region",
             -5: "unable to unlink the shared memory region",
             -6: "unable to munmap the shared memory region",
+            -7: "unable to set the shared memory region",
         }
         self._msg = None
         if type(err) == str: