Skip to content
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

Jump to bottom

[Windows] Support CPU shared memory (Client/Frontend) #551

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
Open

[Windows] Support CPU shared memory (Client/Frontend) #551

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f8852d1
Initial commit
fpetrini15 Mar 15, 2024
d4f9cee
CMakeLists cleanup
fpetrini15 Mar 16, 2024
b69868f
Fix typo
fpetrini15 Mar 20, 2024
d28a780
Use generic wheel
fpetrini15 Mar 20, 2024
875331f
Revert to windows-wheel and moved GPU override
fpetrini15 Mar 27, 2024
dcefeac
Merge Win/Unix functions. Opaque shm file handling
fpetrini15 Apr 4, 2024
caa2ed1
Revert Cmake copyright
fpetrini15 Apr 4, 2024
191be4c
Fix ctypes, don't close fd, use smart pointers
fpetrini15 Apr 6, 2024
bd9bea6
Review comments from Guan
fpetrini15 Apr 11, 2024
f854d4a
Introduce backing file
fpetrini15 Apr 13, 2024
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
79 changes: 61 additions & 18 deletions src/python/library/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -38,9 +38,7 @@ if(${TRITON_ENABLE_PYTHON_HTTP})
file(COPY tritonhttpclient DESTINATION .)
endif() # TRITON_ENABLE_PYTHON_HTTP
file(COPY tritonclientutils DESTINATION .)
if (NOT WIN32)
file(COPY tritonshmutils DESTINATION .)
endif() # NOT WIN32
file(COPY tritonshmutils DESTINATION .)
####################################

file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/TRITON_VERSION ${TRITON_VERSION})
Expand Down Expand Up @@ -91,9 +89,37 @@ add_custom_target(
)

#
# Linux specific Wheel file. Compatible with x86, x64 and aarch64
# Windows-specific Wheel file.
#
if (NOT WIN32)
if(WIN32)
set(WINDOWS_WHEEL_DEPENDS
cshm
${WHEEL_DEPENDS}
)
if (${TRITON_ENABLE_PERF_ANALYZER})
set(perf_analyzer_arg --perf-analyzer ${CMAKE_INSTALL_PREFIX}/bin/perf_analyzer)
endif()
set(windows_wheel_stamp_file "windows_stamp.whl")
add_custom_command(
OUTPUT "${windows_wheel_stamp_file}"
COMMAND python3
ARGS
"${CMAKE_CURRENT_SOURCE_DIR}/build_wheel.py"
--dest-dir "${CMAKE_CURRENT_BINARY_DIR}/windows"
--windows
${perf_analyzer_arg}
DEPENDS ${WINDOWS_WHEEL_DEPENDS}
)

add_custom_target(
windows-client-wheel ALL
DEPENDS
"${windows_wheel_stamp_file}"
)
else()
#
# Linux specific Wheel file. Compatible with x86, x64 and aarch64
#
# Can generate linux specific wheel file on linux systems only.
set(LINUX_WHEEL_DEPENDS
cshm
Expand All @@ -102,7 +128,10 @@ if (NOT WIN32)

if (${TRITON_ENABLE_PERF_ANALYZER})
set(perf_analyzer_arg --perf-analyzer ${CMAKE_INSTALL_PREFIX}/bin/perf_analyzer)
endif()
endif() # TRITON_ENABLE_PERF_ANALYZER
if (${TRITON_ENABLE_GPU})
set(gpu_arg --include-gpu-libs)
endif() # TRITON_ENABLE_GPU
set(linux_wheel_stamp_file "linux_stamp.whl")
add_custom_command(
OUTPUT "${linux_wheel_stamp_file}"
Expand All @@ -112,6 +141,7 @@ if (NOT WIN32)
--dest-dir "${CMAKE_CURRENT_BINARY_DIR}/linux"
--linux
${perf_analyzer_arg}
${gpu_arg}
DEPENDS ${LINUX_WHEEL_DEPENDS}
)

Expand All @@ -120,20 +150,25 @@ if (NOT WIN32)
DEPENDS
"${linux_wheel_stamp_file}"
)
endif() # NOT WIN32
endif() # WIN32

