Merge Win/Unix functions. Opaque shm file handling

fpetrini15 · fpetrini15 · commit 8816dc6dff4a · 2024-04-05T19:42:46.000-07:00
diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt
@@ -1,4 +1,4 @@
-# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -45,13 +45,6 @@ if(NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE Release)
 endif()
 
-# FIXME: Windows client currently does not support GPU tensors.
-# For simplicity, we will override this option here.
-if(WIN32 AND TRITON_ENABLE_GPU)
-  message("GPU shared memory is not currently supported by the Windows python client. Forcing TRITON_ENABLE_GPU to false.")
-  set(TRITON_ENABLE_GPU OFF CACHE BOOL "GPU disabled" FORCE)
-endif()
-
 #
 # Dependencies
 #
diff --git a/src/python/library/CMakeLists.txt b/src/python/library/CMakeLists.txt
@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -180,13 +180,13 @@ endif() # TRITON_ENABLE_PYTHON_GRPC
 
 
 if(WIN32)
-  install(
-    CODE "file(GLOB _Wheel \"${CMAKE_CURRENT_BINARY_DIR}/windows/triton*.whl\")"
-    CODE "file(INSTALL \${_Wheel} DESTINATION \"${CMAKE_INSTALL_PREFIX}/python\")"
-  )
+  set(WHEEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/windows")
 else()
-  install(
-    CODE "file(GLOB _Wheel \"${CMAKE_CURRENT_BINARY_DIR}/linux/triton*.whl\")"
-    CODE "file(INSTALL \${_Wheel} DESTINATION \"${CMAKE_INSTALL_PREFIX}/python\")"
-  )
-endif() # NOT WIN32
+  set(WHEEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/linux")
+endif() # WIN32
+
+install(
+  CODE "file(GLOB _Wheel \"${WHEEL_DIR}/triton*.whl\")"
+  CODE "file(INSTALL \${_Wheel} DESTINATION \"${CMAKE_INSTALL_PREFIX}/python\")"
+)
+
diff --git a/src/python/library/build_wheel.py b/src/python/library/build_wheel.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -78,13 +78,14 @@ def sed(pattern, replace, source, dest=None):
     parser.add_argument(
         "--dest-dir", type=str, required=True, help="Destination directory."
     )
-    parser.add_argument(
+    platform_group = parser.add_mutually_exclusive_group()
+    platform_group.add_argument(
         "--linux",
         action="store_true",
         required=False,
         help="Include linux specific artifacts.",
     )
-    parser.add_argument(
+    platform_group.add_argument(
         "--windows",
         action="store_true",
         required=False,
@@ -210,7 +211,7 @@ def sed(pattern, replace, source, dest=None):
             "tritonclient/utils/Release/cshm.dll",
             os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory/cshm.dll"),
         )
-        # FIXME: Enable when Windows supports GPU tensors
+        # FIXME: Enable when Windows supports GPU tensors DLIS-4169
         # cpdir(
         #     "tritonclient/utils/cuda_shared_memory",
         #     os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"),
diff --git a/src/python/library/setup.py b/src/python/library/setup.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
diff --git a/src/python/library/tritonclient/utils/CMakeLists.txt b/src/python/library/tritonclient/utils/CMakeLists.txt
@@ -1,4 +1,4 @@
-# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -24,6 +24,13 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+# FIXME: Windows client currently does not support GPU tensors.
+# For simplicity, we will override this option here.
+if(WIN32 AND TRITON_ENABLE_GPU)
+  message(FATAL_ERROR  "GPU shared memory is not currently supported by the Windows python client.")
+  set(TRITON_ENABLE_GPU OFF CACHE BOOL "GPU disabled" FORCE)
+endif()
+
 configure_file(__init__.py __init__.py COPYONLY)
 configure_file(_dlpack.py _dlpack.py COPYONLY)
 configure_file(_shared_memory_tensor.py _shared_memory_tensor.py COPYONLY)
@@ -40,7 +47,7 @@ endif() # TRITON_ENABLE_GPU
 
 if(NOT WIN32)
   target_link_libraries(cshm PRIVATE rt)
-endif() # WIN32
+endif() # NOT WIN32
 
 configure_file(shared_memory/__init__.py shared_memory/__init__.py COPYONLY)
 
diff --git a/src/python/library/tritonclient/utils/shared_memory/__init__.py b/src/python/library/tritonclient/utils/shared_memory/__init__.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2019-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -63,7 +63,7 @@ def from_param(cls, value):
     c_void_p,
     POINTER(c_char_p),
     POINTER(c_char_p),
-    POINTER(c_int),
+    POINTER(c_void_p),
     POINTER(c_uint64),
     POINTER(c_uint64),
 ]
@@ -205,7 +205,7 @@ def get_contents_as_numpy(shm_handle, datatype, shape, offset=0):
         The numpy array generated using the contents of the specified shared
         memory region.
     """
-    shm_fd = c_int()
+    shm_file = c_void_p()
     region_offset = c_uint64()
     byte_size = c_uint64()
     shm_addr = c_char_p()
@@ -216,7 +216,7 @@ def get_contents_as_numpy(shm_handle, datatype, shape, offset=0):
                 shm_handle,
                 byref(shm_addr),
                 byref(shm_key),
-                byref(shm_fd),
+                byref(shm_file),
                 byref(region_offset),
                 byref(byte_size),
             )
@@ -285,9 +285,7 @@ def destroy_shared_memory_region(shm_handle):
         If unable to unlink the shared memory region.
     """
 
-    _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle)))
-
-    shm_fd = c_int()
+    shm_file = c_void_p()
     offset = c_uint64()
     byte_size = c_uint64()
     shm_addr = c_char_p()
