chore: update function comment. address comment

triton-inference-server · Oct 23, 2024 · f791cd4 · f791cd4
1 parent bcdfe3e
commit f791cd4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git 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
@@ -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
     ----------