From 4d7a0b88bfdb801556dafb1bae30d0471c3eeb7b Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Oct 2024 16:42:48 -0700 Subject: [PATCH 1/7] test: add unit test for shared memory --- .../library/tests/test_shared_memory.py | 182 ++++++++++++++++++ .../utils/shared_memory/__init__.py | 30 ++- .../utils/shared_memory/shared_memory.cc | 9 +- 3 files changed, 212 insertions(+), 9 deletions(-) create mode 100644 src/python/library/tests/test_shared_memory.py diff --git a/src/python/library/tests/test_shared_memory.py b/src/python/library/tests/test_shared_memory.py new file mode 100644 index 000000000..93a5881eb --- /dev/null +++ b/src/python/library/tests/test_shared_memory.py @@ -0,0 +1,182 @@ +# 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: + # [NOTE] wrapper for old implementation that will fail + try: + shm.destroy_shared_memory_region(shm_handle) + except shm.SharedMemoryException as ex: + if "unlink" in str(ex): + pass + else: + raise ex + + 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_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)) + + # [NOTE] current impl will fail + 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: only allow create by default, flag may be set to 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) + ) + + # 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, create=False) + ) + + self.assertEqual(len(shm.mapped_shared_memory_regions()), 1) + + large_tensor = numpy.ones([4, 4], dtype=numpy.float32) + small_size = 32 + # [NOTE] current impl will fail + with self.assertRaises(shm.SharedMemoryException): + shm.set_shared_memory_region(self.shm_handles[-1], [large_tensor]) + + # [NOTE] current impl will fail + def test_destroy_duplicate(self): + # [NOTE] change in behavior: + # previous: raise exception if underlying shared memory has been unlinked + # now: the exception will be suppressed to align with Windows behavior, unless + # explicitly toggled + 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, create=False) + ) + self.shm_handles.append( + shm.create_shared_memory_region("shm_name", "shm_key", 32, create=False) + ) + 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) + + shm.destroy_shared_memory_region(self.shm_handles.pop(0)) + with self.assertRaises(shm.SharedMemoryException): + shm.destroy_shared_memory_region( + self.shm_handles.pop(0), raise_unlink_exception=True + ) + + 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() diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index fd65191b9..38fd372db 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -90,7 +90,7 @@ def _raise_error(msg): raise ex -def create_shared_memory_region(triton_shm_name, shm_key, byte_size): +def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create=True): """Creates a system shared memory region with the specified name and size. Parameters @@ -113,6 +113,11 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size): If unable to create the shared memory region. """ + if create and shm_key in mapped_shm_regions: + raise SharedMemoryException( + "unable to create the shared memory region, already exists" + ) + shm_handle = c_void_p() _raise_if_error( c_int( @@ -121,7 +126,9 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size): ) ) ) - mapped_shm_regions.append(shm_key) + + if create: + mapped_shm_regions.append(shm_key) return shm_handle @@ -271,7 +278,7 @@ def mapped_shared_memory_regions(): return mapped_shm_regions -def destroy_shared_memory_region(shm_handle): +def destroy_shared_memory_region(shm_handle, raise_unlink_exception=False): """Unlink a system shared memory region with the specified handle. Parameters @@ -306,8 +313,20 @@ def destroy_shared_memory_region(shm_handle): # 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))) + try: + mapped_shm_regions.remove(shm_key.value.decode("utf-8")) + except ValueError: + # okay if mapped_shm_regions doesn't have the key as there can be + # destroy call on handles with the same shared memory key + pass + try: + _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle))) + except SharedMemoryException as ex: + # Suppress unlink exception except when explicitly allow to raise + if not raise_unlink_exception and "unlink" in str(ex): + pass + else: + raise ex return @@ -328,6 +347,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: diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc b/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc index 80720da26..5e68e7068 100644 --- a/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc +++ b/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc @@ -108,10 +108,11 @@ int SharedMemoryRegionSet( void* shm_handle, size_t offset, size_t byte_size, const void* data) { - void* shm_addr = - reinterpret_cast(shm_handle)->base_addr_; - char* shm_addr_offset = reinterpret_cast(shm_addr); - std::memcpy(shm_addr_offset + offset, data, byte_size); + auto shm = reinterpret_cast(shm_handle); + if (shm->byte_size_ < (offset + byte_size)) { + return -7; + } + std::memcpy(shm->base_addr_ + offset, data, byte_size); return 0; } From c28ba19e502219b3cb72c0352126e2ad75053812 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Tue, 22 Oct 2024 18:19:13 -0700 Subject: [PATCH 2/7] refactor: reflect latest comment on create / destroy behavior --- .../library/tests/test_shared_memory.py | 47 ++++--- .../utils/shared_memory/__init__.py | 119 ++++++++++-------- 2 files changed, 90 insertions(+), 76 deletions(-) diff --git a/src/python/library/tests/test_shared_memory.py b/src/python/library/tests/test_shared_memory.py index 93a5881eb..3cbbb44f1 100644 --- a/src/python/library/tests/test_shared_memory.py +++ b/src/python/library/tests/test_shared_memory.py @@ -41,14 +41,7 @@ def setUp(self): def tearDown(self): for shm_handle in self.shm_handles: - # [NOTE] wrapper for old implementation that will fail - try: - shm.destroy_shared_memory_region(shm_handle) - except shm.SharedMemoryException as ex: - if "unlink" in str(ex): - pass - else: - raise ex + shm.destroy_shared_memory_region(shm_handle) def test_lifecycle(self): cpu_tensor = numpy.ones([4, 4], dtype=numpy.float32) @@ -69,6 +62,15 @@ def test_lifecycle(self): 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 @@ -87,7 +89,6 @@ def test_set_region_offset(self): self.assertTrue(numpy.allclose(small_tensor, shm_tensor)) - # [NOTE] current impl will fail def test_set_region_oversize(self): large_tensor = numpy.ones([4, 4], dtype=numpy.float32) small_size = 32 @@ -101,55 +102,51 @@ 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: only allow create by default, flag may be set to return the same handle if - # existed, warning will be print if size is different + # 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) + 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, create=False) + 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) - small_size = 32 - # [NOTE] current impl will fail with self.assertRaises(shm.SharedMemoryException): shm.set_shared_memory_region(self.shm_handles[-1], [large_tensor]) - # [NOTE] current impl will fail def test_destroy_duplicate(self): # [NOTE] change in behavior: # previous: raise exception if underlying shared memory has been unlinked - # now: the exception will be suppressed to align with Windows behavior, unless - # explicitly toggled + # 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, create=False) + shm.create_shared_memory_region("shm_name", "shm_key", 32) ) self.shm_handles.append( - shm.create_shared_memory_region("shm_name", "shm_key", 32, create=False) + 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)) - self.assertEqual(len(shm.mapped_shared_memory_regions()), 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)) - with self.assertRaises(shm.SharedMemoryException): - shm.destroy_shared_memory_region( - self.shm_handles.pop(0), raise_unlink_exception=True - ) + 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) diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index 38fd372db..5b93ba23e 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/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, create=True): +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, create=True 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,21 +131,47 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create=True If unable to create the shared memory region. """ - if create and shm_key in mapped_shm_regions: + if create_only and shm_key in mapped_shm_regions: raise SharedMemoryException( "unable to create the shared memory region, already exists" ) - 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) + 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() + shm_key = c_char_p() + _raise_if_error( + c_int( + _cshm_get_shared_memory_handle_info( + shm_handle._c_handle, + byref(shm_addr), + byref(shm_key), + byref(shm_fd), + byref(region_offset), + byref(shm_byte_size), + ) ) ) - ) - - if create: + 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 @@ -138,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. @@ -167,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), @@ -179,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), @@ -196,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. @@ -220,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), @@ -278,12 +322,12 @@ def mapped_shared_memory_regions(): return mapped_shm_regions -def destroy_shared_memory_region(shm_handle, raise_unlink_exception=False): +def destroy_shared_memory_region(shm_handle): """Unlink a system shared memory region with the specified handle. Parameters ---------- - shm_handle : c_void_p + shm_handle : SharedMemoryRegion The handle for the system shared memory region. Raises @@ -291,43 +335,16 @@ def destroy_shared_memory_region(shm_handle, raise_unlink_exception=False): 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. - try: - mapped_shm_regions.remove(shm_key.value.decode("utf-8")) - except ValueError: - # okay if mapped_shm_regions doesn't have the key as there can be - # destroy call on handles with the same shared memory key - pass - try: - _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle))) - except SharedMemoryException as ex: - # Suppress unlink exception except when explicitly allow to raise - if not raise_unlink_exception and "unlink" in str(ex): - pass - else: - raise ex - 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): From bcdfe3e7ab95afc2ca3cf3d43a1dc6f9d44a8d0b Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Tue, 22 Oct 2024 18:21:33 -0700 Subject: [PATCH 3/7] fix: fix variable shadowing --- .../library/tritonclient/utils/shared_memory/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index 5b93ba23e..1972ccf81 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -146,13 +146,13 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only region_offset = c_uint64() shm_byte_size = c_uint64() shm_addr = c_char_p() - shm_key = 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(shm_key), + byref(c_shm_key), byref(shm_fd), byref(region_offset), byref(shm_byte_size), From f791cd478d2358707113367f138b5ec4f30c2a15 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Tue, 22 Oct 2024 18:44:40 -0700 Subject: [PATCH 4/7] chore: update function comment. address comment --- .../library/tests/test_shared_memory.py | 19 +++++++++++++------ .../utils/shared_memory/__init__.py | 5 +++-- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/python/library/tests/test_shared_memory.py b/src/python/library/tests/test_shared_memory.py index 3cbbb44f1..c349c1eed 100644 --- a/src/python/library/tests/test_shared_memory.py +++ b/src/python/library/tests/test_shared_memory.py @@ -64,12 +64,12 @@ def test_lifecycle(self): def test_invalid_create_shm(self): # Raises error since tried to create invalid system shared memory region - try: + with self.assertRaisesRegex( + shm.SharedMemoryException, "unable to initialize the size" + ): 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) @@ -95,7 +95,9 @@ def test_set_region_oversize(self): self.shm_handles.append( shm.create_shared_memory_region("shm_name", "shm_key", small_size) ) - with self.assertRaises(shm.SharedMemoryException): + with self.assertRaisesRegex( + shm.SharedMemoryException, "unable to set the shared memory region" + ): shm.set_shared_memory_region(self.shm_handles[0], [large_tensor]) def test_duplicate_key(self): @@ -106,7 +108,10 @@ def test_duplicate_key(self): self.shm_handles.append( shm.create_shared_memory_region("shm_name", "shm_key", 32) ) - with self.assertRaises(shm.SharedMemoryException): + with self.assertRaisesRegex( + shm.SharedMemoryException, + "unable to create the shared memory region, already exists", + ): self.shm_handles.append( shm.create_shared_memory_region( "shm_name", "shm_key", 32, create_only=True @@ -122,7 +127,9 @@ def test_duplicate_key(self): self.assertEqual(len(shm.mapped_shared_memory_regions()), 1) large_tensor = numpy.ones([4, 4], dtype=numpy.float32) - with self.assertRaises(shm.SharedMemoryException): + with self.assertRaisesRegex( + shm.SharedMemoryException, "unable to set the shared memory region" + ): shm.set_shared_memory_region(self.shm_handles[-1], [large_tensor]) def test_destroy_duplicate(self): diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index 1972ccf81..d7758a17a 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -104,7 +104,7 @@ def __init__( 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. + """Return a handle of the system shared memory region with the specified name and size. Parameters ---------- @@ -323,7 +323,8 @@ def mapped_shared_memory_regions(): def destroy_shared_memory_region(shm_handle): - """Unlink a system shared memory region with the specified handle. + """Release the handle, unlink a system shared memory region with the specified handle + if it is the last managed handle. Parameters ---------- From 66f3c75a4f464bb396a6cc8d083c20205aa76b62 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 23 Oct 2024 12:49:12 -0700 Subject: [PATCH 5/7] refactor: remove C shared memory shim --- src/python/library/CMakeLists.txt | 1 - src/python/library/build_wheel.py | 4 - src/python/library/setup.py | 2 - .../library/tests/test_shared_memory.py | 4 +- .../library/tritonclient/utils/CMakeLists.txt | 14 - .../utils/shared_memory/__init__.py | 257 +++++------------- .../utils/shared_memory/shared_memory.cc | 152 ----------- .../utils/shared_memory/shared_memory.h | 53 ---- .../shared_memory/shared_memory_handle.h | 44 --- 9 files changed, 66 insertions(+), 465 deletions(-) delete mode 100644 src/python/library/tritonclient/utils/shared_memory/shared_memory.cc delete mode 100644 src/python/library/tritonclient/utils/shared_memory/shared_memory.h delete mode 100644 src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h diff --git a/src/python/library/CMakeLists.txt b/src/python/library/CMakeLists.txt index e552ec2f7..f9a1c5748 100644 --- a/src/python/library/CMakeLists.txt +++ b/src/python/library/CMakeLists.txt @@ -96,7 +96,6 @@ add_custom_target( if (NOT WIN32) # Can generate linux specific wheel file on linux systems only. set(LINUX_WHEEL_DEPENDS - cshm ${WHEEL_DEPENDS} ) diff --git a/src/python/library/build_wheel.py b/src/python/library/build_wheel.py index d32e7732a..3be738086 100755 --- a/src/python/library/build_wheel.py +++ b/src/python/library/build_wheel.py @@ -174,10 +174,6 @@ def sed(pattern, replace, source, dest=None): "tritonclient/utils/shared_memory", os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory"), ) - shutil.copyfile( - "tritonclient/utils/libcshm.so", - os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory/libcshm.so"), - ) cpdir( "tritonclient/utils/cuda_shared_memory", os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"), diff --git a/src/python/library/setup.py b/src/python/library/setup.py index e31f5ddc0..634b8b57e 100755 --- a/src/python/library/setup.py +++ b/src/python/library/setup.py @@ -76,8 +76,6 @@ def req_file(filename, folder="requirements"): extras_require["all"] = list(chain(extras_require.values())) platform_package_data = [] -if PLATFORM_FLAG != "any": - platform_package_data += ["libcshm.so"] data_files = [ ("", ["LICENSE.txt"]), diff --git a/src/python/library/tests/test_shared_memory.py b/src/python/library/tests/test_shared_memory.py index c349c1eed..f765bc781 100644 --- a/src/python/library/tests/test_shared_memory.py +++ b/src/python/library/tests/test_shared_memory.py @@ -65,7 +65,7 @@ def test_lifecycle(self): def test_invalid_create_shm(self): # Raises error since tried to create invalid system shared memory region with self.assertRaisesRegex( - shm.SharedMemoryException, "unable to initialize the size" + shm.SharedMemoryException, "unable to create the shared memory region" ): self.shm_handles.append( shm.create_shared_memory_region("dummy_data", "/dummy_data", -1) @@ -110,7 +110,7 @@ def test_duplicate_key(self): ) with self.assertRaisesRegex( shm.SharedMemoryException, - "unable to create the shared memory region, already exists", + "unable to create the shared memory region", ): self.shm_handles.append( shm.create_shared_memory_region( diff --git a/src/python/library/tritonclient/utils/CMakeLists.txt b/src/python/library/tritonclient/utils/CMakeLists.txt index 7de1acf96..178b0512b 100644 --- a/src/python/library/tritonclient/utils/CMakeLists.txt +++ b/src/python/library/tritonclient/utils/CMakeLists.txt @@ -28,20 +28,6 @@ configure_file(__init__.py __init__.py COPYONLY) configure_file(_dlpack.py _dlpack.py COPYONLY) configure_file(_shared_memory_tensor.py _shared_memory_tensor.py COPYONLY) -if(NOT WIN32) - file(COPY shared_memory DESTINATION .) - - # - # libcshm.so - # - add_library(cshm SHARED shared_memory/shared_memory.cc) - if(${TRITON_ENABLE_GPU}) - target_compile_definitions(cshm PUBLIC TRITON_ENABLE_GPU=1) - target_link_libraries(cshm PUBLIC CUDA::cudart) - endif() # TRITON_ENABLE_GPU - target_link_libraries(cshm PRIVATE rt) -endif() # WIN32 - if(NOT WIN32) configure_file(shared_memory/__init__.py shared_memory/__init__.py COPYONLY) configure_file(cuda_shared_memory/__init__.py cuda_shared_memory/__init__.py COPYONLY) diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index d7758a17a..4c4ca2845 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -29,67 +29,11 @@ import os import struct import warnings -from ctypes import * +from multiprocessing import shared_memory as mpshm import numpy as np -import pkg_resources - - -class _utf8(object): - @classmethod - def from_param(cls, value): - if value is None: - return None - elif isinstance(value, bytes): - return value - else: - return value.encode("utf8") - - -_cshm_lib = "cshm" if os.name == "nt" else "libcshm.so" -_cshm_path = pkg_resources.resource_filename( - "tritonclient.utils.shared_memory", _cshm_lib -) -_cshm = cdll.LoadLibrary(_cshm_path) - -_cshm_shared_memory_region_create = _cshm.SharedMemoryRegionCreate -_cshm_shared_memory_region_create.restype = c_int -_cshm_shared_memory_region_create.argtypes = [_utf8, _utf8, c_uint64, POINTER(c_void_p)] -_cshm_shared_memory_region_set = _cshm.SharedMemoryRegionSet -_cshm_shared_memory_region_set.restype = c_int -_cshm_shared_memory_region_set.argtypes = [c_void_p, c_uint64, c_uint64, c_void_p] -_cshm_get_shared_memory_handle_info = _cshm.GetSharedMemoryHandleInfo -_cshm_get_shared_memory_handle_info.restype = c_int -_cshm_get_shared_memory_handle_info.argtypes = [ - c_void_p, - POINTER(c_char_p), - POINTER(c_char_p), - POINTER(c_int), - POINTER(c_uint64), - POINTER(c_uint64), -] -_cshm_shared_memory_region_destroy = _cshm.SharedMemoryRegionDestroy -_cshm_shared_memory_region_destroy.restype = c_int -_cshm_shared_memory_region_destroy.argtypes = [c_void_p] - -mapped_shm_regions = [] -_key_mapping = {} - - -def _raise_if_error(errno): - """ - Raise SharedMemoryException if 'err' is non-success. - Otherwise return nothing. - """ - if errno.value != 0: - ex = SharedMemoryException(errno) - raise ex - return - -def _raise_error(msg): - ex = SharedMemoryException(msg) - raise ex +_key_mapping = {} class SharedMemoryRegion: @@ -100,7 +44,7 @@ def __init__( ) -> None: self._triton_shm_name = triton_shm_name self._shm_key = shm_key - self._c_handle = c_void_p() + self._mpsm_handle = None def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only=False): @@ -130,49 +74,34 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only SharedMemoryException If unable to create the shared memory region. """ - - 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), - ) - ) - ) - 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) - ) + # Check whether the region exists before creating it + if not create_only: + try: + shm_handle._mpsm_handle = mpshm.SharedMemory(shm_key) + if shm_key not in _key_mapping: + _key_mapping[shm_key] = [False, 0] + _key_mapping[shm_key][1] += 1 + except FileNotFoundError: + pass + if shm_handle._mpsm_handle is None: + try: + shm_handle._mpsm_handle = mpshm.SharedMemory( + shm_key, create=True, size=byte_size ) + except Exception as ex: + raise SharedMemoryException( + "unable to create the shared memory region" + ) from ex + if shm_key not in _key_mapping: + _key_mapping[shm_key] = [False, 0] + _key_mapping[shm_key][0] = True + _key_mapping[shm_key][1] += 1 + + if byte_size > shm_handle._mpsm_handle.size: + warnings.warn( + f"reusing shared memory region with key '{shm_key}', region size is {shm_handle._mpsm_handle.size} instead of requested {byte_size}" ) - _key_mapping[shm_key] = [shm_handle._c_handle, 1] - mapped_shm_regions.append(shm_key) return shm_handle @@ -197,41 +126,33 @@ def set_shared_memory_region(shm_handle, input_values, offset=0): """ if not isinstance(input_values, (list, tuple)): - _raise_error("input_values must be specified as a list/tuple of numpy arrays") + raise SharedMemoryException( + "input_values must be specified as a list/tuple of numpy arrays" + ) for input_value in input_values: if not isinstance(input_value, np.ndarray): - _raise_error("each element of input_values must be a numpy array") + raise SharedMemoryException( + "each element of input_values must be a numpy array" + ) - offset_current = offset - for input_value in input_values: - input_value = np.ascontiguousarray(input_value).flatten() - if input_value.dtype == np.object_: - input_value = input_value.item() - byte_size = np.dtype(np.byte).itemsize * len(input_value) - _raise_if_error( - c_int( - _cshm_shared_memory_region_set( - shm_handle._c_handle, - c_uint64(offset_current), - c_uint64(byte_size), - cast(input_value, c_void_p), - ) + try: + for input_value in input_values: + if input_value.dtype == np.object_: + byte_size = len(input_value.item()) + shm_handle._mpsm_handle.buf[offset : offset + byte_size] = ( + input_value.item() ) - ) - else: - byte_size = input_value.size * input_value.itemsize - _raise_if_error( - c_int( - _cshm_shared_memory_region_set( - shm_handle._c_handle, - c_uint64(offset_current), - c_uint64(byte_size), - input_value.ctypes.data_as(c_void_p), - ) + offset += byte_size + else: + shm_tensor_view = np.ndarray( + input_value.shape, + input_value.dtype, + buffer=shm_handle._mpsm_handle.buf[offset:], ) - ) - offset_current += byte_size - return + shm_tensor_view[:] = input_value[:] + offset += input_value.nbytes + except Exception as ex: + raise SharedMemoryException("unable to set the shared memory region") from ex def get_contents_as_numpy(shm_handle, datatype, shape, offset=0): @@ -256,42 +177,13 @@ def get_contents_as_numpy(shm_handle, datatype, shape, offset=0): The numpy array generated using the contents of the specified shared memory region. """ - shm_fd = c_int() - region_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._c_handle, - byref(shm_addr), - byref(shm_key), - byref(shm_fd), - byref(region_offset), - byref(byte_size), - ) - ) - ) - start_pos = region_offset.value + offset if (datatype != np.object_) and (datatype != np.bytes_): - requested_byte_size = np.prod(shape) * np.dtype(datatype).itemsize - cval_len = start_pos + requested_byte_size - if byte_size.value < cval_len: - _raise_error( - "The size of the shared memory region is insufficient to provide numpy array with requested size" - ) - if cval_len == 0: - result = np.empty(shape, dtype=datatype) - else: - val_buf = cast(shm_addr, POINTER(c_byte * cval_len))[0] - val = np.frombuffer(val_buf, dtype=datatype, offset=start_pos) - - # Reshape the result to the appropriate shape. - result = np.reshape(val, shape) + result = np.ndarray( + shape, datatype, buffer=shm_handle._mpsm_handle.buf[offset:] + ) else: - str_offset = start_pos - val_buf = cast(shm_addr, POINTER(c_byte * byte_size.value))[0] + str_offset = offset + val_buf = shm_handle._mpsm_handle.buf ii = 0 strs = list() while (ii % np.prod(shape) != 0) or (ii == 0): @@ -319,7 +211,7 @@ def mapped_shared_memory_regions(): The list of mapped system shared memory regions. """ - return mapped_shm_regions + return list(_key_mapping.keys()) def destroy_shared_memory_region(shm_handle): @@ -341,38 +233,17 @@ def destroy_shared_memory_region(shm_handle): # 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. + shm_handle._mpsm_handle.close() _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))) + try: + if _key_mapping[shm_handle._shm_key][0]: + shm_handle._mpsm_handle.unlink() + finally: + _key_mapping.pop(shm_handle._shm_key) class SharedMemoryException(Exception): - """Exception indicating non-Success status. - - Parameters - ---------- - err : c_void_p - Pointer to an Error that should be used to initialize the exception. - - """ + """Exception type for shared memory related error.""" - def __init__(self, err): - self.err_code_map = { - -2: "unable to get shared memory descriptor", - -3: "unable to initialize the size", - -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: - self._msg = err - elif err.value != 0 and err.value in self.err_code_map: - self._msg = self.err_code_map[err.value] - - def __str__(self): - msg = super().__str__() if self._msg is None else self._msg - return msg + pass diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc b/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc deleted file mode 100644 index 5e68e7068..000000000 --- a/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc +++ /dev/null @@ -1,152 +0,0 @@ -// Copyright (c) 2019-2020, NVIDIA CORPORATION. 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. - -#include "shared_memory.h" - -#include -#include -#include -#include - -#include -#include -#include - -#include "shared_memory_handle.h" - -//============================================================================== -// SharedMemoryControlContext - -namespace { - -void* -SharedMemoryHandleCreate( - std::string triton_shm_name, void* shm_addr, std::string shm_key, - int shm_fd, size_t offset, size_t byte_size) -{ - SharedMemoryHandle* handle = new SharedMemoryHandle(); - handle->triton_shm_name_ = triton_shm_name; - handle->base_addr_ = shm_addr; - handle->shm_key_ = shm_key; - handle->shm_fd_ = shm_fd; - handle->offset_ = offset; - handle->byte_size_ = byte_size; - return reinterpret_cast(handle); -} - -int -SharedMemoryRegionMap( - int shm_fd, size_t offset, size_t byte_size, void** shm_addr) -{ - // map shared memory to process address space - *shm_addr = mmap(NULL, byte_size, PROT_WRITE, MAP_SHARED, shm_fd, offset); - if (*shm_addr == MAP_FAILED) { - return -1; - } - - // close shared memory descriptor, return 0 if success else return -1 - return close(shm_fd); -} - -} // namespace - -int -SharedMemoryRegionCreate( - const char* triton_shm_name, const char* shm_key, size_t byte_size, - void** shm_handle) -{ - // get shared memory region descriptor - int shm_fd = shm_open(shm_key, O_RDWR | O_CREAT, S_IRUSR | S_IWUSR); - if (shm_fd == -1) { - return -2; - } - - // extend shared memory object as by default it's initialized with size 0 - int res = ftruncate(shm_fd, byte_size); - if (res == -1) { - return -3; - } - - // get base address of shared memory region - void* shm_addr = nullptr; - int err = SharedMemoryRegionMap(shm_fd, 0, byte_size, &shm_addr); - if (err == -1) { - return -4; - } - - // create a handle for the shared memory region - *shm_handle = SharedMemoryHandleCreate( - std::string(triton_shm_name), shm_addr, std::string(shm_key), shm_fd, 0, - byte_size); - return 0; -} - -int -SharedMemoryRegionSet( - void* shm_handle, size_t offset, size_t byte_size, const void* data) -{ - auto shm = reinterpret_cast(shm_handle); - if (shm->byte_size_ < (offset + byte_size)) { - return -7; - } - std::memcpy(shm->base_addr_ + offset, data, byte_size); - return 0; -} - -int -GetSharedMemoryHandleInfo( - void* shm_handle, char** shm_addr, const char** shm_key, int* shm_fd, - size_t* offset, size_t* byte_size) -{ - SharedMemoryHandle* handle = - reinterpret_cast(shm_handle); - *shm_addr = reinterpret_cast(handle->base_addr_); - *shm_key = handle->shm_key_.c_str(); - *shm_fd = handle->shm_fd_; - *offset = handle->offset_; - *byte_size = handle->byte_size_; - return 0; -} - -int -SharedMemoryRegionDestroy(void* shm_handle) -{ - std::unique_ptr handle( - reinterpret_cast(shm_handle)); - void* shm_addr = reinterpret_cast(handle->base_addr_); - int status = munmap(shm_addr, handle->byte_size_); - if (status == -1) { - return -6; - } - - int shm_fd = shm_unlink(handle->shm_key_.c_str()); - if (shm_fd == -1) { - return -5; - } - return 0; -} - -//============================================================================== diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory.h b/src/python/library/tritonclient/utils/shared_memory/shared_memory.h deleted file mode 100644 index 9d3e9519e..000000000 --- a/src/python/library/tritonclient/utils/shared_memory/shared_memory.h +++ /dev/null @@ -1,53 +0,0 @@ -// Copyright (c) 2019-2020, NVIDIA CORPORATION. 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. -#pragma once - -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -//============================================================================== -// SharedMemoryControlContext -int SharedMemoryRegionCreate( - const char* triton_shm_name, const char* shm_key, size_t byte_size, - void** shm_handle); -int SharedMemoryRegionSet( - void* shm_handle, size_t offset, size_t byte_size, const void* data); -int GetSharedMemoryHandleInfo( - void* shm_handle, char** shm_addr, const char** shm_key, int* shm_fd, - size_t* offset, size_t* byte_size); -int SharedMemoryRegionDestroy(void* shm_handle); - -//============================================================================== - -#ifdef __cplusplus -} -#endif diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h b/src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h deleted file mode 100644 index b929ed305..000000000 --- a/src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright (c) 2019-2020, NVIDIA CORPORATION. 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. - -#pragma once - -#ifdef TRITON_ENABLE_GPU -#include -#endif // TRITON_ENABLE_GPU - -struct SharedMemoryHandle { - std::string triton_shm_name_; - std::string shm_key_; -#ifdef TRITON_ENABLE_GPU - cudaIpcMemHandle_t cuda_shm_handle_; - int device_id_; -#endif // TRITON_ENABLE_GPU - void* base_addr_; - int shm_fd_; - size_t offset_; - size_t byte_size_; -}; From 6a87c5ffd5c6adc92ad42f49214e29208490fe26 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 23 Oct 2024 13:27:45 -0700 Subject: [PATCH 6/7] style: add commnet and fix copyright --- src/python/library/build_wheel.py | 2 +- src/python/library/tritonclient/utils/CMakeLists.txt | 2 +- .../library/tritonclient/utils/shared_memory/__init__.py | 4 ++++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/python/library/build_wheel.py b/src/python/library/build_wheel.py index 3be738086..73f727d0d 100755 --- a/src/python/library/build_wheel.py +++ b/src/python/library/build_wheel.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2021-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 diff --git a/src/python/library/tritonclient/utils/CMakeLists.txt b/src/python/library/tritonclient/utils/CMakeLists.txt index 178b0512b..94952efc7 100644 --- a/src/python/library/tritonclient/utils/CMakeLists.txt +++ b/src/python/library/tritonclient/utils/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2020-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 diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index 4c4ca2845..364b58b99 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -83,6 +83,8 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only _key_mapping[shm_key] = [False, 0] _key_mapping[shm_key][1] += 1 except FileNotFoundError: + # File not found means the shared memory region has not been created, + # suppress the exception and attempt to create the region. pass if shm_handle._mpsm_handle is None: try: @@ -137,6 +139,8 @@ def set_shared_memory_region(shm_handle, input_values, offset=0): try: for input_value in input_values: + # numpy array of object type is "syntactic sugar" for the API, should + # be handled by accessing its item and treat as Python object if input_value.dtype == np.object_: byte_size = len(input_value.item()) shm_handle._mpsm_handle.buf[offset : offset + byte_size] = ( From 886d5b2effb2d61faf0aa588e34b0480ecece620 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 30 Oct 2024 19:17:27 -0700 Subject: [PATCH 7/7] fix: address comment --- .../library/tests/test_shared_memory.py | 11 ++++------ .../utils/shared_memory/__init__.py | 22 +++++++++++-------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/python/library/tests/test_shared_memory.py b/src/python/library/tests/test_shared_memory.py index f765bc781..36c64f090 100644 --- a/src/python/library/tests/test_shared_memory.py +++ b/src/python/library/tests/test_shared_memory.py @@ -101,10 +101,8 @@ def test_set_region_oversize(self): 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 + # by default, 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) ) @@ -133,9 +131,8 @@ def test_duplicate_key(self): 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 + # destruction of duplicate shared memory region will occur when the 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) diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py index 364b58b99..12904445e 100755 --- a/src/python/library/tritonclient/utils/shared_memory/__init__.py +++ b/src/python/library/tritonclient/utils/shared_memory/__init__.py @@ -62,7 +62,8 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only 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. + the previously allocated shared memory size can be different from + the size requested. Returns ------- @@ -80,8 +81,11 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only try: shm_handle._mpsm_handle = mpshm.SharedMemory(shm_key) if shm_key not in _key_mapping: - _key_mapping[shm_key] = [False, 0] - _key_mapping[shm_key][1] += 1 + _key_mapping[shm_key] = { + "needs_unlink": False, + "active_handle_count": 0, + } + _key_mapping[shm_key]["active_handle_count"] += 1 except FileNotFoundError: # File not found means the shared memory region has not been created, # suppress the exception and attempt to create the region. @@ -96,9 +100,9 @@ def create_shared_memory_region(triton_shm_name, shm_key, byte_size, create_only "unable to create the shared memory region" ) from ex if shm_key not in _key_mapping: - _key_mapping[shm_key] = [False, 0] - _key_mapping[shm_key][0] = True - _key_mapping[shm_key][1] += 1 + _key_mapping[shm_key] = {"needs_unlink": False, "active_handle_count": 0} + _key_mapping[shm_key]["needs_unlink"] = True + _key_mapping[shm_key]["active_handle_count"] += 1 if byte_size > shm_handle._mpsm_handle.size: warnings.warn( @@ -238,10 +242,10 @@ def destroy_shared_memory_region(shm_handle): # fail to delete a region, we should not report it back to the user # as a valid memory region. shm_handle._mpsm_handle.close() - _key_mapping[shm_handle._shm_key][1] -= 1 - if _key_mapping[shm_handle._shm_key][1] == 0: + _key_mapping[shm_handle._shm_key]["active_handle_count"] -= 1 + if _key_mapping[shm_handle._shm_key]["active_handle_count"] == 0: try: - if _key_mapping[shm_handle._shm_key][0]: + if _key_mapping[shm_handle._shm_key]["needs_unlink"]: shm_handle._mpsm_handle.unlink() finally: _key_mapping.pop(shm_handle._shm_key)