kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

add v1 of python bindings

Browse Source
This commit is contained in:
Linden Li 2023-10-23 09:50:03 -07:00
parent 1c4a0ee24c
commit a10d120e4f
2 changed files with 27 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
cpp/CMakeLists.txt
View File

@ -257,6 +257,8 @@ if(BUILD_PYT)
message(STATUS "Found Python libraries at ${Python3_LIBRARY_DIRS}")
link_directories("${Python3_LIBRARY_DIRS}")
list(APPEND COMMON_HEADER_DIRS ${Python3_INCLUDE_DIRS})
list(APPEND CMAKE_PREFIX_PATH /usr/local/lib/python3.10/dist-packages/pybind11)
find_package(pybind11 REQUIRED)
execute_process(
COMMAND
@ -325,3 +327,7 @@ if(BUILD_BENCHMARKS)
add_subdirectory(${TRT_LLM_ROOT_DIR}/benchmarks/cpp
${CMAKE_BINARY_DIR}/benchmarks)
endif()
pybind11_add_module(pygptmanager include/tensorrt_llm/batch_manager/pygptmanager.cpp)
include_directories(${CMAKE_SOURCE_DIR}/include)
target_link_libraries(pygptmanager PRIVATE pybind11:headers tensorrt_llm)

21
cpp/include/tensorrt_llm/batch_manager/pygptmanager.cpp Normal file
View File

@ -0,0 +1,21 @@
#include "tensorrt_llm/batch_manager/GptManager.h"
#include <pybind11/pybind11.h>
namespace py = pybind11;
PYBIND11_MODULE(pygptmanager, m)
{
m.doc() = "Python bindings for GptManager";
py::class_<tensorrt_llm::batch_manager::GptManager>(m, "GptManager")
.def(py::init<const std::filesystem::path&, tensorrt_llm::batch_manager::TrtGptModelType, int32_t, int32_t,
int32_t, tensorrt_llm::batch_manager::GetInferenceRequestsCallback,
tensorrt_llm::batch_manager::SendResponseCallback,
tensorrt_llm::batch_manager::PollStopSignalCallback>(),
py::arg("trtEnginePath"), py::arg("modelType"), py::arg("maxSeqLen"), py::arg("maxNumRequests"),
py::arg("maxBeamWidth"), py::arg("getInferenceRequestsCb"), py::arg("sendResponseCb"),
py::arg("pollStopSignalCb") = nullptr)
.def("fetchNewRequests", &tensorrt_llm::batch_manager::GptManager::fetchNewRequests)
.def("return_completed_requests", &tensorrt_llm::batch_manager::GptManager::return_completed_requests)
.def("pollStopSignals", &tensorrt_llm::batch_manager::GptManager::pollStopSignals);
}