@@ -298,13 +296,16 @@ def destroy_shared_memory_region(shm_handle):
                 shm_handle,
                 byref(shm_addr),
                 byref(shm_key),
-                byref(shm_fd),
+                byref(shm_file),
                 byref(offset),
                 byref(byte_size),
             )
         )
     )
-    mapped_shm_regions.remove(shm_key.value.decode("utf-8"))
+    shm_key_copy = bytes(shm_key.value)
+    _raise_if_error(c_int(_cshm_shared_memory_region_destroy(shm_handle)))
+
+    mapped_shm_regions.remove(shm_key_copy.decode("utf-8"))
 
     return
 
diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc b/src/python/library/tritonclient/utils/shared_memory/shared_memory.cc
diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory.h b/src/python/library/tritonclient/utils/shared_memory/shared_memory.h
diff --git a/src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h b/src/python/library/tritonclient/utils/shared_memory/shared_memory_handle.h

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.`
	`1`	`+# Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.`
`2`	`2`	`#`
`3`	`3`	`# Redistribution and use in source and binary forms, with or without`
`4`	`4`	`# modification, are permitted provided that the following conditions`
`@@ -45,13 +45,6 @@ if(NOT CMAKE_BUILD_TYPE)`
`45`	`45`	`set(CMAKE_BUILD_TYPE Release)`
`46`	`46`	`endif()`
`47`	`47`
`48`		`-# FIXME: Windows client currently does not support GPU tensors.`
`49`		`-# For simplicity, we will override this option here.`
`50`		`-if(WIN32 AND TRITON_ENABLE_GPU)`
`51`		`- message("GPU shared memory is not currently supported by the Windows python client. Forcing TRITON_ENABLE_GPU to false.")`
`52`		`- set(TRITON_ENABLE_GPU OFF CACHE BOOL "GPU disabled" FORCE)`
`53`		`-endif()`
`54`		`-`
`55`	`48`	`#`
`56`	`49`	`# Dependencies`
`57`	`50`	`#`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`#!/usr/bin/env python3`
`2`	`2`
`3`		`-# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.`
	`3`	`+# Copyright 2020-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.`
`4`	`4`	`#`
`5`	`5`	`# Redistribution and use in source and binary forms, with or without`
`6`	`6`	`# modification, are permitted provided that the following conditions`