Skip to content

fix: Make CUDA shared memory stream cache thread-safe in Python client #915

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jpbhdrey wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: Make CUDA shared memory stream cache thread-safe in Python client #915

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/python/library/tritonclient/utils/cuda_shared_memory/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
import base64
import ctypes
import struct
import threading

import numpy as np

Expand All @@ -58,14 +59,16 @@
# and be reused throughout the process. May revisit for stream pool if
# asynchronous write on CUDA shared memory region is requested
_dlpack_stream = {}
_dlpack_stream_lock = threading.Lock()


# Helper function to retrieve internally managed CUDA stream
def _get_or_create_global_cuda_stream(device_id):
global _dlpack_stream
if device_id not in _dlpack_stream:
_dlpack_stream[device_id] = CudaStream(device_id)
return _dlpack_stream[device_id]._stream
with _dlpack_stream_lock:
if device_id not in _dlpack_stream:
_dlpack_stream[device_id] = CudaStream(device_id)
return _dlpack_stream[device_id]._stream


def _support_uva(shm_device_id, ext_device_id):
Expand Down