if(${TRITON_ENABLE_PYTHON_GRPC})
add_dependencies(
generic-client-wheel
grpc-service-py-library proto-py-library
)

if (NOT WIN32)
if (WIN32)
add_dependencies(
windows-client-wheel
grpc-service-py-library proto-py-library
)
else()
add_dependencies(
linux-client-wheel
grpc-service-py-library proto-py-library
)
endif() # NOT WIN32
)
endif() # WIN32

file(
GLOB generated-py
Expand All @@ -147,14 +182,22 @@ if(${TRITON_ENABLE_PYTHON_GRPC})
)
endif() # TRITON_ENABLE_PYTHON_GRPC

# Generic Wheel
set(WHEEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/generic")
install(
CODE "file(GLOB _Wheel \"${CMAKE_CURRENT_BINARY_DIR}/generic/triton*.whl\")"
CODE "file(GLOB _Wheel \"${WHEEL_DIR}/triton*.whl\")"
CODE "file(INSTALL \${_Wheel} DESTINATION \"${CMAKE_INSTALL_PREFIX}/python\")"
)

# Platform-specific wheels
if(WIN32)
set(WHEEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/windows")
else()
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\")"
)

if (NOT WIN32)
install(
CODE "file(GLOB _Wheel \"${CMAKE_CURRENT_BINARY_DIR}/linux/triton*.whl\")"
CODE "file(INSTALL \${_Wheel} DESTINATION \"${CMAKE_INSTALL_PREFIX}/python\")"
)
endif() # NOT WIN32
48 changes: 41 additions & 7 deletions src/python/library/build_wheel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -28,6 +28,7 @@
import argparse
import os
import pathlib
import platform
import re
import shutil
import subprocess
Expand Down Expand Up @@ -77,12 +78,25 @@
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.",
)
platform_group.add_argument(
"--windows",
action="store_true",
required=False,
help="Include windows specific artifacts.",
)
parser.add_argument(
"--include-gpu-libs",
action="store_true",
required=False,
help="Include gpu specific libraries",
)
parser.add_argument(
"--perf-analyzer",
type=str,
Expand Down Expand Up @@ -118,7 +132,7 @@
cpdir("tritonhttpclient", os.path.join(FLAGS.whl_dir, "tritonhttpclient"))
if os.path.isdir("tritongrpcclient"):
cpdir("tritongrpcclient", os.path.join(FLAGS.whl_dir, "tritongrpcclient"))
if FLAGS.linux:
if FLAGS.linux or FLAGS.windows:
if os.path.isdir("tritonshmutils"):
cpdir("tritonshmutils", os.path.join(FLAGS.whl_dir, "tritonshmutils"))

Expand Down Expand Up @@ -178,10 +192,11 @@
"tritonclient/utils/libcshm.so",
os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory/libcshm.so"),
)
cpdir(
"tritonclient/utils/cuda_shared_memory",
os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"),
)
if FLAGS.include_gpu_libs:
cpdir(
"tritonclient/utils/cuda_shared_memory",
os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"),
)

# Copy the pre-compiled perf_analyzer binary
if FLAGS.perf_analyzer is not None:
Expand All @@ -194,6 +209,22 @@
if not os.path.exists(os.path.join(FLAGS.whl_dir, "perf_client")):
os.symlink("perf_analyzer", os.path.join(FLAGS.whl_dir, "perf_client"))

if FLAGS.windows:
cpdir(
"tritonclient/utils/shared_memory",
os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory"),
)
shutil.copyfile(
"tritonclient/utils/Release/cshm.dll",
os.path.join(FLAGS.whl_dir, "tritonclient/utils/shared_memory/cshm.dll"),
)
# FIXME: Enable when Windows supports GPU tensors DLIS-4169
# if FLAGS.include_gpu_libs:
# cpdir(
# "tritonclient/utils/cuda_shared_memory",
# os.path.join(FLAGS.whl_dir, "tritonclient/utils/cuda_shared_memory"),
Comment on lines +222 to +225

