Skip to content

Commit de8aa9f

Browse files
mc-nvGuanLuo
mc-nv
and
GuanLuo
authored
fix: WAR for Python CUDA library unknown race condition (#8361)
Co-authored-by: GuanLuo <gluo@nvidia.com>
1 parent a514535 commit de8aa9f

1 file changed

Lines changed: 84 additions & 71 deletions

File tree

‎qa/common/shm_util.py‎

Lines changed: 84 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#!/usr/bin/env python3
22

3-
# Copyright 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
44
#
55
# Redistribution and use in source and binary forms, with or without
66
# modification, are permitted provided that the following conditions
@@ -27,6 +27,7 @@
2727
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
2828

2929
import os
30+
import threading
3031
import time
3132
from ctypes import *
3233
from os import listdir
@@ -35,6 +36,8 @@
3536
import tritonclient.http as httpclient
3637
from tritonclient.utils import *
3738

39+
CREATION_LOCK = threading.Lock()
40+
3841
# By default, find tritonserver on "localhost", but can be overridden
3942
# with TRITONSERVER_IPADDR envvar
4043
_tritonserver_ipaddr = os.environ.get("TRITONSERVER_IPADDR", "localhost")
@@ -94,64 +97,69 @@ def create_set_shm_regions(
9497
shm_op0_handle = None
9598
shm_op1_handle = None
9699

97-
if use_system_shared_memory:
98-
shm_ip0_handle = shm.create_shared_memory_region(
99-
shm_region_names[0] + "_data", "/" + shm_region_names[0], input0_byte_size
100-
)
101-
shm_ip1_handle = shm.create_shared_memory_region(
102-
shm_region_names[1] + "_data", "/" + shm_region_names[1], input1_byte_size
103-
)
100+
with CREATION_LOCK:
101+
if use_system_shared_memory:
102+
shm_ip0_handle = shm.create_shared_memory_region(
103+
shm_region_names[0] + "_data",
104+
"/" + shm_region_names[0],
105+
input0_byte_size,
106+
)
107+
shm_ip1_handle = shm.create_shared_memory_region(
108+
shm_region_names[1] + "_data",
109+
"/" + shm_region_names[1],
110+
input1_byte_size,
111+
)
104112

105-
i = 0
106-
if "OUTPUT0" in outputs:
107-
if precreated_shm_regions is None:
108-
shm_op0_handle = shm.create_shared_memory_region(
109-
shm_region_names[2] + "_data",
110-
"/" + shm_region_names[2],
111-
output0_byte_size,
112-
)
113-
else:
114-
shm_op0_handle = precreated_shm_regions[0]
115-
i += 1
116-
if "OUTPUT1" in outputs:
117-
if precreated_shm_regions is None:
118-
shm_op1_handle = shm.create_shared_memory_region(
119-
shm_region_names[2 + i] + "_data",
120-
"/" + shm_region_names[2 + i],
121-
output1_byte_size,
122-
)
123-
else:
124-
shm_op1_handle = precreated_shm_regions[i]
113+
i = 0
114+
if "OUTPUT0" in outputs:
115+
if precreated_shm_regions is None:
116+
shm_op0_handle = shm.create_shared_memory_region(
117+
shm_region_names[2] + "_data",
118+
"/" + shm_region_names[2],
119+
output0_byte_size,
120+
)
121+
else:
122+
shm_op0_handle = precreated_shm_regions[0]
123+
i += 1
124+
if "OUTPUT1" in outputs:
125+
if precreated_shm_regions is None:
126+
shm_op1_handle = shm.create_shared_memory_region(
127+
shm_region_names[2 + i] + "_data",
128+
"/" + shm_region_names[2 + i],
129+
output1_byte_size,
130+
)
131+
else:
132+
shm_op1_handle = precreated_shm_regions[i]
125133

126-
shm.set_shared_memory_region(shm_ip0_handle, input0_list)
127-
shm.set_shared_memory_region(shm_ip1_handle, input1_list)
134+
shm.set_shared_memory_region(shm_ip0_handle, input0_list)
135+
shm.set_shared_memory_region(shm_ip1_handle, input1_list)
128136

129-
if use_cuda_shared_memory:
130-
shm_ip0_handle = cudashm.create_shared_memory_region(
131-
shm_region_names[0] + "_data", input0_byte_size, 0
132-
)
133-
shm_ip1_handle = cudashm.create_shared_memory_region(
134-
shm_region_names[1] + "_data", input1_byte_size, 0
135-
)
136-
i = 0
137-
if "OUTPUT0" in outputs:
138-
if precreated_shm_regions is None:
139-
shm_op0_handle = cudashm.create_shared_memory_region(
140-
shm_region_names[2] + "_data", output0_byte_size, 0
141-
)
142-
else:
143-
shm_op0_handle = precreated_shm_regions[0]
144-
i += 1
145-
if "OUTPUT1" in outputs:
146-
if precreated_shm_regions is None:
147-
shm_op1_handle = cudashm.create_shared_memory_region(
148-
shm_region_names[2 + i] + "_data", output1_byte_size, 0
149-
)
150-
else:
151-
shm_op1_handle = precreated_shm_regions[i]
137+
if use_cuda_shared_memory:
138+
shm_ip0_handle = cudashm.create_shared_memory_region(
139+
shm_region_names[0] + "_data", input0_byte_size, 0
140+
)
141+
shm_ip1_handle = cudashm.create_shared_memory_region(
142+
shm_region_names[1] + "_data", input1_byte_size, 0
143+
)
144+
i = 0
145+
if "OUTPUT0" in outputs:
146+
if precreated_shm_regions is None:
147+
shm_op0_handle = cudashm.create_shared_memory_region(
148+
shm_region_names[2] + "_data", output0_byte_size, 0
149+
)
150+
else:
151+
shm_op0_handle = precreated_shm_regions[0]
152+
i += 1
153+
if "OUTPUT1" in outputs:
154+
if precreated_shm_regions is None:
155+
shm_op1_handle = cudashm.create_shared_memory_region(
156+
shm_region_names[2 + i] + "_data", output1_byte_size, 0
157+
)
158+
else:
159+
shm_op1_handle = precreated_shm_regions[i]
152160

153-
cudashm.set_shared_memory_region(shm_ip0_handle, input0_list)
154-
cudashm.set_shared_memory_region(shm_ip1_handle, input1_list)
161+
cudashm.set_shared_memory_region(shm_ip0_handle, input0_list)
162+
cudashm.set_shared_memory_region(shm_ip1_handle, input1_list)
155163

156164
return shm_region_names, [
157165
shm_ip0_handle,
@@ -337,22 +345,27 @@ def create_set_either_shm_region(
337345
if not (use_system_shared_memory or use_cuda_shared_memory):
338346
return []
339347

340-
if use_cuda_shared_memory:
341-
shm_ip_handle = cudashm.create_shared_memory_region(
342-
shm_region_names[0] + "_data", input_byte_size, 0
343-
)
344-
shm_op_handle = cudashm.create_shared_memory_region(
345-
shm_region_names[1] + "_data", output_byte_size, 0
346-
)
347-
cudashm.set_shared_memory_region(shm_ip_handle, input_list)
348-
elif use_system_shared_memory:
349-
shm_ip_handle = shm.create_shared_memory_region(
350-
shm_region_names[0] + "_data", "/" + shm_region_names[0], input_byte_size
351-
)
352-
shm_op_handle = shm.create_shared_memory_region(
353-
shm_region_names[1] + "_data", "/" + shm_region_names[1], output_byte_size
354-
)
355-
shm.set_shared_memory_region(shm_ip_handle, input_list)
348+
with CREATION_LOCK:
349+
if use_cuda_shared_memory:
350+
shm_ip_handle = cudashm.create_shared_memory_region(
351+
shm_region_names[0] + "_data", input_byte_size, 0
352+
)
353+
shm_op_handle = cudashm.create_shared_memory_region(
354+
shm_region_names[1] + "_data", output_byte_size, 0
355+
)
356+
cudashm.set_shared_memory_region(shm_ip_handle, input_list)
357+
elif use_system_shared_memory:
358+
shm_ip_handle = shm.create_shared_memory_region(
359+
shm_region_names[0] + "_data",
360+
"/" + shm_region_names[0],
361+
input_byte_size,
362+
)
363+
shm_op_handle = shm.create_shared_memory_region(
364+
shm_region_names[1] + "_data",
365+
"/" + shm_region_names[1],
366+
output_byte_size,
367+
)
368+
shm.set_shared_memory_region(shm_ip_handle, input_list)
356369

357370
return [shm_ip_handle, shm_op_handle]
358371

0 commit comments

Comments
 (0)