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 ----------