diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index af877864..99ea7d12 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -129,6 +129,16 @@ target_include_directories(fabric_cli target_link_libraries(fabric_cli PUBLIC transfer_engine gflags::gflags glog::glog pthread) +# Out-of-tree transport plugin: built as a standalone .so that is NOT linked +# into the engine and is dlopened by PluginLoader at runtime +# (fabric_cli plugins ). +add_library(shm_plugin MODULE + ${WORKSPACE}/fabric_conformance/sample_plugin/shm_plugin.cpp) +target_link_libraries(shm_plugin PRIVATE transfer_engine) +set_target_properties(shm_plugin PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/plugins" + PREFIX "") + if (USE_ASCEND_DIRECT) # AscendDirectTransport unit test with mock implementations # Mock implementations are included in the test file via anonymous namespace diff --git a/mooncake-transfer-engine/tests/fabric_conformance/fabric_cli.cpp b/mooncake-transfer-engine/tests/fabric_conformance/fabric_cli.cpp index 88c1b0e7..e5a982c4 100644 --- a/mooncake-transfer-engine/tests/fabric_conformance/fabric_cli.cpp +++ b/mooncake-transfer-engine/tests/fabric_conformance/fabric_cli.cpp @@ -17,6 +17,8 @@ // select run the path selector on representative endpoint pairs and // show the fallback ladder, including a live fault-injection // demo +// cost size-aware cost-model selection sweep +// plugins DIR dlopen transport plugins from DIR and probe them // conformance run the conformance suite over every backend // // probe and conformance drive the real TransferEngine (server-free @@ -30,6 +32,8 @@ #include #include +#include "fabric/plugin_loader.h" + #include "conformance_suite.h" #include "fabric/path_selector.h" @@ -92,28 +96,89 @@ struct DemoKind { fabric::MemoryFabricKind kind; std::string name; bool healthy = true; + fabric::Capability cap; // modelled bandwidth/latency for the cost demo }; -int cmdSelect() { +// Representative bandwidth/latency for a kind, used by the size-aware cost +// model demo. NVLink/RDMA/TCP figures match the order of magnitude measured on +// this class of host (NVLink ~350 GB/s, IB ~25 GB/s, TCP ~5 GB/s); the emulated +// kinds report their host-window throughput. +fabric::Capability modelCap(fabric::MemoryFabricKind kind, + const std::string &name) { + fabric::Capability c; + c.kind = kind; + c.name = name; + using K = fabric::MemoryFabricKind; + // Only NVLink and RDMA (GPUDirect) can address device memory directly; the + // others must stage GPU<->GPU traffic through host memory. + c.supports_device_memory = (kind == K::NVLINK || kind == K::RDMA); + switch (kind) { + case K::NVLINK: + c.max_bandwidth_mbps = 350000; + c.latency_ns = 8000; + break; + case K::CXL: + c.max_bandwidth_mbps = 26000; + c.latency_ns = 300; + break; + case K::RDMA: + c.max_bandwidth_mbps = 25000; + c.latency_ns = 2500; + break; + case K::UALINK: + c.max_bandwidth_mbps = 28000; + c.latency_ns = 250; + break; + case K::NUMA: + c.max_bandwidth_mbps = 18000; + c.latency_ns = 200; + break; + case K::TCP: + c.max_bandwidth_mbps = 5000; + c.latency_ns = 20000; + break; + default: + c.max_bandwidth_mbps = 4000; + c.latency_ns = 5000; + break; + } + return c; +} + +std::vector demoKinds() { std::vector kinds; - // Live-probed kinds. for (auto &proto : builtinBackends()) { conformance::Harness h(proto, 4ull * 1024 * 1024); - if (h.ready()) - kinds.push_back({h.transport()->probe().kind, proto, true}); + if (h.ready()) { + auto cap = h.transport()->probe(); + DemoKind d{cap.kind, proto, true, cap}; + // Backends that do not self-calibrate (TCP/RDMA) get a modelled cap + // so the cost demo has numbers to work with. + if (d.cap.max_bandwidth_mbps == 0) + d.cap = modelCap(cap.kind, proto); + kinds.push_back(d); + } } - // Synthetic high-end intra-node GPU links present in a full deployment. - DemoKind nvlink{fabric::MemoryFabricKind::NVLINK, "nvlink", true}; - kinds.push_back(nvlink); + // Synthetic high-end intra-node GPU link present in a full deployment. + kinds.push_back({fabric::MemoryFabricKind::NVLINK, "nvlink", true, + modelCap(fabric::MemoryFabricKind::NVLINK, "nvlink")}); + return kinds; +} + +int cmdSelect() { + std::vector kinds = demoKinds(); fabric::PathSelector sel; std::vector health(kinds.size(), true); for (size_t i = 0; i < kinds.size(); ++i) { size_t idx = i; - sel.registerKind(kinds[i].name, kinds[i].kind, [&health, idx] { - return health[idx] ? fabric::HealthStatus::HEALTHY - : fabric::HealthStatus::UNREACHABLE; - }); + sel.registerKind( + kinds[i].name, kinds[i].kind, + [&health, idx] { + return health[idx] ? fabric::HealthStatus::HEALTHY + : fabric::HealthStatus::UNREACHABLE; + }, + kinds[i].cap); } Endpoint gpu0{Endpoint::Kind::GPU, 0, 0, "node0", "gpu0"}; @@ -154,7 +219,76 @@ int cmdSelect() { return 0; } -void usage() { std::printf("usage: fabric_cli \n"); } +// Size-aware cost-model demo: for a host_host pair, sweep the transfer size and +// show how the lowest-cost backend changes (small => low latency wins, large => +// high bandwidth wins). +int cmdCost() { + std::vector kinds = demoKinds(); + fabric::PathSelector sel; + for (auto &k : kinds) + sel.registerKind( + k.name, k.kind, [] { return fabric::HealthStatus::HEALTHY; }, + k.cap); + + Endpoint gpu0{Endpoint::Kind::GPU, 0, 0, "node0", "gpu0"}; + Endpoint gpu3{Endpoint::Kind::GPU, 3, 0, "node0", "gpu3"}; + + std::printf( + "--- size-aware cost-model selection (same-node GPU<->GPU) ------\n"); + std::printf("registered candidates and their model:\n"); + for (auto &k : kinds) + std::printf(" %-10s bw=%8lu MB/s lat=%7lu ns\n", k.name.c_str(), + (unsigned long)k.cap.max_bandwidth_mbps, + (unsigned long)k.cap.latency_ns); + std::printf("\n%-12s %-14s %-12s\n", "size", "chosen", "est cost"); + const uint64_t sizes[] = {64, 1024, 64ull * 1024, + 1ull << 20, 16ull << 20, 256ull << 20}; + for (uint64_t s : sizes) { + auto c = sel.selectForSize(gpu0, gpu3, s); + char szbuf[32]; + if (s < (1 << 20)) + std::snprintf(szbuf, sizeof(szbuf), "%lu B", (unsigned long)s); + else + std::snprintf(szbuf, sizeof(szbuf), "%lu MiB", + (unsigned long)(s >> 20)); + std::printf("%-12s %-14s %.0f ns\n", szbuf, + c.found ? c.label.c_str() : "-", c.estimated_cost_ns); + } + std::printf( + "\nSmall transfers favour the lowest-latency fabric; large transfers " + "favour the highest-bandwidth one -- the selector switches " + "automatically.\n"); + return 0; +} + +void usage() { + std::printf( + "usage: fabric_cli \n"); +} + +// Load every transport plugin .so in a directory via the versioned C ABI, then +// probe each one -- demonstrating that a backend can appear without relinking. +int cmdPlugins(const std::string &dir) { + fabric::PluginLoader loader; // declared first => outlives the transports + std::vector infos; + auto plugins = loader.loadDirectory(dir, &infos); + std::printf("--- plugin loading from %s ---\n", dir.c_str()); + for (auto &pi : infos) + std::printf(" %-44s %s%s\n", pi.path.c_str(), + pi.ok ? "loaded" : "FAILED", + pi.ok ? "" : (" : " + pi.error).c_str()); + if (plugins.empty()) { + std::printf(" (no plugins found)\n"); + return 0; + } + std::printf("\n--- probing dynamically-loaded backends ---\n"); + for (auto &t : plugins) { + auto cap = t->probe(); + std::printf(" %s health=%s\n", cap.toLine().c_str(), + fabric::toString(t->health())); + } + return 0; +} } // namespace @@ -165,6 +299,9 @@ int main(int argc, char **argv) { std::string cmd = (argc > 1) ? argv[1] : "probe"; if (cmd == "probe") return cmdProbe(); if (cmd == "select") return cmdSelect(); + if (cmd == "cost") return cmdCost(); + if (cmd == "plugins") + return cmdPlugins(argc > 2 ? argv[2] : "build/plugins"); if (cmd == "conformance") return cmdConformance(); usage(); return 2; diff --git a/mooncake-transfer-engine/tests/fabric_conformance/sample_plugin/shm_plugin.cpp b/mooncake-transfer-engine/tests/fabric_conformance/sample_plugin/shm_plugin.cpp new file mode 100644 index 00000000..8b99f1e7 --- /dev/null +++ b/mooncake-transfer-engine/tests/fabric_conformance/sample_plugin/shm_plugin.cpp @@ -0,0 +1,187 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// A sample out-of-tree transport packaged as a shared object. It is built into +// its own .so and is NOT linked into the Transfer Engine; PluginLoader dlopens +// it at runtime, checks the ABI version, instantiates it, and runs probe() / +// health() -- demonstrating that a backend can appear without relinking +// Mooncake. The transport itself is a host-memory shared-window loopback (the +// same data-plane shape as the CXL backend) so it is conformance-checkable. + +#include + +#include +#include +#include + +#include "fabric/plugin_loader.h" +#include "transport/transport.h" + +namespace mooncake { +namespace { + +class ShmPluginTransport : public Transport { + public: + ShmPluginTransport() { + window_size_ = 64ull * 1024 * 1024; + window_ = mmap(nullptr, window_size_, PROT_READ | PROT_WRITE, + MAP_SHARED | MAP_ANONYMOUS, -1, 0); + if (window_ == MAP_FAILED) window_ = nullptr; + } + ~ShmPluginTransport() override { + if (window_ && window_ != MAP_FAILED) munmap(window_, window_size_); + if (metadata_) metadata_->removeSegmentDesc(local_server_name_); + } + + Status submitTransfer( + BatchID batch_id, + const std::vector &entries) override { + auto &batch_desc = *((BatchDesc *)(batch_id)); + if (batch_desc.task_list.size() + entries.size() > + batch_desc.batch_size) { + return Status::InvalidArgument("shm-plugin: batch capacity"); + } + size_t task_id = batch_desc.task_list.size(); + batch_desc.task_list.resize(task_id + entries.size()); + for (auto &request : entries) { + TransferTask &task = batch_desc.task_list[task_id++]; + task.total_bytes = request.length; + Slice *slice = getSliceCache().allocate(); + slice->source_addr = (char *)request.source; + slice->cxl.dest_addr = (char *)window_ + request.target_offset; + slice->length = request.length; + slice->opcode = request.opcode; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + __sync_fetch_and_add(&task.slice_count, 1); + void *dst = (request.opcode == TransferRequest::READ) + ? slice->source_addr + : (void *)slice->cxl.dest_addr; + const void *src = (request.opcode == TransferRequest::READ) + ? (void *)slice->cxl.dest_addr + : slice->source_addr; + if (request.length) std::memcpy(dst, src, request.length); + __sync_synchronize(); + slice->markSuccess(); + } + return Status::OK(); + } + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus &status) override { + auto &batch_desc = *((BatchDesc *)(batch_id)); + if (task_id >= batch_desc.task_list.size()) + return Status::InvalidArgument("shm-plugin: task id"); + auto &task = batch_desc.task_list[task_id]; + status.transferred_bytes = task.transferred_bytes; + if (task.success_slice_count + task.failed_slice_count == + task.slice_count) { + status.s = task.failed_slice_count ? TransferStatusEnum::FAILED + : TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status.s = TransferStatusEnum::WAITING; + } + return Status::OK(); + } + + fabric::Capability probe() override { + fabric::Capability cap; + cap.kind = fabric::MemoryFabricKind::SHM; + cap.name = "shm-plugin"; + cap.supports_host_memory = true; + cap.supports_p2p = true; + cap.supports_ordered_write = true; + cap.supports_zero_copy = true; + cap.alignment = 1; + cap.max_transfer_size = window_size_; + cap.emulated = true; + cap.notes = "out-of-tree shared-memory loopback plugin (.so)"; + return cap; + } + + fabric::HealthStatus health() override { + return window_ ? fabric::HealthStatus::HEALTHY + : fabric::HealthStatus::UNREACHABLE; + } + + void *windowBase() const { return window_; } + size_t windowSize() const { return window_size_; } + + private: + int install(std::string &local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) override { + (void)topo; + metadata_ = meta; + local_server_name_ = local_server_name; + if (!window_) return -1; + auto desc = metadata_->getSegmentDesc(local_server_name_); + if (!desc) desc = std::make_shared(); + desc->name = local_server_name_; + desc->protocol = "shm-plugin"; + desc->cxl_base_addr = (uint64_t)window_; + desc->cxl_name = "shm-plugin"; + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return metadata_->updateLocalSegmentDesc(); + } + + int registerLocalMemory(void *addr, size_t length, + const std::string &location, bool remote_accessible, + bool update_metadata) override { + (void)location; + (void)remote_accessible; + BufferDesc d; + d.name = local_server_name_; + uintptr_t base = (uintptr_t)window_; + uintptr_t ptr = (uintptr_t)addr; + if (ptr < base || ptr + length > base + window_size_) { + errno = EFAULT; + return -1; + } + d.offset = ptr - base; + d.length = length; + return metadata_->addLocalMemoryBuffer(d, update_metadata); + } + + int unregisterLocalMemory(void *addr, bool update_metadata) override { + return metadata_->removeLocalMemoryBuffer(addr, update_metadata); + } + + int registerLocalMemoryBatch(const std::vector &buffer_list, + const std::string &location) override { + for (auto &b : buffer_list) + registerLocalMemory(b.addr, b.length, location, true, false); + return metadata_->updateLocalSegmentDesc(); + } + + int unregisterLocalMemoryBatch( + const std::vector &addr_list) override { + for (auto &a : addr_list) unregisterLocalMemory(a, false); + return metadata_->updateLocalSegmentDesc(); + } + + const char *getName() const override { return "shm-plugin"; } + + void *window_ = nullptr; + size_t window_size_ = 0; +}; + +} // namespace +} // namespace mooncake + +MOONCAKE_DEFINE_TRANSPORT_PLUGIN(new mooncake::ShmPluginTransport(), + "shm-plugin")