Check notice

Code scanning / CodeQL

Commented-out code Note

This comment appears to contain commented-out code.
# )

shutil.copyfile("LICENSE.txt", os.path.join(FLAGS.whl_dir, "LICENSE.txt"))
shutil.copyfile("setup.py", os.path.join(FLAGS.whl_dir, "setup.py"))
cpdir("requirements", os.path.join(FLAGS.whl_dir, "requirements"))
Expand All @@ -208,6 +239,9 @@
else:
platform_name = "manylinux1_x86_64"
args = ["python3", "setup.py", "bdist_wheel", "--plat-name", platform_name]
elif FLAGS.windows and platform.uname().machine == "AMD64":
platform_name = "win_amd64"
args = ["python3", "setup.py", "bdist_wheel", "--plat-name", platform_name]
else:
args = ["python3", "setup.py", "bdist_wheel"]

Expand Down
6 changes: 4 additions & 2 deletions src/python/library/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -76,8 +76,10 @@ def req_file(filename, folder="requirements"):
extras_require["all"] = list(chain(extras_require.values()))

platform_package_data = []
if PLATFORM_FLAG != "any":
if "linux" in PLATFORM_FLAG:
platform_package_data += ["libcshm.so"]
elif PLATFORM_FLAG == "win_amd64":
platform_package_data += ["cshm.dll"]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we add an else block catching unsupported platform flags?

Copy link
Contributor Author

@fpetrini15 fpetrini15 Apr 2, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The platform flag is currently passed in from build_wheel.py, so, from our perspective, the possible platforms are all handled. If the platform does not contain the substring "linux" or is not "win_amd64", it will be "any". In this case, we do not include either shared library object.

Are you suggesting raising an error in the event the setup.py script is called independently? Like so:

if "linux" in PLATFORM_FLAG:
    platform_package_data += ["libcshm.so"]
elif PLATFORM_FLAG == "win_amd64":
    platform_package_data += ["cshm.dll"]
elif PLATFORM_FLAG != "any":
    raise Exception(f"Unsupported platform flag \"{PLATFORM_FLAG}\" specified.")


data_files = [
("", ["LICENSE.txt"]),
Expand Down
38 changes: 23 additions & 15 deletions src/python/library/tritonclient/utils/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -24,26 +24,34 @@
# (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)

if(NOT WIN32)
file(COPY shared_memory DESTINATION .)
file(COPY shared_memory DESTINATION .)
#
# libcshm.so / cshm.dll
#
add_library(cshm SHARED shared_memory/shared_memory.cc)
if(${TRITON_ENABLE_GPU})
target_compile_definitions(cshm PUBLIC TRITON_ENABLE_GPU=1)
target_link_libraries(cshm PUBLIC CUDA::cudart)
endif() # TRITON_ENABLE_GPU

#
# libcshm.so
#
add_library(cshm SHARED shared_memory/shared_memory.cc)
if(${TRITON_ENABLE_GPU})
target_compile_definitions(cshm PUBLIC TRITON_ENABLE_GPU=1)
target_link_libraries(cshm PUBLIC CUDA::cudart)
endif() # TRITON_ENABLE_GPU
if(NOT WIN32)
target_link_libraries(cshm PRIVATE rt)
endif() # WIN32
endif() # NOT WIN32

if(NOT WIN32)
configure_file(shared_memory/__init__.py shared_memory/__init__.py COPYONLY)
configure_file(shared_memory/__init__.py shared_memory/__init__.py COPYONLY)

if(${TRITON_ENABLE_GPU})
configure_file(cuda_shared_memory/__init__.py cuda_shared_memory/__init__.py COPYONLY)
configure_file(cuda_shared_memory/_utils.py cuda_shared_memory/_utils.py COPYONLY)
endif() # NOT WIN32
endif() # TRITON_ENABLE_GPU
Loading
Loading