Sitelet https://github.com/triton-inference-server/client/commit/9b2b42dc1b1798a6dd9805de22b1b2f042ea53ce
Skip to content

Commit 9b2b42d

Browse files
GuanLuoGuanLuo
authored
refactor: Remove C shared memory shim. Refine shared memory utils behavior. Add unit test for shared memory (#797)
* test: add unit test for shared memory * refactor: reflect latest comment on create / destroy behavior * fix: fix variable shadowing * chore: update function comment. address comment * refactor: remove C shared memory shim * style: add commnet and fix copyright * fix: address comment
1 parent 6962adf commit 9b2b42d

9 files changed

Lines changed: 284 additions & 453 deletions

File tree

‎src/python/library/CMakeLists.txt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,6 @@ add_custom_target(
9696
if (NOT WIN32)
9797
# Can generate linux specific wheel file on linux systems only.
9898
set(LINUX_WHEEL_DEPENDS
99-
cshm
10099
${WHEEL_DEPENDS}
101100
)
102101

‎src/python/library/build_wheel.py‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#!/usr/bin/env python3
2-
# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
#
44
# Redistribution and use in source and binary forms, with or without
55
# modification, are permitted provided that the following conditions
@@ -174,10 +174,6 @@ def sed(pattern, replace, source, dest=None):
174174
"tritonclient/utils/shared_memory",
175175
os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory"),
176176
)
177-
shutil.copyfile(
178-
"tritonclient/utils/libcshm.so",
179-
os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory/libcshm.so"),
180-
)
181177
cpdir(
182178
"tritonclient/utils/cuda_shared_memory",
183179
os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"),

‎src/python/library/setup.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,6 @@ def req_file(filename, folder="requirements"):
7676
extras_require["all"] = list(chain(extras_require.values()))
7777

7878
platform_package_data = []
79-
if PLATFORM_FLAG != "any":
80-
platform_package_data += ["libcshm.so"]
8179

8280
data_files = [
8381
("", ["LICENSE.txt"]),

‎src/python/library/tests/test_shared_memory.py‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# Redistribution and use in source and binary forms, with or without
4+
# modification, are permitted provided that the following conditions
5+
# are met:
6+
# * Redistributions of source code must retain the above copyright
7+
# notice, this list of conditions and the following disclaimer.
8+
# * Redistributions in binary form must reproduce the above copyright
9+
# notice, this list of conditions and the following disclaimer in the
10+
# documentation and/or other materials provided with the distribution.
11+
# * Neither the name of NVIDIA CORPORATION nor the names of its
12+
# contributors may be used to endorse or promote products derived
13+
# from this software without specific prior written permission.
14+
#
15+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
16+
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17+
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18+
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
19+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22+
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
23+
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24+
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25+
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26+
27+
import unittest
28+
29+
import numpy
30+
import tritonclient.utils as utils
31+
import tritonclient.utils.shared_memory as shm
32+
33+
34+
class SharedMemoryTest(unittest.TestCase):
35+
"""
36+
Testing shared memory utilities
37+
"""
38+
39+
def setUp(self):
40+
self.shm_handles = []
41+
42+
def tearDown(self):
43+
for shm_handle in self.shm_handles:
44+
shm.destroy_shared_memory_region(shm_handle)
45+
46+
def test_lifecycle(self):
47+
cpu_tensor = numpy.ones([4, 4], dtype=numpy.float32)
48+
byte_size = 64
49+
self.shm_handles.append(
50+
shm.create_shared_memory_region("shm_name", "shm_key", byte_size)
51+
)
52+
53+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
54+
55+
# Set data from Numpy array
56+
shm.set_shared_memory_region(self.shm_handles[0], [cpu_tensor])
57+
shm_tensor = shm.get_contents_as_numpy(
58+
self.shm_handles[0], numpy.float32, [4, 4]
59+
)
60+
61+
self.assertTrue(numpy.allclose(cpu_tensor, shm_tensor))
62+
63+
shm.destroy_shared_memory_region(self.shm_handles.pop(0))
64+
65+
def test_invalid_create_shm(self):
66+
# Raises error since tried to create invalid system shared memory region
67+
with self.assertRaisesRegex(
68+
shm.SharedMemoryException, "unable to create the shared memory region"
69+
):
70+
self.shm_handles.append(
71+
shm.create_shared_memory_region("dummy_data", "/dummy_data", -1)
72+
)
73+
74+
def test_set_region_offset(self):
75+
large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
76+
large_size = 64
77+
self.shm_handles.append(
78+
shm.create_shared_memory_region("shm_name", "shm_key", large_size)
79+
)
80+
shm.set_shared_memory_region(self.shm_handles[0], [large_tensor])
81+
small_tensor = numpy.zeros([2, 4], dtype=numpy.float32)
82+
small_size = 32
83+
shm.set_shared_memory_region(
84+
self.shm_handles[0], [small_tensor], offset=large_size - small_size
85+
)
86+
shm_tensor = shm.get_contents_as_numpy(
87+
self.shm_handles[0], numpy.float32, [2, 4], offset=large_size - small_size
88+
)
89+
90+
self.assertTrue(numpy.allclose(small_tensor, shm_tensor))
91+
92+
def test_set_region_oversize(self):
93+
large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
94+
small_size = 32
95+
self.shm_handles.append(
96+
shm.create_shared_memory_region("shm_name", "shm_key", small_size)
97+
)
98+
with self.assertRaisesRegex(
99+
shm.SharedMemoryException, "unable to set the shared memory region"
100+
):
101+
shm.set_shared_memory_region(self.shm_handles[0], [large_tensor])
102+
103+
def test_duplicate_key(self):
104+
# by default, return the same handle if existed, warning will be print
105+
# if size is different
106+
self.shm_handles.append(
107+
shm.create_shared_memory_region("shm_name", "shm_key", 32)
108+
)
109+
with self.assertRaisesRegex(
110+
shm.SharedMemoryException,
111+
"unable to create the shared memory region",
112+
):
113+
self.shm_handles.append(
114+
shm.create_shared_memory_region(
115+
"shm_name", "shm_key", 32, create_only=True
116+
)
117+
)
118+
119+
# Get handle to the same shared memory region but with larger size requested,
120+
# check if actual size is checked
121+
self.shm_handles.append(
122+
shm.create_shared_memory_region("shm_name", "shm_key", 64)
123+
)
124+
125+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
126+
127+
large_tensor = numpy.ones([4, 4], dtype=numpy.float32)
128+
with self.assertRaisesRegex(
129+
shm.SharedMemoryException, "unable to set the shared memory region"
130+
):
131+
shm.set_shared_memory_region(self.shm_handles[-1], [large_tensor])
132+
133+
def test_destroy_duplicate(self):
134+
# destruction of duplicate shared memory region will occur when the last
135+
# managed handle is destroyed
136+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 0)
137+
self.shm_handles.append(
138+
shm.create_shared_memory_region("shm_name", "shm_key", 64)
139+
)
140+
self.shm_handles.append(
141+
shm.create_shared_memory_region("shm_name", "shm_key", 32)
142+
)
143+
self.shm_handles.append(
144+
shm.create_shared_memory_region("shm_name", "shm_key", 32)
145+
)
146+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
147+
148+
shm.destroy_shared_memory_region(self.shm_handles.pop(0))
149+
shm.destroy_shared_memory_region(self.shm_handles.pop(0))
150+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 1)
151+
152+
shm.destroy_shared_memory_region(self.shm_handles.pop(0))
153+
self.assertEqual(len(shm.mapped_shared_memory_regions()), 0)
154+
155+
def test_numpy_bytes(self):
156+
int_tensor = numpy.arange(start=0, stop=16, dtype=numpy.int32)
157+
bytes_tensor = numpy.array(
158+
[str(x).encode("utf-8") for x in int_tensor.flatten()], dtype=object
159+
)
160+
bytes_tensor = bytes_tensor.reshape(int_tensor.shape)
161+
bytes_tensor_serialized = utils.serialize_byte_tensor(bytes_tensor)
162+
byte_size = utils.serialized_byte_size(bytes_tensor_serialized)
163+
164+
self.shm_handles.append(
165+
shm.create_shared_memory_region("shm_name", "shm_key", byte_size)
166+
)
167+
168+
# Set data from Numpy array
169+
shm.set_shared_memory_region(self.shm_handles[0], [bytes_tensor_serialized])
170+
171+
shm_tensor = shm.get_contents_as_numpy(
172+
self.shm_handles[0],
173+
numpy.object_,
174+
[
175+
16,
176+
],
177+
)
178+
179+
self.assertTrue(numpy.array_equal(bytes_tensor, shm_tensor))
180+
181+
182+
if __name__ == "__main__":
183+
unittest.main()

‎src/python/library/tritonclient/utils/CMakeLists.txt‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
#
33
# Redistribution and use in source and binary forms, with or without
44
# modification, are permitted provided that the following conditions
@@ -28,20 +28,6 @@ configure_file(__init__.py __init__.py COPYONLY)
2828
configure_file(_dlpack.py _dlpack.py COPYONLY)
2929
configure_file(_shared_memory_tensor.py _shared_memory_tensor.py COPYONLY)
3030

31-
if(NOT WIN32)
32-
file(COPY shared_memory DESTINATION .)
33-
34-
#
35-
# libcshm.so
36-
#
37-
add_library(cshm SHARED shared_memory/shared_memory.cc)
38-
if(${TRITON_ENABLE_GPU})
39-
target_compile_definitions(cshm PUBLIC TRITON_ENABLE_GPU=1)
40-
target_link_libraries(cshm PUBLIC CUDA::cudart)
41-
endif() # TRITON_ENABLE_GPU
42-
target_link_libraries(cshm PRIVATE rt)
43-
endif() # WIN32
44-
4531
if(NOT WIN32)
4632
configure_file(shared_memory/__init__.py shared_memory/__init__.py COPYONLY)
4733
configure_file(cuda_shared_memory/__init__.py cuda_shared_memory/__init__.py COPYONLY)

0 commit comments

Comments
 (0)