@@ -84,7 +84,19 @@ else()
8484endif ()
8585
8686# Python.h needed by torch headers.
87- find_package (Python3 REQUIRED COMPONENTS Development)
87+ find_package (Python3 REQUIRED COMPONENTS Development.Module)
88+
89+ set (RHEL_BUILD OFF )
90+ set (LIB_DIR "lib" )
91+ set (PY_INSTALL_PATH "/usr/local/lib/python3.10/dist-packages" )
92+ if (LINUX)
93+ file (STRINGS "/etc/os-release" DISTRO_ID_LIKE REGEX "ID_LIKE" )
94+ if (${DISTRO_ID_LIKE} MATCHES "rhel|centos" )
95+ set (RHEL_BUILD ON )
96+ set (LIB_DIR "lib64" )
97+ set (PY_INSTALL_PATH "/opt/_internal/cpython-3.10.13/lib/python3.10/site-packages" )
98+ endif (${DISTRO_ID_LIKE} MATCHES "rhel|centos" )
99+ endif (LINUX)
88100
89101#
90102# Dependencies
@@ -220,35 +232,35 @@ if (${TRITON_PYTORCH_DOCKER_BUILD})
220232 COMMAND docker pull ${TRITON_PYTORCH_DOCKER_IMAGE}
221233 COMMAND docker rm pytorch_backend_ptlib || echo "error ignored..." || true
222234 COMMAND docker create --name pytorch_backend_ptlib ${TRITON_PYTORCH_DOCKER_IMAGE}
223- COMMAND /bin/sh -c "for i in ${LIBTORCH_LIBS_STR} ; do echo copying $i && docker cp -L pytorch_backend_ptlib:/ usr/local/lib/$i $i ; done"
224- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libc10.so libc10.so
225- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libc10_cuda.so libc10_cuda.so
226- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libtorch.so libtorch.so
227- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libtorch_cpu.so libtorch_cpu.so
228- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libtorch_cuda.so libtorch_cuda.so
229- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libtorch_cuda_linalg.so libtorch_cuda_linalg.so
230- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libtorch_global_deps.so libtorch_global_deps.so
231- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/lib/libcaffe2_nvrtc.so libcaffe2_nvrtc.so
235+ COMMAND /bin/sh -c "for i in ${LIBTORCH_LIBS_STR} ; do echo copying $i && docker cp -L pytorch_backend_ptlib:$<IF:$<BOOL: ${RHEL_BUILD} >,/opt/_internal/cpython-3.10.13/lib/$i,/ usr/local/lib/$i> $i ; done"
236+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libc10.so libc10.so
237+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libc10_cuda.so libc10_cuda.so
238+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libtorch.so libtorch.so
239+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libtorch_cpu.so libtorch_cpu.so
240+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libtorch_cuda.so libtorch_cuda.so
241+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libtorch_cuda_linalg.so libtorch_cuda_linalg.so
242+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libtorch_global_deps.so libtorch_global_deps.so
243+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/lib/libcaffe2_nvrtc.so libcaffe2_nvrtc.so
232244 # TODO: Revisit when not needed by making it part of cuda base container.
233- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/cuda/lib64/libcusparseLt.so libcusparseLt.so
234- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/libtorchvision.so libtorchvision.so
245+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/cuda/lib64/libcusparseLt.so libcusparseLt.so;
246+ COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then docker cp pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libtorchvision.so libtorchvision.so; fi"
247+ COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then docker cp pytorch_backend_ptlib:/opt/pytorch/vision/torchvision/csrc include/torchvision/torchvision; fi"
235248 COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHTRT} = 'ON' ]; then docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages/torch_tensorrt/lib/libtorchtrt_runtime.so libtorchtrt_runtime.so; fi"
236- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch_tensorrt/bin/torchtrtc torchtrtc || echo "error ignored..." || true
249+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch_tensorrt/bin/torchtrtc torchtrtc || echo "error ignored..." || true
237250 COMMAND docker cp pytorch_backend_ptlib:/opt/pytorch/pytorch/LICENSE LICENSE.pytorch
238- COMMAND docker cp pytorch_backend_ptlib:/usr/local/lib/python3.10/dist-packages /torch/include include /torch
251+ COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH} /torch/include include /torch
239252 COMMAND docker cp pytorch_backend_ptlib:/opt/pytorch/pytorch/torch/csrc/jit/codegen include /torch/torch/csrc/jit/.
240- COMMAND docker cp pytorch_backend_ptlib:/opt/pytorch/vision/torchvision/csrc include /torchvision/torchvision
241- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_videoio.so libopencv_videoio.so
242- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_highgui.so libopencv_highgui.so
243- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_video.so libopencv_video.so
244- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_imgcodecs.so libopencv_imgcodecs.so
245- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_imgproc.so libopencv_imgproc.so
246- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_core.so libopencv_core.so
247- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_calib3d.so libopencv_calib3d.so
248- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_features2d.so libopencv_features2d.so
249- COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/lib/libopencv_flann.so libopencv_flann.so
250- COMMAND docker cp pytorch_backend_ptlib:/usr/lib/${LIBS_ARCH} -linux-gnu/libpng16.so.16.37.0 libpng16.so
251- COMMAND docker cp pytorch_backend_ptlib:/usr/lib/${LIBS_ARCH} -linux-gnu/libjpeg.so.8.2.2 libjpeg.so
253+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_videoio.so libopencv_videoio.so
254+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_highgui.so libopencv_highgui.so
255+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_video.so libopencv_video.so
256+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_imgcodecs.so libopencv_imgcodecs.so
257+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_imgproc.so libopencv_imgproc.so
258+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_core.so libopencv_core.so
259+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_calib3d.so libopencv_calib3d.so
260+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_features2d.so libopencv_features2d.so
261+ COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/${LIB_DIR} /libopencv_flann.so libopencv_flann.so
262+ COMMAND /bin/sh -c "docker cp pytorch_backend_ptlib:$<IF:$<BOOL:${RHEL_BUILD} >,/usr/lib64/libpng16.so.16.34.0,/usr/lib/${LIBS_ARCH} -linux-gnu/libpng16.so.16.37.0> libpng16.so"
263+ COMMAND /bin/sh -c "docker cp pytorch_backend_ptlib:$<IF:$<BOOL:${RHEL_BUILD} >,/usr/lib64/libjpeg.so.62.2.0,/usr/lib/${LIBS_ARCH} -linux-gnu/libjpeg.so.8.2.2> libjpeg.so"
252264 COMMAND /bin/sh -c "if [ -f libmkl_def.so.1 ]; then patchelf --add-needed libmkl_gnu_thread.so.1 libmkl_def.so.1; fi"
253265 COMMAND /bin/sh -c "if [ -f libmkl_def.so.1 ]; then patchelf --add-needed libmkl_core.so.1 libmkl_def.so.1; fi"
254266 COMMAND /bin/sh -c "if [ -f libmkl_avx2.so.1 ]; then patchelf --add-needed libmkl_gnu_thread.so.1 libmkl_avx2.so.1; fi"
0 commit comments