forked from mooncake-track/Mooncake
[TransferEngine] Add out-of-tree transport plugin and fabric diagnostics
Add a sample transport packaged as a standalone shared object (sample_plugin/shm_plugin.cpp, built as shm_plugin.so and not linked into the engine). PluginLoader dlopens it at runtime through the versioned C ABI, checks the version, instantiates it, and probes it -- demonstrating that a backend can appear without relinking Mooncake. Extend fabric_cli with two diagnostics: "cost" sweeps transfer sizes and shows the size-aware selector switching backends, and "plugins <dir>" loads and probes every transport plugin in a directory.
This commit is contained in:
parent
ded8c8751f
commit
0eb607b2bb
|
|
@ -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 <dir>).
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 <string>
|
||||
#include <vector>
|
||||
|
||||
#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<DemoKind> demoKinds() {
|
||||
std::vector<DemoKind> 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<DemoKind> kinds = demoKinds();
|
||||
|
||||
fabric::PathSelector sel;
|
||||
std::vector<bool> 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 <probe|select|conformance>\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<DemoKind> 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 <probe|select|cost|plugins [dir]|conformance>\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<fabric::PluginInfo> 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;
|
||||
|
|
|
|||
|
|
@ -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 <sys/mman.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
|
||||
#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<TransferRequest> &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<TransferMetadata> meta,
|
||||
std::shared_ptr<Topology> 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<SegmentDesc>();
|
||||
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<BufferEntry> &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<void *> &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")
|
||||
Loading…
Reference in New Issue