python api register_memory support location param (#2191)

Co-authored-by: A-Liuhao <liuhao276@hisilicon.com>

Author: A-Liuhao

mooncake-integration/transfer_engine/transfer_engine_py.cpp

@@ -735,15 +735,16 @@ int TransferEnginePy::transferCheckStatus(batch_id_t batch_id) {
 }
 
 int TransferEnginePy::batchRegisterMemory(
-    std::vector<uintptr_t> buffer_addresses, std::vector<size_t> capacities) {
+    std::vector<uintptr_t> buffer_addresses, std::vector<size_t> capacities,
+    const std::string& location) {
     pybind11::gil_scoped_release release;
     auto batch_size = buffer_addresses.size();
     std::vector<BufferEntry> buffers;
     for (size_t i = 0; i < batch_size; i++) {
         buffers.push_back(
             BufferEntry{(void*)buffer_addresses[i], capacities[i]});
     }
-    return engine_->registerLocalMemoryBatch(buffers, kWildcardLocation);
+    return engine_->registerLocalMemoryBatch(buffers, location);
 }
 
 int TransferEnginePy::batchUnregisterMemory(
@@ -757,9 +758,10 @@ int TransferEnginePy::batchUnregisterMemory(
     return engine_->unregisterLocalMemoryBatch(buffers);
 }
 
-int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity) {
+int TransferEnginePy::registerMemory(uintptr_t buffer_addr, size_t capacity,
+                                     const std::string& location) {
     char* buffer = reinterpret_cast<char*>(buffer_addr);
-    return engine_->registerLocalMemory(buffer, capacity);
+    return engine_->registerLocalMemory(buffer, capacity, location);
 }
 
 int TransferEnginePy::unregisterMemory(uintptr_t buffer_addr) {
@@ -1100,10 +1102,14 @@ PYBIND11_MODULE(engine, m) {
             .def("write_bytes_to_buffer", &TransferEnginePy::writeBytesToBuffer)
             .def("read_bytes_from_buffer",
                  &TransferEnginePy::readBytesFromBuffer)
-            .def("register_memory", &TransferEnginePy::registerMemory)
+            .def("register_memory", &TransferEnginePy::registerMemory,
+                 py::arg("buffer_addr"), py::arg("capacity"),
+                 py::arg("location") = kWildcardLocation)
             .def("unregister_memory", &TransferEnginePy::unregisterMemory)
             .def("batch_register_memory",
-                 &TransferEnginePy::batchRegisterMemory)
+                 &TransferEnginePy::batchRegisterMemory,
+                 py::arg("buffer_addresses"), py::arg("capacities"),
+                 py::arg("location") = kWildcardLocation)
             .def("batch_unregister_memory",
                  &TransferEnginePy::batchUnregisterMemory)
             .def("get_local_topology", &TransferEnginePy::getLocalTopology,

mooncake-integration/transfer_engine/transfer_engine_py.h

@@ -169,13 +169,15 @@ class TransferEnginePy {
     }
 
     // FOR EXPERIMENT ONLY
-    int registerMemory(uintptr_t buffer_addr, size_t capacity);
+    int registerMemory(uintptr_t buffer_addr, size_t capacity,
+                       const std::string &location = kWildcardLocation);
 
     // must be called before TransferEnginePy::~TransferEnginePy()
     int unregisterMemory(uintptr_t buffer_addr);
 
     int batchRegisterMemory(std::vector<uintptr_t> buffer_addresses,
-                            std::vector<size_t> capacities);
+                            std::vector<size_t> capacities,
+                            const std::string &location = kWildcardLocation);
 
     int batchUnregisterMemory(std::vector<uintptr_t> buffer_addresses);