From af2ce952f18fd58f7afbc2b78f2cc9add678c4d7 Mon Sep 17 00:00:00 2001 From: Le1zyCatt <2490162471@qq.com> Date: Tue, 28 Jul 2026 11:29:03 +0800 Subject: [PATCH 1/6] [TENT] Add native UB endpoint bootstrap and READ/WRITE data path --- .github/workflows/ci.yml | 11 + .../benchmark/tent_backend.cpp | 2 + mooncake-transfer-engine/benchmark/utils.cpp | 6 +- .../tent/include/tent/rpc/rpc.h | 2 + .../tent/include/tent/runtime/control_plane.h | 77 ++ .../tent/include/tent/transport/ub/endpoint.h | 166 ++++ .../tent/transport/ub/endpoint_store.h | 67 ++ .../tent/include/tent/transport/ub/quota.h | 195 ++++ .../include/tent/transport/ub/rail_monitor.h | 165 ++++ .../tent/include/tent/transport/ub/slice.h | 707 ++++++++++++++ .../include/tent/transport/ub/ub_transport.h | 72 ++ .../tent/include/tent/transport/ub/workers.h | 156 +++ .../tent/src/runtime/control_plane.cpp | 58 ++ .../tent/src/runtime/transport_loader.cpp | 11 + .../tent/src/transport/ub/endpoint.cpp | 472 +++++++++ .../tent/src/transport/ub/endpoint_store.cpp | 189 ++++ .../tent/src/transport/ub/quota.cpp | 371 +++++++ .../tent/src/transport/ub/rail_monitor.cpp | 279 ++++++ .../tent/src/transport/ub/ub_transport.cpp | 718 ++++++++++++++ .../tent/src/transport/ub/workers.cpp | 918 ++++++++++++++++++ .../tent/tests/CMakeLists.txt | 16 + .../tent/tests/transport_selector_test.cpp | 163 ++++ .../tent/tests/ub_core_test.cpp | 304 ++++++ .../tent/tests/ub_native_data_path_test.cpp | 730 ++++++++++++++ 24 files changed, 5853 insertions(+), 2 deletions(-) create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/quota.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp create mode 100644 mooncake-transfer-engine/tent/src/transport/ub/workers.cpp create mode 100644 mooncake-transfer-engine/tent/tests/ub_core_test.cpp create mode 100644 mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9fe555ffe7..7609dca048 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -707,6 +707,17 @@ jobs: run: ./scripts/ci/run_transfer_engine_rust_smoke.sh shell: bash + - name: Smoke test TENT UB benchmark CLI + if: matrix.name == 'ub-mock' + run: | + cd build-tent + help_output="$(./mooncake-transfer-engine/benchmark/tebench \ + --backend=tent --xport_type=ub --tent_transport_hint=ub \ + --help 2>&1)" + grep -q 'iouring|ub|sunrise_link' <<< "${help_output}" + grep -q 'ascend|ub|sunrise_link' <<< "${help_output}" + shell: bash + - name: Run sccache stat for check if: ${{ env.SCCACHE_PATH != '' }} shell: bash diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 8175289bdb..b4e29536f0 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.cpp +++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp @@ -65,6 +65,7 @@ std::shared_ptr loadConfig() { {"gds", "gds"}, {"mnnvl", "mnnvl"}, {"nvlink", "nvlink"}, + {"ub", "ub"}, {"sunrise_link", "sunrise_link"}, {"mpcomm", "mpcomm"}}; @@ -91,6 +92,7 @@ static TransportType getTransportType(const std::string& xport_type) { if (xport_type == "nvlink") return NVLINK; if (xport_type == "tcp") return TCP; if (xport_type == "iouring") return IOURING; + if (xport_type == "ub") return UB; if (xport_type == "sunrise_link") return SUNRISE_LINK; if (xport_type == "mpcomm") return MPCOMM; return UNSPEC; diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp index 25cbb3d417..3f5b5f5e91 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/mooncake-transfer-engine/benchmark/utils.cpp @@ -95,14 +95,16 @@ DEFINE_int32( "RPC server port used for p2p metadata service (0 = auto-select)."); DEFINE_string(xport_type, "", "Transport type: " - "rdma|shm|mnnvl|gds|iouring|sunrise_link|mpcomm|flagcx"); + "rdma|tcp|shm|mnnvl|nvlink|gds|iouring|ub|sunrise_link|mpcomm|" + "flagcx"); DEFINE_string(backend, "tent", "Transport backend: classic|tent"); DEFINE_bool(notifi, false, "Enable RDMA notification for performance measurement."); DEFINE_string( tent_transport_hint, "unspec", "tent only: per-request transport_hint. " - "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|sunrise_link|mpcomm"); + "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|ub|sunrise_link|" + "mpcomm"); DEFINE_string(tent_intent_type, "unspec", "tent only: intent_type attached to every benchmark request. " "unspec|foreground_get|background_prefetch|migration|checkpoint|" diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 8616865454..c706293ae1 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -47,6 +47,8 @@ enum RpcFuncID { Unpin, SubscribeSegmentUpdate, NotifySegmentUpdated, + // Appended to preserve the numeric values of the existing RPCs. + BootstrapUb, }; class ClientPool; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h index 8b691e65d5..3945518df4 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -57,6 +58,65 @@ struct BootstrapDesc { notify_qp_num); }; +// UB/URMA has Jetty, JFC and EID concepts that are not wire-compatible with +// RDMA QPs, CQs and GIDs. Keep a dedicated bootstrap envelope so neither +// transport has to smuggle native identifiers through the other's fields. +struct UbBootstrapDesc { + uint32_t protocol_version = 1; + std::string segment_name; + std::string local_nic_path; + std::string peer_nic_path; + std::string local_device_name; + int local_device_id = -1; + int local_eid_index = -1; + std::string local_eid; + std::vector jetty_ids; + // UASID is part of urma_jetty_id_t on providers that use nonzero address + // spaces. Kept parallel to jetty_ids for protocol-v1 compatibility. + std::vector jetty_uasids; + uint64_t endpoint_generation = 0; + uint64_t segment_generation = 0; + std::vector capabilities; + std::string reply_msg; +}; + +inline void to_json(nlohmann::json& j, const UbBootstrapDesc& desc) { + j = nlohmann::json{{"protocol_version", desc.protocol_version}, + {"segment_name", desc.segment_name}, + {"local_nic_path", desc.local_nic_path}, + {"peer_nic_path", desc.peer_nic_path}, + {"local_device_name", desc.local_device_name}, + {"local_device_id", desc.local_device_id}, + {"local_eid_index", desc.local_eid_index}, + {"local_eid", desc.local_eid}, + {"jetty_ids", desc.jetty_ids}, + {"jetty_uasids", desc.jetty_uasids}, + {"endpoint_generation", desc.endpoint_generation}, + {"segment_generation", desc.segment_generation}, + {"capabilities", desc.capabilities}, + {"reply_msg", desc.reply_msg}}; +} + +inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { + desc.protocol_version = j.value("protocol_version", 0u); + if (desc.protocol_version != 1) { + throw std::invalid_argument("unsupported UB bootstrap version"); + } + desc.segment_name = j.value("segment_name", ""); + desc.local_nic_path = j.value("local_nic_path", ""); + desc.peer_nic_path = j.value("peer_nic_path", ""); + desc.local_device_name = j.value("local_device_name", ""); + desc.local_device_id = j.value("local_device_id", -1); + desc.local_eid_index = j.value("local_eid_index", -1); + desc.local_eid = j.value("local_eid", ""); + desc.jetty_ids = j.value("jetty_ids", std::vector{}); + desc.jetty_uasids = j.value("jetty_uasids", std::vector{}); + desc.endpoint_generation = j.value("endpoint_generation", uint64_t{0}); + desc.segment_generation = j.value("segment_generation", uint64_t{0}); + desc.capabilities = j.value("capabilities", std::vector{}); + desc.reply_msg = j.value("reply_msg", ""); +} + struct XferDataDesc { uint64_t peer_mem_addr; size_t length; @@ -65,6 +125,9 @@ struct XferDataDesc { using OnReceiveBootstrap = std::function; +using OnReceiveUbBootstrap = std::function; + using OnNotify = std::function; class ControlClient { @@ -81,6 +144,10 @@ class ControlClient { const BootstrapDesc& request, BootstrapDesc& response); + static Status bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response); + static Status sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length); @@ -136,6 +203,11 @@ class ControlService { void setBootstrapRdmaCallback(const OnReceiveBootstrap& callback); + void setBootstrapUbCallback(const OnReceiveUbBootstrap& callback) { + std::lock_guard lock(ub_bootstrap_callback_mutex_); + ub_bootstrap_callback_ = callback; + } + void setNotifyCallback(const OnNotify& callback); Status start(uint16_t& port, bool ipv6_ = false, size_t threads = 1); @@ -147,6 +219,8 @@ class ControlService { void onBootstrapRdma(const std::string_view& request, std::string& response); + void onBootstrapUb(const std::string_view& request, std::string& response); + void onSendData(const std::string_view& request, std::string& response); void onRecvData(const std::string_view& request, std::string& response); @@ -184,6 +258,9 @@ class ControlService { OnReceiveBootstrap bootstrap_callback_; static thread_local const ControlService* active_bootstrap_service_; + std::mutex ub_bootstrap_callback_mutex_; + OnReceiveUbBootstrap ub_bootstrap_callback_; + std::mutex notify_cb_mutex_; std::condition_variable notify_cb_cv_; size_t notify_callbacks_in_flight_ = 0; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h new file mode 100644 index 0000000000..3dfe4533ba --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h @@ -0,0 +1,166 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_ENDPOINT_H_ +#define TENT_TRANSPORT_UB_ENDPOINT_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/common/types.h" +#include "tent/runtime/control_plane.h" +#include "tent/runtime/topology.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +// One cache entry represents a local-device to remote-device path. The peer +// NIC path is part of the identity because a remote topology ID may be reused +// after the peer republishes its topology. +struct UbEndpointKey { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + Topology::NicID remote_topology_id{-1}; + std::string peer_nic_path; + + bool operator==(const UbEndpointKey&) const = default; + + [[nodiscard]] bool valid() const noexcept { + return local_topology_id >= 0 && remote_topology_id >= 0 && + !peer_nic_path.empty(); + } +}; + +struct UbEndpointKeyHash { + size_t operator()(const UbEndpointKey& key) const noexcept; +}; + +// A UB endpoint is one immutable incarnation of a Jetty set. Its generation +// is allocated process-wide and is never reused. Lifecycle operations are +// serialized internally; a failed endpoint can only retire, never reconnect. +class UbEndpoint final : public std::enable_shared_from_this { + public: + enum class State : uint8_t { + kUninitialized, + kHandshaking, + kPrepared, + kBinding, + kReady, + kFailed, + kDestroying, + kDestroyed, + }; + + UbEndpoint(UbEndpointKey key, UbContextPtr context, + std::shared_ptr adapter, uint32_t jetty_count, + JettyOptions jetty_options = {}); + ~UbEndpoint(); + + UbEndpoint(const UbEndpoint&) = delete; + UbEndpoint& operator=(const UbEndpoint&) = delete; + + // Creates the local Jetty set. Concurrent calls share the same attempt. + // An error permanently moves this incarnation to kFailed. + Status prepare(); + + // Binds each local Jetty to the peer EID and the corresponding peer Jetty + // ID. The bootstrap must describe exactly one peer Jetty per local Jetty. + Status bind(const UbBootstrapDesc& peer); + + // Builds the local half of the native UB bootstrap after prepare(). + Status makeBootstrapDesc(const std::string& segment_name, + const std::string& local_nic_path, + const std::string& peer_nic_path, + uint64_t segment_generation, + UbBootstrapDesc& output) const; + + // Admission to the posting path is synchronized with retirement: once + // kDestroying is visible no new work can acquire the endpoint. Every + // successful acquire must have one matching release. + [[nodiscard]] bool tryAcquireOutstanding(uint64_t bytes = 0) noexcept; + void releaseOutstanding(uint64_t bytes = 0) noexcept; + + // Stops new posts immediately. With outstanding work, native resources + // remain intact until either completions drain naturally or quiesce() + // establishes an explicit no-more-DMA fence. Idempotent. + Status retire(); + + // Establishes a native drain fence for every Jetty, then resets/unbinds + // them. Returned completions still own the corresponding logical tokens + // and must be dispatched by Workers. A failed fence leaves all resources + // alive so shutdown can be retried safely. + Status quiesce(uint32_t timeout_ms, std::vector& completions); + + [[nodiscard]] const UbEndpointKey& key() const noexcept { return key_; } + [[nodiscard]] const UbContextPtr& context() const noexcept { + return context_; + } + [[nodiscard]] uint64_t generation() const noexcept { return generation_; } + [[nodiscard]] State state() const noexcept { + return state_.load(std::memory_order_acquire); + } + [[nodiscard]] bool ready() const noexcept { + return state() == State::kReady; + } + [[nodiscard]] bool failed() const noexcept { + return state() == State::kFailed; + } + [[nodiscard]] bool reusable() const noexcept; + [[nodiscard]] uint64_t peerGeneration() const noexcept { + return peer_generation_.load(std::memory_order_acquire); + } + [[nodiscard]] uint64_t outstandingWrs() const noexcept { + return outstanding_wrs_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t outstandingBytes() const noexcept { + return outstanding_bytes_.load(std::memory_order_relaxed); + } + [[nodiscard]] size_t jettyCount() const; + [[nodiscard]] JettyPtr jetty(size_t index) const; + [[nodiscard]] size_t jfcIndex(size_t jetty_index) const; + [[nodiscard]] std::vector jetties() const; + [[nodiscard]] Status lifecycleStatus() const; + + private: + static uint64_t allocateGeneration() noexcept; + + Status failLocked(Status status); + Status resetAndUnbindLocked(); + Status deleteJettysLocked(); + Status finishRetireLocked(); + static void rememberFirstError(const Status& candidate, Status& first); + + const UbEndpointKey key_; + const UbContextPtr context_; + const std::shared_ptr adapter_; + const uint32_t jetty_count_; + const JettyOptions jetty_options_; + const uint64_t generation_; + + mutable std::mutex lifecycle_mutex_; + std::vector jetties_; + std::vector jfc_indices_; + UbBootstrapDesc peer_; + Status lifecycle_status_; + Status retire_status_; + bool native_quiesced_{false}; + std::atomic state_{State::kUninitialized}; + std::atomic peer_generation_{0}; + std::atomic outstanding_wrs_{0}; + std::atomic outstanding_bytes_{0}; +}; + +using UbEndpointPtr = std::shared_ptr; +// Keep the spelling used by the design document available to callers. +using UbEndPoint = UbEndpoint; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_ENDPOINT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h new file mode 100644 index 0000000000..06d394c1d9 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint_store.h @@ -0,0 +1,67 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ +#define TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/endpoint.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +// Generation-aware endpoint cache. Failed and retired incarnations are +// unpublished before cleanup, so a subsequent lookup always allocates a new +// generation and can never resurrect a failed Jetty set. +class EndpointStore final { + public: + EndpointStore(std::shared_ptr adapter, size_t max_size, + uint32_t jetty_count, JettyOptions jetty_options = {}); + ~EndpointStore(); + + EndpointStore(const EndpointStore&) = delete; + EndpointStore& operator=(const EndpointStore&) = delete; + + std::shared_ptr get(const UbEndpointKey& key); + Status getOrCreate(const UbEndpointKey& key, const UbContextPtr& context, + std::shared_ptr& endpoint); + + // Only the exact incarnation is removed. A late timeout from generation N + // therefore cannot evict a replacement generation N+1. + bool retire(const UbEndpointKey& key, uint64_t generation); + bool retire(const std::shared_ptr& endpoint); + Status clear(); + [[nodiscard]] size_t size() const; + + private: + struct Entry { + std::shared_ptr endpoint; + uint64_t insertion_order{0}; + }; + + std::shared_ptr adapter_; + const size_t max_size_; + const uint32_t jetty_count_; + const JettyOptions jetty_options_; + mutable std::mutex mutex_; + std::unordered_map endpoints_; + // Unpublished endpoints whose native cleanup failed remain owned until a + // later clear()/shutdown retry. They must never fall through a destructor + // while an ERROR Jetty still lacks its flush fence. + std::vector> quarantined_; + uint64_t next_insertion_order_{1}; +}; + +using UbEndpointStore = EndpointStore; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_ENDPOINT_STORE_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h new file mode 100644 index 0000000000..c135899cdc --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/quota.h @@ -0,0 +1,195 @@ +// Copyright 2026 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. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/transport/ub/rail_monitor.h" + +namespace mooncake::tent::ub { + +struct QuotaLimits { + uint64_t max_inflight_bytes{std::numeric_limits::max()}; + uint64_t max_outstanding_wrs{std::numeric_limits::max()}; + + bool operator==(const QuotaLimits&) const = default; +}; + +struct QuotaUsage { + uint64_t inflight_bytes{0}; + uint64_t outstanding_wrs{0}; + + bool operator==(const QuotaUsage&) const = default; +}; + +struct QuotaStats { + QuotaLimits limits{}; + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; +}; + +// A reservation is a copyable release token. QuotaManager retains the +// authoritative path and charge, so changing those descriptive fields cannot +// decrement the wrong counters, and releasing a copied token cannot do so +// twice. +struct QuotaReservation { + uint64_t id{0}; + UbPostPath path{}; + uint64_t bytes{0}; + uint64_t wrs{0}; + + [[nodiscard]] bool valid() const { return id != 0; } +}; + +struct DeviceQuotaStats : QuotaStats { + Topology::NicID local_topology_id{-1}; +}; + +struct PathQuotaStats : QuotaStats { + UbPostPath path{}; +}; + +struct AggregateQuotaStats { + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + size_t active_reservations{0}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; + uint64_t duplicate_release_attempts{0}; +}; + +// A lock-consistent capacity snapshot used by path selection. Pressure is the +// projected utilization after charging the requested work and is normalized +// to [0, 1]. Unlimited dimensions contribute no pressure. +struct QuotaAvailability { + bool can_acquire{false}; + double normalized_inflight{1.0}; + double normalized_outstanding_wrs{1.0}; +}; + +// Atomically enforces both physical-device and posting-path capacity. This is +// an internal sender-side limit and is intentionally independent of TENT's +// receiver-credit protocol. +class QuotaManager { + public: + explicit QuotaManager(QuotaLimits default_device_limits = {}, + QuotaLimits default_path_limits = {}); + + QuotaManager(const QuotaManager&) = delete; + QuotaManager& operator=(const QuotaManager&) = delete; + + void setDefaultDeviceLimits(const QuotaLimits& limits); + void setDefaultPathLimits(const QuotaLimits& limits); + [[nodiscard]] QuotaLimits defaultDeviceLimits() const; + [[nodiscard]] QuotaLimits defaultPathLimits() const; + + bool setDeviceLimits(Topology::NicID local_topology_id, + const QuotaLimits& limits); + bool clearDeviceLimits(Topology::NicID local_topology_id); + bool setPathLimits(const UbPostPath& path, const QuotaLimits& limits); + bool clearPathLimits(const UbPostPath& path); + + // Acquires device and path charges as one transaction. A zero-byte work + // request is supported, but wrs must be nonzero. + [[nodiscard]] std::optional tryAcquire( + const UbPostPath& path, uint64_t bytes, uint64_t wrs = 1); + + // Tries paths in caller-provided preference order under one lock. This is + // the commit point for multi-rail selection: if a preflight snapshot races + // with another posting worker, later rails are considered before the + // request is deferred. + [[nodiscard]] std::optional tryAcquireFirst( + const std::vector& paths, uint64_t bytes, uint64_t wrs = 1); + + // Returns projected device/path pressure without reserving capacity. + [[nodiscard]] QuotaAvailability availability(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs = 1) const; + + // The first release returns true. Releasing the same token again is a + // harmless no-op and returns false; usage can never underflow. + bool release(const QuotaReservation& reservation); + + [[nodiscard]] DeviceQuotaStats deviceStats( + Topology::NicID local_topology_id) const; + [[nodiscard]] PathQuotaStats pathStats(const UbPostPath& path) const; + [[nodiscard]] std::vector allDeviceStats() const; + [[nodiscard]] std::vector allPathStats() const; + [[nodiscard]] AggregateQuotaStats aggregateStats() const; + [[nodiscard]] size_t activeReservationCount() const; + + private: + struct QuotaRecord { + std::optional override_limits; + // Quota is charged to a physical rail, while this preserves the most + // recent endpoint incarnation for diagnostics. + UbPostPath latest_path{}; + QuotaUsage usage{}; + QuotaUsage peak_usage{}; + uint64_t total_acquisitions{0}; + uint64_t total_releases{0}; + uint64_t rejected_acquisitions{0}; + }; + + struct ActiveReservation { + UbPostPath path{}; + uint64_t bytes{0}; + uint64_t wrs{0}; + }; + + static bool fits(uint64_t current, uint64_t charge, uint64_t limit); + static double normalizedUsage(uint64_t current, uint64_t charge, + uint64_t limit); + static uint64_t saturatingAdd(uint64_t lhs, uint64_t rhs); + static void addUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs); + static void releaseUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs); + static void updatePeak(const QuotaUsage& usage, QuotaUsage& peak); + static QuotaLimits effectiveLimits(const QuotaRecord& record, + const QuotaLimits& defaults); + static QuotaStats makeStats(const QuotaRecord& record, + const QuotaLimits& defaults); + std::optional tryAcquireLocked( + const UbPostPath& path, uint64_t bytes, uint64_t wrs, + bool count_aggregate_reject); + QuotaAvailability availabilityLocked(const UbPostPath& path, uint64_t bytes, + uint64_t wrs) const; + uint64_t nextReservationIdLocked(); + + mutable std::mutex mutex_; + QuotaLimits default_device_limits_; + QuotaLimits default_path_limits_; + std::unordered_map devices_; + std::unordered_map paths_; + std::unordered_map active_reservations_; + AggregateQuotaStats aggregate_stats_{}; + uint64_t next_reservation_id_{1}; +}; + +using UbQuotaManager = QuotaManager; + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_QUOTA_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h new file mode 100644 index 0000000000..1dbf0fd6bf --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/rail_monitor.h @@ -0,0 +1,165 @@ +// Copyright 2026 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. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/transport/ub/slice.h" + +namespace mooncake::tent::ub { + +// Health and learned bandwidth belong to a physical rail, rather than to one +// endpoint incarnation. Endpoint generation is therefore reported in the +// statistics but intentionally excluded from this key. +struct UbRailKey { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + int remote_device_id{-1}; + + bool operator==(const UbRailKey&) const = default; + + [[nodiscard]] bool valid() const { + return local_topology_id >= 0 && remote_device_id >= 0; + } + + static UbRailKey fromPath(const UbPostPath& path) { + return {path.local_topology_id, path.remote_segment_id, + path.remote_device_id}; + } +}; + +struct UbRailKeyHash { + size_t operator()(const UbRailKey& key) const noexcept { + size_t seed = std::hash{}(key.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(key.remote_segment_id)); + combine(std::hash{}(key.remote_device_id)); + return seed; + } +}; + +struct RailMonitorConfig { + uint32_t error_threshold{3}; + uint64_t error_window_ns{10'000'000'000ULL}; + uint64_t cooldown_ns{30'000'000'000ULL}; + // Weight assigned to the newest completion sample. + double ewma_alpha{0.2}; + + [[nodiscard]] bool valid() const { + return error_threshold != 0 && error_window_ns != 0 && + cooldown_ns != 0 && ewma_alpha > 0.0 && ewma_alpha <= 1.0; + } +}; + +struct RailStats { + UbRailKey key{}; + bool paused{false}; + uint32_t errors_in_window{0}; + uint64_t successful_completions{0}; + uint64_t completed_bytes{0}; + uint64_t completion_errors{0}; + uint64_t timeouts{0}; + uint64_t recoveries{0}; + uint64_t endpoint_rebuilds{0}; + uint64_t pauses{0}; + // Bytes/second and nanoseconds respectively. -1 means no valid sample. + double ewma_bandwidth_bytes_per_second{-1.0}; + double ewma_latency_ns{-1.0}; + uint64_t last_success_ns{0}; + uint64_t last_error_ns{0}; + uint64_t pause_started_ns{0}; + uint64_t cooldown_until_ns{0}; + uint64_t latest_endpoint_generation{0}; +}; + +// Thread-safe rolling health and telemetry for UB posting paths. +class RailMonitor { + public: + explicit RailMonitor(RailMonitorConfig config = {}); + + RailMonitor(const RailMonitor&) = delete; + RailMonitor& operator=(const RailMonitor&) = delete; + + // Rejects invalid configurations without changing the active one. + bool configure(const RailMonitorConfig& config); + [[nodiscard]] RailMonitorConfig config() const; + + // Registration is optional; all record operations create the rail lazily. + bool registerPath(const UbPostPath& path); + [[nodiscard]] bool available(const UbPostPath& path, uint64_t now_ns = 0); + + void recordSuccess(const UbPostPath& path, uint64_t bytes, + uint64_t latency_ns, uint64_t now_ns = 0); + void recordError(const UbPostPath& path, uint64_t now_ns = 0); + void recordTimeout(const UbPostPath& path, uint64_t now_ns = 0); + // Records at most one rebuild for each endpoint generation on a physical + // rail. Returns true when telemetry advanced, allowing EndpointStore to + // call this safely from converging/retried rebuild paths. + bool recordEndpointRebuild(const UbPostPath& path, uint64_t now_ns = 0); + + [[nodiscard]] RailStats stats(const UbPostPath& path, uint64_t now_ns = 0); + [[nodiscard]] std::vector allStats(uint64_t now_ns = 0); + + // Adds the best usable remote path sample for each local device. Returns + // -1 until at least one valid completion sample has been observed. + [[nodiscard]] double aggregateBandwidth(uint64_t now_ns = 0); + [[nodiscard]] size_t pathCount() const; + + private: + struct RailState { + RailStats stats{}; + std::deque recent_errors; + // Event timestamps may arrive out of order from different pollers. + // Health decisions never move behind this per-rail watermark. + uint64_t observed_through_ns{0}; + // Late errors at or before a completed recovery epoch must not + // resurrect an already-expired pause. + uint64_t ignore_errors_through_ns{0}; + // Endpoint generations are process-wide monotonic. Keep a separate + // watermark from latest_endpoint_generation because path registration + // may observe the replacement before rebuild telemetry is emitted. + uint64_t recorded_rebuild_generation{0}; + }; + + using RailMap = std::unordered_map; + + static uint64_t normalizedNow(uint64_t now_ns); + static uint64_t deadlineAfter(uint64_t now_ns, uint64_t duration_ns); + static uint64_t observeTimeLocked(RailState& state, uint64_t event_ns); + RailState& getOrCreateLocked(const UbPostPath& path); + static void insertErrorLocked(RailState& state, uint64_t event_ns); + void pruneErrorsLocked(RailState& state, uint64_t now_ns); + void refreshCooldownLocked(RailState& state, uint64_t now_ns); + void recordFailureLocked(const UbPostPath& path, uint64_t now_ns, + bool timeout); + + mutable std::mutex mutex_; + RailMonitorConfig config_; + RailMap rails_; +}; + +using UbRailMonitor = RailMonitor; + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_RAIL_MONITOR_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h new file mode 100644 index 0000000000..3b70c3dc21 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h @@ -0,0 +1,707 @@ +// Copyright 2026 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. + +#ifndef MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ +#define MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/types.h" +#include "tent/runtime/topology.h" + +namespace mooncake::tent::ub { + +inline uint64_t steadyNowNs() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +// Identifies one local-device to remote-device posting path. Endpoint +// generation is part of the identity so a completion from a retired endpoint +// cannot be mistaken for work posted on its replacement. +struct UbPostPath { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + int remote_device_id{-1}; + uint64_t endpoint_generation{0}; + + bool operator==(const UbPostPath&) const = default; + + [[nodiscard]] bool valid() const { + return local_topology_id >= 0 && remote_device_id >= 0 && + endpoint_generation != 0; + } +}; + +struct UbPostPathHash { + size_t operator()(const UbPostPath& path) const noexcept { + size_t seed = std::hash{}(path.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(path.remote_segment_id)); + combine(std::hash{}(path.remote_device_id)); + combine(std::hash{}(path.endpoint_generation)); + return seed; + } +}; + +enum class UbSliceState : uint8_t { + kInitial, + kQueued, + kPosting, + kPosted, + kRetryPending, + kCompleted, + kFailed, + kTimedOut, + kCanceled, + kInvalid, +}; + +inline bool isTerminal(UbSliceState state) { + switch (state) { + case UbSliceState::kCompleted: + case UbSliceState::kFailed: + case UbSliceState::kTimedOut: + case UbSliceState::kCanceled: + case UbSliceState::kInvalid: + return true; + default: + return false; + } +} + +inline TransferStatusEnum transferStatus(UbSliceState state) { + switch (state) { + case UbSliceState::kInitial: + return INITIAL; + case UbSliceState::kCompleted: + return COMPLETED; + case UbSliceState::kFailed: + return FAILED; + case UbSliceState::kTimedOut: + return TIMEOUT; + case UbSliceState::kCanceled: + return CANCELED; + case UbSliceState::kInvalid: + return INVALID; + default: + return PENDING; + } +} + +inline bool isTerminal(TransferStatusEnum status) { + return status == COMPLETED || status == FAILED || status == TIMEOUT || + status == CANCELED || status == INVALID; +} + +inline UbSliceState sliceState(TransferStatusEnum status) { + switch (status) { + case COMPLETED: + return UbSliceState::kCompleted; + case FAILED: + return UbSliceState::kFailed; + case TIMEOUT: + return UbSliceState::kTimedOut; + case CANCELED: + return UbSliceState::kCanceled; + case INVALID: + return UbSliceState::kInvalid; + case INITIAL: + return UbSliceState::kInitial; + case PENDING: + default: + return UbSliceState::kQueued; + } +} + +struct UbSliceSpec { + void* local_address{nullptr}; + uint64_t remote_address{0}; + size_t length{0}; + size_t request_offset{0}; + uint32_t max_retries{0}; +}; + +struct UbAttemptToken { + uint64_t slice_id{0}; + uint32_t attempt{0}; + UbPostPath path{}; + + bool operator==(const UbAttemptToken&) const = default; + + [[nodiscard]] bool valid() const { + return slice_id != 0 && attempt != 0 && path.valid(); + } +}; + +enum class UbAttemptResolution : uint8_t { + kIgnored, + kRetryScheduled, + kTerminal, +}; + +struct UbSliceSnapshot { + uint64_t id{0}; + UbSliceState state{UbSliceState::kInitial}; + UbPostPath path{}; + uint32_t attempt{0}; + uint32_t retry_count{0}; + uint32_t max_retries{0}; + bool cancel_requested{false}; + size_t transferred_bytes{0}; + uint64_t created_ns{0}; + uint64_t queued_ns{0}; + uint64_t attempt_started_ns{0}; + uint64_t posted_ns{0}; + uint64_t terminal_ns{0}; +}; + +struct UbTask; +class UbSlice; + +// Safe user-context payload for a posted operation. Holding this token keeps +// the slice alive until the device completion has been dispatched, while the +// slice itself keeps only a weak reference to its parent task. +struct UbCompletionToken { + std::shared_ptr slice; + UbAttemptToken attempt{}; + + [[nodiscard]] bool valid() const; + bool markPosted(uint64_t now_ns = 0) const; + UbAttemptResolution resolve(TransferStatusEnum outcome, + size_t transferred_bytes, bool retryable, + uint64_t now_ns = 0) const; +}; + +// A logical slice is owned by its UbTask and by any in-flight completion +// token. It keeps only a weak reference back to the task, avoiding a cycle and +// allowing a late completion to be discarded safely after batch teardown. +class UbSlice : public std::enable_shared_from_this { + public: + using Ptr = std::shared_ptr; + + UbSlice(const UbSlice&) = delete; + UbSlice& operator=(const UbSlice&) = delete; + + [[nodiscard]] uint64_t id() const { return id_; } + [[nodiscard]] const UbSliceSpec& spec() const { return spec_; } + + bool markQueued(uint64_t now_ns = 0) { + std::lock_guard lock(mutex_); + if (cancel_requested_.load(std::memory_order_acquire) || + isTerminal(state_)) { + return false; + } + if (state_ != UbSliceState::kInitial && + state_ != UbSliceState::kRetryPending) { + return false; + } + state_ = UbSliceState::kQueued; + queued_ns_ = normalizedNow(now_ns); + return true; + } + + // Claims this slice for one posting attempt. A cancellation racing after + // this point is best-effort: the posting thread owns the device boundary + // and must resolve the attempt instead of reporting an immediate cancel. + std::optional beginAttempt(const UbPostPath& path, + uint64_t now_ns = 0) { + if (!path.valid()) return std::nullopt; + bool canceled = false; + TransferStatusEnum terminal_status = PENDING; + size_t terminal_bytes = 0; + std::optional token; + { + std::lock_guard lock(mutex_); + if (isTerminal(state_)) return std::nullopt; + if (cancel_requested_.load(std::memory_order_acquire)) { + if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + canceled = true; + terminal_status = CANCELED; + } + } else if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + state_ = UbSliceState::kPosting; + path_ = path; + ++attempt_; + attempt_started_ns_ = normalizedNow(now_ns); + posted_ns_ = 0; + token = UbAttemptToken{id_, attempt_, path_}; + } + } + if (canceled) notifyTaskTerminal(terminal_status, terminal_bytes); + return token; + } + + [[nodiscard]] std::optional completionToken( + const UbAttemptToken& token) { + auto self = weak_from_this().lock(); + if (!self) return std::nullopt; + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + (state_ != UbSliceState::kPosting && + state_ != UbSliceState::kPosted)) { + return std::nullopt; + } + return UbCompletionToken{std::move(self), token}; + } + + bool markPosted(const UbAttemptToken& token, uint64_t now_ns = 0) { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + state_ != UbSliceState::kPosting) { + return false; + } + state_ = UbSliceState::kPosted; + posted_ns_ = normalizedNow(now_ns); + return true; + } + + // Linearization point immediately before crossing the native post + // boundary. Cancellation that wins this mutex prevents the WR from being + // posted and terminalizes the slice; cancellation after this point is + // best-effort and waits for completion/drain. + bool tryCommitPost(const UbAttemptToken& token, uint64_t now_ns = 0) { + bool canceled = false; + { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + state_ != UbSliceState::kPosting) { + return false; + } + if (cancel_requested_.load(std::memory_order_acquire)) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + canceled = true; + } else { + state_ = UbSliceState::kPosted; + posted_ns_ = normalizedNow(now_ns); + } + } + if (canceled) notifyTaskTerminal(CANCELED, 0); + return !canceled; + } + + // Resolves exactly one posting attempt. Retriable errors move the logical + // slice back to kRetryPending without notifying the task. Duplicate or + // stale completions are ignored by matching attempt and endpoint + // generation. Only the final resolution contributes task bytes/status. + UbAttemptResolution resolveAttempt(const UbAttemptToken& token, + TransferStatusEnum outcome, + size_t transferred_bytes, bool retryable, + uint64_t now_ns = 0) { + bool terminal = false; + TransferStatusEnum terminal_status = PENDING; + size_t terminal_bytes = 0; + UbAttemptResolution resolution = UbAttemptResolution::kIgnored; + { + std::lock_guard lock(mutex_); + if (!matchesActiveAttemptLocked(token) || + (state_ != UbSliceState::kPosting && + state_ != UbSliceState::kPosted) || + !isTerminal(outcome)) { + return UbAttemptResolution::kIgnored; + } + + const bool canceled = + cancel_requested_.load(std::memory_order_acquire); + if (outcome != COMPLETED && retryable && !canceled && + retry_count_ < spec_.max_retries) { + ++retry_count_; + state_ = UbSliceState::kRetryPending; + queued_ns_ = normalizedNow(now_ns); + resolution = UbAttemptResolution::kRetryScheduled; + } else { + terminal_status = + (canceled && outcome != COMPLETED) ? CANCELED : outcome; + setTerminalLocked(terminal_status, transferred_bytes, + normalizedNow(now_ns)); + terminal_bytes = transferred_bytes_; + terminal = true; + resolution = UbAttemptResolution::kTerminal; + } + } + if (terminal) notifyTaskTerminal(terminal_status, terminal_bytes); + return resolution; + } + + // Resolves work that failed (or completed locally) before any adapter post. + // It deliberately refuses to terminalize kPosting/kPosted work; those must + // drain through resolveAttempt(). + UbAttemptResolution resolveBeforePost(TransferStatusEnum outcome, + size_t transferred_bytes, + bool retryable, uint64_t now_ns = 0) { + bool terminal = false; + TransferStatusEnum terminal_status = outcome; + size_t terminal_bytes = 0; + UbAttemptResolution resolution = UbAttemptResolution::kIgnored; + { + std::lock_guard lock(mutex_); + if (!isTerminal(outcome) || isTerminal(state_) || + state_ == UbSliceState::kPosting || + state_ == UbSliceState::kPosted) { + return UbAttemptResolution::kIgnored; + } + const bool canceled = + cancel_requested_.load(std::memory_order_acquire); + if (outcome != COMPLETED && retryable && !canceled && + retry_count_ < spec_.max_retries) { + ++retry_count_; + state_ = UbSliceState::kRetryPending; + queued_ns_ = normalizedNow(now_ns); + resolution = UbAttemptResolution::kRetryScheduled; + } else { + if (canceled && outcome != COMPLETED) { + terminal_status = CANCELED; + } + setTerminalLocked(terminal_status, transferred_bytes, + normalizedNow(now_ns)); + terminal_bytes = transferred_bytes_; + terminal = true; + resolution = UbAttemptResolution::kTerminal; + } + } + if (terminal) notifyTaskTerminal(terminal_status, terminal_bytes); + return resolution; + } + + bool tryResolveBeforePost(TransferStatusEnum outcome, + size_t transferred_bytes = 0, + uint64_t now_ns = 0) { + return resolveBeforePost(outcome, transferred_bytes, false, now_ns) == + UbAttemptResolution::kTerminal; + } + + // Best-effort cancellation. Unclaimed work becomes terminal immediately; + // posting or posted work only observes the flag and must still be resolved + // by its device completion. + bool requestCancellation(uint64_t now_ns = 0) { + cancel_requested_.store(true, std::memory_order_release); + bool terminal = false; + { + std::lock_guard lock(mutex_); + if (state_ == UbSliceState::kInitial || + state_ == UbSliceState::kQueued || + state_ == UbSliceState::kRetryPending) { + setTerminalLocked(CANCELED, 0, normalizedNow(now_ns)); + terminal = true; + } + } + if (terminal) notifyTaskTerminal(CANCELED, 0); + return terminal; + } + + [[nodiscard]] bool cancellationRequested() const { + return cancel_requested_.load(std::memory_order_acquire); + } + + [[nodiscard]] UbSliceSnapshot snapshot() const { + std::lock_guard lock(mutex_); + return UbSliceSnapshot{ + id_, + state_, + path_, + attempt_, + retry_count_, + spec_.max_retries, + cancel_requested_.load(std::memory_order_acquire), + transferred_bytes_, + created_ns_, + queued_ns_, + attempt_started_ns_, + posted_ns_, + terminal_ns_}; + } + + private: + friend struct UbTask; + + UbSlice(uint64_t id, UbSliceSpec spec, std::weak_ptr task, + uint64_t created_ns) + : id_(id), + spec_(std::move(spec)), + task_(std::move(task)), + created_ns_(normalizedNow(created_ns)) {} + + static uint64_t normalizedNow(uint64_t now_ns) { + return now_ns == 0 ? steadyNowNs() : now_ns; + } + + bool matchesActiveAttemptLocked(const UbAttemptToken& token) const { + return token.slice_id == id_ && token.attempt == attempt_ && + token.path == path_; + } + + void setTerminalLocked(TransferStatusEnum status, size_t bytes, + uint64_t now_ns) { + state_ = sliceState(status); + transferred_bytes_ = std::min(bytes, spec_.length); + terminal_ns_ = now_ns; + } + + void notifyTaskTerminal(TransferStatusEnum status, size_t bytes); + + const uint64_t id_; + const UbSliceSpec spec_; + std::weak_ptr task_; + + mutable std::mutex mutex_; + UbSliceState state_{UbSliceState::kInitial}; + UbPostPath path_{}; + uint32_t attempt_{0}; + uint32_t retry_count_{0}; + std::atomic cancel_requested_{false}; + size_t transferred_bytes_{0}; + uint64_t created_ns_{0}; + uint64_t queued_ns_{0}; + uint64_t attempt_started_ns_{0}; + uint64_t posted_ns_{0}; + uint64_t terminal_ns_{0}; +}; + +struct UbTaskSnapshot { + TransferStatus status{PENDING, 0}; + size_t total_bytes{0}; + size_t total_slices{0}; + size_t resolved_slices{0}; + size_t remaining_slices{0}; + size_t successful_slices{0}; + bool sealed{false}; + bool cancel_requested{false}; + uint64_t created_ns{0}; + uint64_t deadline_ns{0}; + uint64_t terminal_ns{0}; +}; + +struct UbTask : public std::enable_shared_from_this { + public: + using Ptr = std::shared_ptr; + using TerminalCallback = std::function; + + static Ptr create(Request request, TerminalCallback terminal_callback = {}, + uint64_t created_ns = 0) { + return Ptr(new UbTask(std::move(request), std::move(terminal_callback), + created_ns)); + } + + UbTask(const UbTask&) = delete; + UbTask& operator=(const UbTask&) = delete; + + [[nodiscard]] const Request& request() const { return request_; } + + // Slices must be added before seal(). Workers may begin processing only + // after sealing, which prevents an early completion from finalizing a task + // while its remaining slices are still being constructed. + UbSlice::Ptr addSlice(UbSliceSpec spec, uint64_t created_ns = 0) { + UbSlice::Ptr slice; + bool cancel = false; + { + std::lock_guard lock(mutex_); + if (sealed_) return nullptr; + slice = UbSlice::Ptr(new UbSlice(next_slice_id_++, std::move(spec), + weak_from_this(), created_ns)); + slices_.push_back(slice); + cancel = cancel_requested_.load(std::memory_order_acquire); + } + if (cancel) slice->requestCancellation(created_ns); + return slice; + } + + bool seal() { + TerminalCallback callback; + TransferStatus final_status{}; + bool notify = false; + { + std::lock_guard lock(mutex_); + if (sealed_) return false; + sealed_ = true; + notify = maybeFinalizeLocked(final_status, callback); + } + if (notify && callback) callback(final_status); + return true; + } + + // Returns the number of slices canceled before reaching the adapter. Any + // posting/posted slices remain pending until their attempts resolve. + size_t requestCancellation(uint64_t now_ns = 0) { + cancel_requested_.store(true, std::memory_order_release); + std::vector slices; + { + std::lock_guard lock(mutex_); + if (isTerminal(status_.s)) return 0; + slices = slices_; + } + size_t canceled = 0; + for (const auto& slice : slices) { + if (slice->requestCancellation(now_ns)) ++canceled; + } + return canceled; + } + + [[nodiscard]] bool cancellationRequested() const { + return cancel_requested_.load(std::memory_order_acquire); + } + + [[nodiscard]] TransferStatus transferStatus() const { + std::lock_guard lock(mutex_); + return status_; + } + + [[nodiscard]] UbTaskSnapshot snapshot() const { + std::lock_guard lock(mutex_); + return UbTaskSnapshot{status_, + request_.length, + slices_.size(), + resolved_slice_ids_.size(), + slices_.size() - resolved_slice_ids_.size(), + successful_slices_, + sealed_, + cancel_requested_.load(std::memory_order_acquire), + created_ns_, + request_.deadline_ns, + terminal_ns_}; + } + + [[nodiscard]] std::vector slices() const { + std::lock_guard lock(mutex_); + return slices_; + } + + private: + friend class UbSlice; + + UbTask(Request request, TerminalCallback terminal_callback, + uint64_t created_ns) + : request_(std::move(request)), + created_ns_(created_ns == 0 ? steadyNowNs() : created_ns), + terminal_callback_(std::move(terminal_callback)) {} + + static int terminalSeverity(TransferStatusEnum status) { + switch (status) { + case INVALID: + return 4; + case FAILED: + return 3; + case TIMEOUT: + return 2; + case CANCELED: + return 1; + default: + return 0; + } + } + + void onSliceTerminal(uint64_t slice_id, TransferStatusEnum status, + size_t bytes) { + TerminalCallback callback; + TransferStatus final_status{}; + bool notify = false; + { + std::lock_guard lock(mutex_); + if (!resolved_slice_ids_.insert(slice_id).second) return; + if (std::numeric_limits::max() - status_.transferred_bytes < + bytes) { + status_.transferred_bytes = std::numeric_limits::max(); + } else { + status_.transferred_bytes += bytes; + } + if (status == COMPLETED) { + ++successful_slices_; + } else if (terminalSeverity(status) > + terminalSeverity(aggregate_error_)) { + aggregate_error_ = status; + } + notify = maybeFinalizeLocked(final_status, callback); + } + if (notify && callback) callback(final_status); + } + + bool maybeFinalizeLocked(TransferStatus& final_status, + TerminalCallback& callback) { + if (!sealed_ || terminal_notified_ || + resolved_slice_ids_.size() != slices_.size()) { + return false; + } + + if (slices_.empty() || successful_slices_ == slices_.size()) { + status_.s = COMPLETED; + } else { + status_.s = aggregate_error_ == PENDING ? FAILED : aggregate_error_; + } + terminal_notified_ = true; + terminal_ns_ = steadyNowNs(); + final_status = status_; + callback = std::move(terminal_callback_); + return true; + } + + const Request request_; + const uint64_t created_ns_; + mutable std::mutex mutex_; + std::vector slices_; + std::unordered_set resolved_slice_ids_; + uint64_t next_slice_id_{1}; + size_t successful_slices_{0}; + TransferStatusEnum aggregate_error_{PENDING}; + TransferStatus status_{PENDING, 0}; + std::atomic cancel_requested_{false}; + bool sealed_{false}; + bool terminal_notified_{false}; + uint64_t terminal_ns_{0}; + TerminalCallback terminal_callback_; +}; + +inline bool UbCompletionToken::valid() const { + return slice != nullptr && attempt.valid(); +} + +inline bool UbCompletionToken::markPosted(uint64_t now_ns) const { + return valid() && slice->markPosted(attempt, now_ns); +} + +inline UbAttemptResolution UbCompletionToken::resolve( + TransferStatusEnum outcome, size_t transferred_bytes, bool retryable, + uint64_t now_ns) const { + if (!valid()) return UbAttemptResolution::kIgnored; + return slice->resolveAttempt(attempt, outcome, transferred_bytes, retryable, + now_ns); +} + +inline void UbSlice::notifyTaskTerminal(TransferStatusEnum status, + size_t bytes) { + if (auto task = task_.lock()) task->onSliceTerminal(id_, status, bytes); +} + +} // namespace mooncake::tent::ub + +#endif // MOONCAKE_TENT_TRANSPORT_UB_SLICE_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h new file mode 100644 index 0000000000..ea12b9110e --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_transport.h @@ -0,0 +1,72 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_UB_TRANSPORT_H_ +#define TENT_TRANSPORT_UB_UB_TRANSPORT_H_ + +#include +#include + +#include "tent/runtime/control_plane.h" +#include "tent/runtime/transport.h" + +namespace mooncake::tent { + +namespace ub { +class UrmaAdapter; +struct UbTask; +} // namespace ub + +struct UbSubBatch final : public Transport::SubBatch { + std::vector> task_list; + size_t max_size{0}; + + size_t size() const override { return task_list.size(); } +}; + +// TENT-native UB transport. The implementation owns its task/slice scheduler, +// endpoint store and URMA resources; it never converts requests to Classic TE +// types or calls the Classic UbTransport/UbWorkerPool data path. +class UbTransport final : public Transport { + public: + explicit UbTransport(std::shared_ptr adapter = nullptr); + ~UbTransport() override; + + UbTransport(const UbTransport&) = delete; + UbTransport& operator=(const UbTransport&) = delete; + + Status install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr) override; + Status uninstall() override; + + Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override; + Status freeSubBatch(SubBatchRef& batch) override; + Status submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) override; + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override; + + bool supportsCancellation() const override { return true; } + Status cancelTransferTask(SubBatchRef batch, int task_id) override; + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) override; + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override; + Status removeMemoryBuffer(BufferDesc& desc) override; + bool warmupMemory(void* addr, size_t length) override; + + const char* getName() const override { return "ub"; } + double getEstimatedBandwidth() const override; + bool supportNotification() const override { return false; } + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace mooncake::tent + +#endif // TENT_TRANSPORT_UB_UB_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h new file mode 100644 index 0000000000..774169f799 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -0,0 +1,156 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_WORKERS_H_ +#define TENT_TRANSPORT_UB_WORKERS_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/topology.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +class UbEndpoint; + +struct EndpointResolveRequest { + UbContextPtr local_context; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + const SegmentDesc* remote_segment{nullptr}; + Topology::NicID remote_topology_id{-1}; + uint64_t segment_generation{0}; +}; + +using EndpointResolver = std::function&)>; +using EndpointRetirer = std::function&)>; + +// Native UB scheduler. Posting lanes own request selection and URMA post; +// poller lanes own completion dispatch. A monotonic numeric token, never a raw +// UbSlice pointer, crosses the adapter boundary. +class UbWorkers final { + public: + UbWorkers(std::shared_ptr adapter, + std::vector contexts, + std::shared_ptr local_topology, + SegmentManager* segment_manager, UbBufferManager* buffers, + RailMonitor* rail_monitor, QuotaManager* quota, UbParams params, + EndpointResolver endpoint_resolver, + EndpointRetirer endpoint_retirer = {}); + ~UbWorkers(); + + UbWorkers(const UbWorkers&) = delete; + UbWorkers& operator=(const UbWorkers&) = delete; + + Status start(); + Status stop(); + Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL); + Status cancel(const UbTask::Ptr& task); + + [[nodiscard]] bool running() const noexcept { + return accepting_.load(std::memory_order_acquire); + } + [[nodiscard]] size_t queuedCount() const; + [[nodiscard]] size_t inflightCount() const; + + private: + struct PendingSlice { + // Keep the task alive until every queued/in-flight slice reaches a + // terminal state. Callers may release the sub-batch handle before the + // asynchronous data path has drained. + UbTask::Ptr task; + UbSlice::Ptr slice; + uint64_t device_mask{~0ULL}; + int priority{PRIO_HIGH}; + SegmentID target_id{LOCAL_SEGMENT_ID}; + Request::OpCode opcode{Request::READ}; + }; + struct Route; + struct Inflight; + + void postingLoop(size_t worker_index); + void pollingLoop(size_t poller_index); + bool popPending(PendingSlice& pending); + void enqueueRetry(const PendingSlice& pending); + void deferPending(const PendingSlice& pending); + void processPending(const PendingSlice& pending, size_t worker_index); + Status buildRoute(const PendingSlice& pending, Route& route); + Status chooseAndResolveEndpoint(const PendingSlice& pending, Route& route, + std::shared_ptr& endpoint, + UbPostPath& path); + void handleCompletion(const Completion& completion); + void scanTimeouts(); + void releaseInflight(const std::shared_ptr& inflight); + void resolveInflight(const std::shared_ptr& inflight, + TransferStatusEnum outcome, size_t bytes, + bool retryable); + void recordTimeoutOnce(const std::shared_ptr& inflight, + uint64_t now_ns); + void rememberEndpointDrain(const std::shared_ptr& endpoint); + void forgetEndpointDrain(const std::shared_ptr& endpoint); + void failUnposted(const PendingSlice& pending, + TransferStatusEnum outcome = FAILED); + uint64_t nextCompletionToken(); + std::vector orderedLocalDevices( + const PendingSlice& pending) const; + static std::vector orderedRemoteDevices( + const SegmentDesc& segment, const BufferDesc& buffer); + + std::shared_ptr adapter_; + std::vector contexts_; + std::unordered_map context_by_topology_id_; + std::shared_ptr local_topology_; + SegmentManager* segment_manager_; + UbBufferManager* buffers_; + RailMonitor* rail_monitor_; + QuotaManager* quota_; + const UbParams params_; + EndpointResolver endpoint_resolver_; + EndpointRetirer endpoint_retirer_; + + mutable std::mutex queue_mutex_; + std::condition_variable queue_cv_; + std::array, PRIO_LOW + 1> queues_; + + mutable std::mutex inflight_mutex_; + mutable std::mutex endpoint_drain_mutex_; + std::condition_variable inflight_cv_; + std::unordered_map> inflight_; + // Endpoints whose ERROR transition has not yet reached its native flush + // fence remain owned here even if every logical token completes naturally. + std::unordered_map> + draining_endpoints_; + + std::vector> all_jfcs_; + std::unordered_map context_by_jfc_; + std::vector posting_threads_; + std::vector polling_threads_; + std::atomic accepting_{false}; + std::atomic posting_{false}; + std::atomic polling_{false}; + std::atomic timeout_scans_enabled_{false}; + std::atomic next_token_{1}; +}; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_WORKERS_H_ diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 0378593de8..44dc13ea39 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -69,6 +69,24 @@ Status ControlClient::bootstrap(const std::string& server_addr, return Status::OK(); } +Status ControlClient::bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response) { + std::string request_raw, response_raw; + json j = request; + request_raw = j.dump(); + CHECK_STATUS( + tl_rpc_agent.call(server_addr, BootstrapUb, request_raw, response_raw)); + try { + response = json::parse(response_raw).get(); + } catch (const std::exception& e) { + return Status::MalformedJson( + std::string("Malformed UB bootstrap response: ") + e.what() + + LOC_MARK); + } + return Status::OK(); +} + Status ControlClient::sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length) { @@ -249,6 +267,11 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onBootstrapRdma(request, response); }); + rpc_server_->registerFunction( + BootstrapUb, + [this](const std::string_view& request, std::string& response) { + onBootstrapUb(request, response); + }); rpc_server_->registerFunction( SendData, [this](const std::string_view& request, std::string& response) { @@ -298,6 +321,10 @@ ControlService::ControlService(const std::string& type, ControlService::~ControlService() { // Stop RPC workers while callback state and synchronization primitives are // still alive. Member destruction would otherwise tear them down first. + { + std::lock_guard lock(ub_bootstrap_callback_mutex_); + ub_bootstrap_callback_ = {}; + } rpc_server_.reset(); } @@ -410,6 +437,37 @@ void ControlService::onBootstrapRdma(const std::string_view& request, response = j.dump(); } +void ControlService::onBootstrapUb(const std::string_view& request, + std::string& response) { + UbBootstrapDesc response_desc; + try { + auto request_desc = + json::parse(std::string(request)).get(); + int ret = -1; + { + // Callback replacement during uninstall is serialized with + // invocation, so an in-flight bootstrap cannot outlive the UB + // transport object it targets. + std::lock_guard lock(ub_bootstrap_callback_mutex_); + if (ub_bootstrap_callback_) { + ret = ub_bootstrap_callback_(request_desc, response_desc); + } else { + response_desc.reply_msg = + "UB bootstrap callback is not registered"; + } + } + if (ret != 0 && response_desc.reply_msg.empty()) { + response_desc.reply_msg = + "UB bootstrap callback failed, ret=" + std::to_string(ret); + } + } catch (const std::exception& e) { + response_desc.reply_msg = + std::string("Malformed UB bootstrap request: ") + e.what(); + } + json j = response_desc; + response = j.dump(); +} + void ControlService::onSendData(const std::string_view& request, std::string& response) { if (request.size() < sizeof(XferDataDesc)) { diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp index a1444ce5c0..e60c3628b3 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp @@ -20,6 +20,10 @@ #include "tent/transport/rdma/rdma_transport.h" #endif +#ifdef USE_UB +#include "tent/transport/ub/ub_transport.h" +#endif + #ifdef USE_CUDA #include "tent/transport/nvlink/nvlink_transport.h" #include "tent/transport/mnnvl/mnnvl_transport.h" @@ -68,6 +72,13 @@ Status TransferEngineImpl::loadTransports() { } #endif +#ifdef USE_UB + if (conf_->get("transports/ub/enable", false) && + topology_->getNicCount(Topology::NIC_UB)) { + transport_list_[UB] = std::make_shared(); + } +#endif + #ifdef USE_URING if (conf_->get("transports/io_uring/enable", true)) transport_list_[IOURING] = std::make_shared(); diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp new file mode 100644 index 0000000000..a9556dcaae --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp @@ -0,0 +1,472 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/endpoint.h" + +#include +#include +#include +#include +#include +#include + +namespace mooncake::tent::ub { +namespace { + +uint64_t generationSeed() { + const auto wall = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + const auto steady = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + const uint64_t seed = + (wall ^ (steady << 13) ^ (steady >> 7)) & 0x7fffffffffffffffULL; + return seed == 0 ? 1 : seed; +} + +std::atomic g_next_endpoint_generation{generationSeed()}; + +bool samePeer(const UbBootstrapDesc& lhs, const UbBootstrapDesc& rhs) { + return lhs.protocol_version == rhs.protocol_version && + lhs.local_eid == rhs.local_eid && lhs.jetty_ids == rhs.jetty_ids && + lhs.jetty_uasids == rhs.jetty_uasids && + lhs.endpoint_generation == rhs.endpoint_generation; +} + +} // namespace + +size_t UbEndpointKeyHash::operator()(const UbEndpointKey& key) const noexcept { + size_t seed = std::hash{}(key.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(key.remote_segment_id)); + combine(std::hash{}(key.remote_topology_id)); + combine(std::hash{}(key.peer_nic_path)); + return seed; +} + +uint64_t UbEndpoint::allocateGeneration() noexcept { + const uint64_t generation = + g_next_endpoint_generation.fetch_add(1, std::memory_order_relaxed); + // Generation zero is reserved for "no endpoint". Exhausting the 64-bit + // process-wide sequence is not recoverable without reusing generations. + if (generation == 0 || generation == std::numeric_limits::max()) { + std::terminate(); + } + return generation; +} + +UbEndpoint::UbEndpoint(UbEndpointKey key, UbContextPtr context, + std::shared_ptr adapter, + uint32_t jetty_count, JettyOptions jetty_options) + : key_(std::move(key)), + context_(std::move(context)), + adapter_(std::move(adapter)), + jetty_count_(jetty_count), + jetty_options_(jetty_options), + generation_(allocateGeneration()) {} + +UbEndpoint::~UbEndpoint() { + auto status = retire(); + if (!status.ok() || outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + // A caller that drops the final endpoint reference after a failed + // native fence must not let Jetty destructors force RESET/delete. A + // process-lifetime quarantine is safer than DMA-after-free. + static auto* leaked = new std::vector(); + static auto* leaked_mutex = new std::mutex(); + std::scoped_lock lock(lifecycle_mutex_, *leaked_mutex); + leaked->insert(leaked->end(), std::make_move_iterator(jetties_.begin()), + std::make_move_iterator(jetties_.end())); + jetties_.clear(); + } +} + +void UbEndpoint::rememberFirstError(const Status& candidate, Status& first) { + if (first.ok() && !candidate.ok()) first = candidate; +} + +Status UbEndpoint::resetAndUnbindLocked() { + Status first_error = Status::OK(); + if (!adapter_) return first_error; + + for (auto it = jetties_.rbegin(); it != jetties_.rend(); ++it) { + if (!*it) continue; + auto reset_status = adapter_->resetJetty(*it); + rememberFirstError(reset_status, first_error); + // Unimporting a peer before RESET succeeds can sever resources still + // referenced by hardware. Keep this Jetty intact for a retry. + if (!reset_status.ok()) continue; + rememberFirstError(adapter_->unbindJetty(*it), first_error); + } + return first_error; +} + +Status UbEndpoint::deleteJettysLocked() { + Status first_error = Status::OK(); + if (adapter_) { + for (auto it = jetties_.rbegin(); it != jetties_.rend(); ++it) { + if (!*it) continue; + rememberFirstError(adapter_->deleteJetty(*it), first_error); + } + } + jetties_.clear(); + jfc_indices_.clear(); + return first_error; +} + +Status UbEndpoint::failLocked(Status status) { + if (status.ok()) { + status = Status::InternalError( + "UB endpoint entered failed state without an error" LOC_MARK); + } + lifecycle_status_ = status; + state_.store(State::kFailed, std::memory_order_release); + + rememberFirstError(resetAndUnbindLocked(), retire_status_); + rememberFirstError(deleteJettysLocked(), retire_status_); + return lifecycle_status_; +} + +Status UbEndpoint::prepare() { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kPrepared || current == State::kReady) { + return Status::OK(); + } + if (current == State::kFailed) return lifecycle_status_; + if (current != State::kUninitialized) { + return Status::InvalidArgument( + "UB endpoint cannot be prepared in its current state" LOC_MARK); + } + + state_.store(State::kHandshaking, std::memory_order_release); + if (!key_.valid()) { + return failLocked( + Status::InvalidArgument("Invalid UB endpoint key" LOC_MARK)); + } + if (!adapter_ || !context_ || !context_->active() || !context_->handle() || + !context_->handle()->valid()) { + return failLocked(Status::InvalidArgument( + "UB endpoint requires an active context and adapter" LOC_MARK)); + } + if (context_->topologyId() != key_.local_topology_id) { + return failLocked(Status::InvalidArgument( + "UB endpoint key does not match its local context" LOC_MARK)); + } + if (jetty_count_ == 0 || context_->jfcs().empty()) { + return failLocked(Status::InvalidArgument( + "UB endpoint requires at least one Jetty and JFC" LOC_MARK)); + } + const uint32_t max_jetty = context_->deviceInfo().capabilities.max_jetty; + if (max_jetty != 0 && jetty_count_ > max_jetty) { + return failLocked( + Status::InvalidArgument("Requested endpoint Jetty count exceeds " + "device capability" LOC_MARK)); + } + + jetties_.reserve(jetty_count_); + jfc_indices_.reserve(jetty_count_); + for (uint32_t index = 0; index < jetty_count_; ++index) { + const size_t jfc_index = index % context_->jfcs().size(); + auto jfc = context_->jfc(jfc_index); + if (!jfc || !jfc->valid()) { + return failLocked(Status::InvalidArgument( + "UB endpoint selected an inactive JFC" LOC_MARK)); + } + + JettyPtr jetty; + auto status = adapter_->createJetty(context_->handle(), jfc->handle(), + jetty_options_, jetty); + if (!status.ok()) return failLocked(std::move(status)); + if (!jetty || !jetty->valid() || jetty->id() == 0) { + if (jetty) jetties_.push_back(std::move(jetty)); + return failLocked(Status::InternalError( + "URMA adapter returned an invalid Jetty" LOC_MARK)); + } + jetties_.push_back(std::move(jetty)); + jfc_indices_.push_back(jfc_index); + } + + lifecycle_status_ = Status::OK(); + state_.store(State::kPrepared, std::memory_order_release); + return Status::OK(); +} + +Status UbEndpoint::bind(const UbBootstrapDesc& peer) { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kReady) { + if (samePeer(peer_, peer)) return Status::OK(); + return Status::InvalidArgument( + "UB endpoint is already bound to another peer generation" LOC_MARK); + } + if (current == State::kFailed) return lifecycle_status_; + if (current != State::kPrepared) { + return Status::InvalidArgument( + "UB endpoint must be prepared before bind" LOC_MARK); + } + + state_.store(State::kBinding, std::memory_order_release); + if (!peer.reply_msg.empty()) { + return failLocked(Status::RpcServiceError( + std::string("Peer rejected UB bootstrap: ") + peer.reply_msg)); + } + if (peer.protocol_version != 1) { + return failLocked(Status::InvalidArgument( + "Unsupported UB bootstrap protocol version" LOC_MARK)); + } + if (peer.endpoint_generation == 0 || peer.local_eid.empty()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap is missing EID or " + "endpoint generation" LOC_MARK)); + } + if (peer.jetty_ids.size() != jetties_.size()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap Jetty count does not " + "match local endpoint" LOC_MARK)); + } + if (!peer.jetty_uasids.empty() && + peer.jetty_uasids.size() != jetties_.size()) { + return failLocked( + Status::InvalidArgument("Peer UB bootstrap UASID count does not " + "match local endpoint" LOC_MARK)); + } + + for (size_t index = 0; index < jetties_.size(); ++index) { + if (peer.jetty_ids[index] == 0) { + return failLocked(Status::InvalidArgument( + "Peer UB bootstrap contains a zero Jetty ID" LOC_MARK)); + } + RemoteJettyInfo remote; + remote.eid = peer.local_eid; + remote.id = peer.jetty_ids[index]; + if (!peer.jetty_uasids.empty()) { + remote.uasid = peer.jetty_uasids[index]; + } + auto status = adapter_->bindJetty(jetties_[index], remote); + if (!status.ok()) return failLocked(std::move(status)); + } + + peer_ = peer; + peer_generation_.store(peer.endpoint_generation, std::memory_order_release); + lifecycle_status_ = Status::OK(); + state_.store(State::kReady, std::memory_order_release); + return Status::OK(); +} + +Status UbEndpoint::makeBootstrapDesc(const std::string& segment_name, + const std::string& local_nic_path, + const std::string& peer_nic_path, + uint64_t segment_generation, + UbBootstrapDesc& output) const { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current != State::kPrepared && current != State::kReady) { + return Status::InvalidArgument( + "UB endpoint must be prepared before bootstrap" LOC_MARK); + } + + output = UbBootstrapDesc{}; + output.protocol_version = 1; + output.segment_name = segment_name; + output.local_nic_path = local_nic_path.empty() + ? context_->deviceInfo().native_device_path + : local_nic_path; + output.peer_nic_path = + peer_nic_path.empty() ? key_.peer_nic_path : peer_nic_path; + output.local_device_name = context_->deviceInfo().native_device_name; + output.local_device_id = context_->topologyId(); + output.local_eid_index = static_cast(context_->deviceInfo().eid_index); + output.local_eid = context_->deviceInfo().eid; + output.jetty_ids.reserve(jetties_.size()); + output.jetty_uasids.reserve(jetties_.size()); + for (const auto& jetty : jetties_) { + if (!jetty || !jetty->valid() || jetty->id() == 0) { + return Status::InternalError( + "UB endpoint contains an invalid Jetty" LOC_MARK); + } + output.jetty_ids.push_back(jetty->id()); + output.jetty_uasids.push_back(jetty->uasid()); + } + output.endpoint_generation = generation_; + output.segment_generation = segment_generation; + output.capabilities = {"read", "write", "endpoint_generation"}; + return Status::OK(); +} + +bool UbEndpoint::tryAcquireOutstanding(uint64_t bytes) noexcept { + std::lock_guard lock(lifecycle_mutex_); + if (state_.load(std::memory_order_relaxed) != State::kReady || !context_ || + !context_->active()) { + return false; + } + + outstanding_wrs_.fetch_add(1, std::memory_order_relaxed); + outstanding_bytes_.fetch_add(bytes, std::memory_order_relaxed); + context_->addInflight(bytes); + return true; +} + +void UbEndpoint::releaseOutstanding(uint64_t bytes) noexcept { + std::lock_guard lock(lifecycle_mutex_); + const uint64_t current_wrs = + outstanding_wrs_.load(std::memory_order_relaxed); + if (current_wrs == 0) return; + + outstanding_wrs_.store(current_wrs - 1, std::memory_order_relaxed); + const uint64_t current_bytes = + outstanding_bytes_.load(std::memory_order_relaxed); + outstanding_bytes_.store(current_bytes >= bytes ? current_bytes - bytes : 0, + std::memory_order_relaxed); + context_->removeInflight(bytes); + + if (current_wrs == 1 && + state_.load(std::memory_order_relaxed) == State::kDestroying) { + (void)finishRetireLocked(); + } +} + +Status UbEndpoint::finishRetireLocked() { + if (state_.load(std::memory_order_relaxed) == State::kDestroyed) { + return retire_status_; + } + if (outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + return retire_status_; + } + + // With no outstanding WR, RESET itself is a sufficient fence. When + // quiesce() ran earlier this is an idempotent cleanup pass. + auto reset_status = resetAndUnbindLocked(); + if (!reset_status.ok()) { + retire_status_ = reset_status; + return reset_status; + } + native_quiesced_ = true; + auto delete_status = deleteJettysLocked(); + if (!delete_status.ok()) { + retire_status_ = delete_status; + return delete_status; + } + retire_status_ = Status::OK(); + peer_ = UbBootstrapDesc{}; + peer_generation_.store(0, std::memory_order_release); + state_.store(State::kDestroyed, std::memory_order_release); + return retire_status_; +} + +Status UbEndpoint::retire() { + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kDestroyed) return retire_status_; + + if (current != State::kDestroying) { + state_.store(State::kDestroying, std::memory_order_release); + } + if (outstanding_wrs_.load(std::memory_order_relaxed) != 0) { + return retire_status_; + } + return finishRetireLocked(); +} + +Status UbEndpoint::quiesce(uint32_t timeout_ms, + std::vector& completions) { + completions.clear(); + std::lock_guard lock(lifecycle_mutex_); + const State current = state_.load(std::memory_order_relaxed); + if (current == State::kDestroyed) return retire_status_; + if (!adapter_ || timeout_ms == 0) { + return Status::InvalidArgument( + "UB endpoint quiesce requires an adapter and timeout" LOC_MARK); + } + + state_.store(State::kDestroying, std::memory_order_release); + Status fence_error = Status::OK(); + if (!native_quiesced_) { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::milliseconds(timeout_ms); + for (const auto& jetty : jetties_) { + if (!jetty) continue; + const auto now = std::chrono::steady_clock::now(); + if (now >= deadline) { + if (fence_error.ok()) { + fence_error = Status::RdmaError( + "UB endpoint Jetty drain budget exhausted"); + } + break; + } + const auto remaining = + std::chrono::duration_cast(deadline - + now); + const uint32_t remaining_ms = + static_cast(std::max(1, remaining.count())); + std::vector drained; + auto status = adapter_->quiesceJetty(jetty, remaining_ms, drained); + completions.insert(completions.end(), + std::make_move_iterator(drained.begin()), + std::make_move_iterator(drained.end())); + if (!status.ok() && fence_error.ok()) fence_error = status; + } + if (!fence_error.ok()) { + // Do not RESET, unbind, or delete anything without a fence for + // every Jetty. A later shutdown attempt can safely retry. + return fence_error; + } + native_quiesced_ = true; + } + + auto reset_status = resetAndUnbindLocked(); + if (!reset_status.ok()) return reset_status; + if (outstanding_wrs_.load(std::memory_order_relaxed) == 0) { + return finishRetireLocked(); + } + return Status::OK(); +} + +bool UbEndpoint::reusable() const noexcept { + switch (state()) { + case State::kUninitialized: + case State::kHandshaking: + case State::kPrepared: + case State::kBinding: + case State::kReady: + return true; + case State::kFailed: + case State::kDestroying: + case State::kDestroyed: + return false; + } + return false; +} + +size_t UbEndpoint::jettyCount() const { + std::lock_guard lock(lifecycle_mutex_); + return jetties_.size(); +} + +JettyPtr UbEndpoint::jetty(size_t index) const { + std::lock_guard lock(lifecycle_mutex_); + if (jetties_.empty()) return nullptr; + return jetties_[index % jetties_.size()]; +} + +size_t UbEndpoint::jfcIndex(size_t jetty_index) const { + std::lock_guard lock(lifecycle_mutex_); + if (jfc_indices_.empty()) return 0; + return jfc_indices_[jetty_index % jfc_indices_.size()]; +} + +std::vector UbEndpoint::jetties() const { + std::lock_guard lock(lifecycle_mutex_); + return jetties_; +} + +Status UbEndpoint::lifecycleStatus() const { + std::lock_guard lock(lifecycle_mutex_); + return lifecycle_status_; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp new file mode 100644 index 0000000000..295d9553da --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint_store.cpp @@ -0,0 +1,189 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/endpoint_store.h" + +#include +#include +#include +#include + +namespace mooncake::tent::ub { + +EndpointStore::EndpointStore(std::shared_ptr adapter, + size_t max_size, uint32_t jetty_count, + JettyOptions jetty_options) + : adapter_(std::move(adapter)), + max_size_(max_size), + jetty_count_(jetty_count), + jetty_options_(jetty_options) {} + +EndpointStore::~EndpointStore() { + auto status = clear(); + if (!status.ok()) { + // Standalone owners may ignore clear(); preserve unsafe-to-destroy + // endpoints for process lifetime rather than letting native handle + // destructors bypass a failed drain fence. + static auto* leaked = new std::vector>(); + static auto* leaked_mutex = new std::mutex(); + std::scoped_lock lock(mutex_, *leaked_mutex); + leaked->insert(leaked->end(), + std::make_move_iterator(quarantined_.begin()), + std::make_move_iterator(quarantined_.end())); + quarantined_.clear(); + } +} + +std::shared_ptr EndpointStore::get(const UbEndpointKey& key) { + std::shared_ptr retired; + std::shared_ptr result; + { + std::lock_guard lock(mutex_); + auto it = endpoints_.find(key); + if (it == endpoints_.end()) return nullptr; + if (!it->second.endpoint || !it->second.endpoint->reusable()) { + retired = std::move(it->second.endpoint); + endpoints_.erase(it); + } else { + result = it->second.endpoint; + } + } + if (retired) { + auto status = retired->retire(); + if (!status.ok()) { + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(retired)); + } + } + return result; +} + +Status EndpointStore::getOrCreate(const UbEndpointKey& key, + const UbContextPtr& context, + std::shared_ptr& endpoint) { + endpoint.reset(); + if (!key.valid() || !context || + context->topologyId() != key.local_topology_id || max_size_ == 0 || + jetty_count_ == 0) { + return Status::InvalidArgument( + "Invalid UB endpoint store request" LOC_MARK); + } + + std::shared_ptr evicted; + { + std::lock_guard lock(mutex_); + auto existing = endpoints_.find(key); + if (existing != endpoints_.end()) { + if (existing->second.endpoint && + existing->second.endpoint->reusable()) { + endpoint = existing->second.endpoint; + } else { + evicted = std::move(existing->second.endpoint); + endpoints_.erase(existing); + } + } + + if (!endpoint && endpoints_.size() + quarantined_.size() >= max_size_) { + auto victim = endpoints_.end(); + uint64_t oldest = std::numeric_limits::max(); + for (auto it = endpoints_.begin(); it != endpoints_.end(); ++it) { + if (it->second.endpoint && + it->second.endpoint->outstandingWrs() == 0 && + it->second.insertion_order < oldest) { + victim = it; + oldest = it->second.insertion_order; + } + } + if (victim == endpoints_.end()) { + return Status::TooManyRequests( + "All UB endpoint cache entries are in flight" LOC_MARK); + } + if (!evicted) evicted = std::move(victim->second.endpoint); + endpoints_.erase(victim); + } + + if (!endpoint) { + endpoint = std::make_shared( + key, context, adapter_, jetty_count_, jetty_options_); + endpoints_.emplace(key, Entry{endpoint, next_insertion_order_++}); + } + } + if (evicted) { + auto evict_status = evicted->retire(); + if (!evict_status.ok()) { + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(evicted)); + } + } + + auto status = endpoint->prepare(); + if (!status.ok()) { + (void)retire(key, endpoint->generation()); + endpoint.reset(); + return status; + } + return Status::OK(); +} + +bool EndpointStore::retire(const UbEndpointKey& key, uint64_t generation) { + std::shared_ptr endpoint; + { + std::lock_guard lock(mutex_); + auto it = endpoints_.find(key); + if (it == endpoints_.end() || !it->second.endpoint || + it->second.endpoint->generation() != generation) { + return false; + } + endpoint = std::move(it->second.endpoint); + endpoints_.erase(it); + } + auto status = endpoint->retire(); + if (!status.ok()) { + std::lock_guard lock(mutex_); + quarantined_.push_back(std::move(endpoint)); + } + return true; +} + +bool EndpointStore::retire(const std::shared_ptr& endpoint) { + return endpoint && retire(endpoint->key(), endpoint->generation()); +} + +Status EndpointStore::clear() { + std::vector> endpoints; + { + std::lock_guard lock(mutex_); + endpoints.reserve(endpoints_.size()); + for (auto& [_, entry] : endpoints_) { + if (entry.endpoint) endpoints.push_back(std::move(entry.endpoint)); + } + endpoints_.clear(); + for (auto& endpoint : quarantined_) { + if (endpoint) endpoints.push_back(std::move(endpoint)); + } + quarantined_.clear(); + } + Status first_error = Status::OK(); + std::vector> failed; + for (auto& endpoint : endpoints) { + auto status = endpoint->retire(); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + failed.push_back(std::move(endpoint)); + } + } + if (!failed.empty()) { + std::lock_guard lock(mutex_); + quarantined_.insert(quarantined_.end(), + std::make_move_iterator(failed.begin()), + std::make_move_iterator(failed.end())); + } + return first_error; +} + +size_t EndpointStore::size() const { + std::lock_guard lock(mutex_); + return endpoints_.size(); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp b/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp new file mode 100644 index 0000000000..dab6d0da1b --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/quota.cpp @@ -0,0 +1,371 @@ +// Copyright 2026 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. + +#include "tent/transport/ub/quota.h" + +#include + +namespace mooncake::tent::ub { + +QuotaManager::QuotaManager(QuotaLimits default_device_limits, + QuotaLimits default_path_limits) + : default_device_limits_(default_device_limits), + default_path_limits_(default_path_limits) {} + +void QuotaManager::setDefaultDeviceLimits(const QuotaLimits& limits) { + std::lock_guard lock(mutex_); + default_device_limits_ = limits; +} + +void QuotaManager::setDefaultPathLimits(const QuotaLimits& limits) { + std::lock_guard lock(mutex_); + default_path_limits_ = limits; +} + +QuotaLimits QuotaManager::defaultDeviceLimits() const { + std::lock_guard lock(mutex_); + return default_device_limits_; +} + +QuotaLimits QuotaManager::defaultPathLimits() const { + std::lock_guard lock(mutex_); + return default_path_limits_; +} + +bool QuotaManager::setDeviceLimits(Topology::NicID local_topology_id, + const QuotaLimits& limits) { + if (local_topology_id < 0) return false; + std::lock_guard lock(mutex_); + devices_[local_topology_id].override_limits = limits; + return true; +} + +bool QuotaManager::clearDeviceLimits(Topology::NicID local_topology_id) { + if (local_topology_id < 0) return false; + std::lock_guard lock(mutex_); + auto it = devices_.find(local_topology_id); + if (it == devices_.end() || !it->second.override_limits) return false; + it->second.override_limits.reset(); + return true; +} + +bool QuotaManager::setPathLimits(const UbPostPath& path, + const QuotaLimits& limits) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + auto& record = paths_[UbRailKey::fromPath(path)]; + record.latest_path = path; + record.override_limits = limits; + return true; +} + +bool QuotaManager::clearPathLimits(const UbPostPath& path) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + auto it = paths_.find(UbRailKey::fromPath(path)); + if (it == paths_.end() || !it->second.override_limits) return false; + it->second.override_limits.reset(); + return true; +} + +std::optional QuotaManager::tryAcquire(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) { + std::lock_guard lock(mutex_); + return tryAcquireLocked(path, bytes, wrs, true); +} + +std::optional QuotaManager::tryAcquireFirst( + const std::vector& paths, uint64_t bytes, uint64_t wrs) { + std::lock_guard lock(mutex_); + for (const auto& path : paths) { + auto reservation = tryAcquireLocked(path, bytes, wrs, false); + if (reservation) return reservation; + } + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + return std::nullopt; +} + +QuotaAvailability QuotaManager::availability(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) const { + std::lock_guard lock(mutex_); + return availabilityLocked(path, bytes, wrs); +} + +std::optional QuotaManager::tryAcquireLocked( + const UbPostPath& path, uint64_t bytes, uint64_t wrs, + bool count_aggregate_reject) { + if (!path.valid() || wrs == 0) { + if (count_aggregate_reject) { + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + } + return std::nullopt; + } + + auto& device = devices_[path.local_topology_id]; + auto& rail = paths_[UbRailKey::fromPath(path)]; + if (!rail.latest_path.valid() || + path.endpoint_generation > rail.latest_path.endpoint_generation) { + rail.latest_path = path; + } + const auto device_limits = effectiveLimits(device, default_device_limits_); + const auto path_limits = effectiveLimits(rail, default_path_limits_); + const bool device_fits = fits(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes) && + fits(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs); + const bool path_fits = + fits(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes) && + fits(rail.usage.outstanding_wrs, wrs, path_limits.max_outstanding_wrs); + if (!device_fits || !path_fits) { + if (!device_fits) { + device.rejected_acquisitions = + saturatingAdd(device.rejected_acquisitions, 1); + } + if (!path_fits) { + rail.rejected_acquisitions = + saturatingAdd(rail.rejected_acquisitions, 1); + } + if (count_aggregate_reject) { + aggregate_stats_.rejected_acquisitions = + saturatingAdd(aggregate_stats_.rejected_acquisitions, 1); + } + return std::nullopt; + } + + const uint64_t id = nextReservationIdLocked(); + addUsage(device.usage, bytes, wrs); + addUsage(rail.usage, bytes, wrs); + updatePeak(device.usage, device.peak_usage); + updatePeak(rail.usage, rail.peak_usage); + device.total_acquisitions = saturatingAdd(device.total_acquisitions, 1); + rail.total_acquisitions = saturatingAdd(rail.total_acquisitions, 1); + + addUsage(aggregate_stats_.usage, bytes, wrs); + updatePeak(aggregate_stats_.usage, aggregate_stats_.peak_usage); + aggregate_stats_.total_acquisitions = + saturatingAdd(aggregate_stats_.total_acquisitions, 1); + active_reservations_.emplace(id, ActiveReservation{path, bytes, wrs}); + aggregate_stats_.active_reservations = active_reservations_.size(); + return QuotaReservation{id, path, bytes, wrs}; +} + +QuotaAvailability QuotaManager::availabilityLocked(const UbPostPath& path, + uint64_t bytes, + uint64_t wrs) const { + if (!path.valid() || wrs == 0) return {}; + + const auto device_it = devices_.find(path.local_topology_id); + const auto path_it = paths_.find(UbRailKey::fromPath(path)); + const QuotaRecord empty; + const auto& device = + device_it == devices_.end() ? empty : device_it->second; + const auto& rail = path_it == paths_.end() ? empty : path_it->second; + const auto device_limits = effectiveLimits(device, default_device_limits_); + const auto path_limits = effectiveLimits(rail, default_path_limits_); + + QuotaAvailability result; + result.can_acquire = + fits(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes) && + fits(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs) && + fits(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes) && + fits(rail.usage.outstanding_wrs, wrs, path_limits.max_outstanding_wrs); + result.normalized_inflight = + std::max(normalizedUsage(device.usage.inflight_bytes, bytes, + device_limits.max_inflight_bytes), + normalizedUsage(rail.usage.inflight_bytes, bytes, + path_limits.max_inflight_bytes)); + result.normalized_outstanding_wrs = + std::max(normalizedUsage(device.usage.outstanding_wrs, wrs, + device_limits.max_outstanding_wrs), + normalizedUsage(rail.usage.outstanding_wrs, wrs, + path_limits.max_outstanding_wrs)); + return result; +} + +bool QuotaManager::release(const QuotaReservation& reservation) { + std::lock_guard lock(mutex_); + const auto active = active_reservations_.find(reservation.id); + if (reservation.id == 0 || active == active_reservations_.end()) { + aggregate_stats_.duplicate_release_attempts = + saturatingAdd(aggregate_stats_.duplicate_release_attempts, 1); + return false; + } + + // Always release the manager-owned record. The caller's copy may have + // been changed or may refer to an earlier copy of this reservation. + const ActiveReservation charge = active->second; + active_reservations_.erase(active); + + auto device = devices_.find(charge.path.local_topology_id); + if (device != devices_.end()) { + releaseUsage(device->second.usage, charge.bytes, charge.wrs); + device->second.total_releases = + saturatingAdd(device->second.total_releases, 1); + } + auto path = paths_.find(UbRailKey::fromPath(charge.path)); + if (path != paths_.end()) { + releaseUsage(path->second.usage, charge.bytes, charge.wrs); + path->second.total_releases = + saturatingAdd(path->second.total_releases, 1); + } + + releaseUsage(aggregate_stats_.usage, charge.bytes, charge.wrs); + aggregate_stats_.total_releases = + saturatingAdd(aggregate_stats_.total_releases, 1); + aggregate_stats_.active_reservations = active_reservations_.size(); + return true; +} + +DeviceQuotaStats QuotaManager::deviceStats( + Topology::NicID local_topology_id) const { + std::lock_guard lock(mutex_); + DeviceQuotaStats result; + result.local_topology_id = local_topology_id; + auto it = devices_.find(local_topology_id); + if (it == devices_.end()) { + static_cast(result).limits = default_device_limits_; + } else { + static_cast(result) = + makeStats(it->second, default_device_limits_); + } + return result; +} + +PathQuotaStats QuotaManager::pathStats(const UbPostPath& path) const { + std::lock_guard lock(mutex_); + PathQuotaStats result; + result.path = path; + auto it = paths_.find(UbRailKey::fromPath(path)); + if (it == paths_.end()) { + static_cast(result).limits = default_path_limits_; + } else { + static_cast(result) = + makeStats(it->second, default_path_limits_); + } + return result; +} + +std::vector QuotaManager::allDeviceStats() const { + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(devices_.size()); + for (const auto& [id, record] : devices_) { + DeviceQuotaStats stats; + static_cast(stats) = + makeStats(record, default_device_limits_); + stats.local_topology_id = id; + result.push_back(stats); + } + return result; +} + +std::vector QuotaManager::allPathStats() const { + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(paths_.size()); + for (const auto& [_, record] : paths_) { + PathQuotaStats stats; + static_cast(stats) = + makeStats(record, default_path_limits_); + stats.path = record.latest_path; + result.push_back(stats); + } + return result; +} + +AggregateQuotaStats QuotaManager::aggregateStats() const { + std::lock_guard lock(mutex_); + return aggregate_stats_; +} + +size_t QuotaManager::activeReservationCount() const { + std::lock_guard lock(mutex_); + return active_reservations_.size(); +} + +bool QuotaManager::fits(uint64_t current, uint64_t charge, uint64_t limit) { + return current <= limit && charge <= limit - current; +} + +double QuotaManager::normalizedUsage(uint64_t current, uint64_t charge, + uint64_t limit) { + if (limit == std::numeric_limits::max()) return 0.0; + if (limit == 0) return current == 0 && charge == 0 ? 0.0 : 1.0; + const long double projected = std::min( + static_cast(limit), + static_cast(current) + static_cast(charge)); + return static_cast(projected / static_cast(limit)); +} + +uint64_t QuotaManager::saturatingAdd(uint64_t lhs, uint64_t rhs) { + if (std::numeric_limits::max() - lhs < rhs) { + return std::numeric_limits::max(); + } + return lhs + rhs; +} + +void QuotaManager::addUsage(QuotaUsage& usage, uint64_t bytes, uint64_t wrs) { + usage.inflight_bytes = saturatingAdd(usage.inflight_bytes, bytes); + usage.outstanding_wrs = saturatingAdd(usage.outstanding_wrs, wrs); +} + +void QuotaManager::releaseUsage(QuotaUsage& usage, uint64_t bytes, + uint64_t wrs) { + usage.inflight_bytes = + bytes >= usage.inflight_bytes ? 0 : usage.inflight_bytes - bytes; + usage.outstanding_wrs = + wrs >= usage.outstanding_wrs ? 0 : usage.outstanding_wrs - wrs; +} + +void QuotaManager::updatePeak(const QuotaUsage& usage, QuotaUsage& peak) { + peak.inflight_bytes = std::max(peak.inflight_bytes, usage.inflight_bytes); + peak.outstanding_wrs = + std::max(peak.outstanding_wrs, usage.outstanding_wrs); +} + +QuotaLimits QuotaManager::effectiveLimits(const QuotaRecord& record, + const QuotaLimits& defaults) { + return record.override_limits.value_or(defaults); +} + +QuotaStats QuotaManager::makeStats(const QuotaRecord& record, + const QuotaLimits& defaults) { + return QuotaStats{effectiveLimits(record, defaults), + record.usage, + record.peak_usage, + record.total_acquisitions, + record.total_releases, + record.rejected_acquisitions}; +} + +uint64_t QuotaManager::nextReservationIdLocked() { + // IDs are never reused while active, including across uint64_t wrap. + while (next_reservation_id_ == 0 || + active_reservations_.contains(next_reservation_id_)) { + ++next_reservation_id_; + } + const uint64_t result = next_reservation_id_++; + if (next_reservation_id_ == 0) ++next_reservation_id_; + return result; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp b/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp new file mode 100644 index 0000000000..02ae248945 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/rail_monitor.cpp @@ -0,0 +1,279 @@ +// Copyright 2026 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. + +#include "tent/transport/ub/rail_monitor.h" + +#include +#include +#include + +namespace mooncake::tent::ub { + +namespace { + +double updateEwma(double current, double sample, double alpha) { + if (current < 0.0) return sample; + return alpha * sample + (1.0 - alpha) * current; +} + +} // namespace + +RailMonitor::RailMonitor(RailMonitorConfig config) { + if (config.valid()) config_ = config; +} + +bool RailMonitor::configure(const RailMonitorConfig& config) { + if (!config.valid()) return false; + std::lock_guard lock(mutex_); + config_ = config; + return true; +} + +RailMonitorConfig RailMonitor::config() const { + std::lock_guard lock(mutex_); + return config_; +} + +bool RailMonitor::registerPath(const UbPostPath& path) { + if (!path.valid()) return false; + std::lock_guard lock(mutex_); + getOrCreateLocked(path); + return true; +} + +bool RailMonitor::available(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return false; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + return !state.stats.paused; +} + +void RailMonitor::recordSuccess(const UbPostPath& path, uint64_t bytes, + uint64_t latency_ns, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + + ++state.stats.successful_completions; + if (std::numeric_limits::max() - state.stats.completed_bytes < + bytes) { + state.stats.completed_bytes = std::numeric_limits::max(); + } else { + state.stats.completed_bytes += bytes; + } + state.stats.last_success_ns = std::max(state.stats.last_success_ns, now_ns); + + // Zero-byte or zero-latency completions are valid completions but are not + // usable bandwidth observations. + if (bytes == 0 || latency_ns == 0) return; + const double bandwidth = static_cast(bytes) * 1'000'000'000.0 / + static_cast(latency_ns); + state.stats.ewma_bandwidth_bytes_per_second = + updateEwma(state.stats.ewma_bandwidth_bytes_per_second, bandwidth, + config_.ewma_alpha); + state.stats.ewma_latency_ns = + updateEwma(state.stats.ewma_latency_ns, static_cast(latency_ns), + config_.ewma_alpha); +} + +void RailMonitor::recordError(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + recordFailureLocked(path, now_ns, false); +} + +void RailMonitor::recordTimeout(const UbPostPath& path, uint64_t now_ns) { + if (!path.valid()) return; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + recordFailureLocked(path, now_ns, true); +} + +bool RailMonitor::recordEndpointRebuild(const UbPostPath& path, + uint64_t now_ns) { + if (!path.valid()) return false; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + if (path.endpoint_generation <= state.recorded_rebuild_generation) { + return false; + } + state.recorded_rebuild_generation = path.endpoint_generation; + ++state.stats.endpoint_rebuilds; + return true; +} + +RailStats RailMonitor::stats(const UbPostPath& path, uint64_t now_ns) { + RailStats result; + result.key = UbRailKey::fromPath(path); + if (!path.valid()) return result; + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + auto& state = getOrCreateLocked(path); + now_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, now_ns); + pruneErrorsLocked(state, now_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + return state.stats; +} + +std::vector RailMonitor::allStats(uint64_t now_ns) { + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + std::vector result; + result.reserve(rails_.size()); + for (auto& [key, state] : rails_) { + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + pruneErrorsLocked(state, observed_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + result.push_back(state.stats); + } + return result; +} + +double RailMonitor::aggregateBandwidth(uint64_t now_ns) { + now_ns = normalizedNow(now_ns); + std::lock_guard lock(mutex_); + + bool has_sample = false; + std::unordered_map best_per_device; + for (auto& [key, state] : rails_) { + const uint64_t observed_ns = observeTimeLocked(state, now_ns); + refreshCooldownLocked(state, observed_ns); + const double bandwidth = state.stats.ewma_bandwidth_bytes_per_second; + if (bandwidth < 0.0) continue; + has_sample = true; + if (state.stats.paused) continue; + auto [it, inserted] = + best_per_device.emplace(key.local_topology_id, bandwidth); + if (!inserted) it->second = std::max(it->second, bandwidth); + } + + if (!has_sample) return -1.0; + double aggregate = 0.0; + for (const auto& entry : best_per_device) { + aggregate += entry.second; + } + return aggregate; +} + +size_t RailMonitor::pathCount() const { + std::lock_guard lock(mutex_); + return rails_.size(); +} + +uint64_t RailMonitor::normalizedNow(uint64_t now_ns) { + return now_ns == 0 ? steadyNowNs() : now_ns; +} + +uint64_t RailMonitor::deadlineAfter(uint64_t now_ns, uint64_t duration_ns) { + if (std::numeric_limits::max() - now_ns < duration_ns) { + return std::numeric_limits::max(); + } + return now_ns + duration_ns; +} + +uint64_t RailMonitor::observeTimeLocked(RailState& state, uint64_t event_ns) { + state.observed_through_ns = std::max(state.observed_through_ns, event_ns); + return state.observed_through_ns; +} + +RailMonitor::RailState& RailMonitor::getOrCreateLocked(const UbPostPath& path) { + const auto key = UbRailKey::fromPath(path); + auto [it, inserted] = rails_.try_emplace(key); + if (inserted) it->second.stats.key = key; + it->second.stats.latest_endpoint_generation = std::max( + it->second.stats.latest_endpoint_generation, path.endpoint_generation); + return it->second; +} + +void RailMonitor::insertErrorLocked(RailState& state, uint64_t event_ns) { + const auto position = std::upper_bound(state.recent_errors.begin(), + state.recent_errors.end(), event_ns); + state.recent_errors.insert(position, event_ns); +} + +void RailMonitor::pruneErrorsLocked(RailState& state, uint64_t now_ns) { + while (!state.recent_errors.empty()) { + const uint64_t error_ns = state.recent_errors.front(); + if (error_ns > now_ns || now_ns - error_ns < config_.error_window_ns) { + break; + } + state.recent_errors.pop_front(); + } +} + +void RailMonitor::refreshCooldownLocked(RailState& state, uint64_t now_ns) { + pruneErrorsLocked(state, now_ns); + if (!state.stats.paused || now_ns < state.stats.cooldown_until_ns) return; + state.ignore_errors_through_ns = + std::max(state.ignore_errors_through_ns, state.stats.cooldown_until_ns); + state.stats.paused = false; + state.stats.cooldown_until_ns = 0; + state.recent_errors.clear(); + state.stats.errors_in_window = 0; + ++state.stats.recoveries; +} + +void RailMonitor::recordFailureLocked(const UbPostPath& path, uint64_t now_ns, + bool timeout) { + auto& state = getOrCreateLocked(path); + const uint64_t event_ns = now_ns; + const uint64_t observed_ns = observeTimeLocked(state, event_ns); + refreshCooldownLocked(state, observed_ns); + ++state.stats.completion_errors; + if (timeout) ++state.stats.timeouts; + state.stats.last_error_ns = std::max(state.stats.last_error_ns, event_ns); + + if (event_ns <= state.ignore_errors_through_ns) return; + insertErrorLocked(state, event_ns); + pruneErrorsLocked(state, observed_ns); + state.stats.errors_in_window = static_cast(std::min( + state.recent_errors.size(), std::numeric_limits::max())); + + if (state.recent_errors.size() < config_.error_threshold) return; + const uint64_t cooldown_until = + deadlineAfter(state.recent_errors.back(), config_.cooldown_ns); + if (cooldown_until <= observed_ns) { + // This entire failure burst and its cooldown are already in the past + // relative to the watermark. Treat it as a completed health epoch + // instead of resurrecting an expired pause because an event was late. + state.ignore_errors_through_ns = + std::max(state.ignore_errors_through_ns, cooldown_until); + state.recent_errors.clear(); + state.stats.errors_in_window = 0; + return; + } + if (!state.stats.paused) { + state.stats.paused = true; + state.stats.pause_started_ns = state.recent_errors.back(); + ++state.stats.pauses; + } + state.stats.cooldown_until_ns = + std::max(state.stats.cooldown_until_ns, cooldown_until); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp new file mode 100644 index 0000000000..be573db317 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -0,0 +1,718 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/ub_transport.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "tent/common/utils/string_builder.h" +#include "tent/runtime/segment.h" +#include "tent/thirdparty/nlohmann/json.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/endpoint.h" +#include "tent/transport/ub/endpoint_store.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" +#include "tent/transport/ub/urma_adapter.h" +#include "tent/transport/ub/workers.h" + +namespace mooncake::tent { +namespace { + +bool filterAllows(const std::vector& filter, + const ub::DeviceInfo& device) { + if (filter.empty()) return true; + return std::find(filter.begin(), filter.end(), device.topology_name) != + filter.end() || + std::find(filter.begin(), filter.end(), device.native_device_name) != + filter.end(); +} + +uint64_t generationSeed() { + const auto wall = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + const auto steady = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + const uint64_t seed = + (wall ^ (steady << 19) ^ (steady >> 5)) & 0x7fffffffffffffffULL; + return seed == 0 ? 1 : seed; +} + +std::string encodeDeviceMetadata(Topology::NicID topology_id, + const ub::DeviceInfo& device) { + const auto& caps = device.capabilities; + return nlohmann::json{{"schema_version", 1}, + {"topology_id", topology_id}, + {"native_device_name", device.native_device_name}, + {"native_device_path", device.native_device_path}, + {"eid_index", device.eid_index}, + {"eid", device.eid}, + {"active", device.active}, + {"capabilities", + {{"max_jfc", caps.max_jfc}, + {"max_jfc_depth", caps.max_jfc_depth}, + {"max_jfr_depth", caps.max_jfr_depth}, + {"max_jetty", caps.max_jetty}, + {"max_jetty_depth", caps.max_jetty_depth}, + {"max_send_sge", caps.max_send_sge}, + {"max_remote_sge", caps.max_remote_sge}, + {"max_message_size", caps.max_message_size}, + {"max_read_size", caps.max_read_size}, + {"max_write_size", caps.max_write_size}, + {"feature_flags", caps.feature_flags}, + {"transport_modes", caps.transport_modes}}}} + .dump(); +} + +} // namespace + +struct UbTransport::Impl { + explicit Impl(std::shared_ptr injected) + : adapter(injected ? std::move(injected) + : ub::createDefaultUrmaAdapter()) {} + + Status install(const std::string& segment_name, + std::shared_ptr control, + std::shared_ptr topology, + std::shared_ptr config) { + std::lock_guard lock(lifecycle_mutex); + if (installed.load(std::memory_order_acquire)) { + return Status::InvalidArgument( + "UB transport has already been installed" LOC_MARK); + } + if (shutting_down.load(std::memory_order_acquire) || workers || + adapter_initialized) { + return Status::InvalidArgument( + "A previous UB uninstall has not drained yet" LOC_MARK); + } + if (segment_name.empty() || !control || !topology) { + return Status::InvalidArgument( + "UB install requires segment, control service and " + "topology" LOC_MARK); + } + if (!adapter || !adapter->available()) { + return Status::DeviceNotFound( + "A real or injected URMA adapter is unavailable" LOC_MARK); + } + if (!config) config = std::make_shared(); + ub::UbParams parsed; + CHECK_STATUS(ub::UbParams::FromConfig(*config, parsed)); + if (!parsed.enable) { + return Status::InvalidArgument( + "UB transport is disabled by configuration" LOC_MARK); + } + if (parsed.enable_notifications) { + return Status::NotImplemented( + "UB notifications are not supported by protocol version " + "1" LOC_MARK); + } + + shutting_down.store(false, std::memory_order_release); + local_segment_name = segment_name; + metadata = std::move(control); + local_topology = std::move(topology); + conf = std::move(config); + params = parsed; + + auto status = adapter->initialize(); + if (!status.ok()) return failInstall(status); + adapter_initialized = true; + + std::vector discovered; + status = adapter->discoverDevices(discovered); + if (!status.ok()) return failInstall(status); + std::unordered_map by_topology_name; + std::unordered_map jfc_depth_by_device; + for (auto& device : discovered) { + by_topology_name.emplace(device.topology_name, std::move(device)); + } + + for (size_t id = 0; id < local_topology->getNicCount(); ++id) { + const auto* nic = local_topology->getNicEntry(static_cast(id)); + if (!nic || nic->type != Topology::NIC_UB) continue; + auto found = by_topology_name.find(nic->name); + if (found == by_topology_name.end() || !found->second.active) { + continue; + } + if (!filterAllows(params.device_filter, found->second)) continue; + + ub::JfcOptions jfc_options; + const auto& device_caps = found->second.capabilities; + if (device_caps.max_jfc_depth != 0) { + jfc_options.depth = + std::min(jfc_options.depth, device_caps.max_jfc_depth); + jfc_options.receiver_depth = std::min( + jfc_options.receiver_depth, device_caps.max_jfc_depth); + } + if (device_caps.max_jfr_depth != 0) { + jfc_options.receiver_depth = std::min( + jfc_options.receiver_depth, device_caps.max_jfr_depth); + } + auto context = std::make_shared( + static_cast(id), found->second, adapter); + status = context->initialize(params.jfc_per_context, jfc_options); + if (!status.ok()) { + LOG(WARNING) << "Disable UB device " << nic->name << ": " + << status.ToString(); + continue; + } + context_by_topology_id.emplace(static_cast(id), context); + context_by_topology_name.emplace(nic->name, context); + jfc_depth_by_device.emplace(static_cast(id), + jfc_options.depth); + contexts.push_back(std::move(context)); + } + if (contexts.empty()) { + return failInstall(Status::DeviceNotFound( + "No UB context initialized successfully" LOC_MARK)); + } + + size_t safe_slice_size = params.slice_size; + ub::JettyOptions jetty_options; + for (const auto& context : contexts) { + const auto& device_caps = context->deviceInfo().capabilities; + auto clamp_slice = [&safe_slice_size](uint64_t limit) { + if (limit != 0) { + safe_slice_size = static_cast( + std::min(safe_slice_size, limit)); + } + }; + clamp_slice(device_caps.max_message_size); + clamp_slice(device_caps.max_read_size); + clamp_slice(device_caps.max_write_size); + if (device_caps.max_jetty_depth != 0) { + jetty_options.depth = + std::min(jetty_options.depth, device_caps.max_jetty_depth); + } + if (device_caps.max_send_sge != 0) { + jetty_options.max_sge = static_cast(std::min( + jetty_options.max_sge, device_caps.max_send_sge)); + } + if (device_caps.max_remote_sge != 0) { + jetty_options.max_sge = static_cast(std::min( + jetty_options.max_sge, device_caps.max_remote_sge)); + } + } + if (safe_slice_size == 0 || jetty_options.depth == 0 || + jetty_options.max_sge == 0) { + return failInstall(Status::InvalidArgument( + "UB device capabilities cannot support the configured data " + "path" LOC_MARK)); + } + if (safe_slice_size != params.slice_size) { + LOG(WARNING) << "Clamp UB slice_size from " << params.slice_size + << " to device limit " << safe_slice_size; + params.slice_size = safe_slice_size; + } + + buffers = std::make_unique(adapter, contexts); + endpoints = std::make_unique( + adapter, params.max_endpoints, params.jetty_per_endpoint, + jetty_options); + ub::RailMonitorConfig rail_config; + rail_config.cooldown_ns = + static_cast(params.endpoint_cooldown_ms) * 1'000'000ULL; + rails = std::make_unique(rail_config); + auto saturatedProduct = [](uint64_t lhs, uint64_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return std::numeric_limits::max(); + } + return lhs * rhs; + }; + const uint64_t path_wrs = + saturatedProduct(jetty_options.depth, params.jetty_per_endpoint); + ub::QuotaLimits path_limits{ + saturatedProduct(path_wrs, params.slice_size), path_wrs}; + quota = std::make_unique(path_limits, path_limits); + for (const auto& context : contexts) { + const auto depth = jfc_depth_by_device.at(context->topologyId()); + const uint64_t device_wrs = + saturatedProduct(depth, context->jfcs().size()); + (void)quota->setDeviceLimits( + context->topologyId(), + ub::QuotaLimits{saturatedProduct(device_wrs, params.slice_size), + device_wrs}); + } + + status = publishLocalDevices(); + if (!status.ok()) return failInstall(status); + + metadata->setBootstrapUbCallback( + [this](const UbBootstrapDesc& request, UbBootstrapDesc& response) { + return onBootstrap(request, response); + }); + callback_installed = true; + + workers = std::make_unique( + adapter, contexts, local_topology, &metadata->segmentManager(), + buffers.get(), rails.get(), quota.get(), params, + [this](const ub::EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + return resolveEndpoint(request, endpoint); + }, + [this](const std::shared_ptr& endpoint) { + if (endpoints) (void)endpoints->retire(endpoint); + }); + status = workers->start(); + if (!status.ok()) return failInstall(status); + + installed.store(true, std::memory_order_release); + return Status::OK(); + } + + Status uninstall() { + std::lock_guard lock(lifecycle_mutex); + return shutdownUnlocked(); + } + + Status failInstall(Status failure) { + (void)shutdownUnlocked(); + return failure; + } + + Status shutdownUnlocked() { + installed.store(false, std::memory_order_release); + shutting_down.store(true, std::memory_order_release); + Status first_error = Status::OK(); + auto remember = [&first_error](const Status& status) { + if (first_error.ok() && !status.ok()) first_error = status; + }; + + // Callback replacement waits for a currently executing UB bootstrap + // handler, fencing all control-plane access before resources retire. + if (callback_installed && metadata) { + metadata->setBootstrapUbCallback({}); + callback_installed = false; + } + if (workers) { + auto status = workers->stop(); + if (!status.ok()) { + // A failed native drain fence is not permission to destroy + // memory registrations, JFCs, Contexts, or the adapter. Keep + // the complete ownership graph alive so uninstall can be + // retried after a late completion/provider recovery. + return status; + } + workers.reset(); + } + if (endpoints) { + auto status = endpoints->clear(); + if (!status.ok()) return status; + endpoints.reset(); + } + if (buffers) { + auto status = buffers->clear(); + if (!status.ok()) return status; + buffers.reset(); + } + for (auto it = contexts.rbegin(); it != contexts.rend(); ++it) { + if (*it) remember((*it)->shutdown()); + } + contexts.clear(); + context_by_topology_id.clear(); + context_by_topology_name.clear(); + rails.reset(); + quota.reset(); + if (adapter_initialized && adapter) { + remember(adapter->shutdown()); + adapter_initialized = false; + } + metadata.reset(); + local_topology.reset(); + conf.reset(); + local_segment_name.clear(); + shutting_down.store(false, std::memory_order_release); + return first_error; + } + + Status publishLocalDevices() { + auto& manager = metadata->segmentManager(); + CHECK_STATUS( + manager.updateLocal([this](SegmentDesc& segment) -> Status { + if (segment.type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "Local segment is not memory-backed" LOC_MARK); + } + auto& detail = std::get(segment.detail); + detail.transport_attrs[TransportType::UB] = nlohmann::json{ + {"schema_version", 1}, + {"protocol", "urma"}, + {"notifications", false}}.dump(); + std::unordered_set existing; + for (const auto& device : detail.devices) { + existing.insert(device.name); + } + for (const auto& context : contexts) { + if (!context || !context->active()) continue; + if (existing.insert(context->deviceInfo().topology_name) + .second) { + DeviceDesc device; + device.name = context->deviceInfo().topology_name; + device.lid = 0; + device.gid.clear(); + device.transport_attrs[TransportType::UB] = + encodeDeviceMetadata(context->topologyId(), + context->deviceInfo()); + detail.devices.push_back(std::move(device)); + } + } + return Status::OK(); + })); + return manager.synchronizeLocal(); + } + + Status resolveEndpoint(const ub::EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + endpoint.reset(); + if (shutting_down.load(std::memory_order_acquire) || !endpoints || + !request.local_context || !request.remote_segment) { + return Status::InvalidArgument( + "UB transport is shutting down or route is invalid" LOC_MARK); + } + const auto* remote_nic = + request.remote_segment->getMemory().topology.getNicEntry( + request.remote_topology_id); + if (!remote_nic || remote_nic->type != Topology::NIC_UB) { + return Status::DeviceNotFound( + "Remote segment does not advertise a UB topology " + "device" LOC_MARK); + } + const std::string peer_path = + MakeNicPath(request.remote_segment->name, remote_nic->name); + ub::UbEndpointKey key{request.local_context->topologyId(), + request.remote_segment_id, + request.remote_topology_id, peer_path}; + CHECK_STATUS( + endpoints->getOrCreate(key, request.local_context, endpoint)); + if (endpoint->ready()) return Status::OK(); + + const std::string local_path = + MakeNicPath(local_segment_name, + request.local_context->deviceInfo().topology_name); + UbBootstrapDesc bootstrap; + auto status = endpoint->makeBootstrapDesc( + local_segment_name, local_path, peer_path, + request.segment_generation, bootstrap); + if (!status.ok()) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return status; + } + UbBootstrapDesc response; + status = ControlClient::bootstrapUb( + request.remote_segment->rpc_server_addr, bootstrap, response); + if (status.ok()) status = endpoint->bind(response); + if (!status.ok()) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return status; + } + return Status::OK(); + } + + int onBootstrap(const UbBootstrapDesc& request, UbBootstrapDesc& response) { + if (shutting_down.load(std::memory_order_acquire) || !endpoints) { + response.reply_msg = "UB transport is shutting down"; + return -1; + } + const std::string local_name = + getNicNameFromNicPath(request.peer_nic_path); + auto context_it = context_by_topology_name.find(local_name); + if (local_name.empty() || + context_it == context_by_topology_name.end()) { + response.reply_msg = "UB bootstrap selected an unknown local NIC"; + return -1; + } + if (request.local_device_id < 0 || request.local_eid.empty() || + request.jetty_ids.empty() || request.endpoint_generation == 0) { + response.reply_msg = "UB bootstrap request is incomplete"; + return -1; + } + + ub::UbEndpointKey key{context_it->second->topologyId(), + LOCAL_SEGMENT_ID, request.local_device_id, + request.local_nic_path}; + std::shared_ptr endpoint; + auto status = endpoints->getOrCreate(key, context_it->second, endpoint); + if (status.ok()) status = endpoint->bind(request); + if (status.ok()) { + status = endpoint->makeBootstrapDesc( + local_segment_name, request.peer_nic_path, + request.local_nic_path, currentSegmentGeneration(), response); + } + if (!status.ok()) { + if (endpoint) (void)endpoints->retire(endpoint); + response = UbBootstrapDesc{}; + response.reply_msg = status.ToString(); + return -1; + } + return 0; + } + + uint64_t currentSegmentGeneration() const { + // Buffer generations are carried in buffer metadata. The bootstrap + // field is a peer restart/refresh hint and must still be nonzero even + // before the first user buffer is registered. + return segment_generation.load(std::memory_order_relaxed); + } + + mutable std::mutex lifecycle_mutex; + std::shared_ptr adapter; + bool adapter_initialized{false}; + bool callback_installed{false}; + std::atomic installed{false}; + std::atomic shutting_down{false}; + std::atomic segment_generation{generationSeed()}; + ub::UbParams params; + std::string local_segment_name; + std::shared_ptr metadata; + std::shared_ptr local_topology; + std::shared_ptr conf; + std::vector contexts; + std::unordered_map + context_by_topology_id; + std::unordered_map context_by_topology_name; + std::unique_ptr buffers; + std::unique_ptr endpoints; + std::unique_ptr rails; + std::unique_ptr quota; + std::unique_ptr workers; +}; + +UbTransport::UbTransport(std::shared_ptr adapter) + : impl_(std::make_unique(std::move(adapter))) {} + +UbTransport::~UbTransport() { + auto status = uninstall(); + if (!status.ok()) { + // There is no caller left to retry an explicit uninstall. Leaking the + // still-live ownership graph is the only safe failure mode: destroying + // joinable pollers or registered memory after a failed device fence + // would terminate the process or permit DMA-after-free. The OS/provider + // reclaims these resources at process exit. + LOG(ERROR) << "Preserve undrained UB resources during destruction: " + << status.ToString(); + (void)impl_.release(); + } +} + +Status UbTransport::install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf) { + auto status = impl_->install(local_segment_name, std::move(metadata), + std::move(local_topology), std::move(conf)); + if (status.ok()) { + caps = Capabilities{}; + caps.dram_to_dram = true; + } + return status; +} + +Status UbTransport::uninstall() { + auto status = impl_->uninstall(); + caps = Capabilities{}; + return status; +} + +Status UbTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) { + batch = nullptr; + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->installed.load(std::memory_order_acquire)) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto* ub_batch = new (std::nothrow) UbSubBatch(); + if (!ub_batch) { + return Status::InternalError( + "Unable to allocate UB sub-batch" LOC_MARK); + } + ub_batch->max_size = max_size; + ub_batch->task_list.reserve(max_size); + batch = ub_batch; + return Status::OK(); +} + +Status UbTransport::freeSubBatch(SubBatchRef& batch) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument("Invalid UB sub-batch" LOC_MARK); + } + ub_batch->task_list.clear(); + delete ub_batch; + batch = nullptr; + return Status::OK(); +} + +Status UbTransport::submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument("Invalid UB sub-batch" LOC_MARK); + } + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->installed.load(std::memory_order_acquire) || !impl_->workers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + if (request_list.size() > ub_batch->max_size - ub_batch->task_list.size()) { + return Status::TooManyRequests("Exceed UB batch capacity" LOC_MARK); + } + for (const auto& request : request_list) { + const auto local_address = reinterpret_cast(request.source); + if (!request.source || request.length == 0 || + request.length > + std::numeric_limits::max() - request.target_offset || + request.length > + std::numeric_limits::max() - local_address) { + return Status::InvalidArgument( + "UB request range is empty or overflows" LOC_MARK); + } + } + + const auto notify_progress = ub_batch->notify_progress; + const auto progress_batch_id = ub_batch->progress_batch_id; + std::vector new_tasks; + new_tasks.reserve(request_list.size()); + for (const auto& request : request_list) { + auto task = ub::UbTask::create( + request, + [notify_progress, progress_batch_id](const TransferStatus&) { + if (notify_progress) notify_progress(progress_batch_id); + }); + size_t offset = 0; + while (offset < request.length) { + const size_t length = + std::min(impl_->params.slice_size, request.length - offset); + ub::UbSliceSpec spec; + spec.local_address = static_cast(request.source) + offset; + spec.remote_address = request.target_offset + offset; + spec.length = length; + spec.request_offset = offset; + spec.max_retries = impl_->params.max_retries; + if (!task->addSlice(spec)) { + return Status::InternalError( + "Unable to construct UB slice" LOC_MARK); + } + offset += length; + } + (void)task->seal(); + new_tasks.push_back(std::move(task)); + } + + for (auto& task : new_tasks) { + ub_batch->task_list.push_back(task); + auto status = impl_->workers->submit(task, ub_batch->device_mask); + if (!status.ok()) { + for (auto& queued : new_tasks) { + if (queued) queued->requestCancellation(); + } + return status; + } + } + return Status::OK(); +} + +Status UbTransport::getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch || task_id < 0 || + static_cast(task_id) >= ub_batch->task_list.size()) { + return Status::InvalidArgument("Invalid UB task ID" LOC_MARK); + } + status = ub_batch->task_list[task_id]->transferStatus(); + return Status::OK(); +} + +Status UbTransport::cancelTransferTask(SubBatchRef batch, int task_id) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch || task_id < 0 || + static_cast(task_id) >= ub_batch->task_list.size()) { + return Status::InvalidArgument("Invalid UB task ID" LOC_MARK); + } + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->workers) { + return Status::InvalidArgument("UB workers are not running" LOC_MARK); + } + return impl_->workers->cancel(ub_batch->task_list[task_id]); +} + +Status UbTransport::addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto status = impl_->buffers->addBuffer(desc, options); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +Status UbTransport::addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) { + return Status::InvalidArgument( + "UB transport is not installed" LOC_MARK); + } + auto status = impl_->buffers->addBuffers(desc_list, options); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +Status UbTransport::removeMemoryBuffer(BufferDesc& desc) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->buffers) return Status::OK(); + auto status = impl_->buffers->removeBuffer(desc); + if (status.ok()) { + impl_->segment_generation.fetch_add(1, std::memory_order_relaxed); + } + return status; +} + +bool UbTransport::warmupMemory(void* addr, size_t length) { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!addr || length == 0 || !impl_->adapter || impl_->contexts.empty()) { + return false; + } + ub::LocalSegmentPtr segment; + ub::SegmentOptions options; + auto status = impl_->adapter->registerLocalSegment( + impl_->contexts.front()->handle(), reinterpret_cast(addr), + length, options, segment); + if (!status.ok()) return false; + return impl_->adapter->unregisterLocalSegment(segment).ok(); +} + +double UbTransport::getEstimatedBandwidth() const { + std::lock_guard lock(impl_->lifecycle_mutex); + if (!impl_->params.enable_bandwidth_estimation || !impl_->rails) + return -1.0; + return impl_->rails->aggregateBandwidth(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp new file mode 100644 index 0000000000..9e64193d2e --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -0,0 +1,918 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/workers.h" + +#include +#include +#include +#include +#include + +#include + +#include "tent/runtime/platform.h" +#include "tent/transport/ub/endpoint.h" + +namespace mooncake::tent::ub { +namespace { + +uint64_t deadlineAfter(uint64_t now_ns, uint64_t timeout_ns) { + return timeout_ns > std::numeric_limits::max() - now_ns + ? std::numeric_limits::max() + : now_ns + timeout_ns; +} + +bool retryableStatus(const Status& status) { + return status.IsDeviceNotFound() || status.IsNeedsRefreshCache() || + status.IsRpcServiceError() || status.IsInternalError() || + status.IsRdmaError(); +} + +} // namespace + +struct UbWorkers::Route { + SegmentDescRef pin; + BufferDesc* remote_buffer{nullptr}; + UbBufferMetadata metadata; + std::vector remote_devices; +}; + +struct UbWorkers::Inflight { + uint64_t completion_token{0}; + PendingSlice pending; + UbAttemptToken attempt; + UbPostPath path; + std::shared_ptr endpoint; + LocalSegmentPtr local_segment; + RemoteSegmentPtr remote_segment; + QuotaReservation quota; + uint64_t posted_ns{0}; + uint64_t deadline_ns{0}; + std::atomic timed_out{false}; + std::atomic timeout_recorded{false}; + std::atomic resources_released{false}; +}; + +UbWorkers::UbWorkers(std::shared_ptr adapter, + std::vector contexts, + std::shared_ptr local_topology, + SegmentManager* segment_manager, UbBufferManager* buffers, + RailMonitor* rail_monitor, QuotaManager* quota, + UbParams params, EndpointResolver endpoint_resolver, + EndpointRetirer endpoint_retirer) + : adapter_(std::move(adapter)), + contexts_(std::move(contexts)), + local_topology_(std::move(local_topology)), + segment_manager_(segment_manager), + buffers_(buffers), + rail_monitor_(rail_monitor), + quota_(quota), + params_(std::move(params)), + endpoint_resolver_(std::move(endpoint_resolver)), + endpoint_retirer_(std::move(endpoint_retirer)) { + for (const auto& context : contexts_) { + if (!context) continue; + context_by_topology_id_[context->topologyId()] = context; + for (const auto& jfc : context->jfcs()) { + if (jfc) { + all_jfcs_.push_back(jfc); + context_by_jfc_[jfc.get()] = context; + } + } + } +} + +UbWorkers::~UbWorkers() { (void)stop(); } + +Status UbWorkers::start() { + bool expected = false; + if (!accepting_.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + return Status::InvalidArgument("UB workers already started" LOC_MARK); + } + if (!adapter_ || contexts_.empty() || all_jfcs_.empty() || + !segment_manager_ || !buffers_ || !rail_monitor_ || !quota_ || + !endpoint_resolver_ || params_.worker_count == 0 || + params_.poller_count == 0) { + accepting_.store(false, std::memory_order_release); + return Status::InvalidArgument( + "UB workers have incomplete dependencies" LOC_MARK); + } + + posting_.store(true, std::memory_order_release); + polling_.store(true, std::memory_order_release); + timeout_scans_enabled_.store(true, std::memory_order_release); + try { + posting_threads_.reserve(params_.worker_count); + for (uint32_t i = 0; i < params_.worker_count; ++i) { + posting_threads_.emplace_back(&UbWorkers::postingLoop, this, i); + } + const size_t poller_count = + std::min(params_.poller_count, all_jfcs_.size()); + polling_threads_.reserve(poller_count); + for (size_t i = 0; i < poller_count; ++i) { + polling_threads_.emplace_back(&UbWorkers::pollingLoop, this, i); + } + } catch (const std::exception& error) { + accepting_.store(false, std::memory_order_release); + posting_.store(false, std::memory_order_release); + polling_.store(false, std::memory_order_release); + timeout_scans_enabled_.store(false, std::memory_order_release); + queue_cv_.notify_all(); + for (auto& thread : posting_threads_) { + if (thread.joinable()) thread.join(); + } + for (auto& thread : polling_threads_) { + if (thread.joinable()) thread.join(); + } + posting_threads_.clear(); + polling_threads_.clear(); + return Status::InternalError(std::string("Cannot start UB workers: ") + + error.what() + LOC_MARK); + } + return Status::OK(); +} + +Status UbWorkers::stop() { + accepting_.store(false, std::memory_order_release); + // Fence retry enqueue before draining the queues. A completion racing + // with shutdown will turn its retry-pending slice into CANCELED instead of + // leaving a new queue entry behind the drain pass. + posting_.store(false, std::memory_order_release); + timeout_scans_enabled_.store(false, std::memory_order_release); + queue_cv_.notify_all(); + + std::vector abandoned; + { + std::lock_guard lock(queue_mutex_); + for (auto& queue : queues_) { + while (!queue.empty()) { + abandoned.push_back(std::move(queue.front())); + queue.pop_front(); + } + } + } + for (auto& pending : abandoned) { + if (pending.slice) pending.slice->requestCancellation(); + } + + for (auto& thread : posting_threads_) { + if (thread.joinable()) thread.join(); + } + posting_threads_.clear(); + + // Wait for a timeout scan that began before the flag change. scanTimeouts + // rechecks under this mutex, so no new endpoint can enter the drain set + // after this barrier and escape the shutdown snapshot. + { + std::lock_guard lock(endpoint_drain_mutex_); + } + + // Snapshot all native work after posting threads have joined; no new WR + // can cross the adapter boundary beyond this point. Quiesce each endpoint + // before stopping pollers. A failed fence leaves pollers, tokens, native + // segments, quota, and endpoints intact so uninstall can be retried + // safely; fixed sleeps are never treated as proof that DMA stopped. + std::vector> still_inflight; + std::vector> endpoints_to_fence; + { + std::lock_guard lock(inflight_mutex_); + still_inflight.reserve(inflight_.size()); + for (const auto& [_, inflight] : inflight_) { + still_inflight.push_back(inflight); + if (inflight && inflight->endpoint) { + endpoints_to_fence.push_back(inflight->endpoint); + } + } + for (const auto& [_, endpoint] : draining_endpoints_) { + if (endpoint) endpoints_to_fence.push_back(endpoint); + } + } + for (const auto& inflight : still_inflight) { + if (inflight && inflight->pending.slice) { + inflight->pending.slice->requestCancellation(); + } + } + + Status fence_error = Status::OK(); + std::unordered_set fenced; + for (const auto& endpoint : endpoints_to_fence) { + if (!endpoint || !fenced.insert(endpoint.get()).second) { + continue; + } + std::vector drained; + auto status = endpoint->quiesce(params_.slice_timeout_ms, drained); + for (const auto& completion : drained) { + if (completion.token != 0) handleCompletion(completion); + } + if (!status.ok()) { + rememberEndpointDrain(endpoint); + if (fence_error.ok()) fence_error = status; + continue; + } + forgetEndpointDrain(endpoint); + if (endpoint_retirer_) endpoint_retirer_(endpoint); + } + if (!fence_error.ok()) { + return fence_error; + } + { + std::lock_guard lock(inflight_mutex_); + if (!draining_endpoints_.empty()) { + return Status::InternalError( + "UB endpoint drain set changed during shutdown" LOC_MARK); + } + } + + // Successful endpoint fences prove that any token still missing from the + // JFC can no longer touch memory. Let pollers consume already-queued CRs, + // then stop them and safely resolve any provider-lost token. + if (!still_inflight.empty()) { + std::unique_lock lock(inflight_mutex_); + (void)inflight_cv_.wait_for(lock, std::chrono::milliseconds(100), + [this] { return inflight_.empty(); }); + } + polling_.store(false, std::memory_order_release); + for (auto& thread : polling_threads_) { + if (thread.joinable()) thread.join(); + } + polling_threads_.clear(); + + still_inflight.clear(); + { + std::lock_guard lock(inflight_mutex_); + for (auto& [_, inflight] : inflight_) { + still_inflight.push_back(std::move(inflight)); + } + inflight_.clear(); + } + for (const auto& inflight : still_inflight) { + releaseInflight(inflight); + (void)inflight->pending.slice->resolveAttempt(inflight->attempt, FAILED, + 0, false); + } + inflight_cv_.notify_all(); + return Status::OK(); +} + +Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { + if (!task) { + return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK); + } + if (!accepting_.load(std::memory_order_acquire)) { + return Status::InvalidArgument( + "UB workers are not accepting work" LOC_MARK); + } + const auto snapshot = task->snapshot(); + if (!snapshot.sealed) { + return Status::InvalidArgument( + "UB task must be sealed before submission" LOC_MARK); + } + const int priority = + std::clamp(task->request().priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW)); + + std::vector pending; + for (const auto& slice : task->slices()) { + if (slice && slice->markQueued()) { + pending.push_back(PendingSlice{task, slice, device_mask, priority, + task->request().target_id, + task->request().opcode}); + } + } + bool stopped_during_submit = false; + { + std::lock_guard lock(queue_mutex_); + if (!accepting_.load(std::memory_order_relaxed)) { + stopped_during_submit = true; + } else { + for (auto& item : pending) { + queues_[item.priority].push_back(std::move(item)); + } + } + } + if (stopped_during_submit) { + for (auto& item : pending) item.slice->requestCancellation(); + return Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + if (!pending.empty()) queue_cv_.notify_all(); + return Status::OK(); +} + +Status UbWorkers::cancel(const UbTask::Ptr& task) { + if (!task) { + return Status::InvalidArgument("Cannot cancel a null UB task" LOC_MARK); + } + (void)task->requestCancellation(); + queue_cv_.notify_all(); + return Status::OK(); +} + +size_t UbWorkers::queuedCount() const { + std::lock_guard lock(queue_mutex_); + size_t count = 0; + for (const auto& queue : queues_) count += queue.size(); + return count; +} + +size_t UbWorkers::inflightCount() const { + std::lock_guard lock(inflight_mutex_); + return inflight_.size(); +} + +bool UbWorkers::popPending(PendingSlice& pending) { + std::unique_lock lock(queue_mutex_); + queue_cv_.wait(lock, [this] { + if (!posting_.load(std::memory_order_acquire)) return true; + for (const auto& queue : queues_) { + if (!queue.empty()) return true; + } + return false; + }); + if (!posting_.load(std::memory_order_acquire)) return false; + for (auto& queue : queues_) { + if (!queue.empty()) { + pending = std::move(queue.front()); + queue.pop_front(); + return true; + } + } + return false; +} + +void UbWorkers::postingLoop(size_t worker_index) { + while (posting_.load(std::memory_order_acquire)) { + PendingSlice pending; + if (!popPending(pending)) continue; + if (pending.slice) processPending(pending, worker_index); + } +} + +void UbWorkers::pollingLoop(size_t poller_index) { + const size_t poller_count = std::max( + 1, std::min(params_.poller_count, all_jfcs_.size())); + uint64_t last_timeout_scan = 0; + while (polling_.load(std::memory_order_acquire)) { + bool progressed = false; + for (size_t index = poller_index; index < all_jfcs_.size(); + index += poller_count) { + std::vector completions; + auto status = all_jfcs_[index]->poll(64, completions); + if (!status.ok()) { + auto context = context_by_jfc_.find(all_jfcs_[index].get()); + if (context != context_by_jfc_.end() && context->second) { + (void)context->second->markUnavailable(); + } + LOG_EVERY_N(WARNING, 1000) + << "UB JFC poll failed: " << status.ToString(); + continue; + } + progressed = progressed || !completions.empty(); + for (const auto& completion : completions) { + if (completion.token != 0) handleCompletion(completion); + } + } + const uint64_t now = steadyNowNs(); + if (timeout_scans_enabled_.load(std::memory_order_acquire) && + now - last_timeout_scan >= 1'000'000ULL) { + scanTimeouts(); + last_timeout_scan = now; + } + if (!progressed) + std::this_thread::sleep_for(std::chrono::microseconds(20)); + } +} + +void UbWorkers::enqueueRetry(const PendingSlice& pending) { + if (!pending.slice) return; + if (!posting_.load(std::memory_order_acquire) || + !pending.slice->markQueued()) { + pending.slice->requestCancellation(); + return; + } + bool stopped = false; + { + std::lock_guard lock(queue_mutex_); + if (!posting_.load(std::memory_order_relaxed)) { + stopped = true; + } else { + queues_[std::clamp(pending.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))] + .push_back(pending); + } + } + if (stopped) { + pending.slice->requestCancellation(); + return; + } + queue_cv_.notify_one(); +} + +void UbWorkers::deferPending(const PendingSlice& pending) { + if (!pending.slice || pending.slice->cancellationRequested() || + !posting_.load(std::memory_order_acquire)) { + if (pending.slice) pending.slice->requestCancellation(); + return; + } + // Capacity pressure is not a transfer failure. A short bounded backoff + // prevents a posting lane from exhausting CPU while completions release + // quota; releaseInflight() also wakes the queue condition variable. + { + std::unique_lock lock(queue_mutex_); + (void)queue_cv_.wait_for(lock, std::chrono::microseconds(50), [&] { + return !posting_.load(std::memory_order_relaxed) || + pending.slice->cancellationRequested(); + }); + if (posting_.load(std::memory_order_relaxed) && + !pending.slice->cancellationRequested()) { + queues_[std::clamp(pending.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))] + .push_back(pending); + lock.unlock(); + queue_cv_.notify_one(); + return; + } + } + pending.slice->requestCancellation(); +} + +Status UbWorkers::buildRoute(const PendingSlice& pending, Route& route) { + return segment_manager_->withCachedSegment( + pending.target_id, route.pin, [&](SegmentDesc* segment) -> Status { + if (!segment || segment->type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "UB target is not a memory segment" LOC_MARK); + } + auto* buffer = + segment->findBuffer(pending.slice->spec().remote_address, + pending.slice->spec().length); + if (!buffer) { + return Status::NeedsRefreshCache( + "UB target range is not registered" LOC_MARK); + } + auto attr = buffer->transport_attrs.find(TransportType::UB); + if (attr == buffer->transport_attrs.end()) { + return Status::NeedsRefreshCache( + "UB target buffer has no transport metadata" LOC_MARK); + } + UbBufferMetadata metadata; + auto status = decodeBufferMetadata(attr->second, metadata); + if (!status.ok()) return status; + route.remote_buffer = buffer; + route.metadata = std::move(metadata); + route.remote_devices = orderedRemoteDevices(*segment, *buffer); + if (route.remote_devices.empty()) { + return Status::DeviceNotFound( + "UB target has no active advertised device" LOC_MARK); + } + return Status::OK(); + }); +} + +std::vector UbWorkers::orderedLocalDevices( + const PendingSlice& pending) const { + std::vector ordered; + std::unordered_set seen; + std::unordered_map device_rank; + std::string location = kWildcardLocation; + auto locations = Platform::getLoader().getLocation( + pending.slice->spec().local_address, 1, true); + if (!locations.empty()) location = locations.front().location; + + auto append = [&](Topology::NicID id) { + auto found = context_by_topology_id_.find(id); + if (found == context_by_topology_id_.end() || !found->second || + !found->second->active()) { + return; + } + const bool allowed = pending.device_mask == ~0ULL || + (id >= 0 && id < 64 && + (pending.device_mask & (uint64_t{1} << id)) != 0); + if (allowed && seen.insert(id).second) ordered.push_back(id); + }; + if (local_topology_) { + if (const auto* memory = local_topology_->getMemEntry(location)) { + for (size_t rank = 0; rank < Topology::DevicePriorityRanks; + ++rank) { + for (auto id : memory->device_list[rank]) { + auto [it, inserted] = device_rank.emplace(id, rank); + if (!inserted) it->second = std::min(it->second, rank); + append(id); + } + } + } + } + for (const auto& context : contexts_) { + if (context) { + device_rank.try_emplace(context->topologyId(), + Topology::DevicePriorityRanks); + append(context->topologyId()); + } + } + std::stable_sort(ordered.begin(), ordered.end(), [&](auto lhs, auto rhs) { + const auto lhs_rank = device_rank.at(lhs); + const auto rhs_rank = device_rank.at(rhs); + if (lhs_rank != rhs_rank) return lhs_rank < rhs_rank; + return context_by_topology_id_.at(lhs)->inflightBytes() < + context_by_topology_id_.at(rhs)->inflightBytes(); + }); + return ordered; +} + +std::vector UbWorkers::orderedRemoteDevices( + const SegmentDesc& segment, const BufferDesc& buffer) { + std::vector ordered; + auto attr = buffer.transport_attrs.find(TransportType::UB); + if (attr == buffer.transport_attrs.end()) return ordered; + UbBufferMetadata metadata; + if (!decodeBufferMetadata(attr->second, metadata).ok()) return ordered; + std::unordered_set advertised; + for (const auto& item : metadata.segments) { + advertised.insert(item.topology_id); + } + std::unordered_set seen; + const auto& topology = segment.getMemory().topology; + auto append = [&](Topology::NicID id) { + const auto* nic = topology.getNicEntry(id); + if (nic && nic->type == Topology::NIC_UB && advertised.count(id) && + seen.insert(id).second) { + ordered.push_back(id); + } + }; + if (const auto* memory = topology.getMemEntry(buffer.location)) { + for (size_t rank = 0; rank < Topology::DevicePriorityRanks; ++rank) { + for (auto id : memory->device_list[rank]) append(id); + } + } + for (const auto& item : metadata.segments) append(item.topology_id); + return ordered; +} + +Status UbWorkers::chooseAndResolveEndpoint( + const PendingSlice& pending, Route& route, + std::shared_ptr& endpoint, UbPostPath& path) { + const auto local_devices = orderedLocalDevices(pending); + if (local_devices.empty() || route.remote_devices.empty()) { + return Status::DeviceNotFound("No usable UB posting path" LOC_MARK); + } + const auto snapshot = pending.slice->snapshot(); + const size_t combinations = + local_devices.size() * route.remote_devices.size(); + const size_t start = + combinations == 0 ? 0 : snapshot.retry_count % combinations; + Status first_error = Status::DeviceNotFound( + "No ready UB endpoint for any posting path" LOC_MARK); + for (size_t candidate = 0; candidate < combinations; ++candidate) { + const size_t flat = (start + candidate) % combinations; + const auto local_id = local_devices[flat / route.remote_devices.size()]; + const auto remote_id = + route.remote_devices[flat % route.remote_devices.size()]; + auto context = context_by_topology_id_.at(local_id); + EndpointResolveRequest request{context, pending.target_id, + route.pin.get(), remote_id, + route.metadata.generation}; + std::shared_ptr candidate_endpoint; + auto status = endpoint_resolver_(request, candidate_endpoint); + if (!status.ok() || !candidate_endpoint || + !candidate_endpoint->ready()) { + if (first_error.IsDeviceNotFound() && !status.ok()) + first_error = status; + continue; + } + UbPostPath candidate_path{local_id, pending.target_id, remote_id, + candidate_endpoint->generation()}; + if (!rail_monitor_->available(candidate_path)) continue; + endpoint = std::move(candidate_endpoint); + path = candidate_path; + return Status::OK(); + } + return first_error; +} + +void UbWorkers::processPending(const PendingSlice& pending, + size_t worker_index) { + if (pending.slice->cancellationRequested()) { + pending.slice->requestCancellation(); + return; + } + Route route; + auto status = buildRoute(pending, route); + if (!status.ok()) { + const auto resolution = pending.slice->resolveBeforePost( + FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + std::shared_ptr endpoint; + UbPostPath path; + status = chooseAndResolveEndpoint(pending, route, endpoint, path); + if (!status.ok()) { + const auto resolution = pending.slice->resolveBeforePost( + FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + LocalSegmentRef local; + status = buffers_->findLocal( + reinterpret_cast(pending.slice->spec().local_address), + pending.slice->spec().length, path.local_topology_id, local); + if (!status.ok()) { + (void)pending.slice->resolveBeforePost(FAILED, 0, false); + return; + } + ImportedSegmentRef remote; + status = buffers_->importRemote(pending.target_id, path.local_topology_id, + path.remote_device_id, *route.remote_buffer, + pending.opcode, + pending.slice->spec().remote_address, + pending.slice->spec().length, remote); + if (!status.ok()) { + const bool retryable = status.IsNeedsRefreshCache(); + if (retryable && pending.target_id != LOCAL_SEGMENT_ID) { + (void)segment_manager_->invalidateRemote(pending.target_id); + } + const auto resolution = + pending.slice->resolveBeforePost(FAILED, 0, retryable); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } + return; + } + + auto reservation = quota_->tryAcquire(path, pending.slice->spec().length); + if (!reservation) { + deferPending(pending); + return; + } + if (!endpoint->tryAcquireOutstanding(pending.slice->spec().length)) { + (void)quota_->release(*reservation); + deferPending(pending); + return; + } + + auto attempt = pending.slice->beginAttempt(path); + if (!attempt) { + endpoint->releaseOutstanding(pending.slice->spec().length); + (void)quota_->release(*reservation); + return; + } + + const uint64_t completion_token = nextCompletionToken(); + auto inflight = std::make_shared(); + inflight->completion_token = completion_token; + inflight->pending = pending; + inflight->attempt = *attempt; + inflight->path = path; + inflight->endpoint = endpoint; + inflight->local_segment = local.segment; + inflight->remote_segment = remote.segment; + inflight->quota = *reservation; + inflight->posted_ns = steadyNowNs(); + inflight->deadline_ns = deadlineAfter( + inflight->posted_ns, + static_cast(params_.slice_timeout_ms) * 1'000'000ULL); + + if (!pending.slice->tryCommitPost(*attempt, inflight->posted_ns)) { + endpoint->releaseOutstanding(pending.slice->spec().length); + (void)quota_->release(*reservation); + return; + } + { + std::lock_guard lock(inflight_mutex_); + inflight_.emplace(completion_token, inflight); + } + + WorkRequest work; + work.operation = + pending.opcode == Request::READ ? Operation::READ : Operation::WRITE; + work.local_address = + reinterpret_cast(pending.slice->spec().local_address); + work.remote_address = pending.slice->spec().remote_address; + work.length = pending.slice->spec().length; + work.token = completion_token; + work.local_segment = local.segment; + work.remote_segment = remote.segment; + auto jetty = endpoint->jetty(worker_index); + size_t posted_count = 0; + status = jetty ? adapter_->post(jetty, {work}, posted_count) + : Status::InternalError("UB endpoint has no Jetty" LOC_MARK); + if (posted_count == 1) { + rail_monitor_->registerPath(path); + if (!status.ok()) { + LOG_EVERY_N(WARNING, 1000) + << "UB post returned an error after accepting the WR: " + << status.ToString(); + } + return; + } + + if (status.ok()) { + status = Status::InternalError( + "URMA adapter accepted an unexpected WR count" LOC_MARK); + } + + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(completion_token); + if (it != inflight_.end() && it->second == inflight) + inflight_.erase(it); + } + inflight_cv_.notify_all(); + releaseInflight(inflight); + rail_monitor_->recordError(path); + if (endpoint_retirer_) endpoint_retirer_(endpoint); + const auto resolution = pending.slice->resolveAttempt( + *attempt, FAILED, 0, retryableStatus(status)); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(pending); + } +} + +void UbWorkers::handleCompletion(const Completion& completion) { + std::shared_ptr inflight; + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(completion.token); + if (it == inflight_.end()) return; + inflight = std::move(it->second); + inflight_.erase(it); + } + inflight_cv_.notify_all(); + releaseInflight(inflight); + if (inflight->timed_out.load(std::memory_order_acquire)) { + // A natural or flush completion proves that this particular WR can no + // longer touch memory. Whichever side wins the timeout/drain race may + // advance the logical attempt; attempt matching makes the other call + // an idempotent no-op. + const uint64_t now = steadyNowNs(); + recordTimeoutOnce(inflight, now); + resolveInflight(inflight, TIMEOUT, 0, true); + return; + } + + const uint64_t now = steadyNowNs(); + const uint64_t latency = + now >= inflight->posted_ns ? now - inflight->posted_ns : 0; + switch (completion.category) { + case CompletionCategory::SUCCESS: + rail_monitor_->recordSuccess(inflight->path, + inflight->pending.slice->spec().length, + latency, now); + resolveInflight(inflight, COMPLETED, + inflight->pending.slice->spec().length, false); + break; + case CompletionCategory::TIMEOUT: + rail_monitor_->recordTimeout(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, TIMEOUT, 0, true); + break; + case CompletionCategory::LOCAL_DEVICE_ERROR: + if (inflight->endpoint && inflight->endpoint->context()) { + (void)inflight->endpoint->context()->markUnavailable(); + } + rail_monitor_->recordError(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, FAILED, 0, true); + break; + case CompletionCategory::REMOTE_PATH_ERROR: + case CompletionCategory::ENDPOINT_ERROR: + rail_monitor_->recordError(inflight->path, now); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + resolveInflight(inflight, FAILED, 0, true); + break; + case CompletionCategory::MEMORY_ERROR: + case CompletionCategory::UNKNOWN_ERROR: + rail_monitor_->recordError(inflight->path, now); + resolveInflight(inflight, FAILED, 0, false); + break; + } +} + +void UbWorkers::scanTimeouts() { + std::lock_guard drain_lock(endpoint_drain_mutex_); + if (!timeout_scans_enabled_.load(std::memory_order_acquire)) return; + const uint64_t now = steadyNowNs(); + std::vector> expired; + { + std::lock_guard lock(inflight_mutex_); + for (const auto& [_, inflight] : inflight_) { + bool expected = false; + if (inflight->deadline_ns <= now && + inflight->timed_out.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + expired.push_back(inflight); + } + } + } + for (const auto& inflight : expired) { + std::vector drained; + auto status = + inflight->endpoint + ? inflight->endpoint->quiesce(params_.slice_timeout_ms, drained) + : Status::InvalidArgument( + "Timed-out UB WR has no endpoint" LOC_MARK); + // Even a partial/failed fence may have returned valid WR completions; + // they must never be lost. + for (const auto& completion : drained) { + if (completion.token != 0) handleCompletion(completion); + } + if (!status.ok()) { + rememberEndpointDrain(inflight->endpoint); + LOG_EVERY_N(ERROR, 100) + << "Cannot establish UB timeout drain fence: " + << status.ToString(); + // Leave the token and all resource references alive. A natural + // completion can still resolve it safely; otherwise a later scan + // retries the native fence. + bool still_inflight = false; + { + std::lock_guard lock(inflight_mutex_); + auto it = inflight_.find(inflight->completion_token); + still_inflight = + it != inflight_.end() && it->second == inflight; + } + if (still_inflight) { + inflight->timed_out.store(false, std::memory_order_release); + } + continue; + } + forgetEndpointDrain(inflight->endpoint); + if (endpoint_retirer_) endpoint_retirer_(inflight->endpoint); + recordTimeoutOnce(inflight, now); + // quiesce() is the proof that retry cannot overlap old DMA. The old + // token remains in inflight_ only until its returned/polled completion + // is dispatched; a provider-lost token is reclaimed at fenced stop. + resolveInflight(inflight, TIMEOUT, 0, true); + } +} + +void UbWorkers::releaseInflight(const std::shared_ptr& inflight) { + bool expected = false; + if (!inflight || !inflight->resources_released.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + return; + } + if (inflight->endpoint) { + inflight->endpoint->releaseOutstanding( + inflight->pending.slice->spec().length); + } + if (inflight->quota.valid()) (void)quota_->release(inflight->quota); + queue_cv_.notify_all(); +} + +void UbWorkers::recordTimeoutOnce(const std::shared_ptr& inflight, + uint64_t now_ns) { + bool expected = false; + if (inflight && inflight->timeout_recorded.compare_exchange_strong( + expected, true, std::memory_order_acq_rel)) { + rail_monitor_->recordTimeout(inflight->path, now_ns); + } +} + +void UbWorkers::rememberEndpointDrain( + const std::shared_ptr& endpoint) { + if (!endpoint) return; + std::lock_guard lock(inflight_mutex_); + draining_endpoints_[endpoint->generation()] = endpoint; +} + +void UbWorkers::forgetEndpointDrain( + const std::shared_ptr& endpoint) { + if (!endpoint) return; + std::lock_guard lock(inflight_mutex_); + auto it = draining_endpoints_.find(endpoint->generation()); + if (it != draining_endpoints_.end() && it->second == endpoint) { + draining_endpoints_.erase(it); + } +} + +void UbWorkers::resolveInflight(const std::shared_ptr& inflight, + TransferStatusEnum outcome, size_t bytes, + bool retryable) { + const auto resolution = inflight->pending.slice->resolveAttempt( + inflight->attempt, outcome, bytes, retryable); + if (resolution == UbAttemptResolution::kRetryScheduled) { + enqueueRetry(inflight->pending); + } +} + +void UbWorkers::failUnposted(const PendingSlice& pending, + TransferStatusEnum outcome) { + if (pending.slice) (void)pending.slice->tryResolveBeforePost(outcome); +} + +uint64_t UbWorkers::nextCompletionToken() { + uint64_t token = next_token_.fetch_add(1, std::memory_order_relaxed); + if (token == 0) token = next_token_.fetch_add(1, std::memory_order_relaxed); + return token; +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 15f3bd0fc0..973818087c 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -329,6 +329,14 @@ add_test(NAME tent_topology_priority_matrix_test COMMAND tent_topology_priority_matrix_test) if(USE_UB) + add_executable( + tent_ub_core_test ub_core_test.cpp ../src/transport/ub/quota.cpp + ../src/transport/ub/rail_monitor.cpp) + target_link_libraries(tent_ub_core_test PRIVATE tent_common gtest gtest_main) + target_include_directories(tent_ub_core_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_test(NAME tent_ub_core_test COMMAND tent_ub_core_test) + add_executable( tent_ub_teardown_test ub_teardown_test.cpp ../src/transport/ub/buffers.cpp @@ -338,6 +346,14 @@ if(USE_UB) target_include_directories(tent_ub_teardown_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_ub_teardown_test COMMAND tent_ub_teardown_test) + + add_executable(tent_ub_native_data_path_test ub_native_data_path_test.cpp) + target_link_libraries(tent_ub_native_data_path_test PRIVATE gtest gtest_main + tent_link_group) + target_include_directories(tent_ub_native_data_path_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_test(NAME tent_ub_native_data_path_test + COMMAND tent_ub_native_data_path_test) endif() # End-to-end failover test: drives real TransferEngineImpl with diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index f42a974f4f..06b792b916 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -307,6 +307,81 @@ TEST(TopologyTest, UbDeviceAttributesRoundTripThroughJson) { EXPECT_EQ(round_tripped->device_attrs, ub.device_attrs); } +TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { + UbBootstrapDesc source; + source.protocol_version = 1; + source.segment_name = "peer-segment"; + source.local_nic_path = "local/ub-device-0/eid-2"; + source.peer_nic_path = "peer/ub-device-1/eid-3"; + source.local_device_name = "ub-device-0"; + source.local_device_id = 7; + source.local_eid_index = 2; + source.local_eid = "e1:02:03:04:05:06:07:08"; + source.jetty_ids = {11, 12}; + source.jetty_uasids = {21, 22}; + source.endpoint_generation = 41; + source.segment_generation = 42; + source.capabilities = {"read", "write"}; + source.reply_msg = "ok"; + + const auto parsed = json(source).get(); + EXPECT_EQ(parsed.protocol_version, source.protocol_version); + EXPECT_EQ(parsed.segment_name, source.segment_name); + EXPECT_EQ(parsed.local_nic_path, source.local_nic_path); + EXPECT_EQ(parsed.peer_nic_path, source.peer_nic_path); + EXPECT_EQ(parsed.local_device_name, source.local_device_name); + EXPECT_EQ(parsed.local_device_id, source.local_device_id); + EXPECT_EQ(parsed.local_eid_index, source.local_eid_index); + EXPECT_EQ(parsed.local_eid, source.local_eid); + EXPECT_EQ(parsed.jetty_ids, source.jetty_ids); + EXPECT_EQ(parsed.jetty_uasids, source.jetty_uasids); + EXPECT_EQ(parsed.endpoint_generation, source.endpoint_generation); + EXPECT_EQ(parsed.segment_generation, source.segment_generation); + EXPECT_EQ(parsed.capabilities, source.capabilities); + EXPECT_EQ(parsed.reply_msg, source.reply_msg); +} + +TEST(ControlPlaneTest, UbBootstrapJsonDefaultsOptionalFields) { + const auto parsed = json{{"protocol_version", 1}}.get(); + EXPECT_EQ(parsed.protocol_version, 1u); + EXPECT_TRUE(parsed.segment_name.empty()); + EXPECT_TRUE(parsed.local_nic_path.empty()); + EXPECT_TRUE(parsed.peer_nic_path.empty()); + EXPECT_TRUE(parsed.local_device_name.empty()); + EXPECT_EQ(parsed.local_device_id, -1); + EXPECT_EQ(parsed.local_eid_index, -1); + EXPECT_TRUE(parsed.local_eid.empty()); + EXPECT_TRUE(parsed.jetty_ids.empty()); + EXPECT_TRUE(parsed.jetty_uasids.empty()); + EXPECT_EQ(parsed.endpoint_generation, 0u); + EXPECT_EQ(parsed.segment_generation, 0u); + EXPECT_TRUE(parsed.capabilities.empty()); + EXPECT_TRUE(parsed.reply_msg.empty()); +} + +TEST(ControlPlaneTest, UbBootstrapJsonRejectsMissingOrUnknownVersion) { + EXPECT_THROW((void)json::object().get(), + std::invalid_argument); + const json unknown_version{{"protocol_version", 2}}; + EXPECT_THROW((void)unknown_version.get(), + std::invalid_argument); +} + +TEST(ControlPlaneTest, UbBootstrapRpcIdIsAppendedWithoutRenumbering) { + EXPECT_EQ(static_cast(GetSegmentDesc), 1); + EXPECT_EQ(static_cast(BootstrapRdma), 2); + EXPECT_EQ(static_cast(SendData), 3); + EXPECT_EQ(static_cast(RecvData), 4); + EXPECT_EQ(static_cast(Notify), 5); + EXPECT_EQ(static_cast(Probe), 6); + EXPECT_EQ(static_cast(Delegate), 7); + EXPECT_EQ(static_cast(Pin), 8); + EXPECT_EQ(static_cast(Unpin), 9); + EXPECT_EQ(static_cast(SubscribeSegmentUpdate), 10); + EXPECT_EQ(static_cast(NotifySegmentUpdated), 11); + EXPECT_EQ(static_cast(BootstrapUb), 12); +} + // --------------------------------------------------------------------------- // Test legacy mode // --------------------------------------------------------------------------- @@ -653,6 +728,64 @@ TEST(TransportSelectorTest, ConfigBasedPolicySelection) { EXPECT_EQ(result.transport, RDMA); } +TEST(TransportSelectorTest, PolicyCanPreferUb) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "ub-preferred"; + policy["segment_type"] = "memory"; + policy["transports"] = {"ub", "rdma"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[UB] = std::make_shared(UB); + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[UB].get())->setDramToDram(true); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA, UB}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "ub-preferred"; + + EXPECT_EQ(selector.select(ctx, transports).transport, UB); +} + +TEST(TransportSelectorTest, PolicyFallsBackWhenUbIsIncapable) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "ub-with-rdma-fallback"; + policy["segment_type"] = "memory"; + policy["transports"] = {"ub", "rdma"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[UB] = std::make_shared(UB); + transports[RDMA] = std::make_shared(RDMA); + // UB is installed but deliberately lacks dram_to_dram capability. + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {UB, RDMA}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "ub-with-rdma-fallback"; + + EXPECT_EQ(selector.select(ctx, transports).transport, RDMA); +} + // --------------------------------------------------------------------------- // hint: per-request transport_hint hook on select() // --------------------------------------------------------------------------- @@ -719,6 +852,36 @@ TEST(TransportSelectorTest, HintIsPrependedToCandidateList) { EXPECT_EQ(r1.transport, RDMA); } +TEST(TransportSelectorTest, UbHintIsPrependedAndThenFallsBackByIndex) { + auto conf = std::make_shared(); + TransportSelector selector(conf); + + std::array, kSupportedTransportTypes> + transports{}; + transports[UB] = std::make_shared(UB); + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[UB].get())->setDramToDram(true); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA, UB}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/0, /*hint=*/UB).transport, + UB); + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/1, /*hint=*/UB).transport, + RDMA); + EXPECT_EQ( + selector.select(ctx, transports, /*index=*/2, /*hint=*/UB).transport, + UNSPEC); +} + TEST(TransportSelectorTest, HintNotInMatchingPolicyReturnsUnspec) { // Selector mode: the matching policy's transports list is the // authorization whitelist. A hint that is not in it produces UNSPEC diff --git a/mooncake-transfer-engine/tent/tests/ub_core_test.cpp b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp new file mode 100644 index 0000000000..9f5db26cdb --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp @@ -0,0 +1,304 @@ +// Copyright 2026 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. + +#include + +#include +#include +#include +#include + +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/params.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/slice.h" + +namespace mooncake::tent::ub { +namespace { + +TEST(UbParamsTest, ParsesStrictValuesAndAllowsRetryDisable) { + Config config; + config.set("transports/ub/max_retries", 0); + config.set("transports/ub/worker_count", 2); + config.set("transports/ub/device_filter", + std::vector{"ub:fake0:eid0"}); + UbParams params; + ASSERT_TRUE(UbParams::FromConfig(config, params).ok()); + EXPECT_EQ(params.max_retries, 0U); + EXPECT_EQ(params.worker_count, 2U); + ASSERT_EQ(params.device_filter.size(), 1U); + + config.set("transports/ub/worker_count", -1); + EXPECT_TRUE(UbParams::FromConfig(config, params).IsInvalidArgument()); +} + +Request makeRequest(size_t length) { + Request request{}; + request.opcode = Request::WRITE; + request.length = length; + return request; +} + +UbPostPath makePath(Topology::NicID local_id = 0, int remote_id = 1, + uint64_t generation = 1) { + return UbPostPath{local_id, 42, remote_id, generation}; +} + +TEST(UbSliceTest, RetryRejectsStaleCompletionAndNotifiesTerminalOnce) { + int callback_count = 0; + TransferStatus callback_status{INITIAL, 0}; + auto task = UbTask::create( + makeRequest(128), + [&](const TransferStatus& status) { + ++callback_count; + callback_status = status; + }, + 1); + auto slice = task->addSlice(UbSliceSpec{nullptr, 0, 128, 0, 1}, 2); + ASSERT_NE(slice, nullptr); + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(slice->markQueued(3)); + + auto first_attempt = slice->beginAttempt(makePath(0, 1, 10), 4); + ASSERT_TRUE(first_attempt.has_value()); + auto first_completion = slice->completionToken(*first_attempt); + ASSERT_TRUE(first_completion.has_value()); + ASSERT_TRUE(first_completion->markPosted(5)); + EXPECT_EQ(first_completion->resolve(TIMEOUT, 0, true, 6), + UbAttemptResolution::kRetryScheduled); + EXPECT_EQ(task->transferStatus().s, PENDING); + EXPECT_EQ(slice->snapshot().retry_count, 1U); + + ASSERT_TRUE(slice->markQueued(7)); + auto second_attempt = slice->beginAttempt(makePath(0, 2, 11), 8); + ASSERT_TRUE(second_attempt.has_value()); + auto second_completion = slice->completionToken(*second_attempt); + ASSERT_TRUE(second_completion.has_value()); + ASSERT_TRUE(second_completion->markPosted(9)); + + // A completion from the retired endpoint generation cannot resolve the + // replacement attempt. + EXPECT_EQ(first_completion->resolve(COMPLETED, 128, false, 10), + UbAttemptResolution::kIgnored); + EXPECT_EQ(second_completion->resolve(COMPLETED, 128, false, 11), + UbAttemptResolution::kTerminal); + EXPECT_EQ(second_completion->resolve(COMPLETED, 128, false, 12), + UbAttemptResolution::kIgnored); + + EXPECT_EQ(callback_count, 1); + EXPECT_EQ(callback_status.s, COMPLETED); + EXPECT_EQ(callback_status.transferred_bytes, 128U); + const auto task_snapshot = task->snapshot(); + EXPECT_EQ(task_snapshot.remaining_slices, 0U); + EXPECT_NE(task_snapshot.terminal_ns, 0U); +} + +TEST(UbSliceTest, CancellationStopsQueuedWorkAndDrainsPostedWork) { + int callback_count = 0; + TransferStatus callback_status{INITIAL, 0}; + auto task = + UbTask::create(makeRequest(128), [&](const TransferStatus& status) { + ++callback_count; + callback_status = status; + }); + auto posted = task->addSlice(UbSliceSpec{nullptr, 0, 64, 0, 1}); + auto queued = task->addSlice(UbSliceSpec{nullptr, 64, 64, 64, 1}); + ASSERT_NE(posted, nullptr); + ASSERT_NE(queued, nullptr); + ASSERT_TRUE(task->seal()); + + ASSERT_TRUE(posted->markQueued(1)); + auto attempt = posted->beginAttempt(makePath(), 2); + ASSERT_TRUE(attempt.has_value()); + auto completion = posted->completionToken(*attempt); + ASSERT_TRUE(completion.has_value()); + ASSERT_TRUE(completion->markPosted(3)); + ASSERT_TRUE(queued->markQueued(4)); + + EXPECT_EQ(task->requestCancellation(5), 1U); + EXPECT_EQ(task->transferStatus().s, PENDING); + EXPECT_EQ(posted->snapshot().state, UbSliceState::kPosted); + EXPECT_EQ(queued->snapshot().state, UbSliceState::kCanceled); + + // Posted work is not fabricated as immediately canceled. Its real + // completion drains first, then the task reaches its single terminal + // state (CANCELED because another slice never reached the device). + EXPECT_EQ(completion->resolve(COMPLETED, 64, false, 6), + UbAttemptResolution::kTerminal); + EXPECT_EQ(completion->resolve(FAILED, 0, false, 7), + UbAttemptResolution::kIgnored); + EXPECT_EQ(callback_count, 1); + EXPECT_EQ(callback_status.s, CANCELED); + EXPECT_EQ(callback_status.transferred_bytes, 64U); +} + +TEST(UbSliceTest, ConcurrentDuplicateCompletionsChooseOneTerminalWinner) { + std::atomic callback_count{0}; + auto task = UbTask::create(makeRequest(64), [&](const TransferStatus&) { + callback_count.fetch_add(1, std::memory_order_relaxed); + }); + auto slice = task->addSlice(UbSliceSpec{nullptr, 0, 64, 0, 0}); + ASSERT_NE(slice, nullptr); + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(slice->markQueued()); + auto attempt = slice->beginAttempt(makePath()); + ASSERT_TRUE(attempt.has_value()); + auto completion = slice->completionToken(*attempt); + ASSERT_TRUE(completion.has_value()); + ASSERT_TRUE(completion->markPosted()); + + std::atomic terminal_winners{0}; + std::vector threads; + for (int i = 0; i < 8; ++i) { + threads.emplace_back([&] { + if (completion->resolve(COMPLETED, 64, false) == + UbAttemptResolution::kTerminal) { + terminal_winners.fetch_add(1, std::memory_order_relaxed); + } + }); + } + for (auto& thread : threads) thread.join(); + + EXPECT_EQ(terminal_winners.load(std::memory_order_relaxed), 1); + EXPECT_EQ(callback_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(task->transferStatus().s, COMPLETED); + EXPECT_EQ(task->transferStatus().transferred_bytes, 64U); +} + +TEST(UbQuotaTest, EnforcesBothLevelsAndReleasesIdempotently) { + QuotaManager quota(/*default_device_limits=*/{100, 2}, + /*default_path_limits=*/{60, 1}); + const auto first_path = makePath(0, 1, 1); + const auto second_path = makePath(0, 2, 1); + + auto reservation = quota.tryAcquire(first_path, 60); + ASSERT_TRUE(reservation.has_value()); + EXPECT_FALSE(quota.tryAcquire(first_path, 1).has_value()); + // The second path has room, but the physical-device byte cap is shared. + EXPECT_FALSE(quota.tryAcquire(second_path, 50).has_value()); + + auto copied_token = *reservation; + copied_token.path = second_path; + copied_token.bytes = 1; + copied_token.wrs = 99; + EXPECT_TRUE(quota.release(copied_token)); + EXPECT_FALSE(quota.release(*reservation)); + + const auto device = quota.deviceStats(0); + const auto path = quota.pathStats(first_path); + const auto aggregate = quota.aggregateStats(); + EXPECT_EQ(device.usage, QuotaUsage{}); + EXPECT_EQ(path.usage, QuotaUsage{}); + EXPECT_EQ(aggregate.usage, QuotaUsage{}); + EXPECT_EQ(aggregate.active_reservations, 0U); + EXPECT_EQ(aggregate.duplicate_release_attempts, 1U); + EXPECT_TRUE(quota.tryAcquire(first_path, 60).has_value()); +} + +TEST(UbRailMonitorTest, PausesOnErrorWindowAndRecoversAfterCooldown) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/100, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(1), -1.0); + monitor.recordSuccess(path, 100, 10, 50); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(50), 10'000'000'000.0); + + monitor.recordError(path, 100); + monitor.recordTimeout(path, 110); + EXPECT_FALSE(monitor.available(path, 159)); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(159), 0.0); + + EXPECT_TRUE(monitor.available(path, 160)); + const auto stats = monitor.stats(path, 160); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.completion_errors, 2U); + EXPECT_EQ(stats.timeouts, 1U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); + EXPECT_DOUBLE_EQ(monitor.aggregateBandwidth(160), 10'000'000'000.0); +} + +TEST(UbRailMonitorTest, ErrorsOutsideWindowDoNotPauseRail) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/10, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + monitor.recordError(path, 1); + monitor.recordError(path, 11); + EXPECT_TRUE(monitor.available(path, 11)); + EXPECT_EQ(monitor.stats(path, 11).errors_in_window, 1U); +} + +TEST(UbRailMonitorTest, OutOfOrderErrorsUseLatestEventForCooldown) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/100, + /*cooldown_ns=*/50, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + monitor.recordError(path, 120); + monitor.recordError(path, 100); + + auto stats = monitor.stats(path, 120); + EXPECT_TRUE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 2U); + EXPECT_EQ(stats.pause_started_ns, 120U); + EXPECT_EQ(stats.cooldown_until_ns, 170U); + EXPECT_FALSE(monitor.available(path, 169)); + EXPECT_TRUE(monitor.available(path, 170)); + + stats = monitor.stats(path, 170); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); + + // Neither an older query nor an error from the recovered epoch may + // rewind the rail or resurrect its completed pause. + EXPECT_TRUE(monitor.available(path, 130)); + monitor.recordError(path, 110); + stats = monitor.stats(path, 130); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 0U); + EXPECT_EQ(stats.completion_errors, 3U); + EXPECT_EQ(stats.pauses, 1U); + EXPECT_EQ(stats.recoveries, 1U); +} + +TEST(UbRailMonitorTest, OutOfOrderExpiredErrorDoesNotTriggerPause) { + RailMonitor monitor(RailMonitorConfig{/*error_threshold=*/2, + /*error_window_ns=*/50, + /*cooldown_ns=*/25, + /*ewma_alpha=*/0.5}); + const auto path = makePath(); + + monitor.recordError(path, 200); + monitor.recordError(path, 100); + + const auto stats = monitor.stats(path, 200); + EXPECT_FALSE(stats.paused); + EXPECT_EQ(stats.errors_in_window, 1U); + EXPECT_EQ(stats.completion_errors, 2U); + EXPECT_EQ(stats.last_error_ns, 200U); + EXPECT_EQ(stats.pauses, 0U); + EXPECT_EQ(stats.recoveries, 0U); +} + +} // namespace +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp new file mode 100644 index 0000000000..54d22dafc4 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -0,0 +1,730 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/segment_registry.h" +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/endpoint_store.h" +#include "tent/transport/ub/quota.h" +#include "tent/transport/ub/rail_monitor.h" +#include "tent/transport/ub/topology_attrs.h" +#include "tent/transport/ub/ub_transport.h" +#include "tent/transport/ub/workers.h" + +namespace mooncake::tent::ub { +namespace { + +class FakeContext final : public Context { + public: + explicit FakeContext(DeviceInfo info) : info_(std::move(info)) {} + bool valid() const noexcept override { return valid_; } + const DeviceInfo& deviceInfo() const noexcept override { return info_; } + int asyncFd() const noexcept override { return -1; } + void close() { valid_ = false; } + + private: + DeviceInfo info_; + bool valid_{true}; +}; + +class FakeJfc final : public Jfc { + public: + bool valid() const noexcept override { return valid_; } + int eventFd() const noexcept override { return -1; } + + void push(Completion completion) { + std::lock_guard lock(mutex_); + completions_.push_back(completion); + } + void poll(size_t maximum, std::vector& output) { + std::lock_guard lock(mutex_); + while (!completions_.empty() && output.size() < maximum) { + output.push_back(completions_.front()); + completions_.pop_front(); + } + } + void close() { valid_ = false; } + + private: + std::mutex mutex_; + std::deque completions_; + bool valid_{true}; +}; + +SegmentDescriptor makeDescriptor(uint64_t address, uint64_t length) { + return SegmentDescriptor{ + SegmentDescriptor::kSchemaVersion, 1, 16, + std::to_string(address) + ":" + std::to_string(length)}; +} + +bool parseDescriptor(const SegmentDescriptor& descriptor, uint64_t& address, + uint64_t& length) { + const auto colon = descriptor.hex.find(':'); + if (descriptor.schema_version != SegmentDescriptor::kSchemaVersion || + descriptor.urma_api_version != 1 || descriptor.urma_abi_size != 16 || + colon == std::string::npos) { + return false; + } + try { + address = std::stoull(descriptor.hex.substr(0, colon)); + length = std::stoull(descriptor.hex.substr(colon + 1)); + return length != 0; + } catch (...) { + return false; + } +} + +class FakeLocalSegment final : public LocalSegment { + public: + FakeLocalSegment(uint64_t address, uint64_t length) + : address_(address), + length_(length), + descriptor_(makeDescriptor(address, length)) {} + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() { valid_ = false; } + + private: + uint64_t address_; + uint64_t length_; + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeRemoteSegment final : public RemoteSegment { + public: + FakeRemoteSegment(uint64_t address, uint64_t length, + SegmentDescriptor descriptor) + : address_(address), + length_(length), + descriptor_(std::move(descriptor)) {} + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() { valid_ = false; } + + private: + uint64_t address_; + uint64_t length_; + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeJetty final : public Jetty { + public: + FakeJetty(uint32_t id, std::shared_ptr jfc) + : id_(id), jfc_(std::move(jfc)) {} + bool valid() const noexcept override { return valid_; } + uint32_t id() const noexcept override { return id_; } + uint32_t uasid() const noexcept override { return 0; } + std::shared_ptr jfc() const { return jfc_; } + bool bound() const { return bound_; } + void bind() { bound_ = true; } + void unbind() { bound_ = false; } + void close() { + valid_ = false; + bound_ = false; + } + + private: + uint32_t id_; + std::shared_ptr jfc_; + bool valid_{true}; + bool bound_{false}; +}; + +class FakeUrmaAdapter final : public UrmaAdapter { + public: + explicit FakeUrmaAdapter(DeviceInfo device) : device_(std::move(device)) {} + + bool available() const noexcept override { return true; } + uint32_t nativeApiVersion() const noexcept override { return 1; } + size_t nativeSegmentDescriptorSize() const noexcept override { return 16; } + Status initialize() override { + initialized_ = true; + return Status::OK(); + } + Status shutdown() override { + initialized_ = false; + return Status::OK(); + } + Status discoverDevices(std::vector& devices) override { + if (!initialized_) return Status::InvalidArgument("not initialized"); + devices = {device_}; + return Status::OK(); + } + Status openContext(const DeviceInfo& device, ContextPtr& context) override { + if (!initialized_ || device.topology_name != device_.topology_name) { + return Status::DeviceNotFound("fake device missing"); + } + context = std::make_shared(device); + return Status::OK(); + } + Status closeContext(ContextPtr& context) override { + if (auto fake = std::dynamic_pointer_cast(context)) { + fake->close(); + } + context.reset(); + return Status::OK(); + } + Status createJfc(const ContextPtr&, const JfcOptions&, + JfcPtr& jfc) override { + jfc = std::make_shared(); + return Status::OK(); + } + Status deleteJfc(JfcPtr& jfc) override { + if (auto fake = std::dynamic_pointer_cast(jfc)) fake->close(); + jfc.reset(); + return Status::OK(); + } + Status registerLocalSegment(const ContextPtr&, uint64_t address, + size_t length, const SegmentOptions& options, + LocalSegmentPtr& segment) override { + if (length == 0) return Status::InvalidArgument("empty segment"); + last_registered_access_.store(options.access, + std::memory_order_release); + segment = std::make_shared(address, length); + return Status::OK(); + } + Status unregisterLocalSegment(LocalSegmentPtr& segment) override { + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + return Status::OK(); + } + Status importRemoteSegment(const ContextPtr&, + const SegmentDescriptor& descriptor, + const SegmentOptions&, + RemoteSegmentPtr& segment) override { + uint64_t address = 0; + uint64_t length = 0; + if (!parseDescriptor(descriptor, address, length)) { + return Status::InvalidArgument("bad descriptor"); + } + segment = + std::make_shared(address, length, descriptor); + return Status::OK(); + } + Status unimportRemoteSegment(RemoteSegmentPtr& segment) override { + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + return Status::OK(); + } + Status createJetty(const ContextPtr&, const JfcPtr& jfc, + const JettyOptions&, JettyPtr& jetty) override { + auto fake_jfc = std::dynamic_pointer_cast(jfc); + if (!fake_jfc) return Status::InvalidArgument("bad JFC"); + jetty = std::make_shared(next_jetty_id_++, fake_jfc); + return Status::OK(); + } + Status deleteJetty(JettyPtr& jetty) override { + if (auto fake = std::dynamic_pointer_cast(jetty)) { + fake->close(); + } + jetty.reset(); + return Status::OK(); + } + Status bindJetty(const JettyPtr& jetty, const RemoteJettyInfo&) override { + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid()) + return Status::InvalidArgument("bad Jetty"); + fake->bind(); + return Status::OK(); + } + Status unbindJetty(const JettyPtr& jetty) override { + if (auto fake = std::dynamic_pointer_cast(jetty)) { + fake->unbind(); + } + return Status::OK(); + } + Status resetJetty(const JettyPtr&) override { return Status::OK(); } + Status quiesceJetty(const JettyPtr& jetty, uint32_t timeout_ms, + std::vector& completions) override { + completions.clear(); + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid() || timeout_ms == 0) { + return Status::InvalidArgument("bad Jetty quiesce"); + } + if (fail_next_quiesce_.exchange(false, std::memory_order_acq_rel)) { + return Status::RdmaError("injected quiesce failure"); + } + std::lock_guard lock(pending_mutex_); + auto it = pending_.begin(); + while (it != pending_.end()) { + if (it->jetty_id != fake->id()) { + ++it; + continue; + } + completions.push_back(Completion{CompletionCategory::ENDPOINT_ERROR, + 0, it->request.token, 0, + fake->id()}); + it = pending_.erase(it); + } + quiesce_calls_.fetch_add(1, std::memory_order_relaxed); + return Status::OK(); + } + Status post(const JettyPtr& jetty, const std::vector& requests, + size_t& posted_count) override { + posted_count = 0; + auto fake = std::dynamic_pointer_cast(jetty); + if (!fake || !fake->valid() || !fake->bound()) { + return Status::InvalidArgument("unbound Jetty"); + } + for (const auto& request : requests) { + if (!request.local_segment || !request.remote_segment || + request.token == 0 || request.length == 0) { + return Status::InvalidArgument("bad WR"); + } + const auto category = next_completion_.exchange( + CompletionCategory::SUCCESS, std::memory_order_acq_rel); + if (hold_next_completion_.exchange(false, + std::memory_order_acq_rel)) { + std::lock_guard lock(pending_mutex_); + pending_.push_back(Pending{fake->id(), request}); + ++posted_count; + continue; + } + if (category == CompletionCategory::SUCCESS) { + if (request.operation == Operation::WRITE) { + std::memcpy( + reinterpret_cast(request.remote_address), + reinterpret_cast(request.local_address), + request.length); + } else { + std::memcpy( + reinterpret_cast(request.local_address), + reinterpret_cast(request.remote_address), + request.length); + } + } + fake->jfc()->push( + Completion{category, 0, request.token, + category == CompletionCategory::SUCCESS + ? static_cast(request.length) + : 0, + fake->id()}); + ++posted_count; + } + return Status::OK(); + } + Status poll(const JfcPtr& jfc, size_t maximum, + std::vector& completions) override { + completions.clear(); + auto fake = std::dynamic_pointer_cast(jfc); + if (!fake || !fake->valid()) return Status::InvalidArgument("bad JFC"); + fake->poll(maximum, completions); + return Status::OK(); + } + + void failNextCompletion(CompletionCategory category) { + next_completion_.store(category, std::memory_order_release); + } + void holdNextCompletion() { + hold_next_completion_.store(true, std::memory_order_release); + } + void failNextQuiesce() { + fail_next_quiesce_.store(true, std::memory_order_release); + } + size_t pendingCount() const { + std::lock_guard lock(pending_mutex_); + return pending_.size(); + } + uint64_t quiesceCalls() const { + return quiesce_calls_.load(std::memory_order_relaxed); + } + uint32_t lastRegisteredAccess() const { + return last_registered_access_.load(std::memory_order_acquire); + } + + private: + struct Pending { + uint32_t jetty_id{0}; + WorkRequest request; + }; + + DeviceInfo device_; + bool initialized_{false}; + uint32_t next_jetty_id_{1}; + std::atomic next_completion_{ + CompletionCategory::SUCCESS}; + std::atomic hold_next_completion_{false}; + std::atomic fail_next_quiesce_{false}; + mutable std::mutex pending_mutex_; + std::vector pending_; + std::atomic quiesce_calls_{0}; + std::atomic last_registered_access_{0}; +}; + +class NullRegistry final : public SegmentRegistry { + public: + Status getSegmentDesc(SegmentDescRef&, const std::string&) override { + return Status::InvalidEntry("not used"); + } + Status putSegmentDesc(SegmentDescRef&) override { return Status::OK(); } + Status deleteSegmentDesc(const std::string&) override { + return Status::OK(); + } +}; + +DeviceInfo fakeDevice() { + DeviceInfo info; + info.topology_name = "ub:fake0:eid0"; + info.native_device_name = "fake0"; + info.native_device_path = "/fake/fake0"; + info.eid_index = 0; + info.eid = "0001:0000:0000:0000:0000:0000:0000:0000"; + info.active = true; + info.capabilities.max_jfc = 4; + info.capabilities.max_jetty = 64; + return info; +} + +std::shared_ptr fakeTopology(bool discovery_active = true) { + auto topology = std::make_shared(); + Topology::NicEntry nic{.name = "ub:fake0:eid0", + .pci_bus_id = "0000:00:00.0", + .type = Topology::NIC_UB, + .numa_node = 0}; + auto device = fakeDevice(); + device.active = discovery_active; + encodeTopologyDeviceAttributes(device, 0, nic.device_attrs); + topology->nic_list_.push_back(std::move(nic)); + Topology::MemEntry memory; + memory.name = kWildcardLocation; + memory.type = Topology::MEM_HOST; + memory.numa_node = -1; + memory.device_list[0].push_back(0); + topology->mem_list_.push_back(std::move(memory)); + return topology; +} + +TEST(UbNativeDataPathTest, LocalBuffersRequestProviderLocalOnlyAccess) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + UbBufferManager buffers(adapter, {context}); + + std::array storage{}; + BufferDesc descriptor{}; + descriptor.addr = reinterpret_cast(storage.data()); + descriptor.length = storage.size(); + descriptor.location = kWildcardLocation; + MemoryOptions options; + options.perm = kLocalReadWrite; + + ASSERT_TRUE(buffers.addBuffer(descriptor, options).ok()); + EXPECT_EQ(adapter->lastRegisteredAccess(), SEGMENT_ACCESS_LOCAL_ONLY); + + EXPECT_TRUE(buffers.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, + MissingCompletionIsFencedBeforeRetryAndEventuallyCompletes) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + std::vector contexts{context}; + auto topology = fakeTopology(); + UbBufferManager buffers(adapter, contexts); + + std::array source{}; + std::array target{}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i + 1); + } + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = kWildcardLocation; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = kWildcardLocation; + MemoryOptions options; + options.perm = kGlobalReadWrite; + ASSERT_TRUE(buffers.addBuffer(source_desc, options).ok()); + ASSERT_TRUE(buffers.addBuffer(target_desc, options).ok()); + + SegmentManager manager(std::make_unique()); + ASSERT_TRUE(manager + .updateLocal([&](SegmentDesc& segment) { + segment.name = "local"; + segment.type = SegmentType::Memory; + segment.detail = MemorySegmentDesc{}; + auto& memory = + std::get(segment.detail); + memory.topology = *topology; + memory.buffers = {source_desc, target_desc}; + return Status::OK(); + }) + .ok()); + + EndpointStore endpoints(adapter, 16, 1); + RailMonitor rails; + QuotaManager quota; + UbParams params; + params.worker_count = 1; + params.poller_count = 1; + params.slice_size = 16; + params.max_retries = 1; + params.slice_timeout_ms = 20; + + EndpointResolver resolver = [&](const EndpointResolveRequest& request, + std::shared_ptr& endpoint) { + UbEndpointKey key{request.local_context->topologyId(), + request.remote_segment_id, request.remote_topology_id, + "local@ub:fake0:eid0"}; + auto status = + endpoints.getOrCreate(key, request.local_context, endpoint); + if (!status.ok() || endpoint->ready()) return status; + UbBootstrapDesc peer; + peer.local_eid = fakeDevice().eid; + peer.endpoint_generation = 100; + peer.jetty_ids = {777}; + return endpoint->bind(peer); + }; + UbWorkers workers(adapter, contexts, topology, &manager, &buffers, &rails, + "a, params, std::move(resolver), + [&](const std::shared_ptr& endpoint) { + (void)endpoints.retire(endpoint); + }); + ASSERT_TRUE(workers.start().ok()); + adapter->holdNextCompletion(); + + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(target.data()); + request.length = source.size(); + auto task = UbTask::create(request); + for (size_t offset = 0; offset < request.length; offset += 16) { + ASSERT_NE(task->addSlice(UbSliceSpec{ + source.data() + offset, + reinterpret_cast(target.data() + offset), 16, + offset, 1}), + nullptr); + } + ASSERT_TRUE(task->seal()); + ASSERT_TRUE(workers.submit(task).ok()); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (task->transferStatus().s == PENDING && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(task->transferStatus().s, COMPLETED); + EXPECT_EQ(task->transferStatus().transferred_bytes, source.size()); + EXPECT_EQ(source, target); + EXPECT_GT(rails.aggregateBandwidth(), 0.0); + EXPECT_GE(rails.stats(UbPostPath{0, LOCAL_SEGMENT_ID, 0, 1}).timeouts, 1U); + EXPECT_EQ(adapter->pendingCount(), 0U); + EXPECT_GE(adapter->quiesceCalls(), 1U); + + EXPECT_TRUE(workers.stop().ok()); + EXPECT_TRUE(endpoints.clear().ok()); + EXPECT_TRUE(buffers.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, EndpointStoreNeverReusesRetiredGeneration) { + auto adapter = std::make_shared(fakeDevice()); + ASSERT_TRUE(adapter->initialize().ok()); + auto context = std::make_shared(0, fakeDevice(), adapter); + ASSERT_TRUE(context->initialize(1, JfcOptions{}).ok()); + EndpointStore store(adapter, 2, 1); + UbEndpointKey key{0, 9, 0, "peer@ub:fake0:eid0"}; + std::array, 8> concurrent; + std::array statuses; + std::vector threads; + for (size_t i = 0; i < concurrent.size(); ++i) { + threads.emplace_back([&, i] { + statuses[i] = store.getOrCreate(key, context, concurrent[i]); + }); + } + for (auto& thread : threads) thread.join(); + for (size_t i = 0; i < concurrent.size(); ++i) { + ASSERT_TRUE(statuses[i].ok()); + EXPECT_EQ(concurrent[i], concurrent[0]); + } + auto first = concurrent[0]; + const uint64_t old_generation = first->generation(); + EXPECT_TRUE(store.retire(key, old_generation)); + std::shared_ptr replacement; + ASSERT_TRUE(store.getOrCreate(key, context, replacement).ok()); + EXPECT_GT(replacement->generation(), old_generation); + EXPECT_FALSE(store.retire(key, old_generation)); + EXPECT_EQ(store.get(key), replacement); + EXPECT_TRUE(store.clear().ok()); + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_TRUE(adapter->shutdown().ok()); +} + +TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { + auto adapter = std::make_shared(fakeDevice()); + // A serialized discovery snapshot is informational. Current local + // discovery and runtime path health remain the scheduling authorities. + auto topology = fakeTopology(false); + auto control = std::make_shared("p2p", "", nullptr); + uint16_t port = 0; + ASSERT_TRUE(control->start(port).ok()); + const std::string segment_name = + "127.0.0.1:" + std::to_string(static_cast(port)); + ASSERT_TRUE(control->segmentManager() + .updateLocal([&](SegmentDesc& segment) { + segment.name = segment_name; + segment.rpc_server_addr = segment_name; + segment.machine_id = "fake-machine"; + segment.type = SegmentType::Memory; + segment.detail = MemorySegmentDesc{}; + std::get(segment.detail).topology = + *topology; + return Status::OK(); + }) + .ok()); + + auto config = std::make_shared(); + config->set("transports/ub/enable", true); + config->set("transports/ub/worker_count", 1); + config->set("transports/ub/poller_count", 1); + config->set("transports/ub/jfc_per_context", 1); + config->set("transports/ub/jetty_per_endpoint", 1); + config->set("transports/ub/max_endpoints", 16); + config->set("transports/ub/slice_size", 16); + config->set("transports/ub/max_retries", 1); + config->set("transports/ub/slice_timeout_ms", 1000); + config->set("transports/ub/endpoint_cooldown_ms", 100); + + UbTransport transport(adapter); + std::string mutable_segment_name = segment_name; + ASSERT_TRUE( + transport.install(mutable_segment_name, control, topology, config) + .ok()); + EXPECT_STREQ(transport.getName(), "ub"); + EXPECT_TRUE(transport.supportsCancellation()); + EXPECT_FALSE(transport.supportNotification()); + EXPECT_TRUE(transport.capabilities().dram_to_dram); + + std::array source{}; + std::array target{}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(100 - i); + } + const auto expected = source; + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = kWildcardLocation; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = kWildcardLocation; + MemoryOptions options; + options.perm = kGlobalReadWrite; + std::vector descriptors{source_desc, target_desc}; + ASSERT_TRUE(transport.addMemoryBuffer(descriptors, options).ok()); + ASSERT_TRUE(control->segmentManager() + .updateLocal([&](SegmentDesc& segment) { + std::get(segment.detail).buffers = + descriptors; + return Status::OK(); + }) + .ok()); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok()); + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(target.data()); + request.length = source.size(); + ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); + + TransferStatus transfer{PENDING, 0}; + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (transfer.s == PENDING && + std::chrono::steady_clock::now() < deadline) { + ASSERT_TRUE(transport.getTransferStatus(batch, 0, transfer).ok()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(transfer.s, COMPLETED); + EXPECT_EQ(transfer.transferred_bytes, source.size()); + EXPECT_EQ(source, target); + EXPECT_GT(transport.getEstimatedBandwidth(), 0.0); + + source.fill(0); + Request read_request = request; + read_request.opcode = Request::READ; + ASSERT_TRUE(transport.submitTransferTasks(batch, {read_request}).ok()); + transfer = TransferStatus{PENDING, 0}; + const auto read_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (transfer.s == PENDING && + std::chrono::steady_clock::now() < read_deadline) { + ASSERT_TRUE(transport.getTransferStatus(batch, 1, transfer).ok()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(transfer.s, COMPLETED); + EXPECT_EQ(transfer.transferred_bytes, source.size()); + EXPECT_EQ(source, expected); + + ASSERT_TRUE(transport.freeSubBatch(batch).ok()); + EXPECT_EQ(batch, nullptr); + + // A failed device fence must make uninstall retryable without destroying + // pollers, registered memory, or the held WR token. + Transport::SubBatchRef draining_batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(draining_batch, 1).ok()); + adapter->holdNextCompletion(); + ASSERT_TRUE(transport.submitTransferTasks(draining_batch, {request}).ok()); + const auto posted_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(1); + while (adapter->pendingCount() == 0 && + std::chrono::steady_clock::now() < posted_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(adapter->pendingCount(), 1U); + adapter->failNextQuiesce(); + EXPECT_FALSE(transport.uninstall().ok()); + EXPECT_EQ(adapter->pendingCount(), 1U); + EXPECT_TRUE(transport.uninstall().ok()); + EXPECT_EQ(adapter->pendingCount(), 0U); + EXPECT_TRUE(transport.freeSubBatch(draining_batch).ok()); + EXPECT_TRUE(transport.uninstall().ok()); +} + +} // namespace +} // namespace mooncake::tent::ub From 44461942a101b16fabce4b2ac1ca433a96c10944 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:38:57 +0800 Subject: [PATCH 2/6] [TENT] Prefer UB bonding devices when device_filter is empty Avoid initializing underlying udmac* contexts alongside bonding_dev_0 by default, which can bypass UBAGG and cause cross-node LOC_ACCESS_ERR. Co-authored-by: Cursor --- .../tent/transport/ub/device_selection.h | 61 +++++++++++++++++++ .../tent/src/transport/ub/ub_transport.cpp | 46 ++++++++++++++ .../tent/tests/ub_core_test.cpp | 60 ++++++++++++++++++ 3 files changed, 167 insertions(+) create mode 100644 mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h new file mode 100644 index 0000000000..ddd410f961 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/device_selection.h @@ -0,0 +1,61 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ +#define TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ + +#include +#include +#include +#include +#include + +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake { +namespace tent { +namespace ub { + +// Heuristic used when device_filter is empty: prefer UBAGG bonding devices +// over underlying physical ports (e.g. udmac*) that appear alongside them. +inline bool isBondingDeviceName(std::string_view name) { + std::string lower(name); + std::transform( + lower.begin(), lower.end(), lower.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + if (lower.rfind("bonding", 0) == 0) return true; + if (lower.find(":bonding") != std::string::npos) return true; + if (lower.find("_bond") != std::string::npos) return true; + if (lower.find("-bond") != std::string::npos) return true; + return false; +} + +inline bool isBondingDevice(const DeviceInfo& device) { + return isBondingDeviceName(device.native_device_name) || + isBondingDeviceName(device.topology_name); +} + +// When explicit_filter is true, returns devices unchanged (caller already +// applied device_filter). When false and at least one bonding device is +// present, returns only bonding devices; otherwise returns all devices. +inline std::vector preferBondingDevicesIfPresent( + const std::vector& devices, bool explicit_filter) { + if (explicit_filter || devices.empty()) return devices; + const bool has_bonding = + std::any_of(devices.begin(), devices.end(), + [](const DeviceInfo& d) { return isBondingDevice(d); }); + if (!has_bonding) return devices; + + std::vector selected; + selected.reserve(devices.size()); + for (const auto& device : devices) { + if (isBondingDevice(device)) selected.push_back(device); + } + return selected; +} + +} // namespace ub +} // namespace tent +} // namespace mooncake + +#endif // TENT_TRANSPORT_UB_DEVICE_SELECTION_H_ diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index be573db317..81460aabe1 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -22,6 +23,7 @@ #include "tent/thirdparty/nlohmann/json.h" #include "tent/transport/ub/buffers.h" #include "tent/transport/ub/context.h" +#include "tent/transport/ub/device_selection.h" #include "tent/transport/ub/endpoint.h" #include "tent/transport/ub/endpoint_store.h" #include "tent/transport/ub/params.h" @@ -146,6 +148,16 @@ struct UbTransport::Impl { by_topology_name.emplace(device.topology_name, std::move(device)); } + const bool explicit_filter = !params.device_filter.empty(); + const bool prefer_bonding = + !explicit_filter && + std::any_of(by_topology_name.begin(), by_topology_name.end(), + [](const auto& entry) { + return ub::isBondingDevice(entry.second); + }); + std::vector preferred_names; + std::vector skipped_names; + for (size_t id = 0; id < local_topology->getNicCount(); ++id) { const auto* nic = local_topology->getNicEntry(static_cast(id)); if (!nic || nic->type != Topology::NIC_UB) continue; @@ -154,6 +166,10 @@ struct UbTransport::Impl { continue; } if (!filterAllows(params.device_filter, found->second)) continue; + if (prefer_bonding && !ub::isBondingDevice(found->second)) { + skipped_names.push_back(found->second.native_device_name); + continue; + } ub::JfcOptions jfc_options; const auto& device_caps = found->second.capabilities; @@ -175,12 +191,42 @@ struct UbTransport::Impl { << status.ToString(); continue; } + preferred_names.push_back(found->second.native_device_name); context_by_topology_id.emplace(static_cast(id), context); context_by_topology_name.emplace(nic->name, context); jfc_depth_by_device.emplace(static_cast(id), jfc_options.depth); contexts.push_back(std::move(context)); } + + { + std::ostringstream selection_log; + if (explicit_filter) { + selection_log << "UB device selection: mode=explicit-filter"; + } else if (prefer_bonding) { + selection_log << "UB device selection: mode=auto-prefer-bonding"; + } else { + selection_log << "UB device selection: mode=all-devices"; + } + if (!preferred_names.empty()) { + selection_log << " preferred=["; + for (size_t i = 0; i < preferred_names.size(); ++i) { + if (i != 0) selection_log << ", "; + selection_log << preferred_names[i]; + } + selection_log << "]"; + } + if (!skipped_names.empty()) { + selection_log << " skipped=["; + for (size_t i = 0; i < skipped_names.size(); ++i) { + if (i != 0) selection_log << ", "; + selection_log << skipped_names[i]; + } + selection_log << "]"; + } + LOG(INFO) << selection_log.str(); + } + if (contexts.empty()) { return failInstall(Status::DeviceNotFound( "No UB context initialized successfully" LOC_MARK)); diff --git a/mooncake-transfer-engine/tent/tests/ub_core_test.cpp b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp index 9f5db26cdb..4581130273 100644 --- a/mooncake-transfer-engine/tent/tests/ub_core_test.cpp +++ b/mooncake-transfer-engine/tent/tests/ub_core_test.cpp @@ -23,6 +23,7 @@ #include "tent/transport/ub/params.h" #include "tent/transport/ub/rail_monitor.h" #include "tent/transport/ub/slice.h" +#include "tent/transport/ub/device_selection.h" namespace mooncake::tent::ub { namespace { @@ -300,5 +301,64 @@ TEST(UbRailMonitorTest, OutOfOrderExpiredErrorDoesNotTriggerPause) { EXPECT_EQ(stats.recoveries, 0U); } +DeviceInfo makeDevice(std::string native_name, std::string topology_name = "") { + DeviceInfo info; + info.native_device_name = std::move(native_name); + info.topology_name = topology_name.empty() + ? ("ub:" + info.native_device_name + ":eid0") + : std::move(topology_name); + info.active = true; + return info; +} + +TEST(UbDeviceSelectionTest, DetectsBondingNamesCaseInsensitively) { + EXPECT_TRUE(isBondingDeviceName("bonding_dev_0")); + EXPECT_TRUE(isBondingDeviceName("Bonding_Dev_0")); + EXPECT_TRUE(isBondingDeviceName("ub:bonding_dev_0:eid0")); + EXPECT_TRUE(isBondingDeviceName("foo_bond_bar")); + EXPECT_FALSE(isBondingDeviceName("udmac1d1e2")); + EXPECT_FALSE(isBondingDeviceName("ub:udmac1d1e2:eid0")); +} + +TEST(UbDeviceSelectionTest, PrefersBondingWhenPresentAndFilterEmpty) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("bonding_dev_0"), + makeDevice("udmac0d0e1"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/false); + ASSERT_EQ(selected.size(), 1U); + EXPECT_EQ(selected[0].native_device_name, "bonding_dev_0"); +} + +TEST(UbDeviceSelectionTest, KeepsAllDevicesWhenNoBondingPresent) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("udmac0d0e1"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/false); + ASSERT_EQ(selected.size(), 2U); + EXPECT_EQ(selected[0].native_device_name, "udmac1d1e2"); + EXPECT_EQ(selected[1].native_device_name, "udmac0d0e1"); +} + +TEST(UbDeviceSelectionTest, ExplicitFilterSkipsAutoPrefer) { + const std::vector devices = { + makeDevice("udmac1d1e2"), + makeDevice("bonding_dev_0"), + }; + const auto selected = + preferBondingDevicesIfPresent(devices, /*explicit_filter=*/true); + ASSERT_EQ(selected.size(), 2U); +} + +TEST(UbDeviceSelectionTest, EmptyInputStaysEmpty) { + const auto selected = + preferBondingDevicesIfPresent({}, /*explicit_filter=*/false); + EXPECT_TRUE(selected.empty()); +} + } // namespace } // namespace mooncake::tent::ub From 6f4e2fca8979ed5de963f4b33cf03d2471663301 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:59:21 +0800 Subject: [PATCH 3/6] [TENT] Port UB bonding runtime patches into urma_adapter Bring over the local working-tree fixes needed for bonding_dev_0: SET_BONDING_MODE, CTP priority/multi_path jetty config, and token_id segment registration. Also wire topology discover_ub from ub/enable. Co-authored-by: Cursor --- .../tent/src/runtime/topology.cpp | 11 +- .../tent/src/transport/ub/urma_adapter.cpp | 108 ++++++++++++++++-- 2 files changed, 102 insertions(+), 17 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/runtime/topology.cpp b/mooncake-transfer-engine/tent/src/runtime/topology.cpp index dcb9b41d79..82eff44c1e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/topology.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/topology.cpp @@ -354,6 +354,7 @@ Status Topology::parseCustomTopology(const std::string& json_content) { Status Topology::loadFromConfig(const Config& conf, const std::vector& platforms) { + const bool discover_ub = conf.get("transports/ub/enable", false); if (conf.contains("topology/priority_matrix")) { std::string matrix_json; if (conf.dumpSubtree("topology/priority_matrix", &matrix_json)) { @@ -365,7 +366,7 @@ Status Topology::loadFromConfig(const Config& conf, LOG(WARNING) << "Failed to parse topology/priority_matrix: " << status.ToString() << ", falling back to auto-discover"; - return discover(platforms); + return discover(platforms, discover_ub); } } @@ -376,7 +377,7 @@ Status Topology::loadFromConfig(const Config& conf, if (!file.is_open()) { LOG(WARNING) << "Failed to load custom topology from " << path << ", falling back to auto-detect."; - return discover(platforms); + return discover(platforms, discover_ub); } std::stringstream buffer; buffer << file.rdbuf(); @@ -384,16 +385,16 @@ Status Topology::loadFromConfig(const Config& conf, if (content.empty()) { LOG(WARNING) << "Failed to load custom topology from " << path << ", falling back to auto-detect."; - return discover(platforms); + return discover(platforms, discover_ub); } auto status = parseCustomTopology(content); if (status.ok()) return Status::OK(); LOG(WARNING) << "Failed to parse custom topology from " << path << ": " << status.ToString() << ", falling back to auto-detect."; - return discover(platforms); + return discover(platforms, discover_ub); } - return discover(platforms); + return discover(platforms, discover_ub); } size_t Topology::getNicCount(NicType type) const { diff --git a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp index 0c3f4b62b8..a3cf9f6d8a 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp @@ -30,8 +30,11 @@ #include #include +#include + #if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA #include +#include #endif namespace mooncake { @@ -352,10 +355,11 @@ class RuntimeLease { class RealContext final : public Context { public: RealContext(std::shared_ptr runtime, DeviceInfo info, - urma_context_t* native) + urma_context_t* native, uint32_t ctp_priority) : runtime_(std::move(runtime)), info_(std::move(info)), - native_(native) {} + native_(native), + ctp_priority_(ctp_priority) {} ~RealContext() override { (void)close(); } @@ -367,6 +371,12 @@ class RealContext final : public Context { urma_context_t* native() const noexcept { return native_; } + bool isBondingDevice() const noexcept { + return info_.native_device_name.rfind("bonding", 0) == 0; + } + + uint32_t ctpPriority() const noexcept { return ctp_priority_; } + Status close() { if (native_ == nullptr) return Status::OK(); const int rc = urma_delete_context(native_); @@ -379,6 +389,7 @@ class RealContext final : public Context { std::shared_ptr runtime_; DeviceInfo info_; urma_context_t* native_ = nullptr; + uint32_t ctp_priority_ = 15; }; class RealJfc final : public Jfc { @@ -554,12 +565,14 @@ class RealLocalSegment final : public LocalSegment { public: RealLocalSegment(std::shared_ptr context, urma_target_seg_t* native, uint64_t address, - uint64_t length, SegmentDescriptor descriptor) + uint64_t length, SegmentDescriptor descriptor, + urma_token_id_t* token_id) : context_(std::move(context)), native_(native), address_(address), length_(length), - descriptor_(std::move(descriptor)) {} + descriptor_(std::move(descriptor)), + token_id_(token_id) {} ~RealLocalSegment() override { (void)close(); } @@ -576,10 +589,19 @@ class RealLocalSegment final : public LocalSegment { } Status close() { - if (native_ == nullptr) return Status::OK(); - const int rc = urma_unregister_seg(native_); - if (rc != URMA_SUCCESS) return nativeError("urma_unregister_seg", rc); - native_ = nullptr; + if (native_ != nullptr) { + const int rc = urma_unregister_seg(native_); + if (rc != URMA_SUCCESS) { + return nativeError("urma_unregister_seg", rc); + } + native_ = nullptr; + } + + if (token_id_ != nullptr) { + (void)urma_free_token_id(token_id_); + token_id_ = nullptr; + } + return Status::OK(); } @@ -589,6 +611,7 @@ class RealLocalSegment final : public LocalSegment { uint64_t address_ = 0; uint64_t length_ = 0; SegmentDescriptor descriptor_; + urma_token_id_t* token_id_ = nullptr; }; class RealRemoteSegment final : public RemoteSegment { @@ -645,7 +668,12 @@ class RealJetty final : public Jetty { config.flag.bs.share_jfr = 1; config.jfs_cfg.depth = options.depth; config.jfs_cfg.trans_mode = URMA_TM_RC; - config.jfs_cfg.priority = options.priority; + config.jfs_cfg.priority = + context_->isBondingDevice() ? context_->ctpPriority() + : options.priority; + if (context_->isBondingDevice()) { + config.jfs_cfg.flag.bs.multi_path = 1; + } config.jfs_cfg.max_sge = options.max_sge; config.jfs_cfg.max_rsge = options.max_sge; config.jfs_cfg.rnr_retry = options.rnr_retry; @@ -1109,6 +1137,48 @@ class RealUrmaAdapter final : public UrmaAdapter { return nativePointerError("urma_create_context"); } + if (requested.native_device_name.rfind("bonding", 0) == 0) { + bondp_set_bonding_mode_in_t bonding_mode{}; + bonding_mode.bonding_mode = BONDP_BONDING_MODE_STANDALONE; + bonding_mode.bonding_level = BONDP_BONDING_LEVEL_IODIE; + + urma_user_ctl_in_t ctl_in{}; + ctl_in.addr = reinterpret_cast(&bonding_mode); + ctl_in.len = sizeof(bonding_mode); + ctl_in.opcode = BONDP_USER_CTL_SET_BONDING_MODE; + + urma_user_ctl_out_t ctl_out{}; + const int ctl_rc = + urma_user_ctl(native_context, &ctl_in, &ctl_out); + if (ctl_rc != URMA_SUCCESS) { + (void)urma_delete_context(native_context); + return nativeError( + "urma_user_ctl(SET_BONDING_MODE)", ctl_rc); + } + } + + uint32_t ctp_priority = 15; + bool ctp_priority_found = false; + for (uint32_t priority_index = 0; + priority_index < URMA_MAX_PRIORITY_CNT; ++priority_index) { + const auto& priority_info = + attributes.dev_cap.priority_info[priority_index]; + if (priority_info.tp_type.bs.ctp != 0) { + ctp_priority = priority_index; + LOG(INFO) << "UB_PRIO_SELECTED priority=" << ctp_priority + << " SL=" << static_cast(priority_info.SL); + ctp_priority_found = true; + break; + } + } + + if (requested.native_device_name.rfind("bonding", 0) == 0 && + !ctp_priority_found) { + (void)urma_delete_context(native_context); + return Status::InvalidArgument( + "bonding URMA device has no CTP priority"); + } + DeviceInfo current = requested; current.native_device_path = boundedString(native_device->path, URMA_MAX_PATH); @@ -1119,8 +1189,9 @@ class RealUrmaAdapter final : public UrmaAdapter { ":eid" + std::to_string(current.eid_index); } - output = std::make_shared( - std::move(runtime), std::move(current), native_context); + output = std::make_shared(std::move(runtime), + std::move(current), + native_context, ctp_priority); return Status::OK(); } return Status::DeviceNotFound("URMA device not found: " + @@ -1191,20 +1262,29 @@ class RealUrmaAdapter final : public UrmaAdapter { return Status::InvalidArgument("invalid local segment range"); } + urma_token_id_t* token_id = + urma_alloc_token_id(real_context->native()); + if (token_id == nullptr) { + return nativePointerError("urma_alloc_token_id"); + } + urma_reg_seg_flag_t flags{}; flags.bs.token_policy = URMA_TOKEN_NONE; flags.bs.cacheable = options.cacheable ? URMA_CACHEABLE : URMA_NON_CACHEABLE; flags.bs.access = nativeAccess(options.access); + flags.bs.token_id_valid = 1; urma_seg_cfg_t config{}; config.va = address; config.len = length; + config.token_id = token_id; config.token_value.token = options.token; config.flag = flags; urma_target_seg_t* native_segment = urma_register_seg(real_context->native(), &config); if (native_segment == nullptr) { + (void)urma_free_token_id(token_id); return nativePointerError("urma_register_seg"); } @@ -1214,13 +1294,17 @@ class RealUrmaAdapter final : public UrmaAdapter { wire_descriptor.attr = native_segment->seg.attr; wire_descriptor.token_id = native_segment->seg.token_id; + LOG(INFO) << "UB_SEG_TOKEN_ID token_id_valid=1" + << " va=0x" << std::hex << address << " len=0x" << length + << std::dec; + SegmentDescriptor descriptor; descriptor.urma_api_version = URMA_API_VERSION; descriptor.urma_abi_size = sizeof(urma_seg_t); descriptor.hex = encodeHex(&wire_descriptor, sizeof(wire_descriptor)); output = std::make_shared( std::move(real_context), native_segment, address, length, - std::move(descriptor)); + std::move(descriptor), token_id); return Status::OK(); } From fb5b25a9b514e82696f068e1820f551f6061f47f Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:43:47 +0800 Subject: [PATCH 4/6] [TENT] Unify UB bonding detection and harden segment teardown Reuse device_selection helpers in urma_adapter and install path so heuristics cannot drift, and always free token_id when unregister fails. Co-authored-by: Cursor --- .../tent/src/transport/ub/ub_transport.cpp | 26 +++++++++++++------ .../tent/src/transport/ub/urma_adapter.cpp | 17 +++++++----- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index 81460aabe1..d4dcb491d2 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -142,19 +142,28 @@ struct UbTransport::Impl { std::vector discovered; status = adapter->discoverDevices(discovered); if (!status.ok()) return failInstall(status); + + const bool explicit_filter = !params.device_filter.empty(); + const auto preferred_devices = + ub::preferBondingDevicesIfPresent(discovered, explicit_filter); + const bool prefer_bonding = + !explicit_filter && + std::any_of(discovered.begin(), discovered.end(), + [](const ub::DeviceInfo& device) { + return ub::isBondingDevice(device); + }); + std::unordered_set preferred_topology_names; + preferred_topology_names.reserve(preferred_devices.size()); + for (const auto& device : preferred_devices) { + preferred_topology_names.insert(device.topology_name); + } + std::unordered_map by_topology_name; std::unordered_map jfc_depth_by_device; for (auto& device : discovered) { by_topology_name.emplace(device.topology_name, std::move(device)); } - const bool explicit_filter = !params.device_filter.empty(); - const bool prefer_bonding = - !explicit_filter && - std::any_of(by_topology_name.begin(), by_topology_name.end(), - [](const auto& entry) { - return ub::isBondingDevice(entry.second); - }); std::vector preferred_names; std::vector skipped_names; @@ -166,7 +175,8 @@ struct UbTransport::Impl { continue; } if (!filterAllows(params.device_filter, found->second)) continue; - if (prefer_bonding && !ub::isBondingDevice(found->second)) { + if (preferred_topology_names.find(found->second.topology_name) == + preferred_topology_names.end()) { skipped_names.push_back(found->second.native_device_name); continue; } diff --git a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp index a3cf9f6d8a..dec21df008 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp @@ -32,6 +32,8 @@ #include +#include "tent/transport/ub/device_selection.h" + #if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA #include #include @@ -372,7 +374,7 @@ class RealContext final : public Context { urma_context_t* native() const noexcept { return native_; } bool isBondingDevice() const noexcept { - return info_.native_device_name.rfind("bonding", 0) == 0; + return ::mooncake::tent::ub::isBondingDevice(info_); } uint32_t ctpPriority() const noexcept { return ctp_priority_; } @@ -589,12 +591,14 @@ class RealLocalSegment final : public LocalSegment { } Status close() { + Status status = Status::OK(); if (native_ != nullptr) { const int rc = urma_unregister_seg(native_); if (rc != URMA_SUCCESS) { - return nativeError("urma_unregister_seg", rc); + status = nativeError("urma_unregister_seg", rc); + } else { + native_ = nullptr; } - native_ = nullptr; } if (token_id_ != nullptr) { @@ -602,7 +606,7 @@ class RealLocalSegment final : public LocalSegment { token_id_ = nullptr; } - return Status::OK(); + return status; } private: @@ -1137,7 +1141,7 @@ class RealUrmaAdapter final : public UrmaAdapter { return nativePointerError("urma_create_context"); } - if (requested.native_device_name.rfind("bonding", 0) == 0) { + if (isBondingDevice(requested)) { bondp_set_bonding_mode_in_t bonding_mode{}; bonding_mode.bonding_mode = BONDP_BONDING_MODE_STANDALONE; bonding_mode.bonding_level = BONDP_BONDING_LEVEL_IODIE; @@ -1172,8 +1176,7 @@ class RealUrmaAdapter final : public UrmaAdapter { } } - if (requested.native_device_name.rfind("bonding", 0) == 0 && - !ctp_priority_found) { + if (isBondingDevice(requested) && !ctp_priority_found) { (void)urma_delete_context(native_context); return Status::InvalidArgument( "bonding URMA device has no CTP priority"); From 698aab3996e17dce68138165ff23dd2efcbb2c75 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:44:19 +0800 Subject: [PATCH 5/6] [TENT] Fix clang-format for UB transport files Apply clang-format-20 to changed lines in ub_transport.cpp and urma_adapter.cpp so CI code format check passes. Co-authored-by: Cursor --- .../tent/src/transport/ub/ub_transport.cpp | 3 ++- .../tent/src/transport/ub/urma_adapter.cpp | 19 +++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index d4dcb491d2..c572ff1f8f 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -214,7 +214,8 @@ struct UbTransport::Impl { if (explicit_filter) { selection_log << "UB device selection: mode=explicit-filter"; } else if (prefer_bonding) { - selection_log << "UB device selection: mode=auto-prefer-bonding"; + selection_log + << "UB device selection: mode=auto-prefer-bonding"; } else { selection_log << "UB device selection: mode=all-devices"; } diff --git a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp index dec21df008..eaaac17088 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp @@ -672,9 +672,9 @@ class RealJetty final : public Jetty { config.flag.bs.share_jfr = 1; config.jfs_cfg.depth = options.depth; config.jfs_cfg.trans_mode = URMA_TM_RC; - config.jfs_cfg.priority = - context_->isBondingDevice() ? context_->ctpPriority() - : options.priority; + config.jfs_cfg.priority = context_->isBondingDevice() + ? context_->ctpPriority() + : options.priority; if (context_->isBondingDevice()) { config.jfs_cfg.flag.bs.multi_path = 1; } @@ -1156,8 +1156,8 @@ class RealUrmaAdapter final : public UrmaAdapter { urma_user_ctl(native_context, &ctl_in, &ctl_out); if (ctl_rc != URMA_SUCCESS) { (void)urma_delete_context(native_context); - return nativeError( - "urma_user_ctl(SET_BONDING_MODE)", ctl_rc); + return nativeError("urma_user_ctl(SET_BONDING_MODE)", + ctl_rc); } } @@ -1192,9 +1192,9 @@ class RealUrmaAdapter final : public UrmaAdapter { ":eid" + std::to_string(current.eid_index); } - output = std::make_shared(std::move(runtime), - std::move(current), - native_context, ctp_priority); + output = std::make_shared( + std::move(runtime), std::move(current), native_context, + ctp_priority); return Status::OK(); } return Status::DeviceNotFound("URMA device not found: " + @@ -1265,8 +1265,7 @@ class RealUrmaAdapter final : public UrmaAdapter { return Status::InvalidArgument("invalid local segment range"); } - urma_token_id_t* token_id = - urma_alloc_token_id(real_context->native()); + urma_token_id_t* token_id = urma_alloc_token_id(real_context->native()); if (token_id == nullptr) { return nativePointerError("urma_alloc_token_id"); } From 3066465bfb90249cef0580c73431e96b465a6181 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:16:25 +0800 Subject: [PATCH 6/6] [CI/Build] Ignore gflags exit code in UB tebench smoke test tebench --help exits 1 under gflags; tolerate that so the UB benchmark CLI smoke step checks help text instead of failing early under set -e. Co-authored-by: Cursor --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7609dca048..977641e8b2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -713,7 +713,7 @@ jobs: cd build-tent help_output="$(./mooncake-transfer-engine/benchmark/tebench \ --backend=tent --xport_type=ub --tent_transport_hint=ub \ - --help 2>&1)" + --help 2>&1 || true)" grep -q 'iouring|ub|sunrise_link' <<< "${help_output}" grep -q 'ascend|ub|sunrise_link' <<< "${help_output}" shell: bash