tair-kvcache/kv_cache_manager/client/pybind/py_client_binding.cc at bc071b781c04a0985eda2671797808e45e8f6897 · alibaba/tair-kvcache · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
#include <pybind11/gil.h>
#include <pybind11/native_enum.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "kv_cache_manager/client/include/transfer_client.h"

namespace py = pybind11;
namespace kvcm = kv_cache_manager;

namespace {
struct OSException {
    int errcode;
};
} // namespace

PYBIND11_MODULE(kvcm_py_client, module) {
    pybind11::register_local_exception_translator([](std::exception_ptr p) {
        try {
            if (p)
                std::rethrow_exception(p);
        } catch (const OSException &e) {
            errno = e.errcode;
            PyErr_SetFromErrno(PyExc_OSError);
        }
    });

    module.doc() = "kvcm_py_client pybind11 extension";

    // 绑定枚举
    py::native_enum<kvcm::ClientErrorCode>(module, "ClientErrorCode", "enum.Enum")
        .value("ER_OK", kvcm::ClientErrorCode::ER_OK)
        .value("ER_INVALID_STUB", kvcm::ClientErrorCode::ER_INVALID_STUB)
        .value("ER_INVALID_GRPCSTATUS", kvcm::ClientErrorCode::ER_INVALID_GRPCSTATUS)
        .value("ER_INVALID_PARAMS", kvcm::ClientErrorCode::ER_INVALID_PARAMS)
        .value("ER_INVALID_ROLETYPE", kvcm::ClientErrorCode::ER_INVALID_ROLETYPE)
        .value("ER_INVALID_CLIENT_CONFIG", kvcm::ClientErrorCode::ER_INVALID_CLIENT_CONFIG)
        .value("ER_INVALID_STORAGE_CONFIG", kvcm::ClientErrorCode::ER_INVALID_STORAGE_CONFIG)
        .value("ER_INVALID_SDKWRAPPER_CONFIG", kvcm::ClientErrorCode::ER_INVALID_SDKWRAPPER_CONFIG)
        .value("ER_INVALID_SDKBACKEND_CONFIG", kvcm::ClientErrorCode::ER_INVALID_SDKBACKEND_CONFIG)
        .value("ER_CONNECT_FAIL", kvcm::ClientErrorCode::ER_CONNECT_FAIL)
        .value("ER_THREADPOOL_ERROR", kvcm::ClientErrorCode::ER_THREADPOOL_ERROR)
        .value("ER_SKIPINIT", kvcm::ClientErrorCode::ER_SKIPINIT)
        .value("ER_METACLIENT_INIT_ERROR", kvcm::ClientErrorCode::ER_METACLIENT_INIT_ERROR)
        .value("ER_TRANSFERCLIENT_INIT_ERROR", kvcm::ClientErrorCode::ER_TRANSFERCLIENT_INIT_ERROR)
        .value("ER_MANAGERCLIENT_INIT_ERROR", kvcm::ClientErrorCode::ER_MANAGERCLIENT_INIT_ERROR)
        .value("ER_CLIENT_NOT_EXISTS", kvcm::ClientErrorCode::ER_CLIENT_NOT_EXISTS)
        .value("ER_SERVICE_NO_STATUS", kvcm::ClientErrorCode::ER_SERVICE_NO_STATUS)
        .value("ER_SERVICE_INTERNAL_ERROR", kvcm::ClientErrorCode::ER_SERVICE_INTERNAL_ERROR)
        .value("ER_SERVICE_UNSUPPORTED", kvcm::ClientErrorCode::ER_SERVICE_UNSUPPORTED)
        .value("ER_SERVICE_INVALID_ARGUMENT", kvcm::ClientErrorCode::ER_SERVICE_INVALID_ARGUMENT)
        .value("ER_SERVICE_DUPLICATE_ENTITY", kvcm::ClientErrorCode::ER_SERVICE_DUPLICATE_ENTITY)
        .value("ER_SERVICE_INSTANCE_NOT_EXIST", kvcm::ClientErrorCode::ER_SERVICE_INSTANCE_NOT_EXIST)
        .value("ER_SDK_TIMEOUT", kvcm::ClientErrorCode::ER_SDK_TIMEOUT)
        .value("ER_GETSDK_ERROR", kvcm::ClientErrorCode::ER_GETSDK_ERROR)
        .value("ER_CREATESDK_ERROR", kvcm::ClientErrorCode::ER_CREATESDK_ERROR)
        .value("ER_SDKINIT_ERROR", kvcm::ClientErrorCode::ER_SDKINIT_ERROR)
        .value("ER_SDKREAD_ERROR", kvcm::ClientErrorCode::ER_SDKREAD_ERROR)
        .value("ER_SDKWRITE_ERROR", kvcm::ClientErrorCode::ER_SDKWRITE_ERROR)
        .value("ER_SDKALLOC_ERROR", kvcm::ClientErrorCode::ER_SDKALLOC_ERROR)
        .value("ER_INVALID_ADDRESS", kvcm::ClientErrorCode::ER_INVALID_ADDRESS)
        .value("ER_INVALID_LOCAL_BUFFERS", kvcm::ClientErrorCode::ER_INVALID_LOCAL_BUFFERS)
        .value("ER_UNSUPPORTED_MEMORY_TYPE", kvcm::ClientErrorCode::ER_UNSUPPORTED_MEMORY_TYPE)
        .value("ER_UNCONSISTENT_MEMORY_TYPE", kvcm::ClientErrorCode::ER_UNCONSISTENT_MEMORY_TYPE)
        .value("ER_FILE_IO_ERROR", kvcm::ClientErrorCode::ER_FILE_IO_ERROR)
        .value("ER_CUDAMEMCPY_ERROR", kvcm::ClientErrorCode::ER_CUDAMEMCPY_ERROR)
        .value("ER_EXTRACT_SLICES_ERROR", kvcm::ClientErrorCode::ER_EXTRACT_SLICES_ERROR)
        .value("ER_CUDA_STREAM_CREATE_ERROR", kvcm::ClientErrorCode::ER_CUDA_STREAM_CREATE_ERROR)
        .value("ER_CUDA_STREAM_SYNCHRONIZE_ERROR", kvcm::ClientErrorCode::ER_CUDA_STREAM_SYNCHRONIZE_ERROR)
        .value("ER_CUDA_STREAM_DESTROY_ERROR", kvcm::ClientErrorCode::ER_CUDA_STREAM_DESTROY_ERROR)
        .value("ER_CUDA_HOST_REGISTER_ERROR", kvcm::ClientErrorCode::ER_CUDA_HOST_REGISTER_ERROR)
        .finalize();

    py::native_enum<kvcm::MemoryType>(module, "MemoryType", "enum.Enum")
        .value("CPU", kvcm::MemoryType::CPU)
        .value("GPU", kvcm::MemoryType::GPU)
        .finalize();

    py::native_enum<kvcm::RoleType>(module, "RoleType", "enum.Enum")
        .value("UNKNOWN", kvcm::RoleType::UNKNOWN)
        .value("WORKER", kvcm::RoleType::WORKER)
        .value("SCHEDULER", kvcm::RoleType::SCHEDULER)
        .value("HYBRID", kvcm::RoleType::HYBRID)
        .finalize();

    // 绑定结构体
    py::class_<kvcm::Iov, py::smart_holder>(module, "Iov")
        .def(py::init<>())
        .def_readwrite("type", &kvcm::Iov::type)
        .def_property("base", &kvcm::Iov::base_as_uint64, &kvcm::Iov::set_base_as_uint64)
        .def_readwrite("size", &kvcm::Iov::size)
        .def_readwrite("ignore", &kvcm::Iov::ignore);

    py::class_<kvcm::BlockBuffer, py::smart_holder>(module, "BlockBuffer")
        .def(py::init<>())
        .def_readwrite("iovs", &kvcm::BlockBuffer::iovs);

    py::class_<kvcm::RegistSpan, py::smart_holder>(module, "RegistSpan")
        .def(py::init<>())
        .def_property("base", &kvcm::RegistSpan::base_as_uint64, &kvcm::RegistSpan::set_base_as_uint64)
        .def_readwrite("size", &kvcm::RegistSpan::size);

    py::class_<kvcm::InitParams, py::smart_holder>(module, "InitParams")
        .def(py::init<>())
        .def_readwrite("role_type", &kvcm::InitParams::role_type)
        .def_readwrite("regist_span", &kvcm::InitParams::regist_span)
        .def_readwrite("self_location_spec_name", &kvcm::InitParams::self_location_spec_name)
        .def_readwrite("storage_configs", &kvcm::InitParams::storage_configs);

    py::class_<kvcm::ForwardContext, py::smart_holder>(module, "ForwardContext")
        .def(py::init<>())
        .def_readwrite("metas", &kvcm::ForwardContext::metas)
        .def_readwrite("sw_size", &kvcm::ForwardContext::sw_size);

    py::class_<kvcm::TransferTraceInfo, py::smart_holder>(module, "TransferTraceInfo")
        .def(py::init<>())
        .def_readwrite("need_print", &kvcm::TransferTraceInfo::need_print)
        .def_readwrite("block_ids", &kvcm::TransferTraceInfo::block_ids);

    // 由于pybind11会自动处理std::vector和std::map到Python list/dict的转换，
    // 我们不需要显式绑定这些标准容器类型，直接使用Python的list和dict即可
    // 保留这些类型定义以支持C++接口，但使用Python原生类型进行交互

    // 绑定TransferClient类
    py::class_<kvcm::TransferClient, py::smart_holder>(module, "TransferClient")
        .def_static("Create", &kvcm::TransferClient::Create, py::call_guard<py::gil_scoped_release>())
        // C++ interface gained pointer parameters expected_checksums / out_checksums for
        // Scheme A verification. The Python binding intentionally keeps the legacy
        // 3-argument signature via a lambda (default nullptr), letting connectors
        // (vLLM / SGLang / TRT-LLM) keep building unchanged. Extend here once
        // py_connector adopts checksum reporting.
        .def(
            "LoadKvCaches",
            [](kvcm::TransferClient *self,
               const kvcm::UriStrVec &uri_str_vec,
               const kvcm::BlockBuffers &block_buffers,
               std::shared_ptr<kvcm::TransferTraceInfo> trace_info) {
                return self->LoadKvCaches(uri_str_vec, block_buffers, trace_info);
            },
            py::arg("uri_str_vec"),
            py::arg("block_buffers"),
            py::arg("trace_info") = nullptr,
            py::call_guard<py::gil_scoped_release>())
        .def(
            "SaveKvCaches",
            [](kvcm::TransferClient *self,
               const kvcm::UriStrVec &uri_str_vec,
               const kvcm::BlockBuffers &block_buffers,
               std::shared_ptr<kvcm::TransferTraceInfo> trace_info) {
                return self->SaveKvCaches(uri_str_vec, block_buffers, trace_info);
            },
            py::arg("uri_str_vec"),
            py::arg("block_buffers"),
            py::arg("trace_info") = nullptr,
            py::call_guard<py::gil_scoped_release>());

} // namespace kv_cache_manager