From f791cd478d2358707113367f138b5ec4f30c2a15 Mon Sep 17 00:00:00 2001
From: Guan Luo <gluo@nvidia.com>
Date: Tue, 22 Oct 2024 18:44:40 -0700
Subject: [PATCH] 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
     ----------