From 86321020a00179d560db4582e88a72d00cd2b36e Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 13:59:55 +0000 Subject: [PATCH 01/12] io: add TCP backend design doc --- docs/MORI-IO-TCP-BACKEND.md | 181 ++++++++++++++++++++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 docs/MORI-IO-TCP-BACKEND.md diff --git a/docs/MORI-IO-TCP-BACKEND.md b/docs/MORI-IO-TCP-BACKEND.md new file mode 100644 index 000000000..1c0446cc5 --- /dev/null +++ b/docs/MORI-IO-TCP-BACKEND.md @@ -0,0 +1,181 @@ +# MORI-IO TCP Backend Design + +> **Goal**: Provide a high-reliability, high-performance **TCP fallback transport** for distributed KV-Cache transfer when RDMA/XGMI is unavailable. +> The TCP backend must **simulate RDMA-like Read/Write semantics** on top of TCP byte streams, while keeping the public `mori::io::Backend` / `BackendSession` APIs aligned with existing backends. + +--- + +## 1. Requirements Mapping + +### 1) Semantic Mapping (RDMA-like ops on TCP stream) +TCP is a byte stream; MORI-IO upper layers are op/message oriented (Read/Write/BatchRead/BatchWrite). + +**Solution**: a framed protocol: +- **Control channel** carries op metadata (`WriteReq`, `ReadReq`, `Batch*Req`) and `Completion`. +- **Data channel** carries payload with a lightweight data frame header (`op_id`, `payload_len`). + +This avoids classic TCP sticky/half packet issues by always parsing `FixedHeader + VariableBody`. + +### 2) Concurrency Model (non-blocking + multiplexing) +**Baseline**: `epoll`-based Reactor (single IO thread per backend instance). +- Non-blocking sockets, edge-triggered `epoll`. +- `eventfd` used to wake the reactor when new outbound ops are submitted from application threads. +- Optional `timerfd` for op timeout scanning & housekeeping. + +**Evaluation**: Kernel is 5.15+ (io_uring-capable). The design keeps a clean boundary so `io_uring` Proactor can replace the Reactor later without changing public APIs. + +### 3) Zero-Copy & Memory Management +#### CPU memory +- Outbound: `sendmsg` + `iovec` directly referencing user registered buffers (no extra user-space copy). +- Inbound: `recv` directly into destination user buffer. +- Optional: enable `MSG_ZEROCOPY` for large CPU payloads (best-effort). + +#### GPU memory (fallback path) +TCP cannot directly read/write device VRAM. We stage through pinned host buffers: +- GPU→Host (`hipMemcpyDtoHAsync`) before sending payload. +- Host→GPU (`hipMemcpyHtoDAsync`) after receiving payload. +- Pinned host buffers are pooled to avoid per-op allocations. + +### 4) Latency vs Throughput +- **Metadata**: control channel uses `TCP_NODELAY` for low latency. +- **Payload**: data channel uses large `SO_SNDBUF/SO_RCVBUF`, and can optionally enable `MSG_ZEROCOPY` for large CPU payloads. +- Backpressure: bound in-flight bytes per peer to avoid unbounded buffering. + +### 5) Robustness +- NIC selection: bind outbound sockets to `IOEngineConfig.host` before connect; listen on the same host. +- Keep-alive: enable `SO_KEEPALIVE` + reasonable keepalive parameters. +- Disconnect: fail in-flight ops with `ERR_BAD_STATE` and allow lazy reconnect on next submission. +- Timeout: per-op timeout (configurable) -> fail op and clean state. + +--- + +## 2. Architecture + +### 2.1 Components + +- `TcpBackend` (`mori::io::Backend`) + - Owns one `TcpTransport` instance. + - Implements `RegisterRemoteEngine`, `RegisterMemory`, `Read/Write/Batch*`, `CreateSession`. + +- `TcpBackendSession` (`mori::io::BackendSession`) + - Caches local/remote descriptors and routes ops to `TcpTransport` (low overhead fast path). + +- `TcpTransport` + - A single IO thread (reactor) managing: + - Listener socket + - Per-peer connections (control + data) + - Framed parsing / send queue + - Pending outbound op table + - Inbound op handling and inbound status queue + +- `PeerConn` + - `ctrl_fd` + `data_fd` + - control parser state, data parser state + - outbound queues for ctrl and data + +- `PinnedStagingPool` + - Reusable pinned host buffers keyed by `(device_id, size_class)` (GPU staging). + +### 2.2 Threading + +Single IO thread per `TcpBackend`: +- No thread-per-connection. +- All network IO and protocol parsing happens in this thread. +- Application threads only enqueue op descriptors (cheap) and return immediately with `TransferStatus = IN_PROGRESS`. + +--- + +## 3. Wire Protocol + +### 3.1 Channels +For each peer engine we maintain **two TCP connections**: +- `CTRL`: op metadata + completions +- `DATA`: bulk payload frames + +Both connections connect to the same listen port; the first message is a `HELLO` identifying: +- protocol version +- channel type (`CTRL` / `DATA`) +- `EngineKey` string + +### 3.2 Control Frame +All ctrl messages are framed: + +``` +| CtrlHeader (fixed) | CtrlBody (variable, body_len bytes) | +``` + +`CtrlHeader` includes: +- magic/version +- message type +- body length + +`CtrlBody` contains type-specific fields (op_id, mem_id, offsets, sizes, etc). + +### 3.3 Data Frame + +``` +| DataHeader (fixed) | Payload (payload_len bytes) | +``` + +`DataHeader` includes: +- `op_id` +- `payload_len` + +Ordering: +- Write: `CTRL WriteReq` then `DATA (op_id,payload)` then `CTRL Completion` +- Read: `CTRL ReadReq` then `DATA (op_id,payload)` then `CTRL Completion` + +Receiver keeps op state keyed by `op_id` so ctrl/data can be processed independently. + +--- + +## 4. API Semantics + +- `Write(local, remote, ...)`: + - initiator sends payload; target writes into its local memory region. +- `Read(local, remote, ...)`: + - initiator requests; target reads from its local memory region and sends payload back. +- `PopInboundTransferStatus(remote_key, id, ...)`: + - target side may pop completion of inbound ops (mainly useful for write-like semantics), consistent with existing RDMA notification API. + +--- + +## 5. Config Surface (planned) + +`TcpBackendConfig`: +- `num_io_threads` (default 1) +- `max_inflight_bytes_per_peer` +- `op_timeout_ms` +- `sock_sndbuf_bytes`, `sock_rcvbuf_bytes` +- `enable_zerocopy`, `zerocopy_threshold_bytes` +- `enable_keepalive`, keepalive parameters + +--- + +## 6. Planned Code Changes + +### New files +- `src/io/tcp/backend_impl.hpp` +- `src/io/tcp/backend_impl.cpp` +- `src/io/tcp/protocol.hpp` +- `src/io/tcp/protocol.cpp` +- `src/io/tcp/poller_epoll.hpp` (reactor) +- `src/io/tcp/pinned_staging_pool.hpp` + +### Modified files +- `include/mori/io/backend.hpp` (add `TcpBackendConfig`) +- `src/io/engine.cpp` (wire in `BackendType::TCP`) +- `src/io/CMakeLists.txt` (build TCP backend sources) +- `src/pybind/mori.cpp` (bind `TcpBackendConfig`) +- `python/mori/io/engine.py` (enable `BackendType.TCP`) +- `tests/python/io/benchmark.py` (support `--backend tcp` and `--op` alias) + +--- + +## 7. Known Limits / Future Work + +- `io_uring` Proactor implementation (replace epoll reactor in `TcpTransport`). +- Adaptive multi-connection striping per peer to improve throughput on high-BDP networks. +- Full `MSG_ZEROCOPY` completion accounting + buffer lifetime tracking (currently best-effort). +- Optional TLS / auth for multi-tenant environments. + From 8d4855233f6b66060ebff4fbd8ff41e48c6bbd47 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 18:35:17 +0000 Subject: [PATCH 02/12] io: add TCP backend (epoll reactor) --- include/mori/io/backend.hpp | 44 + src/io/CMakeLists.txt | 1 + src/io/engine.cpp | 21 + src/io/tcp/backend_impl.cpp | 2263 +++++++++++++++++++++++++++++++++++ src/io/tcp/backend_impl.hpp | 97 ++ src/io/tcp/protocol.hpp | 268 +++++ 6 files changed, 2694 insertions(+) create mode 100644 src/io/tcp/backend_impl.cpp create mode 100644 src/io/tcp/backend_impl.hpp create mode 100644 src/io/tcp/protocol.hpp diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index a0b072609..d92cc50b3 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -21,6 +21,8 @@ // SOFTWARE. #pragma once +#include + #include "mori/io/common.hpp" #include "mori/io/enum.hpp" @@ -79,6 +81,48 @@ inline std::ostream& operator<<(std::ostream& os, const XgmiBackendConfig& c) { return os << "numStreams[" << c.numStreams << "] numEvents[" << c.numEvents << "]"; } +struct TcpBackendConfig : public BackendConfig { + TcpBackendConfig() : BackendConfig(BackendType::TCP) {} + TcpBackendConfig(int numIoThreads_, int sockSndbufBytes_, int sockRcvbufBytes_, int opTimeoutMs_, + bool enableKeepalive_, int keepaliveIdleSec_, int keepaliveIntvlSec_, + int keepaliveCnt_, bool enableCtrlNodelay_, bool enableDataNodelay_) + : BackendConfig(BackendType::TCP), + numIoThreads(numIoThreads_), + sockSndbufBytes(sockSndbufBytes_), + sockRcvbufBytes(sockRcvbufBytes_), + opTimeoutMs(opTimeoutMs_), + enableKeepalive(enableKeepalive_), + keepaliveIdleSec(keepaliveIdleSec_), + keepaliveIntvlSec(keepaliveIntvlSec_), + keepaliveCnt(keepaliveCnt_), + enableCtrlNodelay(enableCtrlNodelay_), + enableDataNodelay(enableDataNodelay_) {} + + int numIoThreads{1}; + + int sockSndbufBytes{4 * 1024 * 1024}; + int sockRcvbufBytes{4 * 1024 * 1024}; + + int opTimeoutMs{30 * 1000}; + + bool enableKeepalive{true}; + int keepaliveIdleSec{30}; + int keepaliveIntvlSec{10}; + int keepaliveCnt{3}; + + bool enableCtrlNodelay{true}; + bool enableDataNodelay{false}; +}; + +inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { + return os << "numIoThreads[" << c.numIoThreads << "] sockSndbufBytes[" << c.sockSndbufBytes + << "] sockRcvbufBytes[" << c.sockRcvbufBytes << "] opTimeoutMs[" << c.opTimeoutMs + << "] enableKeepalive[" << c.enableKeepalive << "] keepaliveIdleSec[" + << c.keepaliveIdleSec << "] keepaliveIntvlSec[" << c.keepaliveIntvlSec + << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" + << c.enableCtrlNodelay << "] enableDataNodelay[" << c.enableDataNodelay << "]"; +} + /* ---------------------------------------------------------------------------------------------- */ /* BackendSession */ /* ---------------------------------------------------------------------------------------------- */ diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index 7ece166fd..a7b58f98f 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -5,6 +5,7 @@ add_library( rdma/backend_impl.cpp rdma/executor.cpp rdma/common.cpp + tcp/backend_impl.cpp xgmi/backend_impl.cpp xgmi/hip_resource_pool.cpp) diff --git a/src/io/engine.cpp b/src/io/engine.cpp index c78091547..db5bfa003 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -27,6 +27,7 @@ #include "mori/io/logging.hpp" #include "src/io/rdma/backend_impl.hpp" +#include "src/io/tcp/backend_impl.hpp" #include "src/io/xgmi/backend_impl.hpp" namespace mori { @@ -132,6 +133,26 @@ void IOEngine::CreateBackend(BackendType type, const BackendConfig& beConfig) { static_cast(beConfig)); backends.insert({type, std::move(backend)}); InvalidateRouteCache(); + } else if (type == BackendType::TCP) { + assert(backends.find(type) == backends.end()); + auto backend = std::make_unique(desc.key, config, + static_cast(beConfig)); + + if (config.port == 0) { + auto bound_port_opt = backend->GetListenPort(); + if (!bound_port_opt.has_value() || bound_port_opt.value() == 0) { + MORI_IO_ERROR("IOEngine key {} failed to retrieve bound port after TCP backend init", + desc.key); + assert(false && "Failed to retrieve bound port after TCP backend init"); + } else { + uint16_t bound_port = bound_port_opt.value(); + desc.port = bound_port; + this->config.port = bound_port; + MORI_IO_INFO("IOEngine key {} bound ephemeral port {}", desc.key, bound_port); + } + } + + backends.insert({type, std::move(backend)}); } else { assert(false && "not implemented"); } diff --git a/src/io/tcp/backend_impl.cpp b/src/io/tcp/backend_impl.cpp new file mode 100644 index 000000000..2374d9658 --- /dev/null +++ b/src/io/tcp/backend_impl.cpp @@ -0,0 +1,2263 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "src/io/tcp/backend_impl.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/application/utils/check.hpp" +#include "mori/io/logging.hpp" +#include "src/io/tcp/protocol.hpp" +#include "src/io/xgmi/hip_resource_pool.hpp" + +namespace mori { +namespace io { + +namespace { + +using Clock = std::chrono::steady_clock; + +inline bool IsWouldBlock(int err) { return (err == EAGAIN) || (err == EWOULDBLOCK); } + +inline int SetNonBlocking(int fd) { + int flags = fcntl(fd, F_GETFL, 0); + if (flags < 0) return -1; + if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) return -1; + return 0; +} + +inline void SetSockOptOrLog(int fd, int level, int optname, const void* optval, socklen_t optlen, + const char* name) { + if (setsockopt(fd, level, optname, optval, optlen) != 0) { + MORI_IO_WARN("TCP: setsockopt {} failed: {}", name, strerror(errno)); + } +} + +inline void ConfigureSocketCommon(int fd, const TcpBackendConfig& cfg) { + if (cfg.enableKeepalive) { + int on = 1; + SetSockOptOrLog(fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on), "SO_KEEPALIVE"); + + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPIDLE, &cfg.keepaliveIdleSec, + sizeof(cfg.keepaliveIdleSec), "TCP_KEEPIDLE"); + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPINTVL, &cfg.keepaliveIntvlSec, + sizeof(cfg.keepaliveIntvlSec), "TCP_KEEPINTVL"); + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPCNT, &cfg.keepaliveCnt, sizeof(cfg.keepaliveCnt), + "TCP_KEEPCNT"); + } +} + +inline void ConfigureCtrlSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + if (cfg.enableCtrlNodelay) { + int on = 1; + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(ctrl)"); + } +} + +inline void ConfigureDataSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + if (cfg.enableDataNodelay) { + int on = 1; + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(data)"); + } + if (cfg.sockSndbufBytes > 0) { + SetSockOptOrLog(fd, SOL_SOCKET, SO_SNDBUF, &cfg.sockSndbufBytes, sizeof(cfg.sockSndbufBytes), + "SO_SNDBUF"); + } + if (cfg.sockRcvbufBytes > 0) { + SetSockOptOrLog(fd, SOL_SOCKET, SO_RCVBUF, &cfg.sockRcvbufBytes, sizeof(cfg.sockRcvbufBytes), + "SO_RCVBUF"); + } +} + +inline std::optional ParseIpv4(const std::string& host, uint16_t port) { + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) { + return std::nullopt; + } + return addr; +} + +inline uint16_t GetBoundPort(int fd) { + sockaddr_in addr{}; + socklen_t len = sizeof(addr); + if (getsockname(fd, reinterpret_cast(&addr), &len) != 0) { + return 0; + } + return ntohs(addr.sin_port); +} + +inline size_t RoundUpPow2(size_t v) { + if (v <= 1) return 1; + size_t p = 1; + while (p < v) p <<= 1; + return p; +} + +struct PinnedBuf { + void* ptr{nullptr}; + size_t cap{0}; +}; + +class PinnedStagingPool { + public: + PinnedStagingPool() = default; + ~PinnedStagingPool() { Clear(); } + + PinnedStagingPool(const PinnedStagingPool&) = delete; + PinnedStagingPool& operator=(const PinnedStagingPool&) = delete; + + std::shared_ptr Acquire(size_t size) { + const size_t cap = RoundUpPow2(size); + { + std::lock_guard lock(mu); + auto it = free.find(cap); + if (it != free.end() && !it->second.empty()) { + void* p = it->second.back(); + it->second.pop_back(); + return std::shared_ptr(new PinnedBuf{p, cap}, [this](PinnedBuf* b) { + this->Release(b); + }); + } + } + + void* p = nullptr; + hipError_t err = hipHostMalloc(&p, cap, hipHostMallocDefault); + if (err != hipSuccess) { + MORI_IO_ERROR("TCP: hipHostMalloc({}) failed: {}", cap, hipGetErrorString(err)); + return nullptr; + } + return std::shared_ptr(new PinnedBuf{p, cap}, [this](PinnedBuf* b) { this->Release(b); }); + } + + void Clear() { + std::lock_guard lock(mu); + for (auto& kv : free) { + for (void* p : kv.second) { + hipError_t err = hipHostFree(p); + if (err != hipSuccess) { + MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); + } + } + kv.second.clear(); + } + free.clear(); + } + + private: + void Release(PinnedBuf* b) { + if (b == nullptr) return; + const size_t cap = b->cap; + void* p = b->ptr; + delete b; + + constexpr size_t kMaxCachedPerClass = 8; + std::lock_guard lock(mu); + auto& vec = free[cap]; + if (vec.size() < kMaxCachedPerClass) { + vec.push_back(p); + } else { + hipError_t err = hipHostFree(p); + if (err != hipSuccess) { + MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); + } + } + } + + private: + std::mutex mu; + std::unordered_map> free; +}; + +struct Segment { + uint64_t off{0}; + uint64_t len{0}; +}; + +inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, + uint64_t* outLen) { + if (!outOff || !outLen) return false; + if (segs.empty()) return false; + uint64_t off = segs[0].off; + uint64_t end = off + segs[0].len; + for (size_t i = 1; i < segs.size(); ++i) { + if (segs[i].off != end) return false; + end += segs[i].len; + } + *outOff = off; + *outLen = end - off; + return true; +} + +inline uint64_t SumLens(const std::vector& segs) { + uint64_t total = 0; + for (const auto& s : segs) total += s.len; + return total; +} + +struct InboundStatusEntry { + StatusCode code{StatusCode::INIT}; + std::string msg; +}; + +struct OutboundOpState { + EngineKey peer; + TransferUniqueId id{0}; + bool isRead{false}; + TransferStatus* status{nullptr}; + + MemoryDesc local{}; + std::vector localSegs; + + MemoryDesc remote{}; + std::vector remoteSegs; + + uint64_t expectedRxBytes{0}; + uint64_t rxBytes{0}; + bool completionReceived{false}; + StatusCode completionCode{StatusCode::INIT}; + std::string completionMsg; + + bool gpuCopyPending{false}; + + Clock::time_point startTs{Clock::now()}; +}; + +struct InboundWriteState { + EngineKey peer; + TransferUniqueId id{0}; + + MemoryDesc dst{}; + std::vector dstSegs; + + bool discard{false}; +}; + +struct EarlyWriteState { + uint64_t payloadLen{0}; + std::shared_ptr pinned; + bool complete{false}; +}; + +struct ActiveDataRx { + bool active{false}; + TransferUniqueId id{0}; + uint64_t remaining{0}; + + bool discard{false}; + // CPU scatter target: + void* base{nullptr}; + std::vector segs; + size_t segIdx{0}; + uint64_t segOff{0}; + + // GPU staging target: + bool toGpu{false}; + int gpuDevice{-1}; + void* gpuBase{nullptr}; + std::shared_ptr pinned; + uint64_t pinnedWriteOff{0}; +}; + +struct SendItem { + std::vector header; + std::vector iov; + size_t idx{0}; + size_t off{0}; + int flags{0}; + std::shared_ptr keepalive; + std::function onDone; + + bool Done() const { return idx >= iov.size(); } + + void Advance(size_t n) { + while (n > 0 && idx < iov.size()) { + size_t avail = iov[idx].iov_len - off; + if (n < avail) { + off += n; + n = 0; + break; + } + n -= avail; + idx++; + off = 0; + } + } +}; + +struct Connection { + int fd{-1}; + bool isOutgoing{false}; + bool connecting{false}; + bool helloSent{false}; + bool helloReceived{false}; + tcp::Channel ch{tcp::Channel::CTRL}; + EngineKey peerKey{}; + + // Control parsing buffer (also used for HELLO on both channels). + std::vector inbuf; + + // Data RX state (after HELLO on data channel). + std::array dataHdrBuf{}; + size_t dataHdrGot{0}; + tcp::DataHeaderView curDataHdr{}; + ActiveDataRx rx{}; + + std::deque sendq; +}; + +struct PeerLinks { + int ctrlFd{-1}; + int dataFd{-1}; + bool CtrlUp() const { return ctrlFd >= 0; } + bool DataUp() const { return dataFd >= 0; } +}; + +} // namespace + +class TcpTransport { + public: + TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg) + : myEngKey(std::move(myKey)), engConfig(engCfg), config(cfg) {} + + ~TcpTransport() { Shutdown(); } + + TcpTransport(const TcpTransport&) = delete; + TcpTransport& operator=(const TcpTransport&) = delete; + + void Start() { + if (running.load()) return; + + if (config.numIoThreads != 1) { + MORI_IO_WARN("TCP: numIoThreads={} requested but only 1 is supported; forcing 1", + config.numIoThreads); + } + + epfd = epoll_create1(EPOLL_CLOEXEC); + assert(epfd >= 0); + + // listener + listenFd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); + assert(listenFd >= 0); + + int one = 1; + SetSockOptOrLog(listenFd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one), "SO_REUSEADDR"); + + auto addrOpt = ParseIpv4(engConfig.host.empty() ? std::string("0.0.0.0") : engConfig.host, + engConfig.port); + assert(addrOpt.has_value()); + sockaddr_in addr = addrOpt.value(); + if (bind(listenFd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + MORI_IO_ERROR("TCP: bind {}:{} failed: {}", engConfig.host, engConfig.port, strerror(errno)); + assert(false && "bind failed"); + } + if (listen(listenFd, 256) != 0) { + MORI_IO_ERROR("TCP: listen failed: {}", strerror(errno)); + assert(false && "listen failed"); + } + listenPort = GetBoundPort(listenFd); + MORI_IO_INFO("TCP: listen on {}:{} (port={})", engConfig.host, engConfig.port, listenPort); + + // eventfd for submit queue + wakeFd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + assert(wakeFd >= 0); + + AddEpoll(listenFd, true, false); + AddEpoll(wakeFd, true, false); + + running.store(true); + ioThread = std::thread([this] { this->IoLoop(); }); + } + + void Shutdown() { + bool wasRunning = running.exchange(false); + if (!wasRunning) return; + + // Wake epoll. + if (wakeFd >= 0) { + uint64_t one = 1; + ::write(wakeFd, &one, sizeof(one)); + } + + if (ioThread.joinable()) ioThread.join(); + + // Close all connections. + for (auto& kv : conns) { + CloseConnInternal(kv.second.get()); + } + conns.clear(); + peers.clear(); + + if (listenFd >= 0) { + close(listenFd); + listenFd = -1; + } + if (wakeFd >= 0) { + close(wakeFd); + wakeFd = -1; + } + if (epfd >= 0) { + close(epfd); + epfd = -1; + } + } + + std::optional GetListenPort() const { + if (listenPort == 0) return std::nullopt; + return listenPort; + } + + void RegisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteMu); + remoteEngines[desc.key] = desc; + } + + void DeregisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteMu); + remoteEngines.erase(desc.key); + } + + void RegisterMemory(const MemoryDesc& desc) { + std::lock_guard lock(memMu); + localMems[desc.id] = desc; + } + + void DeregisterMemory(const MemoryDesc& desc) { + std::lock_guard lock(memMu); + localMems.erase(desc.id); + } + + bool PopInboundTransferStatus(const EngineKey& remote, TransferUniqueId id, TransferStatus* status) { + std::lock_guard lock(inboundMu); + auto it = inboundStatus.find(remote); + if (it == inboundStatus.end()) return false; + auto it2 = it->second.find(id); + if (it2 == it->second.end()) return false; + status->Update(it2->second.code, it2->second.msg); + it->second.erase(it2); + return true; + } + + void SubmitReadWrite(const MemoryDesc& local, size_t localOffset, const MemoryDesc& remote, + size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, + bool isRead) { + if (status == nullptr) return; + + if (size == 0) { + status->SetCode(StatusCode::SUCCESS); + return; + } + + if ((localOffset + size) > local.size || (remoteOffset + size) > remote.size) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: offset+size out of range"); + return; + } + + auto op = std::make_unique(); + op->peer = remote.engineKey; + op->id = id; + op->isRead = isRead; + op->status = status; + op->local = local; + op->remote = remote; + op->localSegs = {{static_cast(localOffset), static_cast(size)}}; + op->remoteSegs = {{static_cast(remoteOffset), static_cast(size)}}; + op->expectedRxBytes = isRead ? static_cast(size) : 0; + op->startTs = Clock::now(); + + status->SetCode(StatusCode::IN_PROGRESS); + EnqueueOp(std::move(op)); + } + + void SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& localOffsets, + const MemoryDesc& remote, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) { + if (status == nullptr) return; + + const size_t n = sizes.size(); + if (n == 0) { + status->SetCode(StatusCode::SUCCESS); + return; + } + if (localOffsets.size() != n || remoteOffsets.size() != n) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch vector size mismatch"); + return; + } + + std::vector localSegs; + std::vector remoteSegs; + localSegs.reserve(n); + remoteSegs.reserve(n); + + uint64_t total = 0; + for (size_t i = 0; i < n; ++i) { + const size_t lo = localOffsets[i]; + const size_t ro = remoteOffsets[i]; + const size_t sz = sizes[i]; + if (sz == 0) continue; + if ((lo + sz) > local.size || (ro + sz) > remote.size) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch offset+size out of range"); + return; + } + localSegs.push_back({static_cast(lo), static_cast(sz)}); + remoteSegs.push_back({static_cast(ro), static_cast(sz)}); + total += static_cast(sz); + } + + auto op = std::make_unique(); + op->peer = remote.engineKey; + op->id = id; + op->isRead = isRead; + op->status = status; + op->local = local; + op->remote = remote; + op->localSegs = std::move(localSegs); + op->remoteSegs = std::move(remoteSegs); + op->expectedRxBytes = isRead ? total : 0; + op->startTs = Clock::now(); + + status->SetCode(StatusCode::IN_PROGRESS); + EnqueueOp(std::move(op)); + } + + private: + void EnqueueOp(std::unique_ptr op) { + { + std::lock_guard lock(submitMu); + submitQ.push_back(std::move(op)); + } + uint64_t one = 1; + ::write(wakeFd, &one, sizeof(one)); + } + + void AddEpoll(int fd, bool wantRead, bool wantWrite) { + epoll_event ev{}; + ev.data.fd = fd; + ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev)); + } + + void ModEpoll(int fd, bool wantRead, bool wantWrite) { + epoll_event ev{}; + ev.data.fd = fd; + ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev)); + } + + void DelEpoll(int fd) { epoll_ctl(epfd, EPOLL_CTL_DEL, fd, nullptr); } + + void CloseConnInternal(Connection* c) { + if (c == nullptr) return; + if (c->fd >= 0) { + DelEpoll(c->fd); + shutdown(c->fd, SHUT_RDWR); + close(c->fd); + c->fd = -1; + } + } + + bool PreferOutgoingFor(const EngineKey& peerKey) const { return myEngKey < peerKey; } + + void AssignConnToPeer(Connection* c) { + assert(c && c->helloReceived); + PeerLinks& link = peers[c->peerKey]; + const bool preferOutgoing = PreferOutgoingFor(c->peerKey); + + auto replace_if_needed = [&](int& slotFd) { + if (slotFd < 0) { + slotFd = c->fd; + return; + } + // Collision: choose deterministically by key + direction so both sides keep the same TCP conn. + const int existingFd = slotFd; + const int newFd = c->fd; + Connection* existing = conns[existingFd].get(); + if (!existing) { + slotFd = newFd; + return; + } + const bool keepNew = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); + if (keepNew) { + MORI_IO_WARN("TCP: peer {} channel {} replacing existing fd {} with fd {}", c->peerKey, + static_cast(c->ch), existing->fd, c->fd); + CloseConnInternal(existing); + conns.erase(existingFd); + slotFd = newFd; + } else { + MORI_IO_WARN("TCP: peer {} channel {} dropping duplicate fd {}", c->peerKey, + static_cast(c->ch), c->fd); + CloseConnInternal(c); + conns.erase(newFd); + } + }; + + if (c->ch == tcp::Channel::CTRL) { + replace_if_needed(link.ctrlFd); + } else { + replace_if_needed(link.dataFd); + } + } + + void MaybeDispatchQueuedOps(const EngineKey& peerKey) { + auto it = peers.find(peerKey); + if (it == peers.end()) return; + if (!it->second.CtrlUp() || !it->second.DataUp()) return; + Connection* ctrl = conns[it->second.ctrlFd].get(); + Connection* data = conns[it->second.dataFd].get(); + if (!ctrl || !data) return; + if (!ctrl->helloReceived || !data->helloReceived) return; + + auto qit = waitingOps.find(peerKey); + if (qit == waitingOps.end()) return; + + auto ops = std::move(qit->second); + waitingOps.erase(qit); + MORI_IO_TRACE("TCP: peer {} ready, dispatch {} queued ops", peerKey.c_str(), ops.size()); + for (auto& op : ops) { + DispatchOp(std::move(op)); + } + } + + void EnsurePeerChannels(const EngineKey& peerKey) { + PeerLinks& link = peers[peerKey]; + if (!link.CtrlUp()) ConnectChannel(peerKey, tcp::Channel::CTRL); + if (!link.DataUp()) ConnectChannel(peerKey, tcp::Channel::DATA); + } + + void ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { + EngineDesc desc; + { + std::lock_guard lock(remoteMu); + auto it = remoteEngines.find(peerKey); + if (it == remoteEngines.end()) { + MORI_IO_ERROR("TCP: remote engine {} not registered", peerKey.c_str()); + return; + } + desc = it->second; + } + + auto peerAddrOpt = ParseIpv4(desc.host, static_cast(desc.port)); + if (!peerAddrOpt.has_value()) { + MORI_IO_ERROR("TCP: invalid remote host {}:{}", desc.host, desc.port); + return; + } + sockaddr_in peerAddr = peerAddrOpt.value(); + + int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); + if (fd < 0) { + MORI_IO_ERROR("TCP: socket() failed: {}", strerror(errno)); + return; + } + MORI_IO_TRACE("TCP: connect start peer={} ch={} fd={}", peerKey.c_str(), static_cast(ch), + fd); + + // Bind to our selected host to pin NIC choice when multiple interfaces exist. + if (!engConfig.host.empty()) { + auto localAddrOpt = ParseIpv4(engConfig.host, 0); + if (localAddrOpt.has_value()) { + sockaddr_in localAddr = localAddrOpt.value(); + if (bind(fd, reinterpret_cast(&localAddr), sizeof(localAddr)) != 0) { + MORI_IO_WARN("TCP: bind(local) {} failed: {}", engConfig.host, strerror(errno)); + } + } + } + + int rc = connect(fd, reinterpret_cast(&peerAddr), sizeof(peerAddr)); + bool connecting = false; + if (rc != 0) { + if (errno == EINPROGRESS) { + connecting = true; + } else { + MORI_IO_ERROR("TCP: connect to {}:{} failed: {}", desc.host, desc.port, strerror(errno)); + close(fd); + return; + } + } + + auto conn = std::make_unique(); + conn->fd = fd; + conn->isOutgoing = true; + conn->connecting = connecting; + conn->peerKey = peerKey; + conn->ch = ch; + conn->inbuf.reserve(4096); + + // socket opts + if (ch == tcp::Channel::CTRL) { + ConfigureCtrlSocket(fd, config); + } else { + ConfigureDataSocket(fd, config); + } + + const bool wantWrite = connecting || !conn->sendq.empty(); + AddEpoll(fd, true, wantWrite); + conns[fd] = std::move(conn); + + if (!connecting) { + QueueHello(fd); + ModEpoll(fd, true, true); + } + } + + void QueueHello(int fd) { + Connection* c = conns[fd].get(); + if (!c || c->helloSent) return; + c->helloSent = true; + auto hello = tcp::BuildHello(c->ch, myEngKey); + MORI_IO_TRACE("TCP: queue HELLO fd={} ch={} peer={}", fd, static_cast(c->ch), + c->peerKey.c_str()); + + SendItem item; + item.header = std::move(hello); + item.iov.resize(1); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + c->sendq.push_back(std::move(item)); + } + + void AcceptNew() { + while (true) { + sockaddr_in peer{}; + socklen_t len = sizeof(peer); + int fd = accept4(listenFd, reinterpret_cast(&peer), &len, SOCK_NONBLOCK | SOCK_CLOEXEC); + if (fd < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_WARN("TCP: accept failed: {}", strerror(errno)); + break; + } + MORI_IO_TRACE("TCP: accept fd={}", fd); + + auto conn = std::make_unique(); + conn->fd = fd; + conn->isOutgoing = false; + conn->connecting = false; + conn->helloSent = false; + conn->helloReceived = false; + conn->inbuf.reserve(4096); + + AddEpoll(fd, true, false); + conns[fd] = std::move(conn); + } + } + + void DrainWakeFd() { + uint64_t v = 0; + while (true) { + ssize_t n = ::read(wakeFd, &v, sizeof(v)); + if (n < 0) { + if (IsWouldBlock(errno)) break; + break; + } + if (n == 0) break; + } + + std::deque> ops; + { + std::lock_guard lock(submitMu); + ops.swap(submitQ); + } + + for (auto& op : ops) { + EnsurePeerChannels(op->peer); + if (!IsPeerReady(op->peer)) { + waitingOps[op->peer].push_back(std::move(op)); + continue; + } + DispatchOp(std::move(op)); + } + } + + bool IsPeerReady(const EngineKey& peerKey) { + auto it = peers.find(peerKey); + if (it == peers.end()) return false; + if (!it->second.CtrlUp() || !it->second.DataUp()) return false; + Connection* ctrl = conns[it->second.ctrlFd].get(); + Connection* data = conns[it->second.dataFd].get(); + return ctrl && data && ctrl->helloReceived && data->helloReceived; + } + + void DispatchOp(std::unique_ptr op) { + if (!op) { + MORI_IO_ERROR("TCP: DispatchOp got null op"); + return; + } + const EngineKey peerKey = op->peer; + auto it = peers.find(peerKey); + if (it == peers.end() || !it->second.CtrlUp() || !it->second.DataUp()) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer not connected"); + return; + } + Connection* ctrl = conns[it->second.ctrlFd].get(); + Connection* data = conns[it->second.dataFd].get(); + if (!ctrl || !data) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer connection missing"); + return; + } + + const TransferUniqueId opId = op->id; + auto [itIns, inserted] = pendingOutbound.emplace(opId, std::move(op)); + if (!inserted) { + MORI_IO_ERROR("TCP: duplicate outbound op id={} for peer={}", opId, peerKey.c_str()); + itIns->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: duplicate op id"); + pendingOutbound.erase(itIns); + return; + } + OutboundOpState* st = itIns->second.get(); + if (!st) { + MORI_IO_ERROR("TCP: failed to store outbound op id={} (nullptr)", opId); + pendingOutbound.erase(itIns); + return; + } + MORI_IO_TRACE("TCP: dispatch op id={} peer={} isRead={} segs={}", st->id, peerKey.c_str(), + st->isRead, st->localSegs.size()); + + // CTRL request + std::vector ctrlFrame; + if (st->localSegs.size() == 1) { + if (st->isRead) { + ctrlFrame = + tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, st->remoteSegs[0].len); + } else { + ctrlFrame = + tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, st->remoteSegs[0].len); + } + } else { + std::vector roffs; + std::vector szs; + roffs.reserve(st->remoteSegs.size()); + szs.reserve(st->remoteSegs.size()); + for (const auto& s : st->remoteSegs) { + roffs.push_back(s.off); + szs.push_back(s.len); + } + if (st->isRead) { + ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs); + } else { + ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs); + } + } + + QueueSend(ctrl->fd, std::move(ctrlFrame)); + + // DATA payload (writes only) + if (!st->isRead) { + QueueDataSendForWrite(data, *st); + } + + UpdateWriteInterest(ctrl->fd); + UpdateWriteInterest(data->fd); + } + + void QueueSend(int fd, std::vector bytes, std::function onDone = nullptr) { + Connection* c = conns[fd].get(); + if (!c) return; + SendItem item; + item.header = std::move(bytes); + item.iov.resize(1); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.onDone = std::move(onDone); + c->sendq.push_back(std::move(item)); + } + + void QueueDataSendForWrite(Connection* data, OutboundOpState& st) { + const uint64_t total = SumLens(st.localSegs); + auto hdr = tcp::BuildDataHeader(st.id, total, 0); + + if (st.local.loc == MemoryLocationType::GPU) { + QueueGpuToNetSend(data, st.local, st.localSegs, std::move(hdr), /*onDone*/ nullptr); + return; + } + + SendItem item; + item.header = std::move(hdr); + item.iov.reserve(1 + st.localSegs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + + uint8_t* base = reinterpret_cast(st.local.data); + for (const auto& s : st.localSegs) { + item.iov.push_back({base + s.off, static_cast(s.len)}); + } + data->sendq.push_back(std::move(item)); + } + + void QueueGpuToNetSend(Connection* data, const MemoryDesc& src, const std::vector& srcSegs, + std::vector dataHdr, std::function onDone) { + const uint64_t total = SumLens(srcSegs); + auto pinned = staging.Acquire(static_cast(total)); + if (!pinned) { + MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); + return; + } + + hipStream_t stream = streamPool.GetNextStream(src.deviceId); + hipEvent_t ev = eventPool.GetEvent(src.deviceId); + if (stream == nullptr || ev == nullptr) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send"); + if (ev) eventPool.PutEvent(ev, src.deviceId); + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); + uint8_t* dst = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : srcSegs) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); + HIP_RUNTIME_CHECK( + hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + gpuTasks.push_back({src.deviceId, ev, [this, dataFd = data->fd, pinned, hdr = std::move(dataHdr), + total, onDone = std::move(onDone)]() mutable { + Connection* c = conns[dataFd].get(); + if (!c || c->fd < 0) return; + SendItem item; + item.header = std::move(hdr); + item.iov.reserve(2); + item.iov.push_back({item.header.data(), item.header.size()}); + item.iov.push_back({pinned->ptr, static_cast(total)}); + item.keepalive = pinned; + item.onDone = std::move(onDone); + c->sendq.push_back(std::move(item)); + UpdateWriteInterest(dataFd); + }}); + } + + struct GpuTask { + int deviceId{-1}; + hipEvent_t ev{nullptr}; + std::function onReady; + }; + + void PollGpuTasks() { + for (auto it = gpuTasks.begin(); it != gpuTasks.end();) { + hipError_t st = hipEventQuery(it->ev); + if (st == hipSuccess) { + eventPool.PutEvent(it->ev, it->deviceId); + if (it->onReady) it->onReady(); + it = gpuTasks.erase(it); + } else if (st == hipErrorNotReady) { + ++it; + } else { + MORI_IO_ERROR("TCP: hipEventQuery failed: {}", hipGetErrorString(st)); + eventPool.PutEvent(it->ev, it->deviceId); + it = gpuTasks.erase(it); + } + } + } + + void UpdateWriteInterest(int fd) { + auto it = conns.find(fd); + if (it == conns.end()) return; + Connection* c = it->second.get(); + if (!c || c->fd < 0) return; + + // With edge-triggered epoll, don't rely solely on EPOLLOUT transitions to make progress. + // If the socket is already writable, we may not get another EPOLLOUT edge after enqueueing. + if (!c->connecting && !c->sendq.empty()) { + FlushSend(c); + it = conns.find(fd); + if (it == conns.end()) return; + c = it->second.get(); + if (!c || c->fd < 0) return; + } + + bool wantWrite = c->connecting || !c->sendq.empty(); + ModEpoll(fd, true, wantWrite); + } + + void HandleConnWritable(Connection* c) { + if (c->connecting) { + int err = 0; + socklen_t len = sizeof(err); + if (getsockopt(c->fd, SOL_SOCKET, SO_ERROR, &err, &len) != 0 || err != 0) { + MORI_IO_ERROR("TCP: connect failed fd {}: {}", c->fd, strerror(err == 0 ? errno : err)); + ClosePeerByFd(c->fd); + return; + } + c->connecting = false; + QueueHello(c->fd); + } + UpdateWriteInterest(c->fd); + } + + void FlushSend(Connection* c) { + constexpr size_t kMaxIov = 64; + while (!c->sendq.empty()) { + SendItem& item = c->sendq.front(); + if (item.Done()) { + auto cb = std::move(item.onDone); + c->sendq.pop_front(); + if (cb) cb(); + continue; + } + + iovec iov[kMaxIov]; + size_t cnt = 0; + for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { + iov[cnt] = item.iov[i]; + if (i == item.idx && item.off > 0) { + iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; + iov[cnt].iov_len -= item.off; + } + cnt++; + } + msghdr msg{}; + msg.msg_iov = iov; + msg.msg_iovlen = cnt; + ssize_t n = sendmsg(c->fd, &msg, MSG_NOSIGNAL | item.flags); + if (n < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_ERROR("TCP: sendmsg fd {} failed: {}", c->fd, strerror(errno)); + ClosePeerByFd(c->fd); + break; + } + if (n == 0) break; + item.Advance(static_cast(n)); + } + } + + void ClosePeerByFd(int fd) { + MORI_IO_TRACE("TCP: close fd={}", fd); + // Find which peer+channel this fd belongs to, close both channels and fail pending ops. + EngineKey peer{}; + for (auto& kv : peers) { + if (kv.second.ctrlFd == fd || kv.second.dataFd == fd) { + peer = kv.first; + break; + } + } + + Connection* c = conns[fd].get(); + if (c) { + CloseConnInternal(c); + conns.erase(fd); + } + + if (!peer.empty()) { + PeerLinks& link = peers[peer]; + if (link.ctrlFd >= 0) { + Connection* cc = conns[link.ctrlFd].get(); + if (cc) { + CloseConnInternal(cc); + conns.erase(link.ctrlFd); + } + link.ctrlFd = -1; + } + if (link.dataFd >= 0) { + Connection* dc = conns[link.dataFd].get(); + if (dc) { + CloseConnInternal(dc); + conns.erase(link.dataFd); + } + link.dataFd = -1; + } + peers.erase(peer); + + FailPendingOpsForPeer(peer, "TCP: connection lost"); + } + } + + void FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg) { + for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + if (it->second->peer == peer) { + it->second->status->Update(StatusCode::ERR_BAD_STATE, msg); + it = pendingOutbound.erase(it); + } else { + ++it; + } + } + waitingOps.erase(peer); + inboundWrites.erase(peer); + earlyWrites.erase(peer); + } + + void HandleCtrlReadable(Connection* c) { + while (true) { + uint8_t tmp[4096]; + ssize_t n = ::recv(c->fd, tmp, sizeof(tmp), 0); + if (n < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_ERROR("TCP: recv(ctrl) fd {} failed: {}", c->fd, strerror(errno)); + ClosePeerByFd(c->fd); + return; + } + if (n == 0) { + ClosePeerByFd(c->fd); + return; + } + c->inbuf.insert(c->inbuf.end(), tmp, tmp + n); + } + + // Parse frames. + while (true) { + tcp::CtrlHeaderView hv; + if (!tcp::TryParseCtrlHeader(c->inbuf.data(), c->inbuf.size(), &hv)) { + if (c->inbuf.size() >= tcp::kCtrlHeaderSize) { + MORI_IO_ERROR("TCP: bad ctrl header on fd {}, closing", c->fd); + ClosePeerByFd(c->fd); + } + break; + } + if (c->inbuf.size() < tcp::kCtrlHeaderSize + hv.bodyLen) break; + + const uint8_t* body = c->inbuf.data() + tcp::kCtrlHeaderSize; + HandleCtrlFrame(c, hv.type, body, hv.bodyLen); + + c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + tcp::kCtrlHeaderSize + hv.bodyLen); + + // Data channel transitions to a different framing after HELLO. Any bytes already read into + // inbuf may include the beginning of the data stream; handle them immediately to avoid + // edge-triggered epoll stalls. + if (c->helloReceived && c->ch == tcp::Channel::DATA) { + HandleDataReadable(c); + return; + } + } + } + + void HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const uint8_t* body, size_t len) { + if (type == tcp::CtrlMsgType::HELLO) { + HandleHello(c, body, len); + return; + } + + if (!c->helloReceived) { + MORI_IO_WARN("TCP: received ctrl message before HELLO, dropping"); + return; + } + + switch (type) { + case tcp::CtrlMsgType::WRITE_REQ: + HandleWriteReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::READ_REQ: + HandleReadReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::BATCH_WRITE_REQ: + HandleBatchWriteReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::BATCH_READ_REQ: + HandleBatchReadReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::COMPLETION: + HandleCompletion(c->peerKey, body, len); + break; + default: + MORI_IO_WARN("TCP: unknown ctrl msg type {}", static_cast(type)); + } + } + + void HandleHello(Connection* c, const uint8_t* body, size_t len) { + if (len < 1 + 4) { + MORI_IO_WARN("TCP: bad HELLO len {}", len); + ClosePeerByFd(c->fd); + return; + } + size_t off = 0; + const uint8_t chRaw = body[off++]; + uint32_t keyLen = 0; + if (!tcp::ReadU32BE(body, len, &off, &keyLen)) { + ClosePeerByFd(c->fd); + return; + } + if (off + keyLen > len) { + ClosePeerByFd(c->fd); + return; + } + EngineKey peerKey(reinterpret_cast(body + off), keyLen); + off += keyLen; + + c->peerKey = peerKey; + c->ch = (chRaw == static_cast(tcp::Channel::DATA)) ? tcp::Channel::DATA + : tcp::Channel::CTRL; + c->helloReceived = true; + MORI_IO_TRACE("TCP: recv HELLO fd={} peer={} ch={} outgoing={}", c->fd, c->peerKey.c_str(), + static_cast(c->ch), c->isOutgoing); + + // Respond with our hello if we haven't sent it yet. + if (!c->helloSent) { + QueueHello(c->fd); + UpdateWriteInterest(c->fd); + } + + // Post-handshake: configure socket by channel type. + if (c->ch == tcp::Channel::CTRL) { + ConfigureCtrlSocket(c->fd, config); + } else { + ConfigureDataSocket(c->fd, config); + } + + AssignConnToPeer(c); + MaybeDispatchQueuedOps(peerKey); + } + + std::optional LookupLocalMem(MemoryUniqueId id) { + std::lock_guard lock(memMu); + auto it = localMems.find(id); + if (it == localMems.end()) return std::nullopt; + return it->second; + } + + void RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, + const std::string& msg) { + std::lock_guard lock(inboundMu); + inboundStatus[peer][id] = InboundStatusEntry{code, msg}; + } + + Connection* PeerCtrl(const EngineKey& peer) { + auto it = peers.find(peer); + if (it == peers.end()) return nullptr; + if (!it->second.CtrlUp()) return nullptr; + return conns[it->second.ctrlFd].get(); + } + + Connection* PeerData(const EngineKey& peer) { + auto it = peers.find(peer); + if (it == peers.end()) return nullptr; + if (!it->second.DataUp()) return nullptr; + return conns[it->second.dataFd].get(); + } + + void FinalizeBufferedWrite(const EngineKey& peer, TransferUniqueId opId, InboundWriteState ws, + uint64_t payloadLen, const std::shared_ptr& pinned) { + auto ctrl = PeerCtrl(peer); + if (!ctrl) return; + + if (ws.discard) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), + "TCP: write discarded")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); + return; + } + if (!pinned) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: missing pinned staging (early write)")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging"); + return; + } + + const uint64_t expected = SumLens(ws.dstSegs); + if (expected != payloadLen) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: early payload length mismatch")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: early payload length mismatch"); + return; + } + + if (ws.dst.loc == MemoryLocationType::GPU) { + hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); + hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); + if (stream == nullptr || ev == nullptr) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: failed to get HIP stream/event")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: failed to get HIP stream/event"); + if (ev) eventPool.PutEvent(ev, ws.dst.deviceId); + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); + uint8_t* src = reinterpret_cast(pinned->ptr); + const uint64_t total = payloadLen; + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { + void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : ws.dstSegs) { + void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + gpuTasks.push_back({ws.dst.deviceId, ev, [this, peer, opId, ctrlFd = ctrl->fd, pinned]() { + QueueSend(ctrlFd, + tcp::BuildCompletion(opId, + static_cast(StatusCode::SUCCESS), + "")); + UpdateWriteInterest(ctrlFd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + }}); + return; + } + + // CPU destination. + uint8_t* src = reinterpret_cast(pinned->ptr); + uint8_t* dstBase = reinterpret_cast(ws.dst.data); + const uint64_t total = payloadLen; + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { + std::memcpy(dstBase + spanOff, src, static_cast(total)); + } else { + uint64_t off = 0; + for (const auto& s : ws.dstSegs) { + std::memcpy(dstBase + s.off, src + off, static_cast(s.len)); + off += s.len; + } + } + + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + } + + void HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { + size_t off = 0; + uint64_t opId = 0; + uint32_t memId = 0; + uint64_t remoteOff = 0; + uint64_t size = 0; + if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || + !tcp::ReadU64BE(body, len, &off, &remoteOff) || !tcp::ReadU64BE(body, len, &off, &size)) { + MORI_IO_WARN("TCP: malformed WRITE_REQ"); + return; + } + + auto memOpt = LookupLocalMem(memId); + InboundWriteState ws; + ws.peer = peer; + ws.id = opId; + ws.discard = true; + + if (memOpt.has_value()) { + ws.dst = memOpt.value(); + if (remoteOff + size <= ws.dst.size) { + ws.discard = false; + ws.dstSegs = {{remoteOff, size}}; + } + } + + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt != ewPeerIt->second.end()) { + EarlyWriteState early = std::move(ewIt->second); + ewPeerIt->second.erase(ewIt); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + + if (early.complete) { + FinalizeBufferedWrite(peer, opId, std::move(ws), early.payloadLen, early.pinned); + return; + } + // Data is still in flight; keep the write state so data RX can finalize it. + inboundWrites[peer][opId] = std::move(ws); + return; + } + } + + inboundWrites[peer][opId] = std::move(ws); + } + + void HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { + size_t off = 0; + uint64_t opId = 0; + uint32_t memId = 0; + uint32_t n = 0; + if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || + !tcp::ReadU32BE(body, len, &off, &n)) { + MORI_IO_WARN("TCP: malformed BATCH_WRITE_REQ"); + return; + } + + auto memOpt = LookupLocalMem(memId); + InboundWriteState ws; + ws.peer = peer; + ws.id = opId; + ws.discard = true; + + if (memOpt.has_value()) { + ws.dst = memOpt.value(); + ws.dstSegs.reserve(n); + bool ok = true; + for (uint32_t i = 0; i < n; ++i) { + uint64_t ro = 0, sz = 0; + if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { + ok = false; + break; + } + if (ro + sz > ws.dst.size) ok = false; + if (sz > 0) ws.dstSegs.push_back({ro, sz}); + } + if (ok) ws.discard = false; + } + + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt != ewPeerIt->second.end()) { + EarlyWriteState early = std::move(ewIt->second); + ewPeerIt->second.erase(ewIt); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + + if (early.complete) { + FinalizeBufferedWrite(peer, opId, std::move(ws), early.payloadLen, early.pinned); + return; + } + inboundWrites[peer][opId] = std::move(ws); + return; + } + } + + inboundWrites[peer][opId] = std::move(ws); + } + + void HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { + size_t off = 0; + uint64_t opId = 0; + uint32_t memId = 0; + uint64_t srcOff = 0; + uint64_t size = 0; + if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || + !tcp::ReadU64BE(body, len, &off, &srcOff) || !tcp::ReadU64BE(body, len, &off, &size)) { + MORI_IO_WARN("TCP: malformed READ_REQ"); + return; + } + + auto memOpt = LookupLocalMem(memId); + if (!memOpt.has_value() || (srcOff + size > memOpt->size)) { + // No data; completion with error. + auto ctrl = PeerCtrl(peer); + if (ctrl) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_NOT_FOUND), + "TCP: remote mem not found/out of range")); + UpdateWriteInterest(ctrl->fd); + } + RecordInboundStatus(peer, opId, StatusCode::ERR_NOT_FOUND, "TCP: read mem not found"); + return; + } + + MemoryDesc src = memOpt.value(); + std::vector segs = {{srcOff, size}}; + QueueDataSendForRead(peer, opId, src, segs); + } + + void HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { + size_t off = 0; + uint64_t opId = 0; + uint32_t memId = 0; + uint32_t n = 0; + if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || + !tcp::ReadU32BE(body, len, &off, &n)) { + MORI_IO_WARN("TCP: malformed BATCH_READ_REQ"); + return; + } + + auto memOpt = LookupLocalMem(memId); + if (!memOpt.has_value()) { + auto ctrl = PeerCtrl(peer); + if (ctrl) { + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_NOT_FOUND), + "TCP: remote mem not found")); + UpdateWriteInterest(ctrl->fd); + } + RecordInboundStatus(peer, opId, StatusCode::ERR_NOT_FOUND, "TCP: read mem not found"); + return; + } + + MemoryDesc src = memOpt.value(); + std::vector segs; + segs.reserve(n); + bool ok = true; + for (uint32_t i = 0; i < n; ++i) { + uint64_t ro = 0, sz = 0; + if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { + ok = false; + break; + } + if (ro + sz > src.size) ok = false; + if (sz > 0) segs.push_back({ro, sz}); + } + if (!ok) { + auto ctrl = PeerCtrl(peer); + if (ctrl) { + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), + "TCP: batch read out of range")); + UpdateWriteInterest(ctrl->fd); + } + RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: batch read bad args"); + return; + } + + QueueDataSendForRead(peer, opId, src, segs); + } + + void QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, + const std::vector& srcSegs) { + Connection* data = PeerData(peer); + Connection* ctrl = PeerCtrl(peer); + if (!data || !ctrl) return; + + const uint64_t total = SumLens(srcSegs); + auto hdr = tcp::BuildDataHeader(opId, total, 0); + auto onDone = [this, peer, opId, ctrlFd = ctrl->fd]() { + QueueSend(ctrlFd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); + UpdateWriteInterest(ctrlFd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + }; + + if (src.loc == MemoryLocationType::GPU) { + QueueGpuToNetSend(data, src, srcSegs, std::move(hdr), std::move(onDone)); + return; + } + + SendItem item; + item.header = std::move(hdr); + item.iov.reserve(1 + srcSegs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + uint8_t* base = reinterpret_cast(src.data); + for (const auto& s : srcSegs) { + item.iov.push_back({base + s.off, static_cast(s.len)}); + } + item.onDone = std::move(onDone); + data->sendq.push_back(std::move(item)); + UpdateWriteInterest(data->fd); + } + + void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { + size_t off = 0; + uint64_t opId = 0; + uint32_t code = 0; + uint32_t msgLen = 0; + if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &code) || + !tcp::ReadU32BE(body, len, &off, &msgLen)) { + MORI_IO_WARN("TCP: malformed COMPLETION"); + return; + } + if (off + msgLen > len) { + MORI_IO_WARN("TCP: malformed COMPLETION msg len"); + return; + } + std::string msg(reinterpret_cast(body + off), msgLen); + + auto it = pendingOutbound.find(opId); + if (it == pendingOutbound.end()) return; + OutboundOpState& st = *it->second; + st.completionReceived = true; + st.completionCode = static_cast(code); + st.completionMsg = std::move(msg); + + // Fast fail on remote error. + if (st.completionCode != StatusCode::SUCCESS) { + st.status->Update(st.completionCode, st.completionMsg); + pendingOutbound.erase(it); + return; + } + + MaybeCompleteOutbound(st); + } + + void MaybeCompleteOutbound(OutboundOpState& st) { + if (!st.completionReceived) return; + if (st.isRead) { + if (st.rxBytes != st.expectedRxBytes) return; + if (st.gpuCopyPending) return; + } + st.status->Update(StatusCode::SUCCESS, ""); + pendingOutbound.erase(st.id); + } + + void ConsumeBufferedData(Connection* c) { + while (!c->inbuf.empty()) { + if (!c->rx.active) { + const size_t need = tcp::kDataHeaderSize - c->dataHdrGot; + const size_t take = std::min(need, c->inbuf.size()); + std::memcpy(c->dataHdrBuf.data() + c->dataHdrGot, c->inbuf.data(), take); + c->dataHdrGot += take; + c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + take); + if (c->dataHdrGot < tcp::kDataHeaderSize) return; + + tcp::DataHeaderView hv; + if (!tcp::TryParseDataHeader(c->dataHdrBuf.data(), tcp::kDataHeaderSize, &hv)) { + MORI_IO_ERROR("TCP: bad data header during handoff, closing"); + ClosePeerByFd(c->fd); + return; + } + c->dataHdrGot = 0; + BeginDataRx(c, hv.opId, hv.payloadLen); + continue; + } + + if (c->rx.remaining == 0) { + FinishDataRx(c); + continue; + } + + const size_t take = static_cast(std::min(c->rx.remaining, c->inbuf.size())); + if (take == 0) return; + const uint8_t* src = c->inbuf.data(); + + if (c->rx.discard) { + // nothing + } else if (c->rx.toGpu) { + if (!c->rx.pinned) { + MORI_IO_ERROR("TCP: missing pinned buffer for GPU recv"); + ClosePeerByFd(c->fd); + return; + } + uint8_t* dst = reinterpret_cast(c->rx.pinned->ptr) + c->rx.pinnedWriteOff; + std::memcpy(dst, src, take); + c->rx.pinnedWriteOff += take; + } else { + size_t copied = 0; + while (copied < take) { + if (c->rx.segIdx >= c->rx.segs.size()) { + MORI_IO_ERROR("TCP: cpu scatter overflow during buffered consume"); + ClosePeerByFd(c->fd); + return; + } + Segment& seg = c->rx.segs[c->rx.segIdx]; + const uint64_t segRemain = seg.len - c->rx.segOff; + const size_t chunk = + static_cast(std::min(segRemain, static_cast(take - copied))); + uint8_t* dst = reinterpret_cast(c->rx.base) + seg.off + c->rx.segOff; + std::memcpy(dst, src + copied, chunk); + c->rx.segOff += chunk; + copied += chunk; + if (c->rx.segOff >= seg.len) { + c->rx.segIdx++; + c->rx.segOff = 0; + } + } + } + + c->rx.remaining -= static_cast(take); + c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + take); + } + } + + void HandleDataReadable(Connection* c) { + // If we haven't received HELLO on this channel, treat data as ctrl for handshake. + if (!c->helloReceived) { + HandleCtrlReadable(c); + return; + } + + if (!c->inbuf.empty()) { + ConsumeBufferedData(c); + if (c->fd < 0) return; + } + + while (true) { + if (!c->rx.active) { + // Need header. + while (c->dataHdrGot < tcp::kDataHeaderSize) { + ssize_t n = ::recv(c->fd, c->dataHdrBuf.data() + c->dataHdrGot, + tcp::kDataHeaderSize - c->dataHdrGot, 0); + if (n < 0) { + if (IsWouldBlock(errno)) return; + MORI_IO_ERROR("TCP: recv(data hdr) failed: {}", strerror(errno)); + ClosePeerByFd(c->fd); + return; + } + if (n == 0) { + ClosePeerByFd(c->fd); + return; + } + c->dataHdrGot += static_cast(n); + } + tcp::DataHeaderView hv; + if (!tcp::TryParseDataHeader(c->dataHdrBuf.data(), tcp::kDataHeaderSize, &hv)) { + MORI_IO_ERROR("TCP: bad data header, closing"); + ClosePeerByFd(c->fd); + return; + } + c->dataHdrGot = 0; + BeginDataRx(c, hv.opId, hv.payloadLen); + continue; + } + + if (c->rx.remaining == 0) { + FinishDataRx(c); + continue; + } + + ssize_t n = 0; + if (c->rx.discard) { + uint8_t tmp[8192]; + const size_t want = std::min(sizeof(tmp), c->rx.remaining); + n = ::recv(c->fd, tmp, want, 0); + } else if (c->rx.toGpu) { + uint8_t* dst = reinterpret_cast(c->rx.pinned->ptr) + c->rx.pinnedWriteOff; + const size_t want = static_cast(std::min(c->rx.remaining, 1ULL << 20)); + n = ::recv(c->fd, dst, want, 0); + } else { + // CPU scatter: recv into current segment. + Segment& seg = c->rx.segs[c->rx.segIdx]; + uint8_t* dst = reinterpret_cast(c->rx.base) + seg.off + c->rx.segOff; + uint64_t segRemain = seg.len - c->rx.segOff; + const size_t want = static_cast(std::min(c->rx.remaining, segRemain)); + n = ::recv(c->fd, dst, want, 0); + } + + if (n < 0) { + if (IsWouldBlock(errno)) return; + MORI_IO_ERROR("TCP: recv(data) failed: {}", strerror(errno)); + ClosePeerByFd(c->fd); + return; + } + if (n == 0) { + ClosePeerByFd(c->fd); + return; + } + + const uint64_t got = static_cast(n); + c->rx.remaining -= got; + if (c->rx.discard) { + continue; + } + if (c->rx.toGpu) { + c->rx.pinnedWriteOff += got; + continue; + } + // CPU seg advance + c->rx.segOff += got; + Segment& seg = c->rx.segs[c->rx.segIdx]; + if (c->rx.segOff >= seg.len) { + c->rx.segIdx++; + c->rx.segOff = 0; + } + } + } + + void BeginDataRx(Connection* c, TransferUniqueId opId, uint64_t payloadLen) { + c->rx = ActiveDataRx{}; + c->rx.active = true; + c->rx.id = opId; + c->rx.remaining = payloadLen; + + // 1) inbound write: peer->me, stored by peer key. + auto iwIt = inboundWrites.find(c->peerKey); + if (iwIt != inboundWrites.end()) { + auto it = iwIt->second.find(opId); + if (it != iwIt->second.end()) { + InboundWriteState& ws = it->second; + if (!ws.discard) { + const uint64_t expected = SumLens(ws.dstSegs); + if (expected != payloadLen) { + MORI_IO_WARN("TCP: inbound write op {} payloadLen mismatch expected={} got={}", opId, + expected, payloadLen); + ws.discard = true; + } + } + c->rx.discard = ws.discard; + if (!ws.discard) { + if (ws.dst.loc == MemoryLocationType::GPU) { + c->rx.toGpu = true; + c->rx.gpuDevice = ws.dst.deviceId; + c->rx.gpuBase = reinterpret_cast(ws.dst.data); + c->rx.segs = ws.dstSegs; + const uint64_t total = SumLens(ws.dstSegs); + c->rx.pinned = staging.Acquire(static_cast(total)); + c->rx.pinnedWriteOff = 0; + } else { + c->rx.toGpu = false; + c->rx.base = reinterpret_cast(ws.dst.data); + c->rx.segs = ws.dstSegs; + c->rx.segIdx = 0; + c->rx.segOff = 0; + } + } + return; + } + } + + // 2) outbound read response: peer->me for my pending outbound read. + auto obIt = pendingOutbound.find(opId); + if (obIt != pendingOutbound.end()) { + OutboundOpState& st = *obIt->second; + if (!st.isRead) { + c->rx.discard = true; + return; + } + if (payloadLen != st.expectedRxBytes) { + MORI_IO_ERROR("TCP: outbound read op {} payloadLen mismatch expected={} got={}", opId, + st.expectedRxBytes, payloadLen); + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: read payload length mismatch"); + pendingOutbound.erase(obIt); + c->rx.discard = true; + return; + } + c->rx.discard = false; + if (st.local.loc == MemoryLocationType::GPU) { + c->rx.toGpu = true; + c->rx.gpuDevice = st.local.deviceId; + c->rx.gpuBase = reinterpret_cast(st.local.data); + c->rx.segs = st.localSegs; + c->rx.pinned = staging.Acquire(static_cast(payloadLen)); + c->rx.pinnedWriteOff = 0; + } else { + c->rx.toGpu = false; + c->rx.base = reinterpret_cast(st.local.data); + c->rx.segs = st.localSegs; + c->rx.segIdx = 0; + c->rx.segOff = 0; + } + return; + } + + // Control+data are on separate TCP connections; data may arrive before its CTRL write request. + // Buffer it in pinned host memory, and finalize once CTRL arrives. + if (payloadLen > static_cast(std::numeric_limits::max())) { + MORI_IO_WARN("TCP: data payloadLen too large for op {} from peer {}, discarding", opId, + c->peerKey.c_str()); + c->rx.discard = true; + return; + } + auto& perPeer = earlyWrites[c->peerKey]; + if (perPeer.find(opId) != perPeer.end()) { + MORI_IO_WARN("TCP: duplicate early data for op {} from peer {}, discarding", opId, + c->peerKey.c_str()); + c->rx.discard = true; + return; + } + + const size_t alloc = payloadLen == 0 ? 1 : static_cast(payloadLen); + auto pinned = staging.Acquire(alloc); + if (!pinned) { + MORI_IO_WARN("TCP: failed to allocate pinned buffer for early data op {} from peer {}, discarding", + opId, c->peerKey.c_str()); + c->rx.discard = true; + return; + } + + perPeer.emplace(opId, EarlyWriteState{payloadLen, pinned, false}); + c->rx.discard = false; + c->rx.toGpu = true; // receive into pinned + c->rx.pinned = pinned; + c->rx.pinnedWriteOff = 0; + } + + void FinishDataRx(Connection* c) { + const TransferUniqueId opId = c->rx.id; + const EngineKey peer = c->peerKey; + + // Case A: inbound write complete -> copy to GPU if needed, then send completion. + auto iwIt = inboundWrites.find(peer); + if (iwIt != inboundWrites.end()) { + auto it = iwIt->second.find(opId); + if (it != iwIt->second.end()) { + InboundWriteState ws = std::move(it->second); + iwIt->second.erase(it); + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + ewPeerIt->second.erase(opId); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + } + + if (ws.discard) { + auto ctrl = PeerCtrl(peer); + if (ctrl) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), + "TCP: write discarded")); + UpdateWriteInterest(ctrl->fd); + } + RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); + c->rx = ActiveDataRx{}; + return; + } + + auto ctrl = PeerCtrl(peer); + if (!ctrl) { + c->rx = ActiveDataRx{}; + return; + } + + if (ws.dst.loc == MemoryLocationType::GPU) { + hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); + hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); + if (stream == nullptr || ev == nullptr || !c->rx.pinned) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: GPU staging missing")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: GPU staging missing"); + c->rx = ActiveDataRx{}; + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); + uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); + const uint64_t total = SumLens(ws.dstSegs); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { + void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : ws.dstSegs) { + void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); + HIP_RUNTIME_CHECK( + hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + auto pinned = c->rx.pinned; + c->rx = ActiveDataRx{}; + + gpuTasks.push_back({ws.dst.deviceId, ev, [this, peer, opId, ctrlFd = ctrl->fd, pinned]() { + QueueSend(ctrlFd, + tcp::BuildCompletion(opId, + static_cast(StatusCode::SUCCESS), + "")); + UpdateWriteInterest(ctrlFd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + }}); + return; + } + + if (c->rx.pinned) { + uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); + uint8_t* dstBase = reinterpret_cast(ws.dst.data); + const uint64_t total = SumLens(ws.dstSegs); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { + std::memcpy(dstBase + spanOff, src, static_cast(total)); + } else { + uint64_t off = 0; + for (const auto& s : ws.dstSegs) { + std::memcpy(dstBase + s.off, src + off, static_cast(s.len)); + off += s.len; + } + } + } + + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + c->rx = ActiveDataRx{}; + return; + } + } + + // Case B: outbound read response received -> copy to GPU if needed and update status. + auto obIt = pendingOutbound.find(opId); + if (obIt != pendingOutbound.end()) { + OutboundOpState& st = *obIt->second; + st.rxBytes = st.expectedRxBytes; + if (st.local.loc == MemoryLocationType::GPU) { + hipStream_t stream = streamPool.GetNextStream(st.local.deviceId); + hipEvent_t ev = eventPool.GetEvent(st.local.deviceId); + if (stream == nullptr || ev == nullptr || !c->rx.pinned) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: GPU staging missing (read)"); + pendingOutbound.erase(obIt); + c->rx = ActiveDataRx{}; + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(st.local.deviceId)); + uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); + st.gpuCopyPending = true; + const uint64_t total = st.expectedRxBytes; + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(st.localSegs, &spanOff, &spanLen) && spanLen == total) { + void* gpuPtr = reinterpret_cast(st.local.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : st.localSegs) { + void* gpuPtr = reinterpret_cast(st.local.data + s.off); + HIP_RUNTIME_CHECK( + hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + auto pinned = c->rx.pinned; + c->rx = ActiveDataRx{}; + + gpuTasks.push_back({st.local.deviceId, ev, [this, opId, pinned]() { + auto it = pendingOutbound.find(opId); + if (it == pendingOutbound.end()) return; + it->second->gpuCopyPending = false; + MaybeCompleteOutbound(*it->second); + }}); + return; + } + + c->rx = ActiveDataRx{}; + MaybeCompleteOutbound(st); + return; + } + + // Case C: early-arrived write payload (data before CTRL). Mark it complete and finalize once + // the corresponding CTRL write request arrives. + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt != ewPeerIt->second.end()) { + ewIt->second.complete = true; + c->rx = ActiveDataRx{}; + return; + } + } + + c->rx = ActiveDataRx{}; + } + + void HandleReadable(Connection* c) { + if (!c->helloReceived) { + HandleCtrlReadable(c); + return; + } + if (c->ch == tcp::Channel::CTRL) { + HandleCtrlReadable(c); + } else { + HandleDataReadable(c); + } + } + + void ScanTimeouts() { + if (config.opTimeoutMs <= 0) return; + const auto now = Clock::now(); + const auto timeout = std::chrono::milliseconds(config.opTimeoutMs); + for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + if ((now - it->second->startTs) > timeout) { + it->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: op timeout"); + it = pendingOutbound.erase(it); + } else { + ++it; + } + } + } + + void IoLoop() { + constexpr int kMaxEvents = 128; + epoll_event events[kMaxEvents]; + + while (running.load()) { + PollGpuTasks(); + ScanTimeouts(); + + // When GPU staging is in flight, prefer low-latency polling so we can arm the next + // network send / completion promptly after HIP events complete. + const int timeoutMs = gpuTasks.empty() ? 5 /*ms*/ : 0 /*busy*/; + int nfds = epoll_wait(epfd, events, kMaxEvents, timeoutMs); + if (nfds < 0) { + if (errno == EINTR) continue; + MORI_IO_ERROR("TCP: epoll_wait failed: {}", strerror(errno)); + break; + } + + for (int i = 0; i < nfds; ++i) { + int fd = events[i].data.fd; + uint32_t ev = events[i].events; + + if (fd == listenFd) { + AcceptNew(); + continue; + } + if (fd == wakeFd) { + DrainWakeFd(); + continue; + } + + Connection* c = nullptr; + auto it = conns.find(fd); + if (it != conns.end()) c = it->second.get(); + if (!c) continue; + + if (ev & (EPOLLERR | EPOLLHUP)) { + ClosePeerByFd(fd); + continue; + } + + if (ev & EPOLLIN) { + HandleReadable(c); + auto it2 = conns.find(fd); + if (it2 == conns.end()) { + continue; + } + c = it2->second.get(); + if (!c) continue; + } + if (ev & EPOLLOUT) { + HandleConnWritable(c); + } + } + } + } + + private: + EngineKey myEngKey; + IOEngineConfig engConfig; + TcpBackendConfig config; + + int epfd{-1}; + int listenFd{-1}; + int wakeFd{-1}; + uint16_t listenPort{0}; + + std::atomic running{false}; + std::thread ioThread; + + std::mutex submitMu; + std::deque> submitQ; + + std::mutex remoteMu; + std::unordered_map remoteEngines; + + std::mutex memMu; + std::unordered_map localMems; + + std::mutex inboundMu; + std::unordered_map> inboundStatus; + + std::unordered_map> conns; + std::unordered_map peers; + + std::unordered_map>> waitingOps; + + std::unordered_map> pendingOutbound; + + std::unordered_map> inboundWrites; + std::unordered_map> earlyWrites; + + PinnedStagingPool staging; + StreamPool streamPool{8}; + EventPool eventPool{64}; + std::deque gpuTasks; +}; + +/* ---------------------------------------------------------------------------------------------- */ +/* TcpBackendSession */ +/* ---------------------------------------------------------------------------------------------- */ +TcpBackendSession::TcpBackendSession(const TcpBackendConfig& cfg, const MemoryDesc& l, const MemoryDesc& r, + TcpTransport* t) + : config(cfg), local(l), remote(r), transport(t) {} + +void TcpBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status, + TransferUniqueId id, bool isRead) { + MORI_IO_FUNCTION_TIMER; + transport->SubmitReadWrite(local, localOffset, remote, remoteOffset, size, status, id, isRead); +} + +void TcpBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) { + MORI_IO_FUNCTION_TIMER; + transport->SubmitBatchReadWrite(local, localOffsets, remote, remoteOffsets, sizes, status, id, isRead); +} + +bool TcpBackendSession::Alive() const { return true; } + +/* ---------------------------------------------------------------------------------------------- */ +/* TcpBackend */ +/* ---------------------------------------------------------------------------------------------- */ +TcpBackend::TcpBackend(EngineKey k, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg) + : myEngKey(std::move(k)), config(cfg) { + transport = std::make_unique(myEngKey, engCfg, cfg); + transport->Start(); + MORI_IO_INFO("TcpBackend created key={}", myEngKey.c_str()); +} + +TcpBackend::~TcpBackend() { transport->Shutdown(); } + +std::optional TcpBackend::GetListenPort() const { return transport->GetListenPort(); } + +void TcpBackend::RegisterRemoteEngine(const EngineDesc& desc) { transport->RegisterRemoteEngine(desc); } + +void TcpBackend::DeregisterRemoteEngine(const EngineDesc& desc) { transport->DeregisterRemoteEngine(desc); } + +void TcpBackend::RegisterMemory(MemoryDesc& desc) { transport->RegisterMemory(desc); } + +void TcpBackend::DeregisterMemory(const MemoryDesc& desc) { transport->DeregisterMemory(desc); } + +void TcpBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc, + size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, + bool isRead) { + MORI_IO_FUNCTION_TIMER; + transport->SubmitReadWrite(localDest, localOffset, remoteSrc, remoteOffset, size, status, id, isRead); +} + +void TcpBackend::BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, const MemoryDesc& remoteSrc, + const SizeVec& remoteOffsets, const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id, bool isRead) { + MORI_IO_FUNCTION_TIMER; + transport->SubmitBatchReadWrite(localDest, localOffsets, remoteSrc, remoteOffsets, sizes, status, id, isRead); +} + +BackendSession* TcpBackend::CreateSession(const MemoryDesc& local, const MemoryDesc& remote) { + auto* sess = new TcpBackendSession(config, local, remote, transport.get()); + return sess; +} + +bool TcpBackend::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, TransferStatus* status) { + return transport->PopInboundTransferStatus(remote, id, status); +} + +} // namespace io +} // namespace mori diff --git a/src/io/tcp/backend_impl.hpp b/src/io/tcp/backend_impl.hpp new file mode 100644 index 000000000..153fe6169 --- /dev/null +++ b/src/io/tcp/backend_impl.hpp @@ -0,0 +1,97 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include + +#include "mori/io/backend.hpp" +#include "mori/io/common.hpp" +#include "mori/io/engine.hpp" + +namespace mori { +namespace io { + +class TcpTransport; + +/* ---------------------------------------------------------------------------------------------- */ +/* TcpBackendSession */ +/* ---------------------------------------------------------------------------------------------- */ +class TcpBackendSession : public BackendSession { + public: + TcpBackendSession() = default; + TcpBackendSession(const TcpBackendConfig& config, const MemoryDesc& local, const MemoryDesc& remote, + TcpTransport* transport); + ~TcpBackendSession() override = default; + + void ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status, + TransferUniqueId id, bool isRead) override; + + void BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, + TransferStatus* status, TransferUniqueId id, bool isRead) override; + + bool Alive() const override; + + private: + TcpBackendConfig config{}; + MemoryDesc local{}; + MemoryDesc remote{}; + TcpTransport* transport{nullptr}; +}; + +/* ---------------------------------------------------------------------------------------------- */ +/* TcpBackend */ +/* ---------------------------------------------------------------------------------------------- */ +class TcpBackend : public Backend { + public: + TcpBackend(EngineKey, const IOEngineConfig&, const TcpBackendConfig&); + ~TcpBackend() override; + + std::optional GetListenPort() const; + + void RegisterRemoteEngine(const EngineDesc&) override; + void DeregisterRemoteEngine(const EngineDesc&) override; + + void RegisterMemory(MemoryDesc& desc) override; + void DeregisterMemory(const MemoryDesc& desc) override; + + void ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc, + size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, + bool isRead) override; + + void BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, + const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, const SizeVec& sizes, + TransferStatus* status, TransferUniqueId id, bool isRead) override; + + BackendSession* CreateSession(const MemoryDesc& local, const MemoryDesc& remote) override; + + bool PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, TransferStatus* status) override; + + private: + EngineKey myEngKey; + TcpBackendConfig config{}; + std::unique_ptr transport{nullptr}; +}; + +} // namespace io +} // namespace mori diff --git a/src/io/tcp/protocol.hpp b/src/io/tcp/protocol.hpp new file mode 100644 index 000000000..d9e4e41c3 --- /dev/null +++ b/src/io/tcp/protocol.hpp @@ -0,0 +1,268 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +namespace mori { +namespace io { +namespace tcp { + +constexpr uint32_t kCtrlMagic = 0x4D544330; // "MTC0" +constexpr uint32_t kDataMagic = 0x4D544430; // "MTD0" +constexpr uint16_t kProtoVersion = 1; + +enum class Channel : uint8_t { CTRL = 1, DATA = 2 }; + +enum class CtrlMsgType : uint8_t { + HELLO = 1, + WRITE_REQ = 2, + READ_REQ = 3, + BATCH_WRITE_REQ = 4, + BATCH_READ_REQ = 5, + COMPLETION = 6, +}; + +// Fixed framing: +// CtrlHeader (12B) + body +// magic(4) ver(2) type(1) reserved(1) body_len(4) +// DataHeader (24B) + payload +// magic(4) ver(2) flags(2) op_id(8) payload_len(8) +constexpr size_t kCtrlHeaderSize = 12; +constexpr size_t kDataHeaderSize = 24; + +struct CtrlHeaderView { + CtrlMsgType type{CtrlMsgType::HELLO}; + uint32_t bodyLen{0}; +}; + +struct DataHeaderView { + uint16_t flags{0}; + uint64_t opId{0}; + uint64_t payloadLen{0}; +}; + +inline uint64_t HostToBe64(uint64_t v) { return htobe64(v); } +inline uint64_t BeToHost64(uint64_t v) { return be64toh(v); } + +inline void AppendU8(std::vector& out, uint8_t v) { out.push_back(v); } + +inline void AppendU16BE(std::vector& out, uint16_t v) { + uint16_t be = htons(v); + uint8_t* p = reinterpret_cast(&be); + out.insert(out.end(), p, p + sizeof(be)); +} + +inline void AppendU32BE(std::vector& out, uint32_t v) { + uint32_t be = htonl(v); + uint8_t* p = reinterpret_cast(&be); + out.insert(out.end(), p, p + sizeof(be)); +} + +inline void AppendU64BE(std::vector& out, uint64_t v) { + uint64_t be = HostToBe64(v); + uint8_t* p = reinterpret_cast(&be); + out.insert(out.end(), p, p + sizeof(be)); +} + +inline bool ReadU16BE(const uint8_t* p, size_t len, size_t* off, uint16_t* out) { + if (*off + sizeof(uint16_t) > len) return false; + uint16_t be; + std::memcpy(&be, p + *off, sizeof(be)); + *out = ntohs(be); + *off += sizeof(be); + return true; +} + +inline bool ReadU32BE(const uint8_t* p, size_t len, size_t* off, uint32_t* out) { + if (*off + sizeof(uint32_t) > len) return false; + uint32_t be; + std::memcpy(&be, p + *off, sizeof(be)); + *out = ntohl(be); + *off += sizeof(be); + return true; +} + +inline bool ReadU64BE(const uint8_t* p, size_t len, size_t* off, uint64_t* out) { + if (*off + sizeof(uint64_t) > len) return false; + uint64_t be; + std::memcpy(&be, p + *off, sizeof(be)); + *out = BeToHost64(be); + *off += sizeof(be); + return true; +} + +inline bool TryParseCtrlHeader(const uint8_t* buf, size_t len, CtrlHeaderView* out) { + if (len < kCtrlHeaderSize) return false; + size_t off = 0; + + uint32_t magic = 0; + uint16_t ver = 0; + uint8_t type = 0; + uint8_t reserved = 0; + uint32_t bodyLen = 0; + + if (!ReadU32BE(buf, len, &off, &magic)) return false; + if (!ReadU16BE(buf, len, &off, &ver)) return false; + if (off + 2 > len) return false; + type = buf[off++]; + reserved = buf[off++]; + (void)reserved; + if (!ReadU32BE(buf, len, &off, &bodyLen)) return false; + + if (magic != kCtrlMagic || ver != kProtoVersion) return false; + out->type = static_cast(type); + out->bodyLen = bodyLen; + return true; +} + +inline bool TryParseDataHeader(const uint8_t* buf, size_t len, DataHeaderView* out) { + if (len < kDataHeaderSize) return false; + size_t off = 0; + + uint32_t magic = 0; + uint16_t ver = 0; + uint16_t flags = 0; + uint64_t opId = 0; + uint64_t payloadLen = 0; + + if (!ReadU32BE(buf, len, &off, &magic)) return false; + if (!ReadU16BE(buf, len, &off, &ver)) return false; + if (!ReadU16BE(buf, len, &off, &flags)) return false; + if (!ReadU64BE(buf, len, &off, &opId)) return false; + if (!ReadU64BE(buf, len, &off, &payloadLen)) return false; + + if (magic != kDataMagic || ver != kProtoVersion) return false; + out->flags = flags; + out->opId = opId; + out->payloadLen = payloadLen; + return true; +} + +inline std::vector BuildCtrlFrame(CtrlMsgType type, const std::vector& body) { + std::vector out; + out.reserve(kCtrlHeaderSize + body.size()); + AppendU32BE(out, kCtrlMagic); + AppendU16BE(out, kProtoVersion); + AppendU8(out, static_cast(type)); + AppendU8(out, 0 /*reserved*/); + AppendU32BE(out, static_cast(body.size())); + out.insert(out.end(), body.begin(), body.end()); + return out; +} + +inline std::vector BuildHello(Channel ch, const std::string& engineKey) { + std::vector body; + body.reserve(1 + 4 + engineKey.size()); + AppendU8(body, static_cast(ch)); + AppendU32BE(body, static_cast(engineKey.size())); + body.insert(body.end(), engineKey.begin(), engineKey.end()); + return BuildCtrlFrame(CtrlMsgType::HELLO, body); +} + +inline std::vector BuildCompletion(uint64_t opId, uint32_t statusCode, + const std::string& msg) { + std::vector body; + body.reserve(8 + 4 + 4 + msg.size()); + AppendU64BE(body, opId); + AppendU32BE(body, statusCode); + AppendU32BE(body, static_cast(msg.size())); + body.insert(body.end(), msg.begin(), msg.end()); + return BuildCtrlFrame(CtrlMsgType::COMPLETION, body); +} + +inline std::vector BuildWriteReq(uint64_t opId, uint32_t remoteMemId, uint64_t remoteOff, + uint64_t size) { + std::vector body; + body.reserve(8 + 4 + 8 + 8); + AppendU64BE(body, opId); + AppendU32BE(body, remoteMemId); + AppendU64BE(body, remoteOff); + AppendU64BE(body, size); + return BuildCtrlFrame(CtrlMsgType::WRITE_REQ, body); +} + +inline std::vector BuildReadReq(uint64_t opId, uint32_t srcMemId, uint64_t srcOff, + uint64_t size) { + std::vector body; + body.reserve(8 + 4 + 8 + 8); + AppendU64BE(body, opId); + AppendU32BE(body, srcMemId); + AppendU64BE(body, srcOff); + AppendU64BE(body, size); + return BuildCtrlFrame(CtrlMsgType::READ_REQ, body); +} + +inline std::vector BuildBatchWriteReq(uint64_t opId, uint32_t remoteMemId, + const std::vector& remoteOffs, + const std::vector& sizes) { + std::vector body; + const uint32_t n = static_cast(sizes.size()); + body.reserve(8 + 4 + 4 + n * (8 + 8)); + AppendU64BE(body, opId); + AppendU32BE(body, remoteMemId); + AppendU32BE(body, n); + for (uint32_t i = 0; i < n; ++i) { + AppendU64BE(body, remoteOffs[i]); + AppendU64BE(body, sizes[i]); + } + return BuildCtrlFrame(CtrlMsgType::BATCH_WRITE_REQ, body); +} + +inline std::vector BuildBatchReadReq(uint64_t opId, uint32_t srcMemId, + const std::vector& srcOffs, + const std::vector& sizes) { + std::vector body; + const uint32_t n = static_cast(sizes.size()); + body.reserve(8 + 4 + 4 + n * (8 + 8)); + AppendU64BE(body, opId); + AppendU32BE(body, srcMemId); + AppendU32BE(body, n); + for (uint32_t i = 0; i < n; ++i) { + AppendU64BE(body, srcOffs[i]); + AppendU64BE(body, sizes[i]); + } + return BuildCtrlFrame(CtrlMsgType::BATCH_READ_REQ, body); +} + +inline std::vector BuildDataHeader(uint64_t opId, uint64_t payloadLen, uint16_t flags) { + std::vector out; + out.reserve(kDataHeaderSize); + AppendU32BE(out, kDataMagic); + AppendU16BE(out, kProtoVersion); + AppendU16BE(out, flags); + AppendU64BE(out, opId); + AppendU64BE(out, payloadLen); + return out; +} + +} // namespace tcp +} // namespace io +} // namespace mori + From 333f88bf3c7e09121830edb7729b95e0c417a288 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 18:35:37 +0000 Subject: [PATCH 03/12] python: expose TCP backend config --- python/mori/io/__init__.py | 1 + python/mori/io/engine.py | 2 ++ src/pybind/mori.cpp | 18 ++++++++++++++++++ 3 files changed, 21 insertions(+) diff --git a/python/mori/io/__init__.py b/python/mori/io/__init__.py index 8cf0a7733..d246eae2b 100644 --- a/python/mori/io/__init__.py +++ b/python/mori/io/__init__.py @@ -29,6 +29,7 @@ MemoryLocationType, PollCqMode, RdmaBackendConfig, + TcpBackendConfig, XgmiBackendConfig, set_log_level, ) diff --git a/python/mori/io/engine.py b/python/mori/io/engine.py index 00dc04e90..c1d4fe60b 100644 --- a/python/mori/io/engine.py +++ b/python/mori/io/engine.py @@ -91,6 +91,8 @@ def create_backend(self, type: mori_cpp.BackendType, config=None): config = mori_cpp.RdmaBackendConfig() elif type is mori_cpp.BackendType.XGMI: config = mori_cpp.XgmiBackendConfig() + elif type is mori_cpp.BackendType.TCP: + config = mori_cpp.TcpBackendConfig() else: raise NotImplementedError("backend not implemented yet") return self._engine.CreateBackend(type, config) diff --git a/src/pybind/mori.cpp b/src/pybind/mori.cpp index c565d9d80..f01f3f15d 100644 --- a/src/pybind/mori.cpp +++ b/src/pybind/mori.cpp @@ -758,6 +758,24 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("num_streams", &mori::io::XgmiBackendConfig::numStreams) .def_readwrite("num_events", &mori::io::XgmiBackendConfig::numEvents); + py::class_(m, "TcpBackendConfig") + .def(py::init(), + py::arg("num_io_threads") = 1, py::arg("sock_sndbuf_bytes") = 4 * 1024 * 1024, + py::arg("sock_rcvbuf_bytes") = 4 * 1024 * 1024, py::arg("op_timeout_ms") = 30 * 1000, + py::arg("enable_keepalive") = true, py::arg("keepalive_idle_sec") = 30, + py::arg("keepalive_intvl_sec") = 10, py::arg("keepalive_cnt") = 3, + py::arg("enable_ctrl_nodelay") = true, py::arg("enable_data_nodelay") = false) + .def_readwrite("num_io_threads", &mori::io::TcpBackendConfig::numIoThreads) + .def_readwrite("sock_sndbuf_bytes", &mori::io::TcpBackendConfig::sockSndbufBytes) + .def_readwrite("sock_rcvbuf_bytes", &mori::io::TcpBackendConfig::sockRcvbufBytes) + .def_readwrite("op_timeout_ms", &mori::io::TcpBackendConfig::opTimeoutMs) + .def_readwrite("enable_keepalive", &mori::io::TcpBackendConfig::enableKeepalive) + .def_readwrite("keepalive_idle_sec", &mori::io::TcpBackendConfig::keepaliveIdleSec) + .def_readwrite("keepalive_intvl_sec", &mori::io::TcpBackendConfig::keepaliveIntvlSec) + .def_readwrite("keepalive_cnt", &mori::io::TcpBackendConfig::keepaliveCnt) + .def_readwrite("enable_ctrl_nodelay", &mori::io::TcpBackendConfig::enableCtrlNodelay) + .def_readwrite("enable_data_nodelay", &mori::io::TcpBackendConfig::enableDataNodelay); + py::class_(m, "IOEngineConfig") .def(py::init(), py::arg("host") = "", py::arg("port") = 0) .def_readwrite("host", &mori::io::IOEngineConfig::host) From 65d62aa04c30d1a0c965e2de902b20e2be3a95ec Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 18:35:49 +0000 Subject: [PATCH 04/12] tests: add tcp option to io benchmark --- tests/python/io/benchmark.py | 45 +++++++++++++++++++++++------------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index cc461848c..8a879a690 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -45,14 +45,14 @@ def parse_args(): parser.add_argument( "--backend", type=str, - choices=["rdma", "xgmi"], + choices=["rdma", "xgmi", "tcp"], default="rdma", - help="Backend type: 'rdma' for cross-node, 'xgmi' for intra-node GPU-to-GPU (default: rdma)", + help="Backend type: 'rdma' for cross-node, 'tcp' for cross-node fallback, 'xgmi' for intra-node GPU-to-GPU (default: rdma)", ) parser.add_argument( "--host", type=str, - help="Host IP for mori io engine OOB communication (RDMA only)", + help="Host IP for mori io engine OOB communication (RDMA/TCP only)", ) parser.add_argument( "--src-gpu", @@ -90,6 +90,14 @@ def parse_args(): default="read", help="Type of ops, choices [read, write], default to 'read'", ) + # Backward-compatible alias for existing scripts. + parser.add_argument( + "--op", + dest="op_type", + type=str, + choices=["read", "write"], + help="Alias for --op-type", + ) parser.add_argument( "--buffer-size", type=int, @@ -403,13 +411,16 @@ def _initialize_rdma(self): port=self.port, ) self.engine = IOEngine(key=f"{self.role.name}-{self.role_rank}", config=config) - config = RdmaBackendConfig( - qp_per_transfer=self.num_qp_per_transfer, - post_batch_size=-1, - num_worker_threads=self.num_worker_threads, - poll_cq_mode=self.poll_cq_mode, - ) - self.engine.create_backend(BackendType.RDMA, config) + if self.backend_type == "tcp": + self.engine.create_backend(BackendType.TCP) + else: + config = RdmaBackendConfig( + qp_per_transfer=self.num_qp_per_transfer, + post_batch_size=-1, + num_worker_threads=self.num_worker_threads, + poll_cq_mode=self.poll_cq_mode, + ) + self.engine.create_backend(BackendType.RDMA, config) self.engine_desc = self.engine.get_engine_desc() engine_desc_bytes = self.engine_desc.pack() @@ -494,7 +505,7 @@ def _initialize_xgmi(self): def run_single_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "tcp") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: return 0 @@ -546,7 +557,7 @@ def run_single_once(self, buffer_size, transfer_batch_size): def run_batch_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "tcp") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: return 0 @@ -606,7 +617,7 @@ def _run_and_compute(self, buffer_size, transfer_batch_size, iters): latency.append(duration) if self.role is EngineRole.TARGET and ( - self.backend_type == "rdma" + self.backend_type in ("rdma", "tcp") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ): return 0, 0, 0, 0, 0 @@ -627,6 +638,8 @@ def _get_table_title(self): return f"XGMI Multiprocess Benchmark: Rank {self.role_rank} ({self.role.name})" else: return f"XGMI Benchmark: GPU{self.src_gpu} -> GPU{self.dst_gpu}" + elif self.backend_type == "tcp": + return f"TCP Benchmark: Initiator Rank {self.role_rank}" else: return f"RDMA Benchmark: Initiator Rank {self.role_rank}" @@ -650,7 +663,7 @@ def _run_benchmark_loop(self): cur_size = self.sweep_start_size max_size = self.sweep_max_size while cur_size <= max_size: - if self.backend_type == "rdma" or ( + if self.backend_type in ("rdma", "tcp") or ( self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() @@ -675,7 +688,7 @@ def _run_benchmark_loop(self): cur_transfer_batch_size = 1 max_transfer_batch_size = 32768 while cur_transfer_batch_size <= max_transfer_batch_size: - if self.backend_type == "rdma" or ( + if self.backend_type in ("rdma", "tcp") or ( self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() @@ -826,7 +839,7 @@ def benchmark_engine(local_rank, node_rank, args): sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, sweep_max_size=args.sweep_max_size, - backend_type="rdma", + backend_type=args.backend, host=args.host, port=get_free_port(), node_rank=node_rank, From 6dea9bcdfab5f9d031b77304694328d99c0fcb18 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 18:35:59 +0000 Subject: [PATCH 05/12] tests: add basic TCP engine tests --- tests/python/io/test_engine_tcp.py | 122 +++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 tests/python/io/test_engine_tcp.py diff --git a/tests/python/io/test_engine_tcp.py b/tests/python/io/test_engine_tcp.py new file mode 100644 index 000000000..5bfa1c10f --- /dev/null +++ b/tests/python/io/test_engine_tcp.py @@ -0,0 +1,122 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +import time + +import torch +from tests.python.utils import get_free_port + +from mori.io import ( + BackendType, + IOEngine, + IOEngineConfig, + MemoryDesc, + TcpBackendConfig, + set_log_level, +) + + +def _wait_inbound_status(engine, remote_engine_key, remote_transfer_uid, timeout_s=5.0): + deadline = time.time() + timeout_s + while time.time() < deadline: + st = engine.pop_inbound_transfer_status(remote_engine_key, remote_transfer_uid) + if st is not None: + return st + time.sleep(0.001) + raise RuntimeError("Timed out waiting for inbound status") + + +def _create_tcp_engine_pair(name_prefix, port_a=0, port_b=0): + cfg_a = IOEngineConfig(host="127.0.0.1", port=port_a) + cfg_b = IOEngineConfig(host="127.0.0.1", port=port_b) + + a = IOEngine(key=f"{name_prefix}_a", config=cfg_a) + b = IOEngine(key=f"{name_prefix}_b", config=cfg_b) + + a.create_backend(BackendType.TCP, TcpBackendConfig()) + b.create_backend(BackendType.TCP, TcpBackendConfig()) + + a_desc = a.get_engine_desc() + b_desc = b.get_engine_desc() + a.register_remote_engine(b_desc) + b.register_remote_engine(a_desc) + return a, b, a_desc, b_desc + + +def test_tcp_engine_desc_port_zero_auto_bind(): + set_log_level("error") + engine = IOEngine(key="engine_tcp_port0", config=IOEngineConfig(host="127.0.0.1", port=0)) + engine.create_backend(BackendType.TCP, TcpBackendConfig()) + desc = engine.get_engine_desc() + assert desc.port > 0 + + +def test_tcp_cpu_write_read_and_batch(): + set_log_level("error") + a, b, a_desc, b_desc = _create_tcp_engine_pair("tcp_cpu", get_free_port(), get_free_port()) + + # Allocate CPU tensors and register memory. + src = torch.arange(0, 1024 * 4, dtype=torch.uint8) + dst = torch.zeros_like(src) + src_md = a.register_torch_tensor(src) + dst_md = b.register_torch_tensor(dst) + + # MemoryDesc serialization should work for TCP too. + packed = dst_md.pack() + dst_md_remote = MemoryDesc.unpack(packed) + assert dst_md == dst_md_remote + + # Single write + uid = a.allocate_transfer_uid() + st = a.write(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + assert torch.equal(src, dst) + + # Single read (b -> a) + dst.zero_() + uid = a.allocate_transfer_uid() + st = a.read(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + assert torch.equal(src, dst) + + # Batch write via session + sess = a.create_session(src_md, dst_md) + assert sess is not None + offsets = [0, 256, 512, 768] + sizes = [128, 128, 128, 128] + + dst.zero_() + uid = sess.allocate_transfer_uid() + st = sess.batch_write(offsets, offsets, sizes, uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + + for off, sz in zip(offsets, sizes): + assert torch.equal(src[off : off + sz], dst[off : off + sz]) + From fe431fc55813347d2f6adaefde381ed156ebf129 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 26 Feb 2026 18:36:14 +0000 Subject: [PATCH 06/12] docs: clarify ctrl/data reordering for tcp --- docs/MORI-IO-TCP-BACKEND.md | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/docs/MORI-IO-TCP-BACKEND.md b/docs/MORI-IO-TCP-BACKEND.md index 1c0446cc5..d508507ac 100644 --- a/docs/MORI-IO-TCP-BACKEND.md +++ b/docs/MORI-IO-TCP-BACKEND.md @@ -127,6 +127,15 @@ Ordering: Receiver keeps op state keyed by `op_id` so ctrl/data can be processed independently. +#### Cross-channel reordering (CTRL vs DATA) +Because **CTRL** and **DATA** are separate TCP connections, the receiver may observe **DATA arriving before the corresponding CTRL request** (no global ordering across connections). + +To preserve RDMA-like op semantics without adding an extra RTT handshake, the TCP backend must: +- Buffer such early DATA frames by `(peer_key, op_id)` into pinned host memory. +- Finalize the write once the CTRL request arrives (or fail/cleanup on disconnect/timeout). + +For GPU destinations this maps naturally to the existing pinned-staging path; for CPU destinations this fallback path may add an extra copy only in the reordering case. + --- ## 4. API Semantics @@ -178,4 +187,3 @@ Receiver keeps op state keyed by `op_id` so ctrl/data can be processed independe - Adaptive multi-connection striping per peer to improve throughput on high-BDP networks. - Full `MSG_ZEROCOPY` completion accounting + buffer lifetime tracking (currently best-effort). - Optional TLS / auth for multi-tenant environments. - From 0e5d1ad6030bd5af8110a0645892c5f4c64bf028 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 27 Feb 2026 05:08:08 +0000 Subject: [PATCH 07/12] io/tcp: add multi-stream striping config --- include/mori/io/backend.hpp | 17 ++++++++++++++--- src/io/tcp/protocol.hpp | 25 +++++++++++++++---------- src/pybind/mori.cpp | 9 ++++++--- 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index d92cc50b3..14031dce8 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -85,7 +85,8 @@ struct TcpBackendConfig : public BackendConfig { TcpBackendConfig() : BackendConfig(BackendType::TCP) {} TcpBackendConfig(int numIoThreads_, int sockSndbufBytes_, int sockRcvbufBytes_, int opTimeoutMs_, bool enableKeepalive_, int keepaliveIdleSec_, int keepaliveIntvlSec_, - int keepaliveCnt_, bool enableCtrlNodelay_, bool enableDataNodelay_) + int keepaliveCnt_, bool enableCtrlNodelay_, bool enableDataNodelay_, + int numDataConns_, int stripingThresholdBytes_) : BackendConfig(BackendType::TCP), numIoThreads(numIoThreads_), sockSndbufBytes(sockSndbufBytes_), @@ -96,7 +97,9 @@ struct TcpBackendConfig : public BackendConfig { keepaliveIntvlSec(keepaliveIntvlSec_), keepaliveCnt(keepaliveCnt_), enableCtrlNodelay(enableCtrlNodelay_), - enableDataNodelay(enableDataNodelay_) {} + enableDataNodelay(enableDataNodelay_), + numDataConns(numDataConns_), + stripingThresholdBytes(stripingThresholdBytes_) {} int numIoThreads{1}; @@ -112,6 +115,12 @@ struct TcpBackendConfig : public BackendConfig { bool enableCtrlNodelay{true}; bool enableDataNodelay{false}; + + // Number of parallel DATA TCP connections per peer (iperf-like multi-stream striping). + // Effective only when peer has >= numDataConns established and transfer is contiguous. + int numDataConns{4}; + // Stripe large contiguous transfers; keep small transfers on a single stream for latency. + int stripingThresholdBytes{256 * 1024}; }; inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { @@ -120,7 +129,9 @@ inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { << "] enableKeepalive[" << c.enableKeepalive << "] keepaliveIdleSec[" << c.keepaliveIdleSec << "] keepaliveIntvlSec[" << c.keepaliveIntvlSec << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" - << c.enableCtrlNodelay << "] enableDataNodelay[" << c.enableDataNodelay << "]"; + << c.enableCtrlNodelay << "] enableDataNodelay[" << c.enableDataNodelay + << "] numDataConns[" << c.numDataConns << "] stripingThresholdBytes[" + << c.stripingThresholdBytes << "]"; } /* ---------------------------------------------------------------------------------------------- */ diff --git a/src/io/tcp/protocol.hpp b/src/io/tcp/protocol.hpp index d9e4e41c3..aebbaae55 100644 --- a/src/io/tcp/protocol.hpp +++ b/src/io/tcp/protocol.hpp @@ -36,7 +36,7 @@ namespace tcp { constexpr uint32_t kCtrlMagic = 0x4D544330; // "MTC0" constexpr uint32_t kDataMagic = 0x4D544430; // "MTD0" -constexpr uint16_t kProtoVersion = 1; +constexpr uint16_t kProtoVersion = 2; enum class Channel : uint8_t { CTRL = 1, DATA = 2 }; @@ -198,33 +198,36 @@ inline std::vector BuildCompletion(uint64_t opId, uint32_t statusCode, } inline std::vector BuildWriteReq(uint64_t opId, uint32_t remoteMemId, uint64_t remoteOff, - uint64_t size) { + uint64_t size, uint8_t lanesTotal = 1) { std::vector body; - body.reserve(8 + 4 + 8 + 8); + body.reserve(8 + 4 + 8 + 8 + 1); AppendU64BE(body, opId); AppendU32BE(body, remoteMemId); AppendU64BE(body, remoteOff); AppendU64BE(body, size); + AppendU8(body, lanesTotal); return BuildCtrlFrame(CtrlMsgType::WRITE_REQ, body); } inline std::vector BuildReadReq(uint64_t opId, uint32_t srcMemId, uint64_t srcOff, - uint64_t size) { + uint64_t size, uint8_t lanesTotal = 1) { std::vector body; - body.reserve(8 + 4 + 8 + 8); + body.reserve(8 + 4 + 8 + 8 + 1); AppendU64BE(body, opId); AppendU32BE(body, srcMemId); AppendU64BE(body, srcOff); AppendU64BE(body, size); + AppendU8(body, lanesTotal); return BuildCtrlFrame(CtrlMsgType::READ_REQ, body); } inline std::vector BuildBatchWriteReq(uint64_t opId, uint32_t remoteMemId, const std::vector& remoteOffs, - const std::vector& sizes) { + const std::vector& sizes, + uint8_t lanesTotal = 1) { std::vector body; const uint32_t n = static_cast(sizes.size()); - body.reserve(8 + 4 + 4 + n * (8 + 8)); + body.reserve(8 + 4 + 4 + n * (8 + 8) + 1); AppendU64BE(body, opId); AppendU32BE(body, remoteMemId); AppendU32BE(body, n); @@ -232,15 +235,17 @@ inline std::vector BuildBatchWriteReq(uint64_t opId, uint32_t remoteMem AppendU64BE(body, remoteOffs[i]); AppendU64BE(body, sizes[i]); } + AppendU8(body, lanesTotal); return BuildCtrlFrame(CtrlMsgType::BATCH_WRITE_REQ, body); } inline std::vector BuildBatchReadReq(uint64_t opId, uint32_t srcMemId, const std::vector& srcOffs, - const std::vector& sizes) { + const std::vector& sizes, + uint8_t lanesTotal = 1) { std::vector body; const uint32_t n = static_cast(sizes.size()); - body.reserve(8 + 4 + 4 + n * (8 + 8)); + body.reserve(8 + 4 + 4 + n * (8 + 8) + 1); AppendU64BE(body, opId); AppendU32BE(body, srcMemId); AppendU32BE(body, n); @@ -248,6 +253,7 @@ inline std::vector BuildBatchReadReq(uint64_t opId, uint32_t srcMemId, AppendU64BE(body, srcOffs[i]); AppendU64BE(body, sizes[i]); } + AppendU8(body, lanesTotal); return BuildCtrlFrame(CtrlMsgType::BATCH_READ_REQ, body); } @@ -265,4 +271,3 @@ inline std::vector BuildDataHeader(uint64_t opId, uint64_t payloadLen, } // namespace tcp } // namespace io } // namespace mori - diff --git a/src/pybind/mori.cpp b/src/pybind/mori.cpp index f01f3f15d..28b310917 100644 --- a/src/pybind/mori.cpp +++ b/src/pybind/mori.cpp @@ -759,12 +759,13 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("num_events", &mori::io::XgmiBackendConfig::numEvents); py::class_(m, "TcpBackendConfig") - .def(py::init(), + .def(py::init(), py::arg("num_io_threads") = 1, py::arg("sock_sndbuf_bytes") = 4 * 1024 * 1024, py::arg("sock_rcvbuf_bytes") = 4 * 1024 * 1024, py::arg("op_timeout_ms") = 30 * 1000, py::arg("enable_keepalive") = true, py::arg("keepalive_idle_sec") = 30, py::arg("keepalive_intvl_sec") = 10, py::arg("keepalive_cnt") = 3, - py::arg("enable_ctrl_nodelay") = true, py::arg("enable_data_nodelay") = false) + py::arg("enable_ctrl_nodelay") = true, py::arg("enable_data_nodelay") = false, + py::arg("num_data_conns") = 4, py::arg("striping_threshold_bytes") = 256 * 1024) .def_readwrite("num_io_threads", &mori::io::TcpBackendConfig::numIoThreads) .def_readwrite("sock_sndbuf_bytes", &mori::io::TcpBackendConfig::sockSndbufBytes) .def_readwrite("sock_rcvbuf_bytes", &mori::io::TcpBackendConfig::sockRcvbufBytes) @@ -774,7 +775,9 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("keepalive_intvl_sec", &mori::io::TcpBackendConfig::keepaliveIntvlSec) .def_readwrite("keepalive_cnt", &mori::io::TcpBackendConfig::keepaliveCnt) .def_readwrite("enable_ctrl_nodelay", &mori::io::TcpBackendConfig::enableCtrlNodelay) - .def_readwrite("enable_data_nodelay", &mori::io::TcpBackendConfig::enableDataNodelay); + .def_readwrite("enable_data_nodelay", &mori::io::TcpBackendConfig::enableDataNodelay) + .def_readwrite("num_data_conns", &mori::io::TcpBackendConfig::numDataConns) + .def_readwrite("striping_threshold_bytes", &mori::io::TcpBackendConfig::stripingThresholdBytes); py::class_(m, "IOEngineConfig") .def(py::init(), py::arg("host") = "", py::arg("port") = 0) From 04af586850b8c30c8a689024085515bc570a9e0f Mon Sep 17 00:00:00 2001 From: root Date: Fri, 27 Feb 2026 05:08:19 +0000 Subject: [PATCH 08/12] io/tcp: stripe large transfers over multiple data conns --- src/io/tcp/backend_impl.cpp | 1273 ++++++++++++++++++++++++----------- 1 file changed, 871 insertions(+), 402 deletions(-) diff --git a/src/io/tcp/backend_impl.cpp b/src/io/tcp/backend_impl.cpp index 2374d9658..3cba24098 100644 --- a/src/io/tcp/backend_impl.cpp +++ b/src/io/tcp/backend_impl.cpp @@ -222,6 +222,53 @@ struct Segment { uint64_t len{0}; }; +constexpr uint8_t kLaneBits = 3; // up to 8 lanes +constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; + +inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { return (userOpId << kLaneBits) | lane; } + +inline uint64_t ToUserOpId(uint64_t wireOpId) { return wireOpId >> kLaneBits; } + +struct LaneSpan { + uint64_t off{0}; + uint64_t len{0}; +}; + +inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanesTotal, uint8_t lane) { + if (lanesTotal <= 1) return LaneSpan{0, total}; + const uint64_t base = total / lanesTotal; + const uint64_t rem = total % lanesTotal; + const uint64_t off = static_cast(lane) * base + std::min(lane, rem); + const uint64_t len = base + (lane < rem ? 1 : 0); + return LaneSpan{off, len}; +} + +inline uint8_t LanesAllMask(uint8_t lanesTotal) { + if (lanesTotal >= (1U << kLaneBits)) return 0xFF; + return static_cast((1U << lanesTotal) - 1U); +} + +inline std::vector SliceSegments(const std::vector& segs, uint64_t start, + uint64_t len) { + std::vector out; + if (len == 0) return out; + uint64_t skip = start; + uint64_t remaining = len; + for (const auto& s : segs) { + if (remaining == 0) break; + if (skip >= s.len) { + skip -= s.len; + continue; + } + const uint64_t segOff = s.off + skip; + const uint64_t take = std::min(s.len - skip, remaining); + out.push_back({segOff, take}); + remaining -= take; + skip = 0; + } + return out; +} + inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, uint64_t* outLen) { if (!outOff || !outLen) return false; @@ -263,10 +310,13 @@ struct OutboundOpState { uint64_t expectedRxBytes{0}; uint64_t rxBytes{0}; bool completionReceived{false}; - StatusCode completionCode{StatusCode::INIT}; + uint8_t lanesTotal{1}; + uint8_t lanesDoneMask{0}; + StatusCode completionCode{StatusCode::SUCCESS}; std::string completionMsg; bool gpuCopyPending{false}; + std::shared_ptr pinned; Clock::time_point startTs{Clock::now()}; }; @@ -279,17 +329,30 @@ struct InboundWriteState { std::vector dstSegs; bool discard{false}; + uint8_t lanesTotal{1}; + uint8_t lanesDoneMask{0}; + std::shared_ptr pinned; }; -struct EarlyWriteState { +struct EarlyWriteLaneState { uint64_t payloadLen{0}; std::shared_ptr pinned; bool complete{false}; }; +struct EarlyWriteState { + std::unordered_map lanes; +}; + +enum class DataRxKind : uint8_t { NONE = 0, INBOUND_WRITE = 1, OUTBOUND_READ = 2, EARLY_WRITE = 3 }; + struct ActiveDataRx { bool active{false}; TransferUniqueId id{0}; + uint8_t lane{0}; + uint64_t laneOff{0}; + uint64_t laneLen{0}; + DataRxKind kind{DataRxKind::NONE}; uint64_t remaining{0}; bool discard{false}; @@ -356,9 +419,18 @@ struct Connection { struct PeerLinks { int ctrlFd{-1}; - int dataFd{-1}; + std::vector dataFds; + int ctrlPending{0}; + int dataPending{0}; + size_t rr{0}; bool CtrlUp() const { return ctrlFd >= 0; } - bool DataUp() const { return dataFd >= 0; } + bool DataUp() const { return !dataFds.empty(); } + int PickDataFd() { + if (dataFds.empty()) return -1; + int fd = dataFds[rr % dataFds.size()]; + rr = (rr + 1) % dataFds.size(); + return fd; + } }; } // namespace @@ -538,10 +610,10 @@ class TcpTransport { localSegs.reserve(n); remoteSegs.reserve(n); - uint64_t total = 0; - for (size_t i = 0; i < n; ++i) { - const size_t lo = localOffsets[i]; - const size_t ro = remoteOffsets[i]; + uint64_t total = 0; + for (size_t i = 0; i < n; ++i) { + const size_t lo = localOffsets[i]; + const size_t ro = remoteOffsets[i]; const size_t sz = sizes[i]; if (sz == 0) continue; if ((lo + sz) > local.size || (ro + sz) > remote.size) { @@ -549,12 +621,41 @@ class TcpTransport { return; } localSegs.push_back({static_cast(lo), static_cast(sz)}); - remoteSegs.push_back({static_cast(ro), static_cast(sz)}); - total += static_cast(sz); - } + remoteSegs.push_back({static_cast(ro), static_cast(sz)}); + total += static_cast(sz); + } - auto op = std::make_unique(); - op->peer = remote.engineKey; + // Coalesce adjacent segments when both local and remote are contiguous in lockstep. + // This is critical for TCP performance (reduces iov/copy fan-out), and preserves semantics. + if (localSegs.size() > 1) { + std::vector newLocal; + std::vector newRemote; + newLocal.reserve(localSegs.size()); + newRemote.reserve(remoteSegs.size()); + Segment curL = localSegs[0]; + Segment curR = remoteSegs[0]; + for (size_t i = 1; i < localSegs.size(); ++i) { + const Segment& l = localSegs[i]; + const Segment& r = remoteSegs[i]; + if ((curL.off + curL.len == l.off) && (curR.off + curR.len == r.off) && (curL.len == curR.len) && + (l.len == r.len)) { + curL.len += l.len; + curR.len += r.len; + } else { + newLocal.push_back(curL); + newRemote.push_back(curR); + curL = l; + curR = r; + } + } + newLocal.push_back(curL); + newRemote.push_back(curR); + localSegs = std::move(newLocal); + remoteSegs = std::move(newRemote); + } + + auto op = std::make_unique(); + op->peer = remote.engineKey; op->id = id; op->isRead = isRead; op->status = status; @@ -607,14 +708,24 @@ class TcpTransport { bool PreferOutgoingFor(const EngineKey& peerKey) const { return myEngKey < peerKey; } - void AssignConnToPeer(Connection* c) { - assert(c && c->helloReceived); - PeerLinks& link = peers[c->peerKey]; - const bool preferOutgoing = PreferOutgoingFor(c->peerKey); + void AssignConnToPeer(Connection* c) { + assert(c && c->helloReceived); + PeerLinks& link = peers[c->peerKey]; + const bool preferOutgoing = PreferOutgoingFor(c->peerKey); + const bool wasOutgoing = c->isOutgoing; + const tcp::Channel ch = c->ch; + + if (wasOutgoing) { + if (ch == tcp::Channel::CTRL) { + if (link.ctrlPending > 0) link.ctrlPending--; + } else { + if (link.dataPending > 0) link.dataPending--; + } + } - auto replace_if_needed = [&](int& slotFd) { - if (slotFd < 0) { - slotFd = c->fd; + auto replace_if_needed = [&](int& slotFd) { + if (slotFd < 0) { + slotFd = c->fd; return; } // Collision: choose deterministically by key + direction so both sides keep the same TCP conn. @@ -638,26 +749,49 @@ class TcpTransport { CloseConnInternal(c); conns.erase(newFd); } - }; + }; - if (c->ch == tcp::Channel::CTRL) { - replace_if_needed(link.ctrlFd); - } else { - replace_if_needed(link.dataFd); - } - } + if (ch == tcp::Channel::CTRL) { + replace_if_needed(link.ctrlFd); + return; + } - void MaybeDispatchQueuedOps(const EngineKey& peerKey) { - auto it = peers.find(peerKey); - if (it == peers.end()) return; - if (!it->second.CtrlUp() || !it->second.DataUp()) return; - Connection* ctrl = conns[it->second.ctrlFd].get(); - Connection* data = conns[it->second.dataFd].get(); - if (!ctrl || !data) return; - if (!ctrl->helloReceived || !data->helloReceived) return; + // DATA channel: keep up to numDataConns connections, choosing deterministically by key+direction + // so both sides converge to the same physical set. + const bool keepPreferred = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); + if (!keepPreferred) { + MORI_IO_TRACE("TCP: peer {} dropping non-preferred DATA fd {} outgoing={}", c->peerKey, + c->fd, c->isOutgoing); + const int fd = c->fd; + CloseConnInternal(c); + conns.erase(fd); + return; + } + + const size_t want = static_cast(std::max(1, config.numDataConns)); + if (link.dataFds.size() >= want) { + MORI_IO_TRACE("TCP: peer {} dropping extra DATA fd {} (have {} want {})", c->peerKey, c->fd, + link.dataFds.size(), want); + const int fd = c->fd; + CloseConnInternal(c); + conns.erase(fd); + return; + } + link.dataFds.push_back(c->fd); + } - auto qit = waitingOps.find(peerKey); - if (qit == waitingOps.end()) return; + void MaybeDispatchQueuedOps(const EngineKey& peerKey) { + auto it = peers.find(peerKey); + if (it == peers.end()) return; + if (!it->second.CtrlUp() || !it->second.DataUp()) return; + Connection* ctrl = conns[it->second.ctrlFd].get(); + if (!ctrl || !ctrl->helloReceived) return; + int dataFd = it->second.dataFds.empty() ? -1 : it->second.dataFds[0]; + Connection* data = (dataFd >= 0) ? conns[dataFd].get() : nullptr; + if (!data || !data->helloReceived) return; + + auto qit = waitingOps.find(peerKey); + if (qit == waitingOps.end()) return; auto ops = std::move(qit->second); waitingOps.erase(qit); @@ -667,11 +801,15 @@ class TcpTransport { } } - void EnsurePeerChannels(const EngineKey& peerKey) { - PeerLinks& link = peers[peerKey]; - if (!link.CtrlUp()) ConnectChannel(peerKey, tcp::Channel::CTRL); - if (!link.DataUp()) ConnectChannel(peerKey, tcp::Channel::DATA); - } + void EnsurePeerChannels(const EngineKey& peerKey) { + PeerLinks& link = peers[peerKey]; + if (!link.CtrlUp() && link.ctrlPending == 0) ConnectChannel(peerKey, tcp::Channel::CTRL); + + const int want = std::max(1, config.numDataConns); + while (static_cast(link.dataFds.size()) + link.dataPending < want) { + ConnectChannel(peerKey, tcp::Channel::DATA); + } + } void ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { EngineDesc desc; @@ -738,14 +876,23 @@ class TcpTransport { ConfigureDataSocket(fd, config); } - const bool wantWrite = connecting || !conn->sendq.empty(); - AddEpoll(fd, true, wantWrite); - conns[fd] = std::move(conn); + const bool wantWrite = connecting || !conn->sendq.empty(); + AddEpoll(fd, true, wantWrite); + conns[fd] = std::move(conn); - if (!connecting) { - QueueHello(fd); - ModEpoll(fd, true, true); - } + // Track pending outgoing connections so we don't over-connect when many ops are submitted + // before handshake completes. + PeerLinks& link = peers[peerKey]; + if (ch == tcp::Channel::CTRL) { + link.ctrlPending++; + } else { + link.dataPending++; + } + + if (!connecting) { + QueueHello(fd); + ModEpoll(fd, true, true); + } } void QueueHello(int fd) { @@ -821,8 +968,14 @@ class TcpTransport { if (it == peers.end()) return false; if (!it->second.CtrlUp() || !it->second.DataUp()) return false; Connection* ctrl = conns[it->second.ctrlFd].get(); - Connection* data = conns[it->second.dataFd].get(); - return ctrl && data && ctrl->helloReceived && data->helloReceived; + if (!ctrl || !ctrl->helloReceived) return false; + for (int fd : it->second.dataFds) { + auto cit = conns.find(fd); + if (cit == conns.end()) continue; + Connection* data = cit->second.get(); + if (data && data->helloReceived) return true; + } + return false; } void DispatchOp(std::unique_ptr op) { @@ -837,8 +990,21 @@ class TcpTransport { return; } Connection* ctrl = conns[it->second.ctrlFd].get(); - Connection* data = conns[it->second.dataFd].get(); - if (!ctrl || !data) { + if (!ctrl) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer ctrl connection missing"); + return; + } + + std::vector dataFds; + dataFds.reserve(it->second.dataFds.size()); + for (int fd : it->second.dataFds) { + auto dit = conns.find(fd); + if (dit == conns.end()) continue; + Connection* c = dit->second.get(); + if (c && c->helloReceived) dataFds.push_back(fd); + } + + if (dataFds.empty()) { op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer connection missing"); return; } @@ -860,15 +1026,38 @@ class TcpTransport { MORI_IO_TRACE("TCP: dispatch op id={} peer={} isRead={} segs={}", st->id, peerKey.c_str(), st->isRead, st->localSegs.size()); + const uint64_t totalBytes = SumLens(st->localSegs); + int wantLanes = std::max(1, config.numDataConns); + wantLanes = std::min(wantLanes, (1U << kLaneBits)); + uint8_t lanesTotal = 1; + const bool canStripe = (wantLanes > 1) && (config.stripingThresholdBytes > 0) && + (totalBytes >= static_cast(config.stripingThresholdBytes)) && + (st->localSegs.size() == 1) && (st->remoteSegs.size() == 1) && + (dataFds.size() >= 2); + if (canStripe) { + lanesTotal = static_cast(std::min(static_cast(wantLanes), dataFds.size())); + } + st->lanesTotal = lanesTotal; + + if (st->isRead && st->local.loc == MemoryLocationType::GPU) { + // Allocate a single pinned staging buffer for all lanes and do one H2D copy at the end. + st->pinned = staging.Acquire(static_cast(totalBytes)); + if (!st->pinned) { + st->status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging (read)"); + pendingOutbound.erase(opId); + return; + } + } + // CTRL request std::vector ctrlFrame; if (st->localSegs.size() == 1) { if (st->isRead) { - ctrlFrame = - tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, st->remoteSegs[0].len); + ctrlFrame = tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, + st->remoteSegs[0].len, lanesTotal); } else { - ctrlFrame = - tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, st->remoteSegs[0].len); + ctrlFrame = tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, + st->remoteSegs[0].len, lanesTotal); } } else { std::vector roffs; @@ -880,9 +1069,9 @@ class TcpTransport { szs.push_back(s.len); } if (st->isRead) { - ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs); + ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs, lanesTotal); } else { - ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs); + ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs, lanesTotal); } } @@ -890,11 +1079,10 @@ class TcpTransport { // DATA payload (writes only) if (!st->isRead) { - QueueDataSendForWrite(data, *st); + QueueDataSendForWrite(peerKey, dataFds, *st); } UpdateWriteInterest(ctrl->fd); - UpdateWriteInterest(data->fd); } void QueueSend(int fd, std::vector bytes, std::function onDone = nullptr) { @@ -909,25 +1097,72 @@ class TcpTransport { c->sendq.push_back(std::move(item)); } - void QueueDataSendForWrite(Connection* data, OutboundOpState& st) { + void QueueDataSendForWrite(const EngineKey& peerKey, const std::vector& dataFds, + OutboundOpState& st) { + if (dataFds.empty()) return; const uint64_t total = SumLens(st.localSegs); - auto hdr = tcp::BuildDataHeader(st.id, total, 0); + const uint8_t lanesTotal = std::max(1, st.lanesTotal); + + // Unstriped path (can handle iov fan-out). + if (lanesTotal == 1) { + auto hdr = tcp::BuildDataHeader(ToWireOpId(st.id, 0), total, 0); + const int fd = dataFds[0]; + Connection* data = conns[fd].get(); + if (!data) return; + + if (st.local.loc == MemoryLocationType::GPU) { + QueueGpuToNetSend(data, st.local, st.localSegs, std::move(hdr), /*onDone*/ nullptr); + UpdateWriteInterest(fd); + return; + } + + SendItem item; + item.header = std::move(hdr); + item.iov.reserve(1 + st.localSegs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + + uint8_t* base = reinterpret_cast(st.local.data); + for (const auto& s : st.localSegs) { + item.iov.push_back({base + s.off, static_cast(s.len)}); + } + data->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); + return; + } + + // Striped path: requires a single contiguous span. + if (st.localSegs.size() != 1) { + MORI_IO_WARN("TCP: striping requested but localSegs.size={} (expected 1), fallback to 1 lane", + st.localSegs.size()); + st.lanesTotal = 1; + QueueDataSendForWrite(peerKey, dataFds, st); + return; + } + + const uint8_t useLanes = std::min(lanesTotal, static_cast(dataFds.size())); if (st.local.loc == MemoryLocationType::GPU) { - QueueGpuToNetSend(data, st.local, st.localSegs, std::move(hdr), /*onDone*/ nullptr); + QueueGpuToNetSendStriped(peerKey, dataFds, st.id, useLanes, st.local, st.localSegs); return; } - SendItem item; - item.header = std::move(hdr); - item.iov.reserve(1 + st.localSegs.size()); - item.iov.push_back({item.header.data(), item.header.size()}); + uint8_t* base = reinterpret_cast(st.local.data) + st.localSegs[0].off; + for (uint8_t lane = 0; lane < useLanes; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, useLanes, lane); + const int fd = dataFds[lane % dataFds.size()]; + Connection* data = conns[fd].get(); + if (!data) continue; - uint8_t* base = reinterpret_cast(st.local.data); - for (const auto& s : st.localSegs) { - item.iov.push_back({base + s.off, static_cast(s.len)}); + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(st.id, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = base + span.off; + item.iov[1].iov_len = static_cast(span.len); + data->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); } - data->sendq.push_back(std::move(item)); } void QueueGpuToNetSend(Connection* data, const MemoryDesc& src, const std::vector& srcSegs, @@ -981,6 +1216,127 @@ class TcpTransport { }}); } + void QueueGpuToNetSendStriped(const EngineKey& peerKey, const std::vector& dataFds, + uint64_t userOpId, uint8_t lanesTotal, const MemoryDesc& src, + const std::vector& srcSegs) { + (void)peerKey; + if (dataFds.empty()) return; + const uint64_t total = SumLens(srcSegs); + auto pinned = staging.Acquire(static_cast(total)); + if (!pinned) { + MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); + return; + } + + hipStream_t stream = streamPool.GetNextStream(src.deviceId); + hipEvent_t ev = eventPool.GetEvent(src.deviceId); + if (stream == nullptr || ev == nullptr) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send"); + if (ev) eventPool.PutEvent(ev, src.deviceId); + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); + uint8_t* dst = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : srcSegs) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + gpuTasks.push_back({src.deviceId, ev, + [this, dataFds, pinned, userOpId, lanesTotal, total]() mutable { + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + const int fd = dataFds[lane % dataFds.size()]; + Connection* c = conns[fd].get(); + if (!c || c->fd < 0) continue; + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(userOpId, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = + static_cast(pinned->ptr) + static_cast(span.off); + item.iov[1].iov_len = static_cast(span.len); + item.keepalive = pinned; + c->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); + } + }}); + } + + void QueueGpuToNetSendStripedWithOnDone(const EngineKey& peerKey, const std::vector& dataFds, + uint64_t userOpId, uint8_t lanesTotal, const MemoryDesc& src, + const std::vector& srcSegs, + std::function onLaneDone) { + (void)peerKey; + if (dataFds.empty()) return; + const uint64_t total = SumLens(srcSegs); + auto pinned = staging.Acquire(static_cast(total)); + if (!pinned) { + MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send (striped)"); + return; + } + + hipStream_t stream = streamPool.GetNextStream(src.deviceId); + hipEvent_t ev = eventPool.GetEvent(src.deviceId); + if (stream == nullptr || ev == nullptr) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send (striped)"); + if (ev) eventPool.PutEvent(ev, src.deviceId); + return; + } + + HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); + uint8_t* dst = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : srcSegs) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + gpuTasks.push_back({src.deviceId, ev, [this, dataFds, pinned, userOpId, lanesTotal, total, + onLaneDone = std::move(onLaneDone)]() mutable { + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + const int fd = dataFds[lane % dataFds.size()]; + Connection* c = conns[fd].get(); + if (!c || c->fd < 0) continue; + SendItem item; + item.header = + tcp::BuildDataHeader(ToWireOpId(userOpId, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = + static_cast(pinned->ptr) + static_cast(span.off); + item.iov[1].iov_len = static_cast(span.len); + item.keepalive = pinned; + item.onDone = onLaneDone; + c->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); + } + }}); + } + struct GpuTask { int deviceId{-1}; hipEvent_t ev{nullptr}; @@ -1077,43 +1433,42 @@ class TcpTransport { void ClosePeerByFd(int fd) { MORI_IO_TRACE("TCP: close fd={}", fd); - // Find which peer+channel this fd belongs to, close both channels and fail pending ops. + // Find which peer+channel this fd belongs to, close all channels and fail pending ops. EngineKey peer{}; for (auto& kv : peers) { - if (kv.second.ctrlFd == fd || kv.second.dataFd == fd) { + if (kv.second.ctrlFd == fd) { peer = kv.first; break; } + for (int dfd : kv.second.dataFds) { + if (dfd == fd) { + peer = kv.first; + break; + } + } + if (!peer.empty()) break; } - Connection* c = conns[fd].get(); - if (c) { - CloseConnInternal(c); - conns.erase(fd); - } + auto close_fd = [&](int toClose) { + if (toClose < 0) return; + auto it = conns.find(toClose); + if (it == conns.end()) return; + CloseConnInternal(it->second.get()); + conns.erase(it); + }; if (!peer.empty()) { - PeerLinks& link = peers[peer]; - if (link.ctrlFd >= 0) { - Connection* cc = conns[link.ctrlFd].get(); - if (cc) { - CloseConnInternal(cc); - conns.erase(link.ctrlFd); - } - link.ctrlFd = -1; - } - if (link.dataFd >= 0) { - Connection* dc = conns[link.dataFd].get(); - if (dc) { - CloseConnInternal(dc); - conns.erase(link.dataFd); - } - link.dataFd = -1; - } + auto link = peers[peer]; + close_fd(link.ctrlFd); + for (int dfd : link.dataFds) close_fd(dfd); peers.erase(peer); FailPendingOpsForPeer(peer, "TCP: connection lost"); + return; } + + // Unknown fd: close just this connection. + close_fd(fd); } void FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg) { @@ -1274,98 +1629,174 @@ class TcpTransport { auto it = peers.find(peer); if (it == peers.end()) return nullptr; if (!it->second.DataUp()) return nullptr; - return conns[it->second.dataFd].get(); + const int fd = it->second.dataFds.empty() ? -1 : it->second.dataFds[0]; + if (fd < 0) return nullptr; + return conns[fd].get(); } - void FinalizeBufferedWrite(const EngineKey& peer, TransferUniqueId opId, InboundWriteState ws, - uint64_t payloadLen, const std::shared_ptr& pinned) { + static uint8_t ClampLanesTotal(uint8_t lanesTotal) { + if (lanesTotal == 0) return 1; + const uint8_t max = static_cast(1U << kLaneBits); + return std::min(lanesTotal, max); + } + + void MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId) { + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt == inboundWrites.end()) return; + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt == iwPeerIt->second.end()) return; + + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + const uint8_t allMask = LanesAllMask(ws.lanesTotal); + if ((ws.lanesDoneMask & allMask) != allMask) return; + auto ctrl = PeerCtrl(peer); - if (!ctrl) return; + if (!ctrl) { + // Peer is gone; drop state. + iwPeerIt->second.erase(wsIt); + if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); + return; + } if (ws.discard) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), - "TCP: write discarded")); + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), + "TCP: write discarded")); UpdateWriteInterest(ctrl->fd); RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); - return; - } - if (!pinned) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: missing pinned staging (early write)")); + } else if (ws.dst.loc == MemoryLocationType::GPU) { + if (!ws.pinned) { + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: missing pinned staging (write)")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging"); + } else { + hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); + hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); + if (stream == nullptr || ev == nullptr) { + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), + "TCP: failed to get HIP stream/event")); + UpdateWriteInterest(ctrl->fd); + RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, + "TCP: failed to get HIP stream/event"); + if (ev) eventPool.PutEvent(ev, ws.dst.deviceId); + } else { + HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); + uint8_t* src = reinterpret_cast(ws.pinned->ptr); + const uint64_t total = SumLens(ws.dstSegs); + uint64_t spanOff = 0; + uint64_t spanLen = 0; + if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { + void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : ws.dstSegs) { + void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); + HIP_RUNTIME_CHECK( + hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + + auto pinned = ws.pinned; + const int ctrlFd = ctrl->fd; + const int deviceId = ws.dst.deviceId; + gpuTasks.push_back({deviceId, ev, [this, peer, opId, ctrlFd, pinned]() { + QueueSend(ctrlFd, + tcp::BuildCompletion(opId, + static_cast(StatusCode::SUCCESS), + "")); + UpdateWriteInterest(ctrlFd); + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + }}); + } + } + } else { + QueueSend(ctrl->fd, + tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging"); - return; + RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); } - const uint64_t expected = SumLens(ws.dstSegs); - if (expected != payloadLen) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: early payload length mismatch")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: early payload length mismatch"); - return; + // Cleanup state. + iwPeerIt->second.erase(wsIt); + if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + ewPeerIt->second.erase(opId); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); } + } - if (ws.dst.loc == MemoryLocationType::GPU) { - hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); - hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); - if (stream == nullptr || ev == nullptr) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: failed to get HIP stream/event")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: failed to get HIP stream/event"); - if (ev) eventPool.PutEvent(ev, ws.dst.deviceId); - return; + void TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId) { + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt == inboundWrites.end()) return; + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt == iwPeerIt->second.end()) return; + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt == earlyWrites.end()) return; + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt == ewPeerIt->second.end()) return; + + EarlyWriteState& early = ewIt->second; + const uint64_t total = SumLens(ws.dstSegs); + uint8_t* dstBase = reinterpret_cast(ws.dst.data); + + for (auto it = early.lanes.begin(); it != early.lanes.end();) { + const uint8_t lane = it->first; + EarlyWriteLaneState& laneState = it->second; + if (!laneState.complete) { + ++it; + continue; } - HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); - uint8_t* src = reinterpret_cast(pinned->ptr); - const uint64_t total = payloadLen; - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { - void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : ws.dstSegs) { - void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); - off += s.len; + if (lane >= ws.lanesTotal) { + ws.discard = true; + } + const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); + if (span.len != laneState.payloadLen) { + ws.discard = true; + } + + if (!ws.discard && laneState.pinned) { + uint8_t* src = reinterpret_cast(laneState.pinned->ptr); + if (ws.dst.loc == MemoryLocationType::GPU) { + if (!ws.pinned) { + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; + } + if (!ws.discard && ws.pinned) { + std::memcpy(reinterpret_cast(ws.pinned->ptr) + span.off, src, + static_cast(span.len)); + } + } else { + const auto segs = SliceSegments(ws.dstSegs, span.off, span.len); + uint64_t copied = 0; + for (const auto& s : segs) { + std::memcpy(dstBase + s.off, src + copied, static_cast(s.len)); + copied += s.len; + } } } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - gpuTasks.push_back({ws.dst.deviceId, ev, [this, peer, opId, ctrlFd = ctrl->fd, pinned]() { - QueueSend(ctrlFd, - tcp::BuildCompletion(opId, - static_cast(StatusCode::SUCCESS), - "")); - UpdateWriteInterest(ctrlFd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); - }}); - return; + if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + it = early.lanes.erase(it); } - // CPU destination. - uint8_t* src = reinterpret_cast(pinned->ptr); - uint8_t* dstBase = reinterpret_cast(ws.dst.data); - const uint64_t total = payloadLen; - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { - std::memcpy(dstBase + spanOff, src, static_cast(total)); - } else { - uint64_t off = 0; - for (const auto& s : ws.dstSegs) { - std::memcpy(dstBase + s.off, src + off, static_cast(s.len)); - off += s.len; - } + if (early.lanes.empty()) { + ewPeerIt->second.erase(ewIt); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); } - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + MaybeFinalizeInboundWrite(peer, opId); } void HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { @@ -1380,10 +1811,17 @@ class TcpTransport { return; } + uint8_t lanesTotal = 1; + if (off < len) { + lanesTotal = body[off]; + } + lanesTotal = ClampLanesTotal(lanesTotal); + auto memOpt = LookupLocalMem(memId); InboundWriteState ws; ws.peer = peer; ws.id = opId; + ws.lanesTotal = lanesTotal; ws.discard = true; if (memOpt.has_value()) { @@ -1394,25 +1832,14 @@ class TcpTransport { } } - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt != ewPeerIt->second.end()) { - EarlyWriteState early = std::move(ewIt->second); - ewPeerIt->second.erase(ewIt); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); - - if (early.complete) { - FinalizeBufferedWrite(peer, opId, std::move(ws), early.payloadLen, early.pinned); - return; - } - // Data is still in flight; keep the write state so data RX can finalize it. - inboundWrites[peer][opId] = std::move(ws); - return; - } + if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { + const uint64_t total = SumLens(ws.dstSegs); + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; } inboundWrites[peer][opId] = std::move(ws); + TryConsumeEarlyWriteLanes(peer, opId); } void HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { @@ -1432,40 +1859,36 @@ class TcpTransport { ws.id = opId; ws.discard = true; - if (memOpt.has_value()) { - ws.dst = memOpt.value(); - ws.dstSegs.reserve(n); - bool ok = true; - for (uint32_t i = 0; i < n; ++i) { - uint64_t ro = 0, sz = 0; - if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { - ok = false; - break; - } - if (ro + sz > ws.dst.size) ok = false; - if (sz > 0) ws.dstSegs.push_back({ro, sz}); + const bool haveMem = memOpt.has_value(); + if (haveMem) ws.dst = memOpt.value(); + + ws.dstSegs.reserve(n); + bool ok = true; + for (uint32_t i = 0; i < n; ++i) { + uint64_t ro = 0, sz = 0; + if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { + ok = false; + break; } - if (ok) ws.discard = false; + if (haveMem && (ro + sz > ws.dst.size)) ok = false; + if (haveMem && sz > 0) ws.dstSegs.push_back({ro, sz}); } - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt != ewPeerIt->second.end()) { - EarlyWriteState early = std::move(ewIt->second); - ewPeerIt->second.erase(ewIt); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + // `lanesTotal` is appended by the sender; tolerate older senders. + uint8_t lanesTotal = 1; + if (off < len) lanesTotal = body[off]; + ws.lanesTotal = ClampLanesTotal(lanesTotal); - if (early.complete) { - FinalizeBufferedWrite(peer, opId, std::move(ws), early.payloadLen, early.pinned); - return; - } - inboundWrites[peer][opId] = std::move(ws); - return; - } + if (haveMem && ok) ws.discard = false; + + if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { + const uint64_t total = SumLens(ws.dstSegs); + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; } inboundWrites[peer][opId] = std::move(ws); + TryConsumeEarlyWriteLanes(peer, opId); } void HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { @@ -1479,6 +1902,9 @@ class TcpTransport { MORI_IO_WARN("TCP: malformed READ_REQ"); return; } + uint8_t lanesTotal = 1; + if (off < len) lanesTotal = body[off]; + lanesTotal = ClampLanesTotal(lanesTotal); auto memOpt = LookupLocalMem(memId); if (!memOpt.has_value() || (srcOff + size > memOpt->size)) { @@ -1495,7 +1921,7 @@ class TcpTransport { MemoryDesc src = memOpt.value(); std::vector segs = {{srcOff, size}}; - QueueDataSendForRead(peer, opId, src, segs); + QueueDataSendForRead(peer, opId, src, segs, lanesTotal); } void HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { @@ -1535,6 +1961,9 @@ class TcpTransport { if (ro + sz > src.size) ok = false; if (sz > 0) segs.push_back({ro, sz}); } + uint8_t lanesTotal = 1; + if (off < len) lanesTotal = body[off]; + lanesTotal = ClampLanesTotal(lanesTotal); if (!ok) { auto ctrl = PeerCtrl(peer); if (ctrl) { @@ -1547,39 +1976,110 @@ class TcpTransport { return; } - QueueDataSendForRead(peer, opId, src, segs); + QueueDataSendForRead(peer, opId, src, segs, lanesTotal); } void QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, - const std::vector& srcSegs) { - Connection* data = PeerData(peer); + const std::vector& srcSegs, uint8_t lanesTotal) { Connection* ctrl = PeerCtrl(peer); - if (!data || !ctrl) return; + if (!ctrl) return; const uint64_t total = SumLens(srcSegs); - auto hdr = tcp::BuildDataHeader(opId, total, 0); - auto onDone = [this, peer, opId, ctrlFd = ctrl->fd]() { - QueueSend(ctrlFd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); - UpdateWriteInterest(ctrlFd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); + lanesTotal = ClampLanesTotal(lanesTotal); + + // Build list of active data fds. + std::vector dataFds; + auto pit = peers.find(peer); + if (pit != peers.end()) { + for (int fd : pit->second.dataFds) { + auto dit = conns.find(fd); + if (dit == conns.end()) continue; + if (dit->second && dit->second->helloReceived) dataFds.push_back(fd); + } + } + if (dataFds.empty()) return; + lanesTotal = std::min(lanesTotal, static_cast(dataFds.size())); + + struct DoneState { + EngineKey peer; + uint64_t opId{0}; + int ctrlFd{-1}; + int remaining{0}; + }; + auto done = std::make_shared(); + done->peer = peer; + done->opId = opId; + done->ctrlFd = ctrl->fd; + done->remaining = lanesTotal; + auto laneDone = [this, done]() mutable { + done->remaining--; + if (done->remaining > 0) return; + QueueSend(done->ctrlFd, + tcp::BuildCompletion(done->opId, static_cast(StatusCode::SUCCESS), "")); + UpdateWriteInterest(done->ctrlFd); + RecordInboundStatus(done->peer, done->opId, StatusCode::SUCCESS, ""); }; + // Unstriped path. + if (lanesTotal == 1) { + const int fd = dataFds[0]; + Connection* data = conns[fd].get(); + if (!data) return; + auto hdr = tcp::BuildDataHeader(ToWireOpId(opId, 0), total, 0); + if (src.loc == MemoryLocationType::GPU) { + QueueGpuToNetSend(data, src, srcSegs, std::move(hdr), std::move(laneDone)); + UpdateWriteInterest(fd); + return; + } + + SendItem item; + item.header = std::move(hdr); + item.iov.reserve(1 + srcSegs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + uint8_t* base = reinterpret_cast(src.data); + for (const auto& s : srcSegs) { + item.iov.push_back({base + s.off, static_cast(s.len)}); + } + item.onDone = std::move(laneDone); + data->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); + return; + } + + // Striped path: requires a single contiguous span. + if (srcSegs.size() != 1) { + MORI_IO_WARN("TCP: peer {} READ striping requested but srcSegs.size={} (expected 1), fallback to 1 lane", + peer.c_str(), srcSegs.size()); + QueueDataSendForRead(peer, opId, src, srcSegs, /*lanesTotal=*/1); + return; + } + if (src.loc == MemoryLocationType::GPU) { - QueueGpuToNetSend(data, src, srcSegs, std::move(hdr), std::move(onDone)); + // Stage once, then stripe sends from pinned buffer. + const uint8_t useLanes = + std::min(lanesTotal, static_cast(dataFds.size())); + QueueGpuToNetSendStripedWithOnDone(peer, dataFds, opId, useLanes, src, srcSegs, std::move(laneDone)); return; } - SendItem item; - item.header = std::move(hdr); - item.iov.reserve(1 + srcSegs.size()); - item.iov.push_back({item.header.data(), item.header.size()}); - uint8_t* base = reinterpret_cast(src.data); - for (const auto& s : srcSegs) { - item.iov.push_back({base + s.off, static_cast(s.len)}); + const uint8_t useLanes = std::min(lanesTotal, static_cast(dataFds.size())); + uint8_t* base = reinterpret_cast(src.data) + srcSegs[0].off; + for (uint8_t lane = 0; lane < useLanes; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, useLanes, lane); + const int fd = dataFds[lane % dataFds.size()]; + Connection* data = conns[fd].get(); + if (!data) continue; + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = base + span.off; + item.iov[1].iov_len = static_cast(span.len); + item.onDone = laneDone; + data->sendq.push_back(std::move(item)); + UpdateWriteInterest(fd); } - item.onDone = std::move(onDone); - data->sendq.push_back(std::move(item)); - UpdateWriteInterest(data->fd); } void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { @@ -1618,6 +2118,9 @@ class TcpTransport { void MaybeCompleteOutbound(OutboundOpState& st) { if (!st.completionReceived) return; if (st.isRead) { + const uint8_t allMask = + st.lanesTotal >= (1U << kLaneBits) ? 0xFF : static_cast((1U << st.lanesTotal) - 1U); + if (st.lanesDoneMask != allMask) return; if (st.rxBytes != st.expectedRxBytes) return; if (st.gpuCopyPending) return; } @@ -1788,40 +2291,56 @@ class TcpTransport { } } - void BeginDataRx(Connection* c, TransferUniqueId opId, uint64_t payloadLen) { + void BeginDataRx(Connection* c, uint64_t wireOpId, uint64_t payloadLen) { c->rx = ActiveDataRx{}; c->rx.active = true; - c->rx.id = opId; + + const uint8_t lane = static_cast(wireOpId & kLaneMask); + const TransferUniqueId userOpId = static_cast(ToUserOpId(wireOpId)); + + c->rx.id = userOpId; + c->rx.lane = lane; + c->rx.laneLen = payloadLen; c->rx.remaining = payloadLen; - // 1) inbound write: peer->me, stored by peer key. - auto iwIt = inboundWrites.find(c->peerKey); - if (iwIt != inboundWrites.end()) { - auto it = iwIt->second.find(opId); - if (it != iwIt->second.end()) { - InboundWriteState& ws = it->second; - if (!ws.discard) { - const uint64_t expected = SumLens(ws.dstSegs); - if (expected != payloadLen) { - MORI_IO_WARN("TCP: inbound write op {} payloadLen mismatch expected={} got={}", opId, - expected, payloadLen); - ws.discard = true; - } + // 1) inbound write: peer -> me. + auto iwPeerIt = inboundWrites.find(c->peerKey); + if (iwPeerIt != inboundWrites.end()) { + auto wsIt = iwPeerIt->second.find(userOpId); + if (wsIt != iwPeerIt->second.end()) { + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + c->rx.kind = DataRxKind::INBOUND_WRITE; + + const uint64_t total = SumLens(ws.dstSegs); + if (lane >= ws.lanesTotal) ws.discard = true; + const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); + c->rx.laneOff = span.off; + c->rx.laneLen = payloadLen; + + if (!ws.discard && span.len != payloadLen) { + MORI_IO_WARN("TCP: inbound write op {} lane {} payloadLen mismatch expected={} got={}", + userOpId, static_cast(lane), span.len, payloadLen); + ws.discard = true; } + c->rx.discard = ws.discard; - if (!ws.discard) { + if (!c->rx.discard) { if (ws.dst.loc == MemoryLocationType::GPU) { - c->rx.toGpu = true; - c->rx.gpuDevice = ws.dst.deviceId; - c->rx.gpuBase = reinterpret_cast(ws.dst.data); - c->rx.segs = ws.dstSegs; - const uint64_t total = SumLens(ws.dstSegs); - c->rx.pinned = staging.Acquire(static_cast(total)); - c->rx.pinnedWriteOff = 0; + if (!ws.pinned) { + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; + } + c->rx.discard = ws.discard; + if (!c->rx.discard) { + c->rx.toGpu = true; + c->rx.pinned = ws.pinned; + c->rx.pinnedWriteOff = span.off; + } } else { c->rx.toGpu = false; c->rx.base = reinterpret_cast(ws.dst.data); - c->rx.segs = ws.dstSegs; + c->rx.segs = SliceSegments(ws.dstSegs, span.off, span.len); c->rx.segIdx = 0; c->rx.segOff = 0; } @@ -1830,52 +2349,73 @@ class TcpTransport { } } - // 2) outbound read response: peer->me for my pending outbound read. - auto obIt = pendingOutbound.find(opId); + // 2) outbound read response: peer -> me for pending outbound read. + auto obIt = pendingOutbound.find(userOpId); if (obIt != pendingOutbound.end()) { OutboundOpState& st = *obIt->second; if (!st.isRead) { + c->rx.kind = DataRxKind::OUTBOUND_READ; c->rx.discard = true; return; } - if (payloadLen != st.expectedRxBytes) { - MORI_IO_ERROR("TCP: outbound read op {} payloadLen mismatch expected={} got={}", opId, - st.expectedRxBytes, payloadLen); + + st.lanesTotal = ClampLanesTotal(st.lanesTotal); + if (lane >= st.lanesTotal) { + c->rx.kind = DataRxKind::OUTBOUND_READ; + c->rx.discard = true; + return; + } + + const LaneSpan span = ComputeLaneSpan(st.expectedRxBytes, st.lanesTotal, lane); + c->rx.kind = DataRxKind::OUTBOUND_READ; + c->rx.laneOff = span.off; + if (span.len != payloadLen) { + MORI_IO_ERROR("TCP: outbound read op {} lane {} payloadLen mismatch expected={} got={}", + userOpId, static_cast(lane), span.len, payloadLen); st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: read payload length mismatch"); pendingOutbound.erase(obIt); c->rx.discard = true; return; } + c->rx.discard = false; if (st.local.loc == MemoryLocationType::GPU) { + if (!st.pinned) { + st.pinned = staging.Acquire(static_cast(st.expectedRxBytes)); + if (!st.pinned) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging (read)"); + pendingOutbound.erase(obIt); + c->rx.discard = true; + return; + } + } c->rx.toGpu = true; - c->rx.gpuDevice = st.local.deviceId; - c->rx.gpuBase = reinterpret_cast(st.local.data); - c->rx.segs = st.localSegs; - c->rx.pinned = staging.Acquire(static_cast(payloadLen)); - c->rx.pinnedWriteOff = 0; + c->rx.pinned = st.pinned; + c->rx.pinnedWriteOff = span.off; } else { c->rx.toGpu = false; c->rx.base = reinterpret_cast(st.local.data); - c->rx.segs = st.localSegs; + c->rx.segs = SliceSegments(st.localSegs, span.off, span.len); c->rx.segIdx = 0; c->rx.segOff = 0; } return; } - // Control+data are on separate TCP connections; data may arrive before its CTRL write request. - // Buffer it in pinned host memory, and finalize once CTRL arrives. + // 3) early-arrived write payload (data before CTRL write request). + c->rx.kind = DataRxKind::EARLY_WRITE; if (payloadLen > static_cast(std::numeric_limits::max())) { - MORI_IO_WARN("TCP: data payloadLen too large for op {} from peer {}, discarding", opId, + MORI_IO_WARN("TCP: data payloadLen too large for op {} from peer {}, discarding", userOpId, c->peerKey.c_str()); c->rx.discard = true; return; } + auto& perPeer = earlyWrites[c->peerKey]; - if (perPeer.find(opId) != perPeer.end()) { - MORI_IO_WARN("TCP: duplicate early data for op {} from peer {}, discarding", opId, - c->peerKey.c_str()); + EarlyWriteState& early = perPeer[userOpId]; + if (early.lanes.find(lane) != early.lanes.end()) { + MORI_IO_WARN("TCP: duplicate early data for op {} lane {} from peer {}, discarding", userOpId, + static_cast(lane), c->peerKey.c_str()); c->rx.discard = true; return; } @@ -1884,12 +2424,12 @@ class TcpTransport { auto pinned = staging.Acquire(alloc); if (!pinned) { MORI_IO_WARN("TCP: failed to allocate pinned buffer for early data op {} from peer {}, discarding", - opId, c->peerKey.c_str()); + userOpId, c->peerKey.c_str()); c->rx.discard = true; return; } - perPeer.emplace(opId, EarlyWriteState{payloadLen, pinned, false}); + early.lanes.emplace(lane, EarlyWriteLaneState{payloadLen, pinned, false}); c->rx.discard = false; c->rx.toGpu = true; // receive into pinned c->rx.pinned = pinned; @@ -1899,125 +2439,59 @@ class TcpTransport { void FinishDataRx(Connection* c) { const TransferUniqueId opId = c->rx.id; const EngineKey peer = c->peerKey; + const uint8_t lane = c->rx.lane; + const uint64_t laneLen = c->rx.laneLen; + const DataRxKind kind = c->rx.kind; - // Case A: inbound write complete -> copy to GPU if needed, then send completion. - auto iwIt = inboundWrites.find(peer); - if (iwIt != inboundWrites.end()) { - auto it = iwIt->second.find(opId); - if (it != iwIt->second.end()) { - InboundWriteState ws = std::move(it->second); - iwIt->second.erase(it); - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - ewPeerIt->second.erase(opId); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); - } + c->rx = ActiveDataRx{}; - if (ws.discard) { - auto ctrl = PeerCtrl(peer); - if (ctrl) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), - "TCP: write discarded")); - UpdateWriteInterest(ctrl->fd); - } - RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); - c->rx = ActiveDataRx{}; - return; - } + if (kind == DataRxKind::INBOUND_WRITE) { + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt == inboundWrites.end()) return; + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt == iwPeerIt->second.end()) return; + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + MaybeFinalizeInboundWrite(peer, opId); + return; + } - auto ctrl = PeerCtrl(peer); - if (!ctrl) { - c->rx = ActiveDataRx{}; + if (kind == DataRxKind::OUTBOUND_READ) { + auto obIt = pendingOutbound.find(opId); + if (obIt == pendingOutbound.end()) return; + OutboundOpState& st = *obIt->second; + st.lanesTotal = ClampLanesTotal(st.lanesTotal); + const uint8_t bit = static_cast(1U << lane); + if ((st.lanesDoneMask & bit) == 0) { + st.lanesDoneMask |= bit; + st.rxBytes += laneLen; + } + + if (st.local.loc == MemoryLocationType::GPU) { + const uint8_t allMask = LanesAllMask(st.lanesTotal); + if ((st.lanesDoneMask & allMask) != allMask) { + MaybeCompleteOutbound(st); return; } - - if (ws.dst.loc == MemoryLocationType::GPU) { - hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); - hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); - if (stream == nullptr || ev == nullptr || !c->rx.pinned) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: GPU staging missing")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: GPU staging missing"); - c->rx = ActiveDataRx{}; - return; - } - - HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); - uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); - const uint64_t total = SumLens(ws.dstSegs); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { - void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : ws.dstSegs) { - void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); - HIP_RUNTIME_CHECK( - hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - auto pinned = c->rx.pinned; - c->rx = ActiveDataRx{}; - - gpuTasks.push_back({ws.dst.deviceId, ev, [this, peer, opId, ctrlFd = ctrl->fd, pinned]() { - QueueSend(ctrlFd, - tcp::BuildCompletion(opId, - static_cast(StatusCode::SUCCESS), - "")); - UpdateWriteInterest(ctrlFd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); - }}); + if (st.gpuCopyPending) return; + if (!st.pinned) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging (read)"); + pendingOutbound.erase(obIt); return; - } - - if (c->rx.pinned) { - uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); - uint8_t* dstBase = reinterpret_cast(ws.dst.data); - const uint64_t total = SumLens(ws.dstSegs); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { - std::memcpy(dstBase + spanOff, src, static_cast(total)); - } else { - uint64_t off = 0; - for (const auto& s : ws.dstSegs) { - std::memcpy(dstBase + s.off, src + off, static_cast(s.len)); - off += s.len; - } - } - } - - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); - c->rx = ActiveDataRx{}; - return; - } - } + } - // Case B: outbound read response received -> copy to GPU if needed and update status. - auto obIt = pendingOutbound.find(opId); - if (obIt != pendingOutbound.end()) { - OutboundOpState& st = *obIt->second; - st.rxBytes = st.expectedRxBytes; - if (st.local.loc == MemoryLocationType::GPU) { hipStream_t stream = streamPool.GetNextStream(st.local.deviceId); hipEvent_t ev = eventPool.GetEvent(st.local.deviceId); - if (stream == nullptr || ev == nullptr || !c->rx.pinned) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: GPU staging missing (read)"); + if (stream == nullptr || ev == nullptr) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to get HIP stream/event (read)"); pendingOutbound.erase(obIt); - c->rx = ActiveDataRx{}; + if (ev) eventPool.PutEvent(ev, st.local.deviceId); return; } HIP_RUNTIME_CHECK(hipSetDevice(st.local.deviceId)); - uint8_t* src = reinterpret_cast(c->rx.pinned->ptr); + uint8_t* src = reinterpret_cast(st.pinned->ptr); st.gpuCopyPending = true; const uint64_t total = st.expectedRxBytes; uint64_t spanOff = 0; @@ -2029,16 +2503,13 @@ class TcpTransport { uint64_t off = 0; for (const auto& s : st.localSegs) { void* gpuPtr = reinterpret_cast(st.local.data + s.off); - HIP_RUNTIME_CHECK( - hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); off += s.len; } } HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - auto pinned = c->rx.pinned; - c->rx = ActiveDataRx{}; - + auto pinned = st.pinned; gpuTasks.push_back({st.local.deviceId, ev, [this, opId, pinned]() { auto it = pendingOutbound.find(opId); if (it == pendingOutbound.end()) return; @@ -2048,25 +2519,23 @@ class TcpTransport { return; } - c->rx = ActiveDataRx{}; - MaybeCompleteOutbound(st); - return; - } - - // Case C: early-arrived write payload (data before CTRL). Mark it complete and finalize once - // the corresponding CTRL write request arrives. - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt != ewPeerIt->second.end()) { - ewIt->second.complete = true; - c->rx = ActiveDataRx{}; - return; - } - } + MaybeCompleteOutbound(st); + return; + } - c->rx = ActiveDataRx{}; - } + if (kind == DataRxKind::EARLY_WRITE) { + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt == earlyWrites.end()) return; + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt == ewPeerIt->second.end()) return; + auto laneIt = ewIt->second.lanes.find(lane); + if (laneIt != ewIt->second.lanes.end()) { + laneIt->second.complete = true; + } + TryConsumeEarlyWriteLanes(peer, opId); + return; + } + } void HandleReadable(Connection* c) { if (!c->helloReceived) { From 1222b6e09ef086ea9d7cc108af64e55baf0158c3 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 27 Feb 2026 05:39:33 +0000 Subject: [PATCH 09/12] io/tcp: trace lane selection and data conns --- src/io/tcp/backend_impl.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/io/tcp/backend_impl.cpp b/src/io/tcp/backend_impl.cpp index 3cba24098..c361bc30b 100644 --- a/src/io/tcp/backend_impl.cpp +++ b/src/io/tcp/backend_impl.cpp @@ -778,6 +778,8 @@ class TcpTransport { return; } link.dataFds.push_back(c->fd); + MORI_IO_TRACE("TCP: peer {} DATA conn up {}/{}", c->peerKey.c_str(), link.dataFds.size(), + want); } void MaybeDispatchQueuedOps(const EngineKey& peerKey) { @@ -1038,6 +1040,8 @@ class TcpTransport { lanesTotal = static_cast(std::min(static_cast(wantLanes), dataFds.size())); } st->lanesTotal = lanesTotal; + MORI_IO_TRACE("TCP: op {} totalBytes={} dataConns={} lanesTotal={}", st->id, totalBytes, + dataFds.size(), static_cast(lanesTotal)); if (st->isRead && st->local.loc == MemoryLocationType::GPU) { // Allocate a single pinned staging buffer for all lanes and do one H2D copy at the end. From bbc33e62e401844906c68e4915983fc16c2f9f3f Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Fri, 27 Feb 2026 13:00:22 +0000 Subject: [PATCH 10/12] refactor --- docs/MORI-IO-TCP-BACKEND.md | 189 --- include/mori/io/backend.hpp | 15 +- src/io/CMakeLists.txt | 1 + src/io/tcp/backend_impl.cpp | 2693 +---------------------------------- src/io/tcp/data_worker.hpp | 426 ++++++ src/io/tcp/tcp_types.hpp | 421 ++++++ src/io/tcp/transport.cpp | 1657 +++++++++++++++++++++ src/io/tcp/transport.hpp | 210 +++ 8 files changed, 2751 insertions(+), 2861 deletions(-) delete mode 100644 docs/MORI-IO-TCP-BACKEND.md create mode 100644 src/io/tcp/data_worker.hpp create mode 100644 src/io/tcp/tcp_types.hpp create mode 100644 src/io/tcp/transport.cpp create mode 100644 src/io/tcp/transport.hpp diff --git a/docs/MORI-IO-TCP-BACKEND.md b/docs/MORI-IO-TCP-BACKEND.md deleted file mode 100644 index d508507ac..000000000 --- a/docs/MORI-IO-TCP-BACKEND.md +++ /dev/null @@ -1,189 +0,0 @@ -# MORI-IO TCP Backend Design - -> **Goal**: Provide a high-reliability, high-performance **TCP fallback transport** for distributed KV-Cache transfer when RDMA/XGMI is unavailable. -> The TCP backend must **simulate RDMA-like Read/Write semantics** on top of TCP byte streams, while keeping the public `mori::io::Backend` / `BackendSession` APIs aligned with existing backends. - ---- - -## 1. Requirements Mapping - -### 1) Semantic Mapping (RDMA-like ops on TCP stream) -TCP is a byte stream; MORI-IO upper layers are op/message oriented (Read/Write/BatchRead/BatchWrite). - -**Solution**: a framed protocol: -- **Control channel** carries op metadata (`WriteReq`, `ReadReq`, `Batch*Req`) and `Completion`. -- **Data channel** carries payload with a lightweight data frame header (`op_id`, `payload_len`). - -This avoids classic TCP sticky/half packet issues by always parsing `FixedHeader + VariableBody`. - -### 2) Concurrency Model (non-blocking + multiplexing) -**Baseline**: `epoll`-based Reactor (single IO thread per backend instance). -- Non-blocking sockets, edge-triggered `epoll`. -- `eventfd` used to wake the reactor when new outbound ops are submitted from application threads. -- Optional `timerfd` for op timeout scanning & housekeeping. - -**Evaluation**: Kernel is 5.15+ (io_uring-capable). The design keeps a clean boundary so `io_uring` Proactor can replace the Reactor later without changing public APIs. - -### 3) Zero-Copy & Memory Management -#### CPU memory -- Outbound: `sendmsg` + `iovec` directly referencing user registered buffers (no extra user-space copy). -- Inbound: `recv` directly into destination user buffer. -- Optional: enable `MSG_ZEROCOPY` for large CPU payloads (best-effort). - -#### GPU memory (fallback path) -TCP cannot directly read/write device VRAM. We stage through pinned host buffers: -- GPU→Host (`hipMemcpyDtoHAsync`) before sending payload. -- Host→GPU (`hipMemcpyHtoDAsync`) after receiving payload. -- Pinned host buffers are pooled to avoid per-op allocations. - -### 4) Latency vs Throughput -- **Metadata**: control channel uses `TCP_NODELAY` for low latency. -- **Payload**: data channel uses large `SO_SNDBUF/SO_RCVBUF`, and can optionally enable `MSG_ZEROCOPY` for large CPU payloads. -- Backpressure: bound in-flight bytes per peer to avoid unbounded buffering. - -### 5) Robustness -- NIC selection: bind outbound sockets to `IOEngineConfig.host` before connect; listen on the same host. -- Keep-alive: enable `SO_KEEPALIVE` + reasonable keepalive parameters. -- Disconnect: fail in-flight ops with `ERR_BAD_STATE` and allow lazy reconnect on next submission. -- Timeout: per-op timeout (configurable) -> fail op and clean state. - ---- - -## 2. Architecture - -### 2.1 Components - -- `TcpBackend` (`mori::io::Backend`) - - Owns one `TcpTransport` instance. - - Implements `RegisterRemoteEngine`, `RegisterMemory`, `Read/Write/Batch*`, `CreateSession`. - -- `TcpBackendSession` (`mori::io::BackendSession`) - - Caches local/remote descriptors and routes ops to `TcpTransport` (low overhead fast path). - -- `TcpTransport` - - A single IO thread (reactor) managing: - - Listener socket - - Per-peer connections (control + data) - - Framed parsing / send queue - - Pending outbound op table - - Inbound op handling and inbound status queue - -- `PeerConn` - - `ctrl_fd` + `data_fd` - - control parser state, data parser state - - outbound queues for ctrl and data - -- `PinnedStagingPool` - - Reusable pinned host buffers keyed by `(device_id, size_class)` (GPU staging). - -### 2.2 Threading - -Single IO thread per `TcpBackend`: -- No thread-per-connection. -- All network IO and protocol parsing happens in this thread. -- Application threads only enqueue op descriptors (cheap) and return immediately with `TransferStatus = IN_PROGRESS`. - ---- - -## 3. Wire Protocol - -### 3.1 Channels -For each peer engine we maintain **two TCP connections**: -- `CTRL`: op metadata + completions -- `DATA`: bulk payload frames - -Both connections connect to the same listen port; the first message is a `HELLO` identifying: -- protocol version -- channel type (`CTRL` / `DATA`) -- `EngineKey` string - -### 3.2 Control Frame -All ctrl messages are framed: - -``` -| CtrlHeader (fixed) | CtrlBody (variable, body_len bytes) | -``` - -`CtrlHeader` includes: -- magic/version -- message type -- body length - -`CtrlBody` contains type-specific fields (op_id, mem_id, offsets, sizes, etc). - -### 3.3 Data Frame - -``` -| DataHeader (fixed) | Payload (payload_len bytes) | -``` - -`DataHeader` includes: -- `op_id` -- `payload_len` - -Ordering: -- Write: `CTRL WriteReq` then `DATA (op_id,payload)` then `CTRL Completion` -- Read: `CTRL ReadReq` then `DATA (op_id,payload)` then `CTRL Completion` - -Receiver keeps op state keyed by `op_id` so ctrl/data can be processed independently. - -#### Cross-channel reordering (CTRL vs DATA) -Because **CTRL** and **DATA** are separate TCP connections, the receiver may observe **DATA arriving before the corresponding CTRL request** (no global ordering across connections). - -To preserve RDMA-like op semantics without adding an extra RTT handshake, the TCP backend must: -- Buffer such early DATA frames by `(peer_key, op_id)` into pinned host memory. -- Finalize the write once the CTRL request arrives (or fail/cleanup on disconnect/timeout). - -For GPU destinations this maps naturally to the existing pinned-staging path; for CPU destinations this fallback path may add an extra copy only in the reordering case. - ---- - -## 4. API Semantics - -- `Write(local, remote, ...)`: - - initiator sends payload; target writes into its local memory region. -- `Read(local, remote, ...)`: - - initiator requests; target reads from its local memory region and sends payload back. -- `PopInboundTransferStatus(remote_key, id, ...)`: - - target side may pop completion of inbound ops (mainly useful for write-like semantics), consistent with existing RDMA notification API. - ---- - -## 5. Config Surface (planned) - -`TcpBackendConfig`: -- `num_io_threads` (default 1) -- `max_inflight_bytes_per_peer` -- `op_timeout_ms` -- `sock_sndbuf_bytes`, `sock_rcvbuf_bytes` -- `enable_zerocopy`, `zerocopy_threshold_bytes` -- `enable_keepalive`, keepalive parameters - ---- - -## 6. Planned Code Changes - -### New files -- `src/io/tcp/backend_impl.hpp` -- `src/io/tcp/backend_impl.cpp` -- `src/io/tcp/protocol.hpp` -- `src/io/tcp/protocol.cpp` -- `src/io/tcp/poller_epoll.hpp` (reactor) -- `src/io/tcp/pinned_staging_pool.hpp` - -### Modified files -- `include/mori/io/backend.hpp` (add `TcpBackendConfig`) -- `src/io/engine.cpp` (wire in `BackendType::TCP`) -- `src/io/CMakeLists.txt` (build TCP backend sources) -- `src/pybind/mori.cpp` (bind `TcpBackendConfig`) -- `python/mori/io/engine.py` (enable `BackendType.TCP`) -- `tests/python/io/benchmark.py` (support `--backend tcp` and `--op` alias) - ---- - -## 7. Known Limits / Future Work - -- `io_uring` Proactor implementation (replace epoll reactor in `TcpTransport`). -- Adaptive multi-connection striping per peer to improve throughput on high-BDP networks. -- Full `MSG_ZEROCOPY` completion accounting + buffer lifetime tracking (currently best-effort). -- Optional TLS / auth for multi-tenant environments. diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index 14031dce8..99cf72711 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -103,8 +103,8 @@ struct TcpBackendConfig : public BackendConfig { int numIoThreads{1}; - int sockSndbufBytes{4 * 1024 * 1024}; - int sockRcvbufBytes{4 * 1024 * 1024}; + int sockSndbufBytes{32 * 1024 * 1024}; + int sockRcvbufBytes{32 * 1024 * 1024}; int opTimeoutMs{30 * 1000}; @@ -114,13 +114,13 @@ struct TcpBackendConfig : public BackendConfig { int keepaliveCnt{3}; bool enableCtrlNodelay{true}; - bool enableDataNodelay{false}; + bool enableDataNodelay{true}; // Number of parallel DATA TCP connections per peer (iperf-like multi-stream striping). // Effective only when peer has >= numDataConns established and transfer is contiguous. int numDataConns{4}; // Stripe large contiguous transfers; keep small transfers on a single stream for latency. - int stripingThresholdBytes{256 * 1024}; + int stripingThresholdBytes{64 * 1024}; }; inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { @@ -128,10 +128,9 @@ inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { << "] sockRcvbufBytes[" << c.sockRcvbufBytes << "] opTimeoutMs[" << c.opTimeoutMs << "] enableKeepalive[" << c.enableKeepalive << "] keepaliveIdleSec[" << c.keepaliveIdleSec << "] keepaliveIntvlSec[" << c.keepaliveIntvlSec - << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" - << c.enableCtrlNodelay << "] enableDataNodelay[" << c.enableDataNodelay - << "] numDataConns[" << c.numDataConns << "] stripingThresholdBytes[" - << c.stripingThresholdBytes << "]"; + << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" << c.enableCtrlNodelay + << "] enableDataNodelay[" << c.enableDataNodelay << "] numDataConns[" << c.numDataConns + << "] stripingThresholdBytes[" << c.stripingThresholdBytes << "]"; } /* ---------------------------------------------------------------------------------------------- */ diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index a7b58f98f..342aa7c62 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -6,6 +6,7 @@ add_library( rdma/executor.cpp rdma/common.cpp tcp/backend_impl.cpp + tcp/transport.cpp xgmi/backend_impl.cpp xgmi/hip_resource_pool.cpp) diff --git a/src/io/tcp/backend_impl.cpp b/src/io/tcp/backend_impl.cpp index c361bc30b..309527cac 100644 --- a/src/io/tcp/backend_impl.cpp +++ b/src/io/tcp/backend_impl.cpp @@ -19,2677 +19,34 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. +// Copyright © Advanced Micro Devices, Inc. All rights reserved. #include "src/io/tcp/backend_impl.hpp" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "mori/application/utils/check.hpp" -#include "mori/io/logging.hpp" -#include "src/io/tcp/protocol.hpp" -#include "src/io/xgmi/hip_resource_pool.hpp" +#include "src/io/tcp/transport.hpp" namespace mori { namespace io { -namespace { - -using Clock = std::chrono::steady_clock; - -inline bool IsWouldBlock(int err) { return (err == EAGAIN) || (err == EWOULDBLOCK); } - -inline int SetNonBlocking(int fd) { - int flags = fcntl(fd, F_GETFL, 0); - if (flags < 0) return -1; - if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) return -1; - return 0; -} - -inline void SetSockOptOrLog(int fd, int level, int optname, const void* optval, socklen_t optlen, - const char* name) { - if (setsockopt(fd, level, optname, optval, optlen) != 0) { - MORI_IO_WARN("TCP: setsockopt {} failed: {}", name, strerror(errno)); - } -} - -inline void ConfigureSocketCommon(int fd, const TcpBackendConfig& cfg) { - if (cfg.enableKeepalive) { - int on = 1; - SetSockOptOrLog(fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on), "SO_KEEPALIVE"); - - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPIDLE, &cfg.keepaliveIdleSec, - sizeof(cfg.keepaliveIdleSec), "TCP_KEEPIDLE"); - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPINTVL, &cfg.keepaliveIntvlSec, - sizeof(cfg.keepaliveIntvlSec), "TCP_KEEPINTVL"); - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPCNT, &cfg.keepaliveCnt, sizeof(cfg.keepaliveCnt), - "TCP_KEEPCNT"); - } -} - -inline void ConfigureCtrlSocket(int fd, const TcpBackendConfig& cfg) { - ConfigureSocketCommon(fd, cfg); - if (cfg.enableCtrlNodelay) { - int on = 1; - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(ctrl)"); - } -} - -inline void ConfigureDataSocket(int fd, const TcpBackendConfig& cfg) { - ConfigureSocketCommon(fd, cfg); - if (cfg.enableDataNodelay) { - int on = 1; - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(data)"); - } - if (cfg.sockSndbufBytes > 0) { - SetSockOptOrLog(fd, SOL_SOCKET, SO_SNDBUF, &cfg.sockSndbufBytes, sizeof(cfg.sockSndbufBytes), - "SO_SNDBUF"); - } - if (cfg.sockRcvbufBytes > 0) { - SetSockOptOrLog(fd, SOL_SOCKET, SO_RCVBUF, &cfg.sockRcvbufBytes, sizeof(cfg.sockRcvbufBytes), - "SO_RCVBUF"); - } -} - -inline std::optional ParseIpv4(const std::string& host, uint16_t port) { - sockaddr_in addr{}; - addr.sin_family = AF_INET; - addr.sin_port = htons(port); - if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) { - return std::nullopt; - } - return addr; -} - -inline uint16_t GetBoundPort(int fd) { - sockaddr_in addr{}; - socklen_t len = sizeof(addr); - if (getsockname(fd, reinterpret_cast(&addr), &len) != 0) { - return 0; - } - return ntohs(addr.sin_port); -} - -inline size_t RoundUpPow2(size_t v) { - if (v <= 1) return 1; - size_t p = 1; - while (p < v) p <<= 1; - return p; -} - -struct PinnedBuf { - void* ptr{nullptr}; - size_t cap{0}; -}; - -class PinnedStagingPool { - public: - PinnedStagingPool() = default; - ~PinnedStagingPool() { Clear(); } - - PinnedStagingPool(const PinnedStagingPool&) = delete; - PinnedStagingPool& operator=(const PinnedStagingPool&) = delete; - - std::shared_ptr Acquire(size_t size) { - const size_t cap = RoundUpPow2(size); - { - std::lock_guard lock(mu); - auto it = free.find(cap); - if (it != free.end() && !it->second.empty()) { - void* p = it->second.back(); - it->second.pop_back(); - return std::shared_ptr(new PinnedBuf{p, cap}, [this](PinnedBuf* b) { - this->Release(b); - }); - } - } - - void* p = nullptr; - hipError_t err = hipHostMalloc(&p, cap, hipHostMallocDefault); - if (err != hipSuccess) { - MORI_IO_ERROR("TCP: hipHostMalloc({}) failed: {}", cap, hipGetErrorString(err)); - return nullptr; - } - return std::shared_ptr(new PinnedBuf{p, cap}, [this](PinnedBuf* b) { this->Release(b); }); - } - - void Clear() { - std::lock_guard lock(mu); - for (auto& kv : free) { - for (void* p : kv.second) { - hipError_t err = hipHostFree(p); - if (err != hipSuccess) { - MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); - } - } - kv.second.clear(); - } - free.clear(); - } - - private: - void Release(PinnedBuf* b) { - if (b == nullptr) return; - const size_t cap = b->cap; - void* p = b->ptr; - delete b; - - constexpr size_t kMaxCachedPerClass = 8; - std::lock_guard lock(mu); - auto& vec = free[cap]; - if (vec.size() < kMaxCachedPerClass) { - vec.push_back(p); - } else { - hipError_t err = hipHostFree(p); - if (err != hipSuccess) { - MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); - } - } - } - - private: - std::mutex mu; - std::unordered_map> free; -}; - -struct Segment { - uint64_t off{0}; - uint64_t len{0}; -}; - -constexpr uint8_t kLaneBits = 3; // up to 8 lanes -constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; - -inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { return (userOpId << kLaneBits) | lane; } - -inline uint64_t ToUserOpId(uint64_t wireOpId) { return wireOpId >> kLaneBits; } - -struct LaneSpan { - uint64_t off{0}; - uint64_t len{0}; -}; - -inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanesTotal, uint8_t lane) { - if (lanesTotal <= 1) return LaneSpan{0, total}; - const uint64_t base = total / lanesTotal; - const uint64_t rem = total % lanesTotal; - const uint64_t off = static_cast(lane) * base + std::min(lane, rem); - const uint64_t len = base + (lane < rem ? 1 : 0); - return LaneSpan{off, len}; -} - -inline uint8_t LanesAllMask(uint8_t lanesTotal) { - if (lanesTotal >= (1U << kLaneBits)) return 0xFF; - return static_cast((1U << lanesTotal) - 1U); -} - -inline std::vector SliceSegments(const std::vector& segs, uint64_t start, - uint64_t len) { - std::vector out; - if (len == 0) return out; - uint64_t skip = start; - uint64_t remaining = len; - for (const auto& s : segs) { - if (remaining == 0) break; - if (skip >= s.len) { - skip -= s.len; - continue; - } - const uint64_t segOff = s.off + skip; - const uint64_t take = std::min(s.len - skip, remaining); - out.push_back({segOff, take}); - remaining -= take; - skip = 0; - } - return out; -} - -inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, - uint64_t* outLen) { - if (!outOff || !outLen) return false; - if (segs.empty()) return false; - uint64_t off = segs[0].off; - uint64_t end = off + segs[0].len; - for (size_t i = 1; i < segs.size(); ++i) { - if (segs[i].off != end) return false; - end += segs[i].len; - } - *outOff = off; - *outLen = end - off; - return true; -} - -inline uint64_t SumLens(const std::vector& segs) { - uint64_t total = 0; - for (const auto& s : segs) total += s.len; - return total; -} - -struct InboundStatusEntry { - StatusCode code{StatusCode::INIT}; - std::string msg; -}; - -struct OutboundOpState { - EngineKey peer; - TransferUniqueId id{0}; - bool isRead{false}; - TransferStatus* status{nullptr}; - - MemoryDesc local{}; - std::vector localSegs; - - MemoryDesc remote{}; - std::vector remoteSegs; - - uint64_t expectedRxBytes{0}; - uint64_t rxBytes{0}; - bool completionReceived{false}; - uint8_t lanesTotal{1}; - uint8_t lanesDoneMask{0}; - StatusCode completionCode{StatusCode::SUCCESS}; - std::string completionMsg; - - bool gpuCopyPending{false}; - std::shared_ptr pinned; - - Clock::time_point startTs{Clock::now()}; -}; - -struct InboundWriteState { - EngineKey peer; - TransferUniqueId id{0}; - - MemoryDesc dst{}; - std::vector dstSegs; - - bool discard{false}; - uint8_t lanesTotal{1}; - uint8_t lanesDoneMask{0}; - std::shared_ptr pinned; -}; - -struct EarlyWriteLaneState { - uint64_t payloadLen{0}; - std::shared_ptr pinned; - bool complete{false}; -}; - -struct EarlyWriteState { - std::unordered_map lanes; -}; - -enum class DataRxKind : uint8_t { NONE = 0, INBOUND_WRITE = 1, OUTBOUND_READ = 2, EARLY_WRITE = 3 }; - -struct ActiveDataRx { - bool active{false}; - TransferUniqueId id{0}; - uint8_t lane{0}; - uint64_t laneOff{0}; - uint64_t laneLen{0}; - DataRxKind kind{DataRxKind::NONE}; - uint64_t remaining{0}; - - bool discard{false}; - // CPU scatter target: - void* base{nullptr}; - std::vector segs; - size_t segIdx{0}; - uint64_t segOff{0}; - - // GPU staging target: - bool toGpu{false}; - int gpuDevice{-1}; - void* gpuBase{nullptr}; - std::shared_ptr pinned; - uint64_t pinnedWriteOff{0}; -}; - -struct SendItem { - std::vector header; - std::vector iov; - size_t idx{0}; - size_t off{0}; - int flags{0}; - std::shared_ptr keepalive; - std::function onDone; - - bool Done() const { return idx >= iov.size(); } - - void Advance(size_t n) { - while (n > 0 && idx < iov.size()) { - size_t avail = iov[idx].iov_len - off; - if (n < avail) { - off += n; - n = 0; - break; - } - n -= avail; - idx++; - off = 0; - } - } -}; - -struct Connection { - int fd{-1}; - bool isOutgoing{false}; - bool connecting{false}; - bool helloSent{false}; - bool helloReceived{false}; - tcp::Channel ch{tcp::Channel::CTRL}; - EngineKey peerKey{}; - - // Control parsing buffer (also used for HELLO on both channels). - std::vector inbuf; - - // Data RX state (after HELLO on data channel). - std::array dataHdrBuf{}; - size_t dataHdrGot{0}; - tcp::DataHeaderView curDataHdr{}; - ActiveDataRx rx{}; - - std::deque sendq; -}; - -struct PeerLinks { - int ctrlFd{-1}; - std::vector dataFds; - int ctrlPending{0}; - int dataPending{0}; - size_t rr{0}; - bool CtrlUp() const { return ctrlFd >= 0; } - bool DataUp() const { return !dataFds.empty(); } - int PickDataFd() { - if (dataFds.empty()) return -1; - int fd = dataFds[rr % dataFds.size()]; - rr = (rr + 1) % dataFds.size(); - return fd; - } -}; - -} // namespace - -class TcpTransport { - public: - TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg) - : myEngKey(std::move(myKey)), engConfig(engCfg), config(cfg) {} - - ~TcpTransport() { Shutdown(); } - - TcpTransport(const TcpTransport&) = delete; - TcpTransport& operator=(const TcpTransport&) = delete; - - void Start() { - if (running.load()) return; - - if (config.numIoThreads != 1) { - MORI_IO_WARN("TCP: numIoThreads={} requested but only 1 is supported; forcing 1", - config.numIoThreads); - } - - epfd = epoll_create1(EPOLL_CLOEXEC); - assert(epfd >= 0); - - // listener - listenFd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); - assert(listenFd >= 0); - - int one = 1; - SetSockOptOrLog(listenFd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one), "SO_REUSEADDR"); - - auto addrOpt = ParseIpv4(engConfig.host.empty() ? std::string("0.0.0.0") : engConfig.host, - engConfig.port); - assert(addrOpt.has_value()); - sockaddr_in addr = addrOpt.value(); - if (bind(listenFd, reinterpret_cast(&addr), sizeof(addr)) != 0) { - MORI_IO_ERROR("TCP: bind {}:{} failed: {}", engConfig.host, engConfig.port, strerror(errno)); - assert(false && "bind failed"); - } - if (listen(listenFd, 256) != 0) { - MORI_IO_ERROR("TCP: listen failed: {}", strerror(errno)); - assert(false && "listen failed"); - } - listenPort = GetBoundPort(listenFd); - MORI_IO_INFO("TCP: listen on {}:{} (port={})", engConfig.host, engConfig.port, listenPort); - - // eventfd for submit queue - wakeFd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); - assert(wakeFd >= 0); - - AddEpoll(listenFd, true, false); - AddEpoll(wakeFd, true, false); - - running.store(true); - ioThread = std::thread([this] { this->IoLoop(); }); - } - - void Shutdown() { - bool wasRunning = running.exchange(false); - if (!wasRunning) return; - - // Wake epoll. - if (wakeFd >= 0) { - uint64_t one = 1; - ::write(wakeFd, &one, sizeof(one)); - } - - if (ioThread.joinable()) ioThread.join(); - - // Close all connections. - for (auto& kv : conns) { - CloseConnInternal(kv.second.get()); - } - conns.clear(); - peers.clear(); - - if (listenFd >= 0) { - close(listenFd); - listenFd = -1; - } - if (wakeFd >= 0) { - close(wakeFd); - wakeFd = -1; - } - if (epfd >= 0) { - close(epfd); - epfd = -1; - } - } - - std::optional GetListenPort() const { - if (listenPort == 0) return std::nullopt; - return listenPort; - } - - void RegisterRemoteEngine(const EngineDesc& desc) { - std::lock_guard lock(remoteMu); - remoteEngines[desc.key] = desc; - } - - void DeregisterRemoteEngine(const EngineDesc& desc) { - std::lock_guard lock(remoteMu); - remoteEngines.erase(desc.key); - } - - void RegisterMemory(const MemoryDesc& desc) { - std::lock_guard lock(memMu); - localMems[desc.id] = desc; - } - - void DeregisterMemory(const MemoryDesc& desc) { - std::lock_guard lock(memMu); - localMems.erase(desc.id); - } - - bool PopInboundTransferStatus(const EngineKey& remote, TransferUniqueId id, TransferStatus* status) { - std::lock_guard lock(inboundMu); - auto it = inboundStatus.find(remote); - if (it == inboundStatus.end()) return false; - auto it2 = it->second.find(id); - if (it2 == it->second.end()) return false; - status->Update(it2->second.code, it2->second.msg); - it->second.erase(it2); - return true; - } - - void SubmitReadWrite(const MemoryDesc& local, size_t localOffset, const MemoryDesc& remote, - size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, - bool isRead) { - if (status == nullptr) return; - - if (size == 0) { - status->SetCode(StatusCode::SUCCESS); - return; - } - - if ((localOffset + size) > local.size || (remoteOffset + size) > remote.size) { - status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: offset+size out of range"); - return; - } - - auto op = std::make_unique(); - op->peer = remote.engineKey; - op->id = id; - op->isRead = isRead; - op->status = status; - op->local = local; - op->remote = remote; - op->localSegs = {{static_cast(localOffset), static_cast(size)}}; - op->remoteSegs = {{static_cast(remoteOffset), static_cast(size)}}; - op->expectedRxBytes = isRead ? static_cast(size) : 0; - op->startTs = Clock::now(); - - status->SetCode(StatusCode::IN_PROGRESS); - EnqueueOp(std::move(op)); - } - - void SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& localOffsets, - const MemoryDesc& remote, const SizeVec& remoteOffsets, - const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, - bool isRead) { - if (status == nullptr) return; - - const size_t n = sizes.size(); - if (n == 0) { - status->SetCode(StatusCode::SUCCESS); - return; - } - if (localOffsets.size() != n || remoteOffsets.size() != n) { - status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch vector size mismatch"); - return; - } - - std::vector localSegs; - std::vector remoteSegs; - localSegs.reserve(n); - remoteSegs.reserve(n); - - uint64_t total = 0; - for (size_t i = 0; i < n; ++i) { - const size_t lo = localOffsets[i]; - const size_t ro = remoteOffsets[i]; - const size_t sz = sizes[i]; - if (sz == 0) continue; - if ((lo + sz) > local.size || (ro + sz) > remote.size) { - status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch offset+size out of range"); - return; - } - localSegs.push_back({static_cast(lo), static_cast(sz)}); - remoteSegs.push_back({static_cast(ro), static_cast(sz)}); - total += static_cast(sz); - } - - // Coalesce adjacent segments when both local and remote are contiguous in lockstep. - // This is critical for TCP performance (reduces iov/copy fan-out), and preserves semantics. - if (localSegs.size() > 1) { - std::vector newLocal; - std::vector newRemote; - newLocal.reserve(localSegs.size()); - newRemote.reserve(remoteSegs.size()); - Segment curL = localSegs[0]; - Segment curR = remoteSegs[0]; - for (size_t i = 1; i < localSegs.size(); ++i) { - const Segment& l = localSegs[i]; - const Segment& r = remoteSegs[i]; - if ((curL.off + curL.len == l.off) && (curR.off + curR.len == r.off) && (curL.len == curR.len) && - (l.len == r.len)) { - curL.len += l.len; - curR.len += r.len; - } else { - newLocal.push_back(curL); - newRemote.push_back(curR); - curL = l; - curR = r; - } - } - newLocal.push_back(curL); - newRemote.push_back(curR); - localSegs = std::move(newLocal); - remoteSegs = std::move(newRemote); - } - - auto op = std::make_unique(); - op->peer = remote.engineKey; - op->id = id; - op->isRead = isRead; - op->status = status; - op->local = local; - op->remote = remote; - op->localSegs = std::move(localSegs); - op->remoteSegs = std::move(remoteSegs); - op->expectedRxBytes = isRead ? total : 0; - op->startTs = Clock::now(); - - status->SetCode(StatusCode::IN_PROGRESS); - EnqueueOp(std::move(op)); - } - - private: - void EnqueueOp(std::unique_ptr op) { - { - std::lock_guard lock(submitMu); - submitQ.push_back(std::move(op)); - } - uint64_t one = 1; - ::write(wakeFd, &one, sizeof(one)); - } - - void AddEpoll(int fd, bool wantRead, bool wantWrite) { - epoll_event ev{}; - ev.data.fd = fd; - ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); - SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev)); - } - - void ModEpoll(int fd, bool wantRead, bool wantWrite) { - epoll_event ev{}; - ev.data.fd = fd; - ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); - SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev)); - } - - void DelEpoll(int fd) { epoll_ctl(epfd, EPOLL_CTL_DEL, fd, nullptr); } - - void CloseConnInternal(Connection* c) { - if (c == nullptr) return; - if (c->fd >= 0) { - DelEpoll(c->fd); - shutdown(c->fd, SHUT_RDWR); - close(c->fd); - c->fd = -1; - } - } - - bool PreferOutgoingFor(const EngineKey& peerKey) const { return myEngKey < peerKey; } - - void AssignConnToPeer(Connection* c) { - assert(c && c->helloReceived); - PeerLinks& link = peers[c->peerKey]; - const bool preferOutgoing = PreferOutgoingFor(c->peerKey); - const bool wasOutgoing = c->isOutgoing; - const tcp::Channel ch = c->ch; - - if (wasOutgoing) { - if (ch == tcp::Channel::CTRL) { - if (link.ctrlPending > 0) link.ctrlPending--; - } else { - if (link.dataPending > 0) link.dataPending--; - } - } - - auto replace_if_needed = [&](int& slotFd) { - if (slotFd < 0) { - slotFd = c->fd; - return; - } - // Collision: choose deterministically by key + direction so both sides keep the same TCP conn. - const int existingFd = slotFd; - const int newFd = c->fd; - Connection* existing = conns[existingFd].get(); - if (!existing) { - slotFd = newFd; - return; - } - const bool keepNew = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); - if (keepNew) { - MORI_IO_WARN("TCP: peer {} channel {} replacing existing fd {} with fd {}", c->peerKey, - static_cast(c->ch), existing->fd, c->fd); - CloseConnInternal(existing); - conns.erase(existingFd); - slotFd = newFd; - } else { - MORI_IO_WARN("TCP: peer {} channel {} dropping duplicate fd {}", c->peerKey, - static_cast(c->ch), c->fd); - CloseConnInternal(c); - conns.erase(newFd); - } - }; - - if (ch == tcp::Channel::CTRL) { - replace_if_needed(link.ctrlFd); - return; - } - - // DATA channel: keep up to numDataConns connections, choosing deterministically by key+direction - // so both sides converge to the same physical set. - const bool keepPreferred = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); - if (!keepPreferred) { - MORI_IO_TRACE("TCP: peer {} dropping non-preferred DATA fd {} outgoing={}", c->peerKey, - c->fd, c->isOutgoing); - const int fd = c->fd; - CloseConnInternal(c); - conns.erase(fd); - return; - } - - const size_t want = static_cast(std::max(1, config.numDataConns)); - if (link.dataFds.size() >= want) { - MORI_IO_TRACE("TCP: peer {} dropping extra DATA fd {} (have {} want {})", c->peerKey, c->fd, - link.dataFds.size(), want); - const int fd = c->fd; - CloseConnInternal(c); - conns.erase(fd); - return; - } - link.dataFds.push_back(c->fd); - MORI_IO_TRACE("TCP: peer {} DATA conn up {}/{}", c->peerKey.c_str(), link.dataFds.size(), - want); - } - - void MaybeDispatchQueuedOps(const EngineKey& peerKey) { - auto it = peers.find(peerKey); - if (it == peers.end()) return; - if (!it->second.CtrlUp() || !it->second.DataUp()) return; - Connection* ctrl = conns[it->second.ctrlFd].get(); - if (!ctrl || !ctrl->helloReceived) return; - int dataFd = it->second.dataFds.empty() ? -1 : it->second.dataFds[0]; - Connection* data = (dataFd >= 0) ? conns[dataFd].get() : nullptr; - if (!data || !data->helloReceived) return; - - auto qit = waitingOps.find(peerKey); - if (qit == waitingOps.end()) return; - - auto ops = std::move(qit->second); - waitingOps.erase(qit); - MORI_IO_TRACE("TCP: peer {} ready, dispatch {} queued ops", peerKey.c_str(), ops.size()); - for (auto& op : ops) { - DispatchOp(std::move(op)); - } - } - - void EnsurePeerChannels(const EngineKey& peerKey) { - PeerLinks& link = peers[peerKey]; - if (!link.CtrlUp() && link.ctrlPending == 0) ConnectChannel(peerKey, tcp::Channel::CTRL); - - const int want = std::max(1, config.numDataConns); - while (static_cast(link.dataFds.size()) + link.dataPending < want) { - ConnectChannel(peerKey, tcp::Channel::DATA); - } - } - - void ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { - EngineDesc desc; - { - std::lock_guard lock(remoteMu); - auto it = remoteEngines.find(peerKey); - if (it == remoteEngines.end()) { - MORI_IO_ERROR("TCP: remote engine {} not registered", peerKey.c_str()); - return; - } - desc = it->second; - } - - auto peerAddrOpt = ParseIpv4(desc.host, static_cast(desc.port)); - if (!peerAddrOpt.has_value()) { - MORI_IO_ERROR("TCP: invalid remote host {}:{}", desc.host, desc.port); - return; - } - sockaddr_in peerAddr = peerAddrOpt.value(); - - int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); - if (fd < 0) { - MORI_IO_ERROR("TCP: socket() failed: {}", strerror(errno)); - return; - } - MORI_IO_TRACE("TCP: connect start peer={} ch={} fd={}", peerKey.c_str(), static_cast(ch), - fd); - - // Bind to our selected host to pin NIC choice when multiple interfaces exist. - if (!engConfig.host.empty()) { - auto localAddrOpt = ParseIpv4(engConfig.host, 0); - if (localAddrOpt.has_value()) { - sockaddr_in localAddr = localAddrOpt.value(); - if (bind(fd, reinterpret_cast(&localAddr), sizeof(localAddr)) != 0) { - MORI_IO_WARN("TCP: bind(local) {} failed: {}", engConfig.host, strerror(errno)); - } - } - } - - int rc = connect(fd, reinterpret_cast(&peerAddr), sizeof(peerAddr)); - bool connecting = false; - if (rc != 0) { - if (errno == EINPROGRESS) { - connecting = true; - } else { - MORI_IO_ERROR("TCP: connect to {}:{} failed: {}", desc.host, desc.port, strerror(errno)); - close(fd); - return; - } - } - - auto conn = std::make_unique(); - conn->fd = fd; - conn->isOutgoing = true; - conn->connecting = connecting; - conn->peerKey = peerKey; - conn->ch = ch; - conn->inbuf.reserve(4096); - - // socket opts - if (ch == tcp::Channel::CTRL) { - ConfigureCtrlSocket(fd, config); - } else { - ConfigureDataSocket(fd, config); - } - - const bool wantWrite = connecting || !conn->sendq.empty(); - AddEpoll(fd, true, wantWrite); - conns[fd] = std::move(conn); - - // Track pending outgoing connections so we don't over-connect when many ops are submitted - // before handshake completes. - PeerLinks& link = peers[peerKey]; - if (ch == tcp::Channel::CTRL) { - link.ctrlPending++; - } else { - link.dataPending++; - } - - if (!connecting) { - QueueHello(fd); - ModEpoll(fd, true, true); - } - } - - void QueueHello(int fd) { - Connection* c = conns[fd].get(); - if (!c || c->helloSent) return; - c->helloSent = true; - auto hello = tcp::BuildHello(c->ch, myEngKey); - MORI_IO_TRACE("TCP: queue HELLO fd={} ch={} peer={}", fd, static_cast(c->ch), - c->peerKey.c_str()); - - SendItem item; - item.header = std::move(hello); - item.iov.resize(1); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - c->sendq.push_back(std::move(item)); - } - - void AcceptNew() { - while (true) { - sockaddr_in peer{}; - socklen_t len = sizeof(peer); - int fd = accept4(listenFd, reinterpret_cast(&peer), &len, SOCK_NONBLOCK | SOCK_CLOEXEC); - if (fd < 0) { - if (IsWouldBlock(errno)) break; - MORI_IO_WARN("TCP: accept failed: {}", strerror(errno)); - break; - } - MORI_IO_TRACE("TCP: accept fd={}", fd); - - auto conn = std::make_unique(); - conn->fd = fd; - conn->isOutgoing = false; - conn->connecting = false; - conn->helloSent = false; - conn->helloReceived = false; - conn->inbuf.reserve(4096); - - AddEpoll(fd, true, false); - conns[fd] = std::move(conn); - } - } - - void DrainWakeFd() { - uint64_t v = 0; - while (true) { - ssize_t n = ::read(wakeFd, &v, sizeof(v)); - if (n < 0) { - if (IsWouldBlock(errno)) break; - break; - } - if (n == 0) break; - } - - std::deque> ops; - { - std::lock_guard lock(submitMu); - ops.swap(submitQ); - } - - for (auto& op : ops) { - EnsurePeerChannels(op->peer); - if (!IsPeerReady(op->peer)) { - waitingOps[op->peer].push_back(std::move(op)); - continue; - } - DispatchOp(std::move(op)); - } - } - - bool IsPeerReady(const EngineKey& peerKey) { - auto it = peers.find(peerKey); - if (it == peers.end()) return false; - if (!it->second.CtrlUp() || !it->second.DataUp()) return false; - Connection* ctrl = conns[it->second.ctrlFd].get(); - if (!ctrl || !ctrl->helloReceived) return false; - for (int fd : it->second.dataFds) { - auto cit = conns.find(fd); - if (cit == conns.end()) continue; - Connection* data = cit->second.get(); - if (data && data->helloReceived) return true; - } - return false; - } - - void DispatchOp(std::unique_ptr op) { - if (!op) { - MORI_IO_ERROR("TCP: DispatchOp got null op"); - return; - } - const EngineKey peerKey = op->peer; - auto it = peers.find(peerKey); - if (it == peers.end() || !it->second.CtrlUp() || !it->second.DataUp()) { - op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer not connected"); - return; - } - Connection* ctrl = conns[it->second.ctrlFd].get(); - if (!ctrl) { - op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer ctrl connection missing"); - return; - } - - std::vector dataFds; - dataFds.reserve(it->second.dataFds.size()); - for (int fd : it->second.dataFds) { - auto dit = conns.find(fd); - if (dit == conns.end()) continue; - Connection* c = dit->second.get(); - if (c && c->helloReceived) dataFds.push_back(fd); - } - - if (dataFds.empty()) { - op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer connection missing"); - return; - } - - const TransferUniqueId opId = op->id; - auto [itIns, inserted] = pendingOutbound.emplace(opId, std::move(op)); - if (!inserted) { - MORI_IO_ERROR("TCP: duplicate outbound op id={} for peer={}", opId, peerKey.c_str()); - itIns->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: duplicate op id"); - pendingOutbound.erase(itIns); - return; - } - OutboundOpState* st = itIns->second.get(); - if (!st) { - MORI_IO_ERROR("TCP: failed to store outbound op id={} (nullptr)", opId); - pendingOutbound.erase(itIns); - return; - } - MORI_IO_TRACE("TCP: dispatch op id={} peer={} isRead={} segs={}", st->id, peerKey.c_str(), - st->isRead, st->localSegs.size()); - - const uint64_t totalBytes = SumLens(st->localSegs); - int wantLanes = std::max(1, config.numDataConns); - wantLanes = std::min(wantLanes, (1U << kLaneBits)); - uint8_t lanesTotal = 1; - const bool canStripe = (wantLanes > 1) && (config.stripingThresholdBytes > 0) && - (totalBytes >= static_cast(config.stripingThresholdBytes)) && - (st->localSegs.size() == 1) && (st->remoteSegs.size() == 1) && - (dataFds.size() >= 2); - if (canStripe) { - lanesTotal = static_cast(std::min(static_cast(wantLanes), dataFds.size())); - } - st->lanesTotal = lanesTotal; - MORI_IO_TRACE("TCP: op {} totalBytes={} dataConns={} lanesTotal={}", st->id, totalBytes, - dataFds.size(), static_cast(lanesTotal)); - - if (st->isRead && st->local.loc == MemoryLocationType::GPU) { - // Allocate a single pinned staging buffer for all lanes and do one H2D copy at the end. - st->pinned = staging.Acquire(static_cast(totalBytes)); - if (!st->pinned) { - st->status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging (read)"); - pendingOutbound.erase(opId); - return; - } - } - - // CTRL request - std::vector ctrlFrame; - if (st->localSegs.size() == 1) { - if (st->isRead) { - ctrlFrame = tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, - st->remoteSegs[0].len, lanesTotal); - } else { - ctrlFrame = tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, - st->remoteSegs[0].len, lanesTotal); - } - } else { - std::vector roffs; - std::vector szs; - roffs.reserve(st->remoteSegs.size()); - szs.reserve(st->remoteSegs.size()); - for (const auto& s : st->remoteSegs) { - roffs.push_back(s.off); - szs.push_back(s.len); - } - if (st->isRead) { - ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs, lanesTotal); - } else { - ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs, lanesTotal); - } - } - - QueueSend(ctrl->fd, std::move(ctrlFrame)); - - // DATA payload (writes only) - if (!st->isRead) { - QueueDataSendForWrite(peerKey, dataFds, *st); - } - - UpdateWriteInterest(ctrl->fd); - } - - void QueueSend(int fd, std::vector bytes, std::function onDone = nullptr) { - Connection* c = conns[fd].get(); - if (!c) return; - SendItem item; - item.header = std::move(bytes); - item.iov.resize(1); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.onDone = std::move(onDone); - c->sendq.push_back(std::move(item)); - } - - void QueueDataSendForWrite(const EngineKey& peerKey, const std::vector& dataFds, - OutboundOpState& st) { - if (dataFds.empty()) return; - const uint64_t total = SumLens(st.localSegs); - const uint8_t lanesTotal = std::max(1, st.lanesTotal); - - // Unstriped path (can handle iov fan-out). - if (lanesTotal == 1) { - auto hdr = tcp::BuildDataHeader(ToWireOpId(st.id, 0), total, 0); - const int fd = dataFds[0]; - Connection* data = conns[fd].get(); - if (!data) return; - - if (st.local.loc == MemoryLocationType::GPU) { - QueueGpuToNetSend(data, st.local, st.localSegs, std::move(hdr), /*onDone*/ nullptr); - UpdateWriteInterest(fd); - return; - } - - SendItem item; - item.header = std::move(hdr); - item.iov.reserve(1 + st.localSegs.size()); - item.iov.push_back({item.header.data(), item.header.size()}); - - uint8_t* base = reinterpret_cast(st.local.data); - for (const auto& s : st.localSegs) { - item.iov.push_back({base + s.off, static_cast(s.len)}); - } - data->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - return; - } - - // Striped path: requires a single contiguous span. - if (st.localSegs.size() != 1) { - MORI_IO_WARN("TCP: striping requested but localSegs.size={} (expected 1), fallback to 1 lane", - st.localSegs.size()); - st.lanesTotal = 1; - QueueDataSendForWrite(peerKey, dataFds, st); - return; - } - - const uint8_t useLanes = std::min(lanesTotal, static_cast(dataFds.size())); - - if (st.local.loc == MemoryLocationType::GPU) { - QueueGpuToNetSendStriped(peerKey, dataFds, st.id, useLanes, st.local, st.localSegs); - return; - } - - uint8_t* base = reinterpret_cast(st.local.data) + st.localSegs[0].off; - for (uint8_t lane = 0; lane < useLanes; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, useLanes, lane); - const int fd = dataFds[lane % dataFds.size()]; - Connection* data = conns[fd].get(); - if (!data) continue; - - SendItem item; - item.header = tcp::BuildDataHeader(ToWireOpId(st.id, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = base + span.off; - item.iov[1].iov_len = static_cast(span.len); - data->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - } - } - - void QueueGpuToNetSend(Connection* data, const MemoryDesc& src, const std::vector& srcSegs, - std::vector dataHdr, std::function onDone) { - const uint64_t total = SumLens(srcSegs); - auto pinned = staging.Acquire(static_cast(total)); - if (!pinned) { - MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); - return; - } - - hipStream_t stream = streamPool.GetNextStream(src.deviceId); - hipEvent_t ev = eventPool.GetEvent(src.deviceId); - if (stream == nullptr || ev == nullptr) { - MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send"); - if (ev) eventPool.PutEvent(ev, src.deviceId); - return; - } - - HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); - uint8_t* dst = reinterpret_cast(pinned->ptr); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : srcSegs) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); - HIP_RUNTIME_CHECK( - hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - gpuTasks.push_back({src.deviceId, ev, [this, dataFd = data->fd, pinned, hdr = std::move(dataHdr), - total, onDone = std::move(onDone)]() mutable { - Connection* c = conns[dataFd].get(); - if (!c || c->fd < 0) return; - SendItem item; - item.header = std::move(hdr); - item.iov.reserve(2); - item.iov.push_back({item.header.data(), item.header.size()}); - item.iov.push_back({pinned->ptr, static_cast(total)}); - item.keepalive = pinned; - item.onDone = std::move(onDone); - c->sendq.push_back(std::move(item)); - UpdateWriteInterest(dataFd); - }}); - } - - void QueueGpuToNetSendStriped(const EngineKey& peerKey, const std::vector& dataFds, - uint64_t userOpId, uint8_t lanesTotal, const MemoryDesc& src, - const std::vector& srcSegs) { - (void)peerKey; - if (dataFds.empty()) return; - const uint64_t total = SumLens(srcSegs); - auto pinned = staging.Acquire(static_cast(total)); - if (!pinned) { - MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); - return; - } - - hipStream_t stream = streamPool.GetNextStream(src.deviceId); - hipEvent_t ev = eventPool.GetEvent(src.deviceId); - if (stream == nullptr || ev == nullptr) { - MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send"); - if (ev) eventPool.PutEvent(ev, src.deviceId); - return; - } - - HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); - uint8_t* dst = reinterpret_cast(pinned->ptr); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : srcSegs) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - gpuTasks.push_back({src.deviceId, ev, - [this, dataFds, pinned, userOpId, lanesTotal, total]() mutable { - for (uint8_t lane = 0; lane < lanesTotal; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); - const int fd = dataFds[lane % dataFds.size()]; - Connection* c = conns[fd].get(); - if (!c || c->fd < 0) continue; - SendItem item; - item.header = tcp::BuildDataHeader(ToWireOpId(userOpId, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = - static_cast(pinned->ptr) + static_cast(span.off); - item.iov[1].iov_len = static_cast(span.len); - item.keepalive = pinned; - c->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - } - }}); - } - - void QueueGpuToNetSendStripedWithOnDone(const EngineKey& peerKey, const std::vector& dataFds, - uint64_t userOpId, uint8_t lanesTotal, const MemoryDesc& src, - const std::vector& srcSegs, - std::function onLaneDone) { - (void)peerKey; - if (dataFds.empty()) return; - const uint64_t total = SumLens(srcSegs); - auto pinned = staging.Acquire(static_cast(total)); - if (!pinned) { - MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send (striped)"); - return; - } - - hipStream_t stream = streamPool.GetNextStream(src.deviceId); - hipEvent_t ev = eventPool.GetEvent(src.deviceId); - if (stream == nullptr || ev == nullptr) { - MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU send (striped)"); - if (ev) eventPool.PutEvent(ev, src.deviceId); - return; - } - - HIP_RUNTIME_CHECK(hipSetDevice(src.deviceId)); - uint8_t* dst = reinterpret_cast(pinned->ptr); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : srcSegs) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - gpuTasks.push_back({src.deviceId, ev, [this, dataFds, pinned, userOpId, lanesTotal, total, - onLaneDone = std::move(onLaneDone)]() mutable { - for (uint8_t lane = 0; lane < lanesTotal; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); - const int fd = dataFds[lane % dataFds.size()]; - Connection* c = conns[fd].get(); - if (!c || c->fd < 0) continue; - SendItem item; - item.header = - tcp::BuildDataHeader(ToWireOpId(userOpId, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = - static_cast(pinned->ptr) + static_cast(span.off); - item.iov[1].iov_len = static_cast(span.len); - item.keepalive = pinned; - item.onDone = onLaneDone; - c->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - } - }}); - } - - struct GpuTask { - int deviceId{-1}; - hipEvent_t ev{nullptr}; - std::function onReady; - }; - - void PollGpuTasks() { - for (auto it = gpuTasks.begin(); it != gpuTasks.end();) { - hipError_t st = hipEventQuery(it->ev); - if (st == hipSuccess) { - eventPool.PutEvent(it->ev, it->deviceId); - if (it->onReady) it->onReady(); - it = gpuTasks.erase(it); - } else if (st == hipErrorNotReady) { - ++it; - } else { - MORI_IO_ERROR("TCP: hipEventQuery failed: {}", hipGetErrorString(st)); - eventPool.PutEvent(it->ev, it->deviceId); - it = gpuTasks.erase(it); - } - } - } - - void UpdateWriteInterest(int fd) { - auto it = conns.find(fd); - if (it == conns.end()) return; - Connection* c = it->second.get(); - if (!c || c->fd < 0) return; - - // With edge-triggered epoll, don't rely solely on EPOLLOUT transitions to make progress. - // If the socket is already writable, we may not get another EPOLLOUT edge after enqueueing. - if (!c->connecting && !c->sendq.empty()) { - FlushSend(c); - it = conns.find(fd); - if (it == conns.end()) return; - c = it->second.get(); - if (!c || c->fd < 0) return; - } - - bool wantWrite = c->connecting || !c->sendq.empty(); - ModEpoll(fd, true, wantWrite); - } - - void HandleConnWritable(Connection* c) { - if (c->connecting) { - int err = 0; - socklen_t len = sizeof(err); - if (getsockopt(c->fd, SOL_SOCKET, SO_ERROR, &err, &len) != 0 || err != 0) { - MORI_IO_ERROR("TCP: connect failed fd {}: {}", c->fd, strerror(err == 0 ? errno : err)); - ClosePeerByFd(c->fd); - return; - } - c->connecting = false; - QueueHello(c->fd); - } - UpdateWriteInterest(c->fd); - } - - void FlushSend(Connection* c) { - constexpr size_t kMaxIov = 64; - while (!c->sendq.empty()) { - SendItem& item = c->sendq.front(); - if (item.Done()) { - auto cb = std::move(item.onDone); - c->sendq.pop_front(); - if (cb) cb(); - continue; - } - - iovec iov[kMaxIov]; - size_t cnt = 0; - for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { - iov[cnt] = item.iov[i]; - if (i == item.idx && item.off > 0) { - iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; - iov[cnt].iov_len -= item.off; - } - cnt++; - } - msghdr msg{}; - msg.msg_iov = iov; - msg.msg_iovlen = cnt; - ssize_t n = sendmsg(c->fd, &msg, MSG_NOSIGNAL | item.flags); - if (n < 0) { - if (IsWouldBlock(errno)) break; - MORI_IO_ERROR("TCP: sendmsg fd {} failed: {}", c->fd, strerror(errno)); - ClosePeerByFd(c->fd); - break; - } - if (n == 0) break; - item.Advance(static_cast(n)); - } - } - - void ClosePeerByFd(int fd) { - MORI_IO_TRACE("TCP: close fd={}", fd); - // Find which peer+channel this fd belongs to, close all channels and fail pending ops. - EngineKey peer{}; - for (auto& kv : peers) { - if (kv.second.ctrlFd == fd) { - peer = kv.first; - break; - } - for (int dfd : kv.second.dataFds) { - if (dfd == fd) { - peer = kv.first; - break; - } - } - if (!peer.empty()) break; - } - - auto close_fd = [&](int toClose) { - if (toClose < 0) return; - auto it = conns.find(toClose); - if (it == conns.end()) return; - CloseConnInternal(it->second.get()); - conns.erase(it); - }; - - if (!peer.empty()) { - auto link = peers[peer]; - close_fd(link.ctrlFd); - for (int dfd : link.dataFds) close_fd(dfd); - peers.erase(peer); - - FailPendingOpsForPeer(peer, "TCP: connection lost"); - return; - } - - // Unknown fd: close just this connection. - close_fd(fd); - } - - void FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg) { - for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { - if (it->second->peer == peer) { - it->second->status->Update(StatusCode::ERR_BAD_STATE, msg); - it = pendingOutbound.erase(it); - } else { - ++it; - } - } - waitingOps.erase(peer); - inboundWrites.erase(peer); - earlyWrites.erase(peer); - } - - void HandleCtrlReadable(Connection* c) { - while (true) { - uint8_t tmp[4096]; - ssize_t n = ::recv(c->fd, tmp, sizeof(tmp), 0); - if (n < 0) { - if (IsWouldBlock(errno)) break; - MORI_IO_ERROR("TCP: recv(ctrl) fd {} failed: {}", c->fd, strerror(errno)); - ClosePeerByFd(c->fd); - return; - } - if (n == 0) { - ClosePeerByFd(c->fd); - return; - } - c->inbuf.insert(c->inbuf.end(), tmp, tmp + n); - } - - // Parse frames. - while (true) { - tcp::CtrlHeaderView hv; - if (!tcp::TryParseCtrlHeader(c->inbuf.data(), c->inbuf.size(), &hv)) { - if (c->inbuf.size() >= tcp::kCtrlHeaderSize) { - MORI_IO_ERROR("TCP: bad ctrl header on fd {}, closing", c->fd); - ClosePeerByFd(c->fd); - } - break; - } - if (c->inbuf.size() < tcp::kCtrlHeaderSize + hv.bodyLen) break; - - const uint8_t* body = c->inbuf.data() + tcp::kCtrlHeaderSize; - HandleCtrlFrame(c, hv.type, body, hv.bodyLen); - - c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + tcp::kCtrlHeaderSize + hv.bodyLen); - - // Data channel transitions to a different framing after HELLO. Any bytes already read into - // inbuf may include the beginning of the data stream; handle them immediately to avoid - // edge-triggered epoll stalls. - if (c->helloReceived && c->ch == tcp::Channel::DATA) { - HandleDataReadable(c); - return; - } - } - } - - void HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const uint8_t* body, size_t len) { - if (type == tcp::CtrlMsgType::HELLO) { - HandleHello(c, body, len); - return; - } - - if (!c->helloReceived) { - MORI_IO_WARN("TCP: received ctrl message before HELLO, dropping"); - return; - } - - switch (type) { - case tcp::CtrlMsgType::WRITE_REQ: - HandleWriteReq(c->peerKey, body, len); - break; - case tcp::CtrlMsgType::READ_REQ: - HandleReadReq(c->peerKey, body, len); - break; - case tcp::CtrlMsgType::BATCH_WRITE_REQ: - HandleBatchWriteReq(c->peerKey, body, len); - break; - case tcp::CtrlMsgType::BATCH_READ_REQ: - HandleBatchReadReq(c->peerKey, body, len); - break; - case tcp::CtrlMsgType::COMPLETION: - HandleCompletion(c->peerKey, body, len); - break; - default: - MORI_IO_WARN("TCP: unknown ctrl msg type {}", static_cast(type)); - } - } - - void HandleHello(Connection* c, const uint8_t* body, size_t len) { - if (len < 1 + 4) { - MORI_IO_WARN("TCP: bad HELLO len {}", len); - ClosePeerByFd(c->fd); - return; - } - size_t off = 0; - const uint8_t chRaw = body[off++]; - uint32_t keyLen = 0; - if (!tcp::ReadU32BE(body, len, &off, &keyLen)) { - ClosePeerByFd(c->fd); - return; - } - if (off + keyLen > len) { - ClosePeerByFd(c->fd); - return; - } - EngineKey peerKey(reinterpret_cast(body + off), keyLen); - off += keyLen; - - c->peerKey = peerKey; - c->ch = (chRaw == static_cast(tcp::Channel::DATA)) ? tcp::Channel::DATA - : tcp::Channel::CTRL; - c->helloReceived = true; - MORI_IO_TRACE("TCP: recv HELLO fd={} peer={} ch={} outgoing={}", c->fd, c->peerKey.c_str(), - static_cast(c->ch), c->isOutgoing); - - // Respond with our hello if we haven't sent it yet. - if (!c->helloSent) { - QueueHello(c->fd); - UpdateWriteInterest(c->fd); - } - - // Post-handshake: configure socket by channel type. - if (c->ch == tcp::Channel::CTRL) { - ConfigureCtrlSocket(c->fd, config); - } else { - ConfigureDataSocket(c->fd, config); - } - - AssignConnToPeer(c); - MaybeDispatchQueuedOps(peerKey); - } - - std::optional LookupLocalMem(MemoryUniqueId id) { - std::lock_guard lock(memMu); - auto it = localMems.find(id); - if (it == localMems.end()) return std::nullopt; - return it->second; - } - - void RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, - const std::string& msg) { - std::lock_guard lock(inboundMu); - inboundStatus[peer][id] = InboundStatusEntry{code, msg}; - } - - Connection* PeerCtrl(const EngineKey& peer) { - auto it = peers.find(peer); - if (it == peers.end()) return nullptr; - if (!it->second.CtrlUp()) return nullptr; - return conns[it->second.ctrlFd].get(); - } - - Connection* PeerData(const EngineKey& peer) { - auto it = peers.find(peer); - if (it == peers.end()) return nullptr; - if (!it->second.DataUp()) return nullptr; - const int fd = it->second.dataFds.empty() ? -1 : it->second.dataFds[0]; - if (fd < 0) return nullptr; - return conns[fd].get(); - } - - static uint8_t ClampLanesTotal(uint8_t lanesTotal) { - if (lanesTotal == 0) return 1; - const uint8_t max = static_cast(1U << kLaneBits); - return std::min(lanesTotal, max); - } - - void MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId) { - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt == inboundWrites.end()) return; - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt == iwPeerIt->second.end()) return; - - InboundWriteState& ws = wsIt->second; - ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - const uint8_t allMask = LanesAllMask(ws.lanesTotal); - if ((ws.lanesDoneMask & allMask) != allMask) return; - - auto ctrl = PeerCtrl(peer); - if (!ctrl) { - // Peer is gone; drop state. - iwPeerIt->second.erase(wsIt); - if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); - return; - } - - if (ws.discard) { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), - "TCP: write discarded")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); - } else if (ws.dst.loc == MemoryLocationType::GPU) { - if (!ws.pinned) { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: missing pinned staging (write)")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging"); - } else { - hipStream_t stream = streamPool.GetNextStream(ws.dst.deviceId); - hipEvent_t ev = eventPool.GetEvent(ws.dst.deviceId); - if (stream == nullptr || ev == nullptr) { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_BAD_STATE), - "TCP: failed to get HIP stream/event")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::ERR_BAD_STATE, - "TCP: failed to get HIP stream/event"); - if (ev) eventPool.PutEvent(ev, ws.dst.deviceId); - } else { - HIP_RUNTIME_CHECK(hipSetDevice(ws.dst.deviceId)); - uint8_t* src = reinterpret_cast(ws.pinned->ptr); - const uint64_t total = SumLens(ws.dstSegs); - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(ws.dstSegs, &spanOff, &spanLen) && spanLen == total) { - void* gpuPtr = reinterpret_cast(ws.dst.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : ws.dstSegs) { - void* gpuPtr = reinterpret_cast(ws.dst.data + s.off); - HIP_RUNTIME_CHECK( - hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - auto pinned = ws.pinned; - const int ctrlFd = ctrl->fd; - const int deviceId = ws.dst.deviceId; - gpuTasks.push_back({deviceId, ev, [this, peer, opId, ctrlFd, pinned]() { - QueueSend(ctrlFd, - tcp::BuildCompletion(opId, - static_cast(StatusCode::SUCCESS), - "")); - UpdateWriteInterest(ctrlFd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); - }}); - } - } - } else { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::SUCCESS), "")); - UpdateWriteInterest(ctrl->fd); - RecordInboundStatus(peer, opId, StatusCode::SUCCESS, ""); - } - - // Cleanup state. - iwPeerIt->second.erase(wsIt); - if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - ewPeerIt->second.erase(opId); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); - } - } - - void TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId) { - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt == inboundWrites.end()) return; - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt == iwPeerIt->second.end()) return; - InboundWriteState& ws = wsIt->second; - ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt == earlyWrites.end()) return; - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt == ewPeerIt->second.end()) return; - - EarlyWriteState& early = ewIt->second; - const uint64_t total = SumLens(ws.dstSegs); - uint8_t* dstBase = reinterpret_cast(ws.dst.data); - - for (auto it = early.lanes.begin(); it != early.lanes.end();) { - const uint8_t lane = it->first; - EarlyWriteLaneState& laneState = it->second; - if (!laneState.complete) { - ++it; - continue; - } - - if (lane >= ws.lanesTotal) { - ws.discard = true; - } - const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); - if (span.len != laneState.payloadLen) { - ws.discard = true; - } - - if (!ws.discard && laneState.pinned) { - uint8_t* src = reinterpret_cast(laneState.pinned->ptr); - if (ws.dst.loc == MemoryLocationType::GPU) { - if (!ws.pinned) { - ws.pinned = staging.Acquire(static_cast(total)); - if (!ws.pinned) ws.discard = true; - } - if (!ws.discard && ws.pinned) { - std::memcpy(reinterpret_cast(ws.pinned->ptr) + span.off, src, - static_cast(span.len)); - } - } else { - const auto segs = SliceSegments(ws.dstSegs, span.off, span.len); - uint64_t copied = 0; - for (const auto& s : segs) { - std::memcpy(dstBase + s.off, src + copied, static_cast(s.len)); - copied += s.len; - } - } - } - - if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); - it = early.lanes.erase(it); - } - - if (early.lanes.empty()) { - ewPeerIt->second.erase(ewIt); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); - } - - MaybeFinalizeInboundWrite(peer, opId); - } - - void HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { - size_t off = 0; - uint64_t opId = 0; - uint32_t memId = 0; - uint64_t remoteOff = 0; - uint64_t size = 0; - if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || - !tcp::ReadU64BE(body, len, &off, &remoteOff) || !tcp::ReadU64BE(body, len, &off, &size)) { - MORI_IO_WARN("TCP: malformed WRITE_REQ"); - return; - } - - uint8_t lanesTotal = 1; - if (off < len) { - lanesTotal = body[off]; - } - lanesTotal = ClampLanesTotal(lanesTotal); - - auto memOpt = LookupLocalMem(memId); - InboundWriteState ws; - ws.peer = peer; - ws.id = opId; - ws.lanesTotal = lanesTotal; - ws.discard = true; - - if (memOpt.has_value()) { - ws.dst = memOpt.value(); - if (remoteOff + size <= ws.dst.size) { - ws.discard = false; - ws.dstSegs = {{remoteOff, size}}; - } - } - - if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { - const uint64_t total = SumLens(ws.dstSegs); - ws.pinned = staging.Acquire(static_cast(total)); - if (!ws.pinned) ws.discard = true; - } - - inboundWrites[peer][opId] = std::move(ws); - TryConsumeEarlyWriteLanes(peer, opId); - } - - void HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { - size_t off = 0; - uint64_t opId = 0; - uint32_t memId = 0; - uint32_t n = 0; - if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || - !tcp::ReadU32BE(body, len, &off, &n)) { - MORI_IO_WARN("TCP: malformed BATCH_WRITE_REQ"); - return; - } - - auto memOpt = LookupLocalMem(memId); - InboundWriteState ws; - ws.peer = peer; - ws.id = opId; - ws.discard = true; - - const bool haveMem = memOpt.has_value(); - if (haveMem) ws.dst = memOpt.value(); - - ws.dstSegs.reserve(n); - bool ok = true; - for (uint32_t i = 0; i < n; ++i) { - uint64_t ro = 0, sz = 0; - if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { - ok = false; - break; - } - if (haveMem && (ro + sz > ws.dst.size)) ok = false; - if (haveMem && sz > 0) ws.dstSegs.push_back({ro, sz}); - } - - // `lanesTotal` is appended by the sender; tolerate older senders. - uint8_t lanesTotal = 1; - if (off < len) lanesTotal = body[off]; - ws.lanesTotal = ClampLanesTotal(lanesTotal); - - if (haveMem && ok) ws.discard = false; - - if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { - const uint64_t total = SumLens(ws.dstSegs); - ws.pinned = staging.Acquire(static_cast(total)); - if (!ws.pinned) ws.discard = true; - } - - inboundWrites[peer][opId] = std::move(ws); - TryConsumeEarlyWriteLanes(peer, opId); - } - - void HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { - size_t off = 0; - uint64_t opId = 0; - uint32_t memId = 0; - uint64_t srcOff = 0; - uint64_t size = 0; - if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || - !tcp::ReadU64BE(body, len, &off, &srcOff) || !tcp::ReadU64BE(body, len, &off, &size)) { - MORI_IO_WARN("TCP: malformed READ_REQ"); - return; - } - uint8_t lanesTotal = 1; - if (off < len) lanesTotal = body[off]; - lanesTotal = ClampLanesTotal(lanesTotal); - - auto memOpt = LookupLocalMem(memId); - if (!memOpt.has_value() || (srcOff + size > memOpt->size)) { - // No data; completion with error. - auto ctrl = PeerCtrl(peer); - if (ctrl) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_NOT_FOUND), - "TCP: remote mem not found/out of range")); - UpdateWriteInterest(ctrl->fd); - } - RecordInboundStatus(peer, opId, StatusCode::ERR_NOT_FOUND, "TCP: read mem not found"); - return; - } - - MemoryDesc src = memOpt.value(); - std::vector segs = {{srcOff, size}}; - QueueDataSendForRead(peer, opId, src, segs, lanesTotal); - } - - void HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { - size_t off = 0; - uint64_t opId = 0; - uint32_t memId = 0; - uint32_t n = 0; - if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &memId) || - !tcp::ReadU32BE(body, len, &off, &n)) { - MORI_IO_WARN("TCP: malformed BATCH_READ_REQ"); - return; - } - - auto memOpt = LookupLocalMem(memId); - if (!memOpt.has_value()) { - auto ctrl = PeerCtrl(peer); - if (ctrl) { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_NOT_FOUND), - "TCP: remote mem not found")); - UpdateWriteInterest(ctrl->fd); - } - RecordInboundStatus(peer, opId, StatusCode::ERR_NOT_FOUND, "TCP: read mem not found"); - return; - } - - MemoryDesc src = memOpt.value(); - std::vector segs; - segs.reserve(n); - bool ok = true; - for (uint32_t i = 0; i < n; ++i) { - uint64_t ro = 0, sz = 0; - if (!tcp::ReadU64BE(body, len, &off, &ro) || !tcp::ReadU64BE(body, len, &off, &sz)) { - ok = false; - break; - } - if (ro + sz > src.size) ok = false; - if (sz > 0) segs.push_back({ro, sz}); - } - uint8_t lanesTotal = 1; - if (off < len) lanesTotal = body[off]; - lanesTotal = ClampLanesTotal(lanesTotal); - if (!ok) { - auto ctrl = PeerCtrl(peer); - if (ctrl) { - QueueSend(ctrl->fd, - tcp::BuildCompletion(opId, static_cast(StatusCode::ERR_INVALID_ARGS), - "TCP: batch read out of range")); - UpdateWriteInterest(ctrl->fd); - } - RecordInboundStatus(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: batch read bad args"); - return; - } - - QueueDataSendForRead(peer, opId, src, segs, lanesTotal); - } - - void QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, - const std::vector& srcSegs, uint8_t lanesTotal) { - Connection* ctrl = PeerCtrl(peer); - if (!ctrl) return; - - const uint64_t total = SumLens(srcSegs); - lanesTotal = ClampLanesTotal(lanesTotal); - - // Build list of active data fds. - std::vector dataFds; - auto pit = peers.find(peer); - if (pit != peers.end()) { - for (int fd : pit->second.dataFds) { - auto dit = conns.find(fd); - if (dit == conns.end()) continue; - if (dit->second && dit->second->helloReceived) dataFds.push_back(fd); - } - } - if (dataFds.empty()) return; - lanesTotal = std::min(lanesTotal, static_cast(dataFds.size())); - - struct DoneState { - EngineKey peer; - uint64_t opId{0}; - int ctrlFd{-1}; - int remaining{0}; - }; - auto done = std::make_shared(); - done->peer = peer; - done->opId = opId; - done->ctrlFd = ctrl->fd; - done->remaining = lanesTotal; - auto laneDone = [this, done]() mutable { - done->remaining--; - if (done->remaining > 0) return; - QueueSend(done->ctrlFd, - tcp::BuildCompletion(done->opId, static_cast(StatusCode::SUCCESS), "")); - UpdateWriteInterest(done->ctrlFd); - RecordInboundStatus(done->peer, done->opId, StatusCode::SUCCESS, ""); - }; - - // Unstriped path. - if (lanesTotal == 1) { - const int fd = dataFds[0]; - Connection* data = conns[fd].get(); - if (!data) return; - auto hdr = tcp::BuildDataHeader(ToWireOpId(opId, 0), total, 0); - if (src.loc == MemoryLocationType::GPU) { - QueueGpuToNetSend(data, src, srcSegs, std::move(hdr), std::move(laneDone)); - UpdateWriteInterest(fd); - return; - } - - SendItem item; - item.header = std::move(hdr); - item.iov.reserve(1 + srcSegs.size()); - item.iov.push_back({item.header.data(), item.header.size()}); - uint8_t* base = reinterpret_cast(src.data); - for (const auto& s : srcSegs) { - item.iov.push_back({base + s.off, static_cast(s.len)}); - } - item.onDone = std::move(laneDone); - data->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - return; - } - - // Striped path: requires a single contiguous span. - if (srcSegs.size() != 1) { - MORI_IO_WARN("TCP: peer {} READ striping requested but srcSegs.size={} (expected 1), fallback to 1 lane", - peer.c_str(), srcSegs.size()); - QueueDataSendForRead(peer, opId, src, srcSegs, /*lanesTotal=*/1); - return; - } - - if (src.loc == MemoryLocationType::GPU) { - // Stage once, then stripe sends from pinned buffer. - const uint8_t useLanes = - std::min(lanesTotal, static_cast(dataFds.size())); - QueueGpuToNetSendStripedWithOnDone(peer, dataFds, opId, useLanes, src, srcSegs, std::move(laneDone)); - return; - } - - const uint8_t useLanes = std::min(lanesTotal, static_cast(dataFds.size())); - uint8_t* base = reinterpret_cast(src.data) + srcSegs[0].off; - for (uint8_t lane = 0; lane < useLanes; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, useLanes, lane); - const int fd = dataFds[lane % dataFds.size()]; - Connection* data = conns[fd].get(); - if (!data) continue; - SendItem item; - item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = base + span.off; - item.iov[1].iov_len = static_cast(span.len); - item.onDone = laneDone; - data->sendq.push_back(std::move(item)); - UpdateWriteInterest(fd); - } - } - - void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { - size_t off = 0; - uint64_t opId = 0; - uint32_t code = 0; - uint32_t msgLen = 0; - if (!tcp::ReadU64BE(body, len, &off, &opId) || !tcp::ReadU32BE(body, len, &off, &code) || - !tcp::ReadU32BE(body, len, &off, &msgLen)) { - MORI_IO_WARN("TCP: malformed COMPLETION"); - return; - } - if (off + msgLen > len) { - MORI_IO_WARN("TCP: malformed COMPLETION msg len"); - return; - } - std::string msg(reinterpret_cast(body + off), msgLen); - - auto it = pendingOutbound.find(opId); - if (it == pendingOutbound.end()) return; - OutboundOpState& st = *it->second; - st.completionReceived = true; - st.completionCode = static_cast(code); - st.completionMsg = std::move(msg); - - // Fast fail on remote error. - if (st.completionCode != StatusCode::SUCCESS) { - st.status->Update(st.completionCode, st.completionMsg); - pendingOutbound.erase(it); - return; - } - - MaybeCompleteOutbound(st); - } - - void MaybeCompleteOutbound(OutboundOpState& st) { - if (!st.completionReceived) return; - if (st.isRead) { - const uint8_t allMask = - st.lanesTotal >= (1U << kLaneBits) ? 0xFF : static_cast((1U << st.lanesTotal) - 1U); - if (st.lanesDoneMask != allMask) return; - if (st.rxBytes != st.expectedRxBytes) return; - if (st.gpuCopyPending) return; - } - st.status->Update(StatusCode::SUCCESS, ""); - pendingOutbound.erase(st.id); - } - - void ConsumeBufferedData(Connection* c) { - while (!c->inbuf.empty()) { - if (!c->rx.active) { - const size_t need = tcp::kDataHeaderSize - c->dataHdrGot; - const size_t take = std::min(need, c->inbuf.size()); - std::memcpy(c->dataHdrBuf.data() + c->dataHdrGot, c->inbuf.data(), take); - c->dataHdrGot += take; - c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + take); - if (c->dataHdrGot < tcp::kDataHeaderSize) return; - - tcp::DataHeaderView hv; - if (!tcp::TryParseDataHeader(c->dataHdrBuf.data(), tcp::kDataHeaderSize, &hv)) { - MORI_IO_ERROR("TCP: bad data header during handoff, closing"); - ClosePeerByFd(c->fd); - return; - } - c->dataHdrGot = 0; - BeginDataRx(c, hv.opId, hv.payloadLen); - continue; - } - - if (c->rx.remaining == 0) { - FinishDataRx(c); - continue; - } - - const size_t take = static_cast(std::min(c->rx.remaining, c->inbuf.size())); - if (take == 0) return; - const uint8_t* src = c->inbuf.data(); - - if (c->rx.discard) { - // nothing - } else if (c->rx.toGpu) { - if (!c->rx.pinned) { - MORI_IO_ERROR("TCP: missing pinned buffer for GPU recv"); - ClosePeerByFd(c->fd); - return; - } - uint8_t* dst = reinterpret_cast(c->rx.pinned->ptr) + c->rx.pinnedWriteOff; - std::memcpy(dst, src, take); - c->rx.pinnedWriteOff += take; - } else { - size_t copied = 0; - while (copied < take) { - if (c->rx.segIdx >= c->rx.segs.size()) { - MORI_IO_ERROR("TCP: cpu scatter overflow during buffered consume"); - ClosePeerByFd(c->fd); - return; - } - Segment& seg = c->rx.segs[c->rx.segIdx]; - const uint64_t segRemain = seg.len - c->rx.segOff; - const size_t chunk = - static_cast(std::min(segRemain, static_cast(take - copied))); - uint8_t* dst = reinterpret_cast(c->rx.base) + seg.off + c->rx.segOff; - std::memcpy(dst, src + copied, chunk); - c->rx.segOff += chunk; - copied += chunk; - if (c->rx.segOff >= seg.len) { - c->rx.segIdx++; - c->rx.segOff = 0; - } - } - } - - c->rx.remaining -= static_cast(take); - c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + take); - } - } - - void HandleDataReadable(Connection* c) { - // If we haven't received HELLO on this channel, treat data as ctrl for handshake. - if (!c->helloReceived) { - HandleCtrlReadable(c); - return; - } - - if (!c->inbuf.empty()) { - ConsumeBufferedData(c); - if (c->fd < 0) return; - } - - while (true) { - if (!c->rx.active) { - // Need header. - while (c->dataHdrGot < tcp::kDataHeaderSize) { - ssize_t n = ::recv(c->fd, c->dataHdrBuf.data() + c->dataHdrGot, - tcp::kDataHeaderSize - c->dataHdrGot, 0); - if (n < 0) { - if (IsWouldBlock(errno)) return; - MORI_IO_ERROR("TCP: recv(data hdr) failed: {}", strerror(errno)); - ClosePeerByFd(c->fd); - return; - } - if (n == 0) { - ClosePeerByFd(c->fd); - return; - } - c->dataHdrGot += static_cast(n); - } - tcp::DataHeaderView hv; - if (!tcp::TryParseDataHeader(c->dataHdrBuf.data(), tcp::kDataHeaderSize, &hv)) { - MORI_IO_ERROR("TCP: bad data header, closing"); - ClosePeerByFd(c->fd); - return; - } - c->dataHdrGot = 0; - BeginDataRx(c, hv.opId, hv.payloadLen); - continue; - } - - if (c->rx.remaining == 0) { - FinishDataRx(c); - continue; - } - - ssize_t n = 0; - if (c->rx.discard) { - uint8_t tmp[8192]; - const size_t want = std::min(sizeof(tmp), c->rx.remaining); - n = ::recv(c->fd, tmp, want, 0); - } else if (c->rx.toGpu) { - uint8_t* dst = reinterpret_cast(c->rx.pinned->ptr) + c->rx.pinnedWriteOff; - const size_t want = static_cast(std::min(c->rx.remaining, 1ULL << 20)); - n = ::recv(c->fd, dst, want, 0); - } else { - // CPU scatter: recv into current segment. - Segment& seg = c->rx.segs[c->rx.segIdx]; - uint8_t* dst = reinterpret_cast(c->rx.base) + seg.off + c->rx.segOff; - uint64_t segRemain = seg.len - c->rx.segOff; - const size_t want = static_cast(std::min(c->rx.remaining, segRemain)); - n = ::recv(c->fd, dst, want, 0); - } - - if (n < 0) { - if (IsWouldBlock(errno)) return; - MORI_IO_ERROR("TCP: recv(data) failed: {}", strerror(errno)); - ClosePeerByFd(c->fd); - return; - } - if (n == 0) { - ClosePeerByFd(c->fd); - return; - } - - const uint64_t got = static_cast(n); - c->rx.remaining -= got; - if (c->rx.discard) { - continue; - } - if (c->rx.toGpu) { - c->rx.pinnedWriteOff += got; - continue; - } - // CPU seg advance - c->rx.segOff += got; - Segment& seg = c->rx.segs[c->rx.segIdx]; - if (c->rx.segOff >= seg.len) { - c->rx.segIdx++; - c->rx.segOff = 0; - } - } - } - - void BeginDataRx(Connection* c, uint64_t wireOpId, uint64_t payloadLen) { - c->rx = ActiveDataRx{}; - c->rx.active = true; - - const uint8_t lane = static_cast(wireOpId & kLaneMask); - const TransferUniqueId userOpId = static_cast(ToUserOpId(wireOpId)); - - c->rx.id = userOpId; - c->rx.lane = lane; - c->rx.laneLen = payloadLen; - c->rx.remaining = payloadLen; - - // 1) inbound write: peer -> me. - auto iwPeerIt = inboundWrites.find(c->peerKey); - if (iwPeerIt != inboundWrites.end()) { - auto wsIt = iwPeerIt->second.find(userOpId); - if (wsIt != iwPeerIt->second.end()) { - InboundWriteState& ws = wsIt->second; - ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - c->rx.kind = DataRxKind::INBOUND_WRITE; - - const uint64_t total = SumLens(ws.dstSegs); - if (lane >= ws.lanesTotal) ws.discard = true; - const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); - c->rx.laneOff = span.off; - c->rx.laneLen = payloadLen; - - if (!ws.discard && span.len != payloadLen) { - MORI_IO_WARN("TCP: inbound write op {} lane {} payloadLen mismatch expected={} got={}", - userOpId, static_cast(lane), span.len, payloadLen); - ws.discard = true; - } - - c->rx.discard = ws.discard; - if (!c->rx.discard) { - if (ws.dst.loc == MemoryLocationType::GPU) { - if (!ws.pinned) { - ws.pinned = staging.Acquire(static_cast(total)); - if (!ws.pinned) ws.discard = true; - } - c->rx.discard = ws.discard; - if (!c->rx.discard) { - c->rx.toGpu = true; - c->rx.pinned = ws.pinned; - c->rx.pinnedWriteOff = span.off; - } - } else { - c->rx.toGpu = false; - c->rx.base = reinterpret_cast(ws.dst.data); - c->rx.segs = SliceSegments(ws.dstSegs, span.off, span.len); - c->rx.segIdx = 0; - c->rx.segOff = 0; - } - } - return; - } - } - - // 2) outbound read response: peer -> me for pending outbound read. - auto obIt = pendingOutbound.find(userOpId); - if (obIt != pendingOutbound.end()) { - OutboundOpState& st = *obIt->second; - if (!st.isRead) { - c->rx.kind = DataRxKind::OUTBOUND_READ; - c->rx.discard = true; - return; - } - - st.lanesTotal = ClampLanesTotal(st.lanesTotal); - if (lane >= st.lanesTotal) { - c->rx.kind = DataRxKind::OUTBOUND_READ; - c->rx.discard = true; - return; - } - - const LaneSpan span = ComputeLaneSpan(st.expectedRxBytes, st.lanesTotal, lane); - c->rx.kind = DataRxKind::OUTBOUND_READ; - c->rx.laneOff = span.off; - if (span.len != payloadLen) { - MORI_IO_ERROR("TCP: outbound read op {} lane {} payloadLen mismatch expected={} got={}", - userOpId, static_cast(lane), span.len, payloadLen); - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: read payload length mismatch"); - pendingOutbound.erase(obIt); - c->rx.discard = true; - return; - } - - c->rx.discard = false; - if (st.local.loc == MemoryLocationType::GPU) { - if (!st.pinned) { - st.pinned = staging.Acquire(static_cast(st.expectedRxBytes)); - if (!st.pinned) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging (read)"); - pendingOutbound.erase(obIt); - c->rx.discard = true; - return; - } - } - c->rx.toGpu = true; - c->rx.pinned = st.pinned; - c->rx.pinnedWriteOff = span.off; - } else { - c->rx.toGpu = false; - c->rx.base = reinterpret_cast(st.local.data); - c->rx.segs = SliceSegments(st.localSegs, span.off, span.len); - c->rx.segIdx = 0; - c->rx.segOff = 0; - } - return; - } - - // 3) early-arrived write payload (data before CTRL write request). - c->rx.kind = DataRxKind::EARLY_WRITE; - if (payloadLen > static_cast(std::numeric_limits::max())) { - MORI_IO_WARN("TCP: data payloadLen too large for op {} from peer {}, discarding", userOpId, - c->peerKey.c_str()); - c->rx.discard = true; - return; - } - - auto& perPeer = earlyWrites[c->peerKey]; - EarlyWriteState& early = perPeer[userOpId]; - if (early.lanes.find(lane) != early.lanes.end()) { - MORI_IO_WARN("TCP: duplicate early data for op {} lane {} from peer {}, discarding", userOpId, - static_cast(lane), c->peerKey.c_str()); - c->rx.discard = true; - return; - } - - const size_t alloc = payloadLen == 0 ? 1 : static_cast(payloadLen); - auto pinned = staging.Acquire(alloc); - if (!pinned) { - MORI_IO_WARN("TCP: failed to allocate pinned buffer for early data op {} from peer {}, discarding", - userOpId, c->peerKey.c_str()); - c->rx.discard = true; - return; - } - - early.lanes.emplace(lane, EarlyWriteLaneState{payloadLen, pinned, false}); - c->rx.discard = false; - c->rx.toGpu = true; // receive into pinned - c->rx.pinned = pinned; - c->rx.pinnedWriteOff = 0; - } - - void FinishDataRx(Connection* c) { - const TransferUniqueId opId = c->rx.id; - const EngineKey peer = c->peerKey; - const uint8_t lane = c->rx.lane; - const uint64_t laneLen = c->rx.laneLen; - const DataRxKind kind = c->rx.kind; - - c->rx = ActiveDataRx{}; - - if (kind == DataRxKind::INBOUND_WRITE) { - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt == inboundWrites.end()) return; - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt == iwPeerIt->second.end()) return; - InboundWriteState& ws = wsIt->second; - ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); - MaybeFinalizeInboundWrite(peer, opId); - return; - } - - if (kind == DataRxKind::OUTBOUND_READ) { - auto obIt = pendingOutbound.find(opId); - if (obIt == pendingOutbound.end()) return; - OutboundOpState& st = *obIt->second; - st.lanesTotal = ClampLanesTotal(st.lanesTotal); - const uint8_t bit = static_cast(1U << lane); - if ((st.lanesDoneMask & bit) == 0) { - st.lanesDoneMask |= bit; - st.rxBytes += laneLen; - } - - if (st.local.loc == MemoryLocationType::GPU) { - const uint8_t allMask = LanesAllMask(st.lanesTotal); - if ((st.lanesDoneMask & allMask) != allMask) { - MaybeCompleteOutbound(st); - return; - } - if (st.gpuCopyPending) return; - if (!st.pinned) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging (read)"); - pendingOutbound.erase(obIt); - return; - } - - hipStream_t stream = streamPool.GetNextStream(st.local.deviceId); - hipEvent_t ev = eventPool.GetEvent(st.local.deviceId); - if (stream == nullptr || ev == nullptr) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to get HIP stream/event (read)"); - pendingOutbound.erase(obIt); - if (ev) eventPool.PutEvent(ev, st.local.deviceId); - return; - } - - HIP_RUNTIME_CHECK(hipSetDevice(st.local.deviceId)); - uint8_t* src = reinterpret_cast(st.pinned->ptr); - st.gpuCopyPending = true; - const uint64_t total = st.expectedRxBytes; - uint64_t spanOff = 0; - uint64_t spanLen = 0; - if (IsSingleContiguousSpan(st.localSegs, &spanOff, &spanLen) && spanLen == total) { - void* gpuPtr = reinterpret_cast(st.local.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); - } else { - uint64_t off = 0; - for (const auto& s : st.localSegs) { - void* gpuPtr = reinterpret_cast(st.local.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); - off += s.len; - } - } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - - auto pinned = st.pinned; - gpuTasks.push_back({st.local.deviceId, ev, [this, opId, pinned]() { - auto it = pendingOutbound.find(opId); - if (it == pendingOutbound.end()) return; - it->second->gpuCopyPending = false; - MaybeCompleteOutbound(*it->second); - }}); - return; - } - - MaybeCompleteOutbound(st); - return; - } - - if (kind == DataRxKind::EARLY_WRITE) { - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt == earlyWrites.end()) return; - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt == ewPeerIt->second.end()) return; - auto laneIt = ewIt->second.lanes.find(lane); - if (laneIt != ewIt->second.lanes.end()) { - laneIt->second.complete = true; - } - TryConsumeEarlyWriteLanes(peer, opId); - return; - } - } - - void HandleReadable(Connection* c) { - if (!c->helloReceived) { - HandleCtrlReadable(c); - return; - } - if (c->ch == tcp::Channel::CTRL) { - HandleCtrlReadable(c); - } else { - HandleDataReadable(c); - } - } - - void ScanTimeouts() { - if (config.opTimeoutMs <= 0) return; - const auto now = Clock::now(); - const auto timeout = std::chrono::milliseconds(config.opTimeoutMs); - for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { - if ((now - it->second->startTs) > timeout) { - it->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: op timeout"); - it = pendingOutbound.erase(it); - } else { - ++it; - } - } - } - - void IoLoop() { - constexpr int kMaxEvents = 128; - epoll_event events[kMaxEvents]; - - while (running.load()) { - PollGpuTasks(); - ScanTimeouts(); - - // When GPU staging is in flight, prefer low-latency polling so we can arm the next - // network send / completion promptly after HIP events complete. - const int timeoutMs = gpuTasks.empty() ? 5 /*ms*/ : 0 /*busy*/; - int nfds = epoll_wait(epfd, events, kMaxEvents, timeoutMs); - if (nfds < 0) { - if (errno == EINTR) continue; - MORI_IO_ERROR("TCP: epoll_wait failed: {}", strerror(errno)); - break; - } - - for (int i = 0; i < nfds; ++i) { - int fd = events[i].data.fd; - uint32_t ev = events[i].events; - - if (fd == listenFd) { - AcceptNew(); - continue; - } - if (fd == wakeFd) { - DrainWakeFd(); - continue; - } - - Connection* c = nullptr; - auto it = conns.find(fd); - if (it != conns.end()) c = it->second.get(); - if (!c) continue; - - if (ev & (EPOLLERR | EPOLLHUP)) { - ClosePeerByFd(fd); - continue; - } - - if (ev & EPOLLIN) { - HandleReadable(c); - auto it2 = conns.find(fd); - if (it2 == conns.end()) { - continue; - } - c = it2->second.get(); - if (!c) continue; - } - if (ev & EPOLLOUT) { - HandleConnWritable(c); - } - } - } - } - - private: - EngineKey myEngKey; - IOEngineConfig engConfig; - TcpBackendConfig config; - - int epfd{-1}; - int listenFd{-1}; - int wakeFd{-1}; - uint16_t listenPort{0}; - - std::atomic running{false}; - std::thread ioThread; - - std::mutex submitMu; - std::deque> submitQ; - - std::mutex remoteMu; - std::unordered_map remoteEngines; - - std::mutex memMu; - std::unordered_map localMems; - - std::mutex inboundMu; - std::unordered_map> inboundStatus; - - std::unordered_map> conns; - std::unordered_map peers; - - std::unordered_map>> waitingOps; - - std::unordered_map> pendingOutbound; - - std::unordered_map> inboundWrites; - std::unordered_map> earlyWrites; - - PinnedStagingPool staging; - StreamPool streamPool{8}; - EventPool eventPool{64}; - std::deque gpuTasks; -}; - -/* ---------------------------------------------------------------------------------------------- */ -/* TcpBackendSession */ -/* ---------------------------------------------------------------------------------------------- */ -TcpBackendSession::TcpBackendSession(const TcpBackendConfig& cfg, const MemoryDesc& l, const MemoryDesc& r, - TcpTransport* t) +TcpBackendSession::TcpBackendSession(const TcpBackendConfig& cfg, const MemoryDesc& l, + const MemoryDesc& r, TcpTransport* t) : config(cfg), local(l), remote(r), transport(t) {} -void TcpBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status, - TransferUniqueId id, bool isRead) { +void TcpBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, + TransferStatus* status, TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; transport->SubmitReadWrite(local, localOffset, remote, remoteOffset, size, status, id, isRead); } void TcpBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, - const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, - bool isRead) { + const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; - transport->SubmitBatchReadWrite(local, localOffsets, remote, remoteOffsets, sizes, status, id, isRead); + transport->SubmitBatchReadWrite(local, localOffsets, remote, remoteOffsets, sizes, status, id, + isRead); } bool TcpBackendSession::Alive() const { return true; } -/* ---------------------------------------------------------------------------------------------- */ -/* TcpBackend */ -/* ---------------------------------------------------------------------------------------------- */ TcpBackend::TcpBackend(EngineKey k, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg) : myEngKey(std::move(k)), config(cfg) { transport = std::make_unique(myEngKey, engCfg, cfg); @@ -2701,26 +58,33 @@ TcpBackend::~TcpBackend() { transport->Shutdown(); } std::optional TcpBackend::GetListenPort() const { return transport->GetListenPort(); } -void TcpBackend::RegisterRemoteEngine(const EngineDesc& desc) { transport->RegisterRemoteEngine(desc); } +void TcpBackend::RegisterRemoteEngine(const EngineDesc& desc) { + transport->RegisterRemoteEngine(desc); +} -void TcpBackend::DeregisterRemoteEngine(const EngineDesc& desc) { transport->DeregisterRemoteEngine(desc); } +void TcpBackend::DeregisterRemoteEngine(const EngineDesc& desc) { + transport->DeregisterRemoteEngine(desc); +} void TcpBackend::RegisterMemory(MemoryDesc& desc) { transport->RegisterMemory(desc); } void TcpBackend::DeregisterMemory(const MemoryDesc& desc) { transport->DeregisterMemory(desc); } -void TcpBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc, - size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, - bool isRead) { +void TcpBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, + const MemoryDesc& remoteSrc, size_t remoteOffset, size_t size, + TransferStatus* status, TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; - transport->SubmitReadWrite(localDest, localOffset, remoteSrc, remoteOffset, size, status, id, isRead); + transport->SubmitReadWrite(localDest, localOffset, remoteSrc, remoteOffset, size, status, id, + isRead); } -void TcpBackend::BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, const MemoryDesc& remoteSrc, - const SizeVec& remoteOffsets, const SizeVec& sizes, TransferStatus* status, - TransferUniqueId id, bool isRead) { +void TcpBackend::BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, + const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) { MORI_IO_FUNCTION_TIMER; - transport->SubmitBatchReadWrite(localDest, localOffsets, remoteSrc, remoteOffsets, sizes, status, id, isRead); + transport->SubmitBatchReadWrite(localDest, localOffsets, remoteSrc, remoteOffsets, sizes, status, + id, isRead); } BackendSession* TcpBackend::CreateSession(const MemoryDesc& local, const MemoryDesc& remote) { @@ -2728,7 +92,8 @@ BackendSession* TcpBackend::CreateSession(const MemoryDesc& local, const MemoryD return sess; } -bool TcpBackend::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, TransferStatus* status) { +bool TcpBackend::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, + TransferStatus* status) { return transport->PopInboundTransferStatus(remote, id, status); } diff --git a/src/io/tcp/data_worker.hpp b/src/io/tcp/data_worker.hpp new file mode 100644 index 000000000..58c4999d8 --- /dev/null +++ b/src/io/tcp/data_worker.hpp @@ -0,0 +1,426 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include + +#include + +#include "src/io/tcp/tcp_types.hpp" + +namespace mori { +namespace io { + +class DataConnectionWorker { + public: + DataConnectionWorker(int fd, EngineKey peer, PinnedStagingPool* staging) + : fd_(fd), peerKey_(std::move(peer)), staging_(staging) { + notifyFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + wakeFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + } + + ~DataConnectionWorker() { + Stop(); + if (notifyFd_ >= 0) { + close(notifyFd_); + notifyFd_ = -1; + } + if (wakeFd_ >= 0) { + close(wakeFd_); + wakeFd_ = -1; + } + } + + DataConnectionWorker(const DataConnectionWorker&) = delete; + DataConnectionWorker& operator=(const DataConnectionWorker&) = delete; + + void Start() { + if (running_.load()) return; + running_.store(true); + thread_ = std::thread(&DataConnectionWorker::Run, this); + } + + void Stop() { + bool was = running_.exchange(false); + if (!was) return; + WakeWorker(); + if (thread_.joinable()) thread_.join(); + } + + int NotifyFd() const { return notifyFd_; } + int Fd() const { return fd_; } + + void SubmitSend(SendItem item) { + { + std::lock_guard lk(sendMu_); + sendQ_.push_back(std::move(item)); + } + WakeWorker(); + } + + void RegisterRecvTarget(TransferUniqueId opId, const WorkerRecvTarget& target) { + std::lock_guard lk(targetMu_); + recvTargets_[opId] = target; + } + + void RemoveRecvTarget(TransferUniqueId opId) { + std::lock_guard lk(targetMu_); + recvTargets_.erase(opId); + } + + void DrainEvents(std::deque& out) { + uint64_t v; + while (::read(notifyFd_, &v, sizeof(v)) > 0) { + } + std::lock_guard lk(eventMu_); + while (!eventQ_.empty()) { + out.push_back(std::move(eventQ_.front())); + eventQ_.pop_front(); + } + } + + private: + void WakeWorker() { + uint64_t one = 1; + ::write(wakeFd_, &one, sizeof(one)); + } + + void NotifyMain() { + uint64_t one = 1; + ::write(notifyFd_, &one, sizeof(one)); + } + + void PostEvent(WorkerEvent ev) { + { + std::lock_guard lk(eventMu_); + eventQ_.push_back(std::move(ev)); + } + NotifyMain(); + } + + void PostRecvDone(TransferUniqueId opId, uint8_t lane, uint64_t laneLen, bool discarded = false) { + WorkerEvent ev; + ev.type = WorkerEventType::RECV_DONE; + ev.peerKey = peerKey_; + ev.opId = opId; + ev.lane = lane; + ev.laneLen = laneLen; + ev.discarded = discarded; + PostEvent(std::move(ev)); + } + + void PostEarlyData(TransferUniqueId opId, uint8_t lane, uint64_t laneLen, + std::shared_ptr buf) { + WorkerEvent ev; + ev.type = WorkerEventType::EARLY_DATA; + ev.peerKey = peerKey_; + ev.opId = opId; + ev.lane = lane; + ev.laneLen = laneLen; + ev.earlyBuf = std::move(buf); + PostEvent(std::move(ev)); + } + + void PostCallback(std::function cb) { + WorkerEvent ev; + ev.type = WorkerEventType::SEND_CALLBACK; + ev.callback = std::move(cb); + PostEvent(std::move(ev)); + } + + void PostError(const std::string& msg) { + WorkerEvent ev; + ev.type = WorkerEventType::CONN_ERROR; + ev.peerKey = peerKey_; + ev.errorMsg = msg; + PostEvent(std::move(ev)); + } + + void Run() { + MORI_IO_TRACE("TCP: DataWorker fd={} peer={} started", fd_, peerKey_); + struct pollfd pfds[2]; + pfds[0].fd = fd_; + pfds[1].fd = wakeFd_; + pfds[1].events = POLLIN; + + while (running_.load()) { + bool hasSend; + { + std::lock_guard lk(sendMu_); + hasSend = !sendQ_.empty(); + } + + pfds[0].events = POLLIN | (hasSend ? POLLOUT : 0); + pfds[0].revents = 0; + pfds[1].revents = 0; + + int n = ::poll(pfds, 2, hasSend ? 0 : 1); + if (n < 0) { + if (errno == EINTR) continue; + PostError(std::string("poll failed: ") + strerror(errno)); + break; + } + + if (pfds[1].revents & POLLIN) { + uint64_t v; + while (::read(wakeFd_, &v, sizeof(v)) > 0) { + } + } + + if (pfds[0].revents & (POLLERR | POLLHUP | POLLNVAL)) { + PostError("data connection error/hangup"); + break; + } + + if (pfds[0].revents & POLLOUT) { + if (!ProcessSend()) break; + } + + if (pfds[0].revents & POLLIN) { + if (!ProcessRecv()) break; + } + } + MORI_IO_TRACE("TCP: DataWorker fd={} peer={} exiting", fd_, peerKey_); + } + + bool ProcessSend() { + std::deque batch; + { + std::lock_guard lk(sendMu_); + batch.swap(sendQ_); + } + + for (auto& item : batch) { + while (!item.Done()) { + constexpr size_t kMaxIov = 64; + iovec iov[kMaxIov]; + size_t cnt = 0; + for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { + iov[cnt] = item.iov[i]; + if (i == item.idx && item.off > 0) { + iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; + iov[cnt].iov_len -= item.off; + } + cnt++; + } + msghdr msg{}; + msg.msg_iov = iov; + msg.msg_iovlen = cnt; + ssize_t n = ::sendmsg(fd_, &msg, MSG_NOSIGNAL | item.flags); + if (n < 0) { + if (IsWouldBlock(errno)) { + goto requeue_batch; + } + PostError(std::string("sendmsg failed: ") + strerror(errno)); + return false; + } + if (n == 0) { + goto requeue_batch; + } + item.Advance(static_cast(n)); + } + if (item.onDone) { + PostCallback(std::move(item.onDone)); + } + } + return true; + + requeue_batch: { + std::lock_guard lk(sendMu_); + for (auto rit = batch.rbegin(); rit != batch.rend(); ++rit) { + if (!rit->Done()) { + sendQ_.push_front(std::move(*rit)); + } + } + } + return true; + } + + bool ProcessRecv() { + while (true) { + while (hdrGot_ < tcp::kDataHeaderSize) { + ssize_t n = ::recv(fd_, hdrBuf_ + hdrGot_, tcp::kDataHeaderSize - hdrGot_, 0); + if (n < 0) { + if (IsWouldBlock(errno)) return true; + PostError(std::string("recv header failed: ") + strerror(errno)); + return false; + } + if (n == 0) { + PostError("data connection closed by peer"); + return false; + } + hdrGot_ += static_cast(n); + } + hdrGot_ = 0; + + tcp::DataHeaderView hv; + if (!tcp::TryParseDataHeader(hdrBuf_, tcp::kDataHeaderSize, &hv)) { + PostError("bad data header"); + return false; + } + + const uint8_t lane = static_cast(hv.opId & kLaneMask); + const TransferUniqueId userOpId = static_cast(ToUserOpId(hv.opId)); + const uint64_t payloadLen = hv.payloadLen; + + WorkerRecvTarget target; + bool hasTarget = false; + { + std::lock_guard lk(targetMu_); + auto it = recvTargets_.find(userOpId); + if (it != recvTargets_.end()) { + target = it->second; + hasTarget = true; + } + } + + if (hasTarget && !target.discard) { + const LaneSpan span = ComputeLaneSpan(target.totalLen, target.lanesTotal, lane); + if (span.len != payloadLen) { + MORI_IO_WARN("TCP: worker recv op {} lane {} len mismatch expected={} got={}", userOpId, + (uint32_t)lane, span.len, payloadLen); + if (!DiscardPayload(payloadLen)) return false; + PostRecvDone(userOpId, lane, payloadLen, true); + } else if (target.toGpu) { + uint8_t* dst = reinterpret_cast(target.pinned->ptr) + span.off; + if (!RecvExact(dst, payloadLen)) return false; + PostRecvDone(userOpId, lane, payloadLen); + } else { + auto segs = SliceSegments(target.segs, span.off, span.len); + if (!RecvIntoSegments(reinterpret_cast(target.cpuBase), segs, payloadLen)) + return false; + PostRecvDone(userOpId, lane, payloadLen); + } + } else if (hasTarget && target.discard) { + if (!DiscardPayload(payloadLen)) return false; + PostRecvDone(userOpId, lane, payloadLen, true); + } else { + if (payloadLen == 0) { + PostEarlyData(userOpId, lane, 0, nullptr); + } else { + auto buf = staging_->Acquire(static_cast(payloadLen)); + if (!buf) { + if (!DiscardPayload(payloadLen)) return false; + PostRecvDone(userOpId, lane, payloadLen, true); + } else { + if (!RecvExact(reinterpret_cast(buf->ptr), payloadLen)) return false; + PostEarlyData(userOpId, lane, payloadLen, std::move(buf)); + } + } + } + } + return true; + } + + bool RecvExact(uint8_t* dst, uint64_t len) { + uint64_t got = 0; + while (got < len) { + const size_t want = static_cast(std::min(len - got, 16ULL * 1024 * 1024)); + ssize_t n = ::recv(fd_, dst + got, want, 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostError(std::string("recv payload failed: ") + strerror(errno)); + return false; + } + if (n == 0) { + PostError("data connection closed during recv"); + return false; + } + got += static_cast(n); + } + return true; + } + + bool RecvIntoSegments(uint8_t* base, const std::vector& segs, uint64_t totalLen) { + uint64_t remaining = totalLen; + size_t segIdx = 0; + uint64_t segOff = 0; + while (remaining > 0 && segIdx < segs.size()) { + const Segment& seg = segs[segIdx]; + const uint64_t segRemain = seg.len - segOff; + const size_t want = static_cast( + std::min(remaining, std::min(segRemain, 16ULL * 1024 * 1024))); + uint8_t* dst = base + seg.off + segOff; + ssize_t n = ::recv(fd_, dst, want, 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostError(std::string("recv seg failed: ") + strerror(errno)); + return false; + } + if (n == 0) { + PostError("data connection closed during seg recv"); + return false; + } + remaining -= static_cast(n); + segOff += static_cast(n); + if (segOff >= seg.len) { + segIdx++; + segOff = 0; + } + } + return (remaining == 0); + } + + bool DiscardPayload(uint64_t len) { + uint8_t tmp[65536]; + uint64_t remaining = len; + while (remaining > 0) { + const size_t want = static_cast(std::min(remaining, sizeof(tmp))); + ssize_t n = ::recv(fd_, tmp, want, 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostError(std::string("recv discard failed: ") + strerror(errno)); + return false; + } + if (n == 0) { + PostError("data connection closed during discard"); + return false; + } + remaining -= static_cast(n); + } + return true; + } + + int fd_; + EngineKey peerKey_; + PinnedStagingPool* staging_; + std::atomic running_{false}; + std::thread thread_; + int notifyFd_{-1}; + int wakeFd_{-1}; + + std::mutex sendMu_; + std::deque sendQ_; + + std::mutex targetMu_; + std::unordered_map recvTargets_; + + std::mutex eventMu_; + std::deque eventQ_; + + uint8_t hdrBuf_[tcp::kDataHeaderSize]{}; + size_t hdrGot_{0}; +}; + +} // namespace io +} // namespace mori diff --git a/src/io/tcp/tcp_types.hpp b/src/io/tcp/tcp_types.hpp new file mode 100644 index 000000000..49f06f5e0 --- /dev/null +++ b/src/io/tcp/tcp_types.hpp @@ -0,0 +1,421 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/io/backend.hpp" +#include "mori/io/common.hpp" +#include "mori/io/logging.hpp" +#include "src/io/tcp/protocol.hpp" + +namespace mori { +namespace io { + +class DataConnectionWorker; + +using Clock = std::chrono::steady_clock; + +inline bool IsWouldBlock(int err) { return (err == EAGAIN) || (err == EWOULDBLOCK); } + +inline int SetNonBlocking(int fd) { + int flags = fcntl(fd, F_GETFL, 0); + if (flags < 0) return -1; + if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) return -1; + return 0; +} + +inline int SetBlocking(int fd) { + int flags = fcntl(fd, F_GETFL, 0); + if (flags < 0) return -1; + if (fcntl(fd, F_SETFL, flags & ~O_NONBLOCK) < 0) return -1; + return 0; +} + +inline void SetSockOptOrLog(int fd, int level, int optname, const void* optval, socklen_t optlen, + const char* name) { + if (setsockopt(fd, level, optname, optval, optlen) != 0) { + MORI_IO_WARN("TCP: setsockopt {} failed: {}", name, strerror(errno)); + } +} + +inline void ConfigureSocketCommon(int fd, const TcpBackendConfig& cfg) { + if (cfg.enableKeepalive) { + int on = 1; + SetSockOptOrLog(fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on), "SO_KEEPALIVE"); + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPIDLE, &cfg.keepaliveIdleSec, + sizeof(cfg.keepaliveIdleSec), "TCP_KEEPIDLE"); + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPINTVL, &cfg.keepaliveIntvlSec, + sizeof(cfg.keepaliveIntvlSec), "TCP_KEEPINTVL"); + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPCNT, &cfg.keepaliveCnt, sizeof(cfg.keepaliveCnt), + "TCP_KEEPCNT"); + } +} + +inline void ConfigureCtrlSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + if (cfg.enableCtrlNodelay) { + int on = 1; + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(ctrl)"); + } +} + +inline void ConfigureDataSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + { + int on = 1; + SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(data)"); + } + if (cfg.sockSndbufBytes > 0) { + SetSockOptOrLog(fd, SOL_SOCKET, SO_SNDBUF, &cfg.sockSndbufBytes, sizeof(cfg.sockSndbufBytes), + "SO_SNDBUF"); + } + if (cfg.sockRcvbufBytes > 0) { + SetSockOptOrLog(fd, SOL_SOCKET, SO_RCVBUF, &cfg.sockRcvbufBytes, sizeof(cfg.sockRcvbufBytes), + "SO_RCVBUF"); + } +} + +inline std::optional ParseIpv4(const std::string& host, uint16_t port) { + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) return std::nullopt; + return addr; +} + +inline uint16_t GetBoundPort(int fd) { + sockaddr_in addr{}; + socklen_t len = sizeof(addr); + if (getsockname(fd, reinterpret_cast(&addr), &len) != 0) return 0; + return ntohs(addr.sin_port); +} + +inline size_t RoundUpPow2(size_t v) { + if (v <= 1) return 1; + size_t p = 1; + while (p < v) p <<= 1; + return p; +} + +struct PinnedBuf { + void* ptr{nullptr}; + size_t cap{0}; +}; + +class PinnedStagingPool { + public: + PinnedStagingPool() = default; + ~PinnedStagingPool() { Clear(); } + + PinnedStagingPool(const PinnedStagingPool&) = delete; + PinnedStagingPool& operator=(const PinnedStagingPool&) = delete; + + std::shared_ptr Acquire(size_t size) { + const size_t cap = RoundUpPow2(size); + { + std::lock_guard lock(mu); + auto it = free.find(cap); + if (it != free.end() && !it->second.empty()) { + void* p = it->second.back(); + it->second.pop_back(); + return std::shared_ptr(new PinnedBuf{p, cap}, + [this](PinnedBuf* b) { this->Release(b); }); + } + } + void* p = nullptr; + hipError_t err = hipHostMalloc(&p, cap, hipHostMallocDefault); + if (err != hipSuccess) { + MORI_IO_ERROR("TCP: hipHostMalloc({}) failed: {}", cap, hipGetErrorString(err)); + return nullptr; + } + return std::shared_ptr(new PinnedBuf{p, cap}, + [this](PinnedBuf* b) { this->Release(b); }); + } + + void Clear() { + std::lock_guard lock(mu); + for (auto& kv : free) { + for (void* p : kv.second) { + hipError_t err = hipHostFree(p); + if (err != hipSuccess) { + MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); + } + } + kv.second.clear(); + } + free.clear(); + } + + private: + void Release(PinnedBuf* b) { + if (b == nullptr) return; + const size_t cap = b->cap; + void* p = b->ptr; + delete b; + constexpr size_t kMaxCachedPerClass = 8; + std::lock_guard lock(mu); + auto& vec = free[cap]; + if (vec.size() < kMaxCachedPerClass) { + vec.push_back(p); + } else { + hipError_t err = hipHostFree(p); + if (err != hipSuccess) { + MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); + } + } + } + + std::mutex mu; + std::unordered_map> free; +}; + +struct Segment { + uint64_t off{0}; + uint64_t len{0}; +}; + +constexpr uint8_t kLaneBits = 3; +constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; + +inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { + return (userOpId << kLaneBits) | lane; +} +inline uint64_t ToUserOpId(uint64_t wireOpId) { return wireOpId >> kLaneBits; } + +struct LaneSpan { + uint64_t off{0}; + uint64_t len{0}; +}; + +inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanesTotal, uint8_t lane) { + if (lanesTotal <= 1) return LaneSpan{0, total}; + const uint64_t base = total / lanesTotal; + const uint64_t rem = total % lanesTotal; + return LaneSpan{static_cast(lane) * base + std::min(lane, rem), + base + (lane < rem ? 1 : 0)}; +} + +inline uint8_t LanesAllMask(uint8_t lanesTotal) { + if (lanesTotal >= (1U << kLaneBits)) return 0xFF; + return static_cast((1U << lanesTotal) - 1U); +} + +inline uint8_t ClampLanesTotal(uint8_t lanesTotal) { + if (lanesTotal == 0) return 1; + return std::min(lanesTotal, static_cast(1U << kLaneBits)); +} + +inline std::vector SliceSegments(const std::vector& segs, uint64_t start, + uint64_t len) { + std::vector out; + if (len == 0) return out; + uint64_t skip = start; + uint64_t remaining = len; + for (const auto& s : segs) { + if (remaining == 0) break; + if (skip >= s.len) { + skip -= s.len; + continue; + } + const uint64_t take = std::min(s.len - skip, remaining); + out.push_back({s.off + skip, take}); + remaining -= take; + skip = 0; + } + return out; +} + +inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, + uint64_t* outLen) { + if (!outOff || !outLen || segs.empty()) return false; + uint64_t off = segs[0].off; + uint64_t end = off + segs[0].len; + for (size_t i = 1; i < segs.size(); ++i) { + if (segs[i].off != end) return false; + end += segs[i].len; + } + *outOff = off; + *outLen = end - off; + return true; +} + +inline uint64_t SumLens(const std::vector& segs) { + uint64_t total = 0; + for (const auto& s : segs) total += s.len; + return total; +} + +struct SendItem { + std::vector header; + std::vector iov; + size_t idx{0}; + size_t off{0}; + int flags{0}; + std::shared_ptr keepalive; + std::function onDone; + + bool Done() const { return idx >= iov.size(); } + + void Advance(size_t n) { + while (n > 0 && idx < iov.size()) { + size_t avail = iov[idx].iov_len - off; + if (n < avail) { + off += n; + n = 0; + break; + } + n -= avail; + idx++; + off = 0; + } + } +}; + +struct Connection { + int fd{-1}; + bool isOutgoing{false}; + bool connecting{false}; + bool helloSent{false}; + bool helloReceived{false}; + tcp::Channel ch{tcp::Channel::CTRL}; + EngineKey peerKey{}; + std::vector inbuf; + std::deque sendq; +}; + +struct PeerLinks { + int ctrlFd{-1}; + std::vector dataFds; + std::vector workers; + int ctrlPending{0}; + int dataPending{0}; + size_t rr{0}; + bool CtrlUp() const { return ctrlFd >= 0; } + bool DataUp() const { return !dataFds.empty(); } +}; + +struct InboundStatusEntry { + StatusCode code{StatusCode::INIT}; + std::string msg; +}; + +struct OutboundOpState { + EngineKey peer; + TransferUniqueId id{0}; + bool isRead{false}; + TransferStatus* status{nullptr}; + MemoryDesc local{}; + std::vector localSegs; + MemoryDesc remote{}; + std::vector remoteSegs; + uint64_t expectedRxBytes{0}; + uint64_t rxBytes{0}; + bool completionReceived{false}; + uint8_t lanesTotal{1}; + uint8_t lanesDoneMask{0}; + StatusCode completionCode{StatusCode::SUCCESS}; + std::string completionMsg; + bool gpuCopyPending{false}; + std::shared_ptr pinned; + Clock::time_point startTs{Clock::now()}; +}; + +struct InboundWriteState { + EngineKey peer; + TransferUniqueId id{0}; + MemoryDesc dst{}; + std::vector dstSegs; + bool discard{false}; + uint8_t lanesTotal{1}; + uint8_t lanesDoneMask{0}; + std::shared_ptr pinned; +}; + +struct EarlyWriteLaneState { + uint64_t payloadLen{0}; + std::shared_ptr pinned; + bool complete{false}; +}; + +struct EarlyWriteState { + std::unordered_map lanes; +}; + +struct WorkerRecvTarget { + uint8_t lanesTotal{1}; + uint64_t totalLen{0}; + bool discard{false}; + bool toGpu{false}; + void* cpuBase{nullptr}; + std::vector segs; + std::shared_ptr pinned; +}; + +enum class WorkerEventType : uint8_t { + RECV_DONE = 0, + EARLY_DATA = 1, + SEND_CALLBACK = 2, + CONN_ERROR = 3, +}; + +struct WorkerEvent { + WorkerEventType type{WorkerEventType::RECV_DONE}; + EngineKey peerKey; + TransferUniqueId opId{0}; + uint8_t lane{0}; + uint64_t laneLen{0}; + bool discarded{false}; + std::shared_ptr earlyBuf; + std::function callback; + std::string errorMsg; +}; + +struct GpuTask { + int deviceId{-1}; + hipEvent_t ev{nullptr}; + std::function onReady; +}; + +} // namespace io +} // namespace mori diff --git a/src/io/tcp/transport.cpp b/src/io/tcp/transport.cpp new file mode 100644 index 000000000..6f7ba61ec --- /dev/null +++ b/src/io/tcp/transport.cpp @@ -0,0 +1,1657 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "src/io/tcp/transport.hpp" + +#include + +namespace mori { +namespace io { + +namespace { + +struct LinearReqView { + uint64_t opId{0}; + uint32_t memId{0}; + uint64_t remoteOff{0}; + uint64_t size{0}; + uint8_t lanesTotal{1}; +}; + +bool ParseLinearReq(const uint8_t* body, size_t len, LinearReqView* out) { + if (out == nullptr) return false; + size_t off = 0; + if (!tcp::ReadU64BE(body, len, &off, &out->opId) || + !tcp::ReadU32BE(body, len, &off, &out->memId) || + !tcp::ReadU64BE(body, len, &off, &out->remoteOff) || + !tcp::ReadU64BE(body, len, &off, &out->size)) { + return false; + } + if (off < len) out->lanesTotal = body[off]; + out->lanesTotal = ClampLanesTotal(out->lanesTotal); + return true; +} + +struct BatchReqView { + uint64_t opId{0}; + uint32_t memId{0}; + std::vector segs; + uint8_t lanesTotal{1}; +}; + +bool ParseBatchReq(const uint8_t* body, size_t len, BatchReqView* out) { + if (out == nullptr) return false; + size_t off = 0; + uint32_t n = 0; + if (!tcp::ReadU64BE(body, len, &off, &out->opId) || + !tcp::ReadU32BE(body, len, &off, &out->memId) || !tcp::ReadU32BE(body, len, &off, &n)) { + return false; + } + + out->segs.clear(); + out->segs.reserve(n); + for (uint32_t i = 0; i < n; ++i) { + uint64_t remoteOff = 0; + uint64_t size = 0; + if (!tcp::ReadU64BE(body, len, &off, &remoteOff) || !tcp::ReadU64BE(body, len, &off, &size)) { + return false; + } + if (size > 0) out->segs.push_back({remoteOff, size}); + } + + if (off < len) out->lanesTotal = body[off]; + out->lanesTotal = ClampLanesTotal(out->lanesTotal); + return true; +} + +struct CompletionView { + uint64_t opId{0}; + uint32_t statusCode{0}; + std::string msg; +}; + +bool ParseCompletion(const uint8_t* body, size_t len, CompletionView* out) { + if (out == nullptr) return false; + size_t off = 0; + uint32_t msgLen = 0; + if (!tcp::ReadU64BE(body, len, &off, &out->opId) || + !tcp::ReadU32BE(body, len, &off, &out->statusCode) || + !tcp::ReadU32BE(body, len, &off, &msgLen)) { + return false; + } + if (off + msgLen > len) return false; + out->msg.assign(reinterpret_cast(body + off), msgLen); + return true; +} + +bool SegmentsInRange(const std::vector& segs, uint64_t memSize) { + for (const auto& seg : segs) { + if (seg.off + seg.len > memSize) return false; + } + return true; +} + +} // namespace + +TcpTransport::TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, + const TcpBackendConfig& cfg) + : myEngKey(std::move(myKey)), engConfig(engCfg), config(cfg) {} + +TcpTransport::~TcpTransport() { Shutdown(); } + +void TcpTransport::Start() { + if (running.load()) return; + + epfd = epoll_create1(EPOLL_CLOEXEC); + assert(epfd >= 0); + + listenFd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); + assert(listenFd >= 0); + + int one = 1; + SetSockOptOrLog(listenFd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one), "SO_REUSEADDR"); + + auto addrOpt = + ParseIpv4(engConfig.host.empty() ? std::string("0.0.0.0") : engConfig.host, engConfig.port); + assert(addrOpt.has_value()); + sockaddr_in addr = addrOpt.value(); + if (bind(listenFd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + MORI_IO_ERROR("TCP: bind {}:{} failed: {}", engConfig.host, engConfig.port, strerror(errno)); + assert(false && "bind failed"); + } + if (listen(listenFd, 256) != 0) { + MORI_IO_ERROR("TCP: listen failed: {}", strerror(errno)); + assert(false && "listen failed"); + } + listenPort = GetBoundPort(listenFd); + MORI_IO_INFO("TCP: listen on {}:{} (port={})", engConfig.host, engConfig.port, listenPort); + + wakeFd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + assert(wakeFd >= 0); + + AddEpoll(listenFd, true, false); + AddEpoll(wakeFd, true, false); + + running.store(true); + ioThread = std::thread([this] { this->IoLoop(); }); +} + +void TcpTransport::Shutdown() { + bool wasRunning = running.exchange(false); + if (!wasRunning) return; + + if (wakeFd >= 0) { + uint64_t one = 1; + ::write(wakeFd, &one, sizeof(one)); + } + if (ioThread.joinable()) ioThread.join(); + + for (auto& kv : dataWorkers) kv.second->Stop(); + dataWorkers.clear(); + workerNotifyMap.clear(); + + for (auto& kv : conns) CloseConnInternal(kv.second.get()); + conns.clear(); + peers.clear(); + + if (listenFd >= 0) { + close(listenFd); + listenFd = -1; + } + if (wakeFd >= 0) { + close(wakeFd); + wakeFd = -1; + } + if (epfd >= 0) { + close(epfd); + epfd = -1; + } +} + +std::optional TcpTransport::GetListenPort() const { + if (listenPort == 0) return std::nullopt; + return listenPort; +} + +void TcpTransport::RegisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteMu); + remoteEngines[desc.key] = desc; +} + +void TcpTransport::DeregisterRemoteEngine(const EngineDesc& desc) { + std::lock_guard lock(remoteMu); + remoteEngines.erase(desc.key); +} + +void TcpTransport::RegisterMemory(const MemoryDesc& desc) { + std::lock_guard lock(memMu); + localMems[desc.id] = desc; +} + +void TcpTransport::DeregisterMemory(const MemoryDesc& desc) { + std::lock_guard lock(memMu); + localMems.erase(desc.id); +} + +bool TcpTransport::PopInboundTransferStatus(const EngineKey& remote, TransferUniqueId id, + TransferStatus* status) { + std::lock_guard lock(inboundMu); + auto it = inboundStatus.find(remote); + if (it == inboundStatus.end()) return false; + auto it2 = it->second.find(id); + if (it2 == it->second.end()) return false; + status->Update(it2->second.code, it2->second.msg); + it->second.erase(it2); + return true; +} + +void TcpTransport::SubmitReadWrite(const MemoryDesc& local, size_t localOffset, + const MemoryDesc& remote, size_t remoteOffset, size_t size, + TransferStatus* status, TransferUniqueId id, bool isRead) { + if (status == nullptr) return; + if (size == 0) { + status->SetCode(StatusCode::SUCCESS); + return; + } + if ((localOffset + size) > local.size || (remoteOffset + size) > remote.size) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: offset+size out of range"); + return; + } + + auto op = std::make_unique(); + op->peer = remote.engineKey; + op->id = id; + op->isRead = isRead; + op->status = status; + op->local = local; + op->remote = remote; + op->localSegs = {{static_cast(localOffset), static_cast(size)}}; + op->remoteSegs = {{static_cast(remoteOffset), static_cast(size)}}; + op->expectedRxBytes = isRead ? static_cast(size) : 0; + op->startTs = Clock::now(); + + status->SetCode(StatusCode::IN_PROGRESS); + EnqueueOp(std::move(op)); +} + +void TcpTransport::SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& localOffsets, + const MemoryDesc& remote, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, + TransferUniqueId id, bool isRead) { + if (status == nullptr) return; + const size_t n = sizes.size(); + if (n == 0) { + status->SetCode(StatusCode::SUCCESS); + return; + } + if (localOffsets.size() != n || remoteOffsets.size() != n) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch vector size mismatch"); + return; + } + + std::vector localSegs, remoteSegs; + localSegs.reserve(n); + remoteSegs.reserve(n); + uint64_t total = 0; + for (size_t i = 0; i < n; ++i) { + const size_t lo = localOffsets[i], ro = remoteOffsets[i], sz = sizes[i]; + if (sz == 0) continue; + if ((lo + sz) > local.size || (ro + sz) > remote.size) { + status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch offset+size out of range"); + return; + } + localSegs.push_back({static_cast(lo), static_cast(sz)}); + remoteSegs.push_back({static_cast(ro), static_cast(sz)}); + total += static_cast(sz); + } + + if (localSegs.size() > 1) { + std::vector newLocal, newRemote; + newLocal.reserve(localSegs.size()); + newRemote.reserve(remoteSegs.size()); + Segment curL = localSegs[0], curR = remoteSegs[0]; + for (size_t i = 1; i < localSegs.size(); ++i) { + const Segment& l = localSegs[i]; + const Segment& r = remoteSegs[i]; + if ((curL.off + curL.len == l.off) && (curR.off + curR.len == r.off) && + (curL.len == curR.len) && (l.len == r.len)) { + curL.len += l.len; + curR.len += r.len; + } else { + newLocal.push_back(curL); + newRemote.push_back(curR); + curL = l; + curR = r; + } + } + newLocal.push_back(curL); + newRemote.push_back(curR); + localSegs = std::move(newLocal); + remoteSegs = std::move(newRemote); + } + + auto op = std::make_unique(); + op->peer = remote.engineKey; + op->id = id; + op->isRead = isRead; + op->status = status; + op->local = local; + op->remote = remote; + op->localSegs = std::move(localSegs); + op->remoteSegs = std::move(remoteSegs); + op->expectedRxBytes = isRead ? total : 0; + op->startTs = Clock::now(); + + status->SetCode(StatusCode::IN_PROGRESS); + EnqueueOp(std::move(op)); +} + +void TcpTransport::EnqueueOp(std::unique_ptr op) { + { + std::lock_guard lock(submitMu); + submitQ.push_back(std::move(op)); + } + uint64_t one = 1; + ::write(wakeFd, &one, sizeof(one)); +} + +void TcpTransport::AddEpoll(int fd, bool wantRead, bool wantWrite) { + epoll_event ev{}; + ev.data.fd = fd; + ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev)); +} + +void TcpTransport::ModEpoll(int fd, bool wantRead, bool wantWrite) { + epoll_event ev{}; + ev.data.fd = fd; + ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev)); +} + +void TcpTransport::DelEpoll(int fd) { epoll_ctl(epfd, EPOLL_CTL_DEL, fd, nullptr); } + +void TcpTransport::CloseConnInternal(Connection* c) { + if (c == nullptr) return; + if (c->fd >= 0) { + DelEpoll(c->fd); + shutdown(c->fd, SHUT_RDWR); + close(c->fd); + c->fd = -1; + } +} + +bool TcpTransport::PreferOutgoingFor(const EngineKey& peerKey) const { return myEngKey < peerKey; } + +void TcpTransport::AssignConnToPeer(Connection* c) { + assert(c && c->helloReceived); + PeerLinks& link = peers[c->peerKey]; + const bool preferOutgoing = PreferOutgoingFor(c->peerKey); + const bool wasOutgoing = c->isOutgoing; + const tcp::Channel ch = c->ch; + + if (wasOutgoing) { + if (ch == tcp::Channel::CTRL) { + if (link.ctrlPending > 0) link.ctrlPending--; + } else { + if (link.dataPending > 0) link.dataPending--; + } + } + + auto replace_if_needed = [&](int& slotFd) { + if (slotFd < 0) { + slotFd = c->fd; + return; + } + const int existingFd = slotFd; + const int newFd = c->fd; + Connection* existing = conns[existingFd].get(); + if (!existing) { + slotFd = newFd; + return; + } + const bool keepNew = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); + if (keepNew) { + MORI_IO_WARN("TCP: peer {} channel {} replacing fd {} with fd {}", c->peerKey, + static_cast(c->ch), existing->fd, c->fd); + CloseConnInternal(existing); + conns.erase(existingFd); + slotFd = newFd; + } else { + MORI_IO_WARN("TCP: peer {} channel {} dropping duplicate fd {}", c->peerKey, + static_cast(c->ch), c->fd); + CloseConnInternal(c); + conns.erase(newFd); + } + }; + + if (ch == tcp::Channel::CTRL) { + replace_if_needed(link.ctrlFd); + return; + } + + const bool keepPreferred = + (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); + if (!keepPreferred) { + MORI_IO_TRACE("TCP: peer {} dropping non-preferred DATA fd {} outgoing={}", c->peerKey, c->fd, + c->isOutgoing); + const int fd = c->fd; + CloseConnInternal(c); + conns.erase(fd); + return; + } + + const size_t want = static_cast(std::max(1, config.numDataConns)); + if (link.dataFds.size() >= want) { + MORI_IO_TRACE("TCP: peer {} dropping extra DATA fd {} (have {} want {})", c->peerKey, c->fd, + link.dataFds.size(), want); + const int fd = c->fd; + CloseConnInternal(c); + conns.erase(fd); + return; + } + + const int dataFd = c->fd; + link.dataFds.push_back(dataFd); + MORI_IO_TRACE("TCP: peer {} DATA conn up {}/{}", c->peerKey.c_str(), link.dataFds.size(), want); + + DelEpoll(dataFd); + SetNonBlocking(dataFd); + ConfigureDataSocket(dataFd, config); + + auto worker = std::make_unique(dataFd, c->peerKey, &staging); + worker->Start(); + AddEpoll(worker->NotifyFd(), true, false); + workerNotifyMap[worker->NotifyFd()] = worker.get(); + link.workers.push_back(worker.get()); + dataWorkers[dataFd] = std::move(worker); +} + +void TcpTransport::MaybeDispatchQueuedOps(const EngineKey& peerKey) { + auto it = peers.find(peerKey); + if (it == peers.end()) return; + if (!it->second.CtrlUp() || !it->second.DataUp()) return; + Connection* ctrl = conns[it->second.ctrlFd].get(); + if (!ctrl || !ctrl->helloReceived) return; + if (it->second.workers.empty()) return; + + auto qit = waitingOps.find(peerKey); + if (qit == waitingOps.end()) return; + + auto ops = std::move(qit->second); + waitingOps.erase(qit); + MORI_IO_TRACE("TCP: peer {} ready, dispatch {} queued ops", peerKey.c_str(), ops.size()); + for (auto& op : ops) DispatchOp(std::move(op)); +} + +void TcpTransport::EnsurePeerChannels(const EngineKey& peerKey) { + PeerLinks& link = peers[peerKey]; + if (!link.CtrlUp() && link.ctrlPending == 0) ConnectChannel(peerKey, tcp::Channel::CTRL); + const int want = std::max(1, config.numDataConns); + while (static_cast(link.dataFds.size()) + link.dataPending < want) + ConnectChannel(peerKey, tcp::Channel::DATA); +} + +void TcpTransport::ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { + EngineDesc desc; + { + std::lock_guard lock(remoteMu); + auto it = remoteEngines.find(peerKey); + if (it == remoteEngines.end()) { + MORI_IO_ERROR("TCP: remote engine {} not registered", peerKey.c_str()); + return; + } + desc = it->second; + } + + auto peerAddrOpt = ParseIpv4(desc.host, static_cast(desc.port)); + if (!peerAddrOpt.has_value()) { + MORI_IO_ERROR("TCP: invalid remote host {}:{}", desc.host, desc.port); + return; + } + sockaddr_in peerAddr = peerAddrOpt.value(); + + int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); + if (fd < 0) { + MORI_IO_ERROR("TCP: socket() failed: {}", strerror(errno)); + return; + } + MORI_IO_TRACE("TCP: connect start peer={} ch={} fd={}", peerKey.c_str(), static_cast(ch), + fd); + + if (!engConfig.host.empty()) { + auto localAddrOpt = ParseIpv4(engConfig.host, 0); + if (localAddrOpt.has_value()) { + sockaddr_in localAddr = localAddrOpt.value(); + if (bind(fd, reinterpret_cast(&localAddr), sizeof(localAddr)) != 0) { + MORI_IO_WARN("TCP: bind(local) {} failed: {}", engConfig.host, strerror(errno)); + } + } + } + + int rc = connect(fd, reinterpret_cast(&peerAddr), sizeof(peerAddr)); + bool connecting = false; + if (rc != 0) { + if (errno == EINPROGRESS) { + connecting = true; + } else { + MORI_IO_ERROR("TCP: connect to {}:{} failed: {}", desc.host, desc.port, strerror(errno)); + close(fd); + return; + } + } + + auto conn = std::make_unique(); + conn->fd = fd; + conn->isOutgoing = true; + conn->connecting = connecting; + conn->peerKey = peerKey; + conn->ch = ch; + conn->inbuf.reserve(4096); + + if (ch == tcp::Channel::CTRL) ConfigureCtrlSocket(fd, config); + + const bool wantWrite = connecting || !conn->sendq.empty(); + AddEpoll(fd, true, wantWrite); + conns[fd] = std::move(conn); + + PeerLinks& link = peers[peerKey]; + if (ch == tcp::Channel::CTRL) + link.ctrlPending++; + else + link.dataPending++; + + if (!connecting) { + QueueHello(fd); + ModEpoll(fd, true, true); + } +} + +void TcpTransport::QueueHello(int fd) { + Connection* c = conns[fd].get(); + if (!c || c->helloSent) return; + c->helloSent = true; + auto hello = tcp::BuildHello(c->ch, myEngKey); + MORI_IO_TRACE("TCP: queue HELLO fd={} ch={} peer={}", fd, static_cast(c->ch), + c->peerKey.c_str()); + + SendItem item; + item.header = std::move(hello); + item.iov.resize(1); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + c->sendq.push_back(std::move(item)); +} + +void TcpTransport::AcceptNew() { + while (true) { + sockaddr_in peer{}; + socklen_t len = sizeof(peer); + int fd = + accept4(listenFd, reinterpret_cast(&peer), &len, SOCK_NONBLOCK | SOCK_CLOEXEC); + if (fd < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_WARN("TCP: accept failed: {}", strerror(errno)); + break; + } + MORI_IO_TRACE("TCP: accept fd={}", fd); + + auto conn = std::make_unique(); + conn->fd = fd; + conn->isOutgoing = false; + conn->inbuf.reserve(4096); + AddEpoll(fd, true, false); + conns[fd] = std::move(conn); + } +} + +void TcpTransport::DrainWakeFd() { + uint64_t v = 0; + while (true) { + ssize_t n = ::read(wakeFd, &v, sizeof(v)); + if (n <= 0) break; + } + + std::deque> ops; + { + std::lock_guard lock(submitMu); + ops.swap(submitQ); + } + + for (auto& op : ops) { + EnsurePeerChannels(op->peer); + if (!IsPeerReady(op->peer)) { + waitingOps[op->peer].push_back(std::move(op)); + continue; + } + DispatchOp(std::move(op)); + } +} + +bool TcpTransport::IsPeerReady(const EngineKey& peerKey) { + auto it = peers.find(peerKey); + if (it == peers.end()) return false; + if (!it->second.CtrlUp() || !it->second.DataUp()) return false; + Connection* ctrl = conns[it->second.ctrlFd].get(); + if (!ctrl || !ctrl->helloReceived) return false; + return !it->second.workers.empty(); +} + +void TcpTransport::RegisterRecvTargetWithWorkers(const EngineKey& peerKey, TransferUniqueId opId, + const WorkerRecvTarget& target) { + auto pit = peers.find(peerKey); + if (pit == peers.end()) return; + for (auto* w : pit->second.workers) w->RegisterRecvTarget(opId, target); +} + +void TcpTransport::RemoveRecvTargetFromWorkers(const EngineKey& peerKey, TransferUniqueId opId) { + auto pit = peers.find(peerKey); + if (pit == peers.end()) return; + for (auto* w : pit->second.workers) w->RemoveRecvTarget(opId); +} + +void TcpTransport::DispatchOp(std::unique_ptr op) { + if (!op) { + MORI_IO_ERROR("TCP: DispatchOp got null op"); + return; + } + const EngineKey peerKey = op->peer; + auto it = peers.find(peerKey); + if (it == peers.end() || !it->second.CtrlUp() || !it->second.DataUp()) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer not connected"); + return; + } + Connection* ctrl = conns[it->second.ctrlFd].get(); + if (!ctrl) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer ctrl connection missing"); + return; + } + auto& workerList = it->second.workers; + if (workerList.empty()) { + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: no data workers for peer"); + return; + } + + const TransferUniqueId opId = op->id; + auto [itIns, inserted] = pendingOutbound.emplace(opId, std::move(op)); + if (!inserted) { + MORI_IO_ERROR("TCP: duplicate outbound op id={} for peer={}", opId, peerKey.c_str()); + itIns->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: duplicate op id"); + pendingOutbound.erase(itIns); + return; + } + OutboundOpState* st = itIns->second.get(); + if (!st) { + MORI_IO_ERROR("TCP: failed to store outbound op id={}", opId); + pendingOutbound.erase(itIns); + return; + } + + const uint64_t totalBytes = SumLens(st->localSegs); + int wantLanes = std::max(1, config.numDataConns); + wantLanes = std::min(wantLanes, (1U << kLaneBits)); + uint8_t lanesTotal = 1; + const bool canStripe = (wantLanes > 1) && (config.stripingThresholdBytes > 0) && + (totalBytes >= static_cast(config.stripingThresholdBytes)) && + (st->localSegs.size() == 1) && (st->remoteSegs.size() == 1) && + (workerList.size() >= 2); + if (canStripe) { + lanesTotal = + static_cast(std::min(static_cast(wantLanes), workerList.size())); + } + st->lanesTotal = lanesTotal; + + if (st->isRead && st->local.loc == MemoryLocationType::GPU) { + st->pinned = staging.Acquire(static_cast(totalBytes)); + if (!st->pinned) { + st->status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging"); + pendingOutbound.erase(opId); + return; + } + } + + if (st->isRead) { + WorkerRecvTarget target; + target.lanesTotal = lanesTotal; + target.totalLen = totalBytes; + target.discard = false; + if (st->local.loc == MemoryLocationType::GPU) { + target.toGpu = true; + target.pinned = st->pinned; + } else { + target.toGpu = false; + target.cpuBase = reinterpret_cast(st->local.data); + target.segs = st->localSegs; + } + RegisterRecvTargetWithWorkers(peerKey, opId, target); + } + + std::vector ctrlFrame; + if (st->localSegs.size() == 1) { + if (st->isRead) + ctrlFrame = tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, + st->remoteSegs[0].len, lanesTotal); + else + ctrlFrame = tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, + st->remoteSegs[0].len, lanesTotal); + } else { + std::vector roffs, szs; + roffs.reserve(st->remoteSegs.size()); + szs.reserve(st->remoteSegs.size()); + for (const auto& s : st->remoteSegs) { + roffs.push_back(s.off); + szs.push_back(s.len); + } + if (st->isRead) + ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs, lanesTotal); + else + ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs, lanesTotal); + } + + QueueSend(ctrl->fd, std::move(ctrlFrame)); + if (!st->isRead) QueueDataSendForWrite(workerList, *st); + UpdateWriteInterest(ctrl->fd); +} + +void TcpTransport::QueueSend(int fd, std::vector bytes, std::function onDone) { + Connection* c = conns[fd].get(); + if (!c) return; + SendItem item; + item.header = std::move(bytes); + item.iov.resize(1); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.onDone = std::move(onDone); + c->sendq.push_back(std::move(item)); +} + +void TcpTransport::QueueSegmentSend(DataConnectionWorker* worker, uint64_t wireOpId, uint8_t* base, + const std::vector& segs, uint64_t totalLen, + std::function onDone) { + SendItem item; + item.header = tcp::BuildDataHeader(wireOpId, totalLen, 0); + item.iov.reserve(1 + segs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + for (const auto& s : segs) item.iov.push_back({base + s.off, static_cast(s.len)}); + item.onDone = std::move(onDone); + worker->SubmitSend(std::move(item)); +} + +void TcpTransport::QueueStripedCpuSend(const std::vector& workers, + uint64_t opId, uint8_t lanesTotal, uint8_t* base, + uint64_t baseOff, uint64_t total, + std::function onLaneDone) { + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + auto* worker = workers[lane % workers.size()]; + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = base + baseOff + span.off; + item.iov[1].iov_len = static_cast(span.len); + item.onDone = onLaneDone; + worker->SubmitSend(std::move(item)); + } +} + +bool TcpTransport::ScheduleGpuDtoH(int deviceId, const MemoryDesc& src, + const std::vector& srcSegs, + std::shared_ptr pinned, + std::function onComplete) { + const uint64_t total = SumLens(srcSegs); + hipStream_t stream = streamPool.GetNextStream(deviceId); + hipEvent_t ev = eventPool.GetEvent(deviceId); + if (stream == nullptr || ev == nullptr) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU DtoH"); + if (ev) eventPool.PutEvent(ev, deviceId); + return false; + } + + HIP_RUNTIME_CHECK(hipSetDevice(deviceId)); + uint8_t* dst = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0, spanLen = 0; + if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : srcSegs) { + hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + gpuTasks.push_back({deviceId, ev, std::move(onComplete)}); + return true; +} + +bool TcpTransport::ScheduleGpuHtoD(int deviceId, const MemoryDesc& dst, + const std::vector& dstSegs, + std::shared_ptr pinned, + std::function onComplete) { + const uint64_t total = SumLens(dstSegs); + hipStream_t stream = streamPool.GetNextStream(deviceId); + hipEvent_t ev = eventPool.GetEvent(deviceId); + if (stream == nullptr || ev == nullptr) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU HtoD"); + if (ev) eventPool.PutEvent(ev, deviceId); + return false; + } + + HIP_RUNTIME_CHECK(hipSetDevice(deviceId)); + uint8_t* src = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0, spanLen = 0; + if (IsSingleContiguousSpan(dstSegs, &spanOff, &spanLen) && spanLen == total) { + void* gpuPtr = reinterpret_cast(dst.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + } else { + uint64_t off = 0; + for (const auto& s : dstSegs) { + void* gpuPtr = reinterpret_cast(dst.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + off += s.len; + } + } + HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); + gpuTasks.push_back({deviceId, ev, std::move(onComplete)}); + return true; +} + +void TcpTransport::QueueGpuSend(const std::vector& workers, uint64_t opId, + uint8_t lanesTotal, const MemoryDesc& src, + const std::vector& srcSegs, + std::function onLaneDone) { + const uint64_t total = SumLens(srcSegs); + auto pinned = staging.Acquire(static_cast(total)); + if (!pinned) { + MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); + return; + } + + auto sendCallback = [workers, pinned, opId, lanesTotal, total, + onLaneDone = std::move(onLaneDone)]() { + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + auto* worker = workers[lane % workers.size()]; + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); + item.iov.resize(2); + item.iov[0].iov_base = item.header.data(); + item.iov[0].iov_len = item.header.size(); + item.iov[1].iov_base = static_cast(pinned->ptr) + static_cast(span.off); + item.iov[1].iov_len = static_cast(span.len); + item.keepalive = pinned; + item.onDone = onLaneDone; + worker->SubmitSend(std::move(item)); + } + }; + + ScheduleGpuDtoH(src.deviceId, src, srcSegs, pinned, std::move(sendCallback)); +} + +void TcpTransport::QueueDataSendForWrite(const std::vector& workerList, + OutboundOpState& st) { + if (workerList.empty()) return; + uint8_t lanesTotal = std::max(1, st.lanesTotal); + + if (lanesTotal > 1 && st.localSegs.size() != 1) { + MORI_IO_WARN("TCP: striping requested but localSegs.size={}, fallback to 1 lane", + st.localSegs.size()); + lanesTotal = 1; + st.lanesTotal = 1; + } + + QueueDataSendCommon(workerList, st.local, st.localSegs, st.id, lanesTotal); +} + +void TcpTransport::QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, + const std::vector& srcSegs, uint8_t lanesTotal) { + auto pit = peers.find(peer); + if (pit == peers.end()) return; + auto& workerList = pit->second.workers; + if (workerList.empty()) return; + + struct DoneState { + EngineKey peer; + uint64_t opId{0}; + std::atomic remaining{0}; + }; + auto done = std::make_shared(); + done->peer = peer; + done->opId = opId; + uint8_t useLanes = std::min( + ClampLanesTotal(lanesTotal), static_cast(std::max(1, workerList.size()))); + if (useLanes > 1 && srcSegs.size() != 1) useLanes = 1; + done->remaining.store(useLanes); + auto laneDone = [this, done]() mutable { + if (done->remaining.fetch_sub(1) > 1) return; + SendCompletionAndRecord(done->peer, done->opId, StatusCode::SUCCESS, ""); + }; + + QueueDataSendCommon(workerList, src, srcSegs, opId, useLanes, std::move(laneDone)); +} + +void TcpTransport::QueueDataSendCommon(const std::vector& workerList, + const MemoryDesc& src, const std::vector& srcSegs, + uint64_t opId, uint8_t lanesTotal, + std::function onLaneDone) { + if (workerList.empty()) return; + const uint64_t total = SumLens(srcSegs); + lanesTotal = ClampLanesTotal(lanesTotal); + lanesTotal = std::min(lanesTotal, static_cast(workerList.size())); + if (lanesTotal > 1 && srcSegs.size() != 1) { + MORI_IO_WARN("TCP: striping requested for {} segments, fallback to 1 lane", srcSegs.size()); + lanesTotal = 1; + } + + if (src.loc == MemoryLocationType::GPU) { + std::vector workers(workerList.begin(), workerList.begin() + lanesTotal); + QueueGpuSend(workers, opId, lanesTotal, src, srcSegs, std::move(onLaneDone)); + return; + } + + uint8_t* base = reinterpret_cast(src.data); + if (lanesTotal == 1) { + QueueSegmentSend(workerList[0], ToWireOpId(opId, 0), base, srcSegs, total, + std::move(onLaneDone)); + } else { + QueueStripedCpuSend(workerList, opId, lanesTotal, base, srcSegs[0].off, total, + std::move(onLaneDone)); + } +} + +void TcpTransport::PollGpuTasks() { + for (auto it = gpuTasks.begin(); it != gpuTasks.end();) { + hipError_t st = hipEventQuery(it->ev); + if (st == hipSuccess) { + eventPool.PutEvent(it->ev, it->deviceId); + if (it->onReady) it->onReady(); + it = gpuTasks.erase(it); + } else if (st == hipErrorNotReady) { + ++it; + } else { + MORI_IO_ERROR("TCP: hipEventQuery failed: {}", hipGetErrorString(st)); + eventPool.PutEvent(it->ev, it->deviceId); + it = gpuTasks.erase(it); + } + } +} + +void TcpTransport::UpdateWriteInterest(int fd) { + auto it = conns.find(fd); + if (it == conns.end()) return; + Connection* c = it->second.get(); + if (!c || c->fd < 0) return; + + if (!c->connecting && !c->sendq.empty()) { + FlushSend(c); + it = conns.find(fd); + if (it == conns.end()) return; + c = it->second.get(); + if (!c || c->fd < 0) return; + } + ModEpoll(fd, true, c->connecting || !c->sendq.empty()); +} + +void TcpTransport::HandleConnWritable(Connection* c) { + if (c->connecting) { + int err = 0; + socklen_t len = sizeof(err); + if (getsockopt(c->fd, SOL_SOCKET, SO_ERROR, &err, &len) != 0 || err != 0) { + MORI_IO_ERROR("TCP: connect failed fd {}: {}", c->fd, strerror(err == 0 ? errno : err)); + ClosePeerByFd(c->fd); + return; + } + c->connecting = false; + QueueHello(c->fd); + } + UpdateWriteInterest(c->fd); +} + +void TcpTransport::FlushSend(Connection* c) { + constexpr size_t kMaxIov = 64; + while (!c->sendq.empty()) { + SendItem& item = c->sendq.front(); + if (item.Done()) { + auto cb = std::move(item.onDone); + c->sendq.pop_front(); + if (cb) cb(); + continue; + } + + iovec iov[kMaxIov]; + size_t cnt = 0; + for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { + iov[cnt] = item.iov[i]; + if (i == item.idx && item.off > 0) { + iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; + iov[cnt].iov_len -= item.off; + } + cnt++; + } + msghdr msg{}; + msg.msg_iov = iov; + msg.msg_iovlen = cnt; + ssize_t n = sendmsg(c->fd, &msg, MSG_NOSIGNAL | item.flags); + if (n < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_ERROR("TCP: sendmsg fd {} failed: {}", c->fd, strerror(errno)); + ClosePeerByFd(c->fd); + break; + } + if (n == 0) break; + item.Advance(static_cast(n)); + } +} + +void TcpTransport::CloseAndRemoveFd(int fd) { + if (fd < 0) return; + auto wit = dataWorkers.find(fd); + if (wit != dataWorkers.end()) { + wit->second->Stop(); + auto nit = workerNotifyMap.find(wit->second->NotifyFd()); + if (nit != workerNotifyMap.end()) { + DelEpoll(nit->first); + workerNotifyMap.erase(nit); + } + dataWorkers.erase(wit); + } + auto it = conns.find(fd); + if (it != conns.end()) { + CloseConnInternal(it->second.get()); + conns.erase(it); + } +} + +EngineKey TcpTransport::FindPeerByFd(int fd) { + for (auto& kv : peers) { + if (kv.second.ctrlFd == fd) return kv.first; + for (int dfd : kv.second.dataFds) { + if (dfd == fd) return kv.first; + } + } + return {}; +} + +void TcpTransport::ClosePeerByFd(int fd) { + MORI_IO_TRACE("TCP: close fd={}", fd); + EngineKey peer = FindPeerByFd(fd); + if (!peer.empty()) { + ClosePeerByKey(peer, "TCP: connection lost"); + } else { + CloseAndRemoveFd(fd); + } +} + +void TcpTransport::ClosePeerByKey(const EngineKey& peer, const std::string& reason) { + auto pit = peers.find(peer); + if (pit == peers.end()) return; + auto link = pit->second; + CloseAndRemoveFd(link.ctrlFd); + for (int dfd : link.dataFds) CloseAndRemoveFd(dfd); + peers.erase(peer); + FailPendingOpsForPeer(peer, reason); +} + +void TcpTransport::FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg) { + for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + if (it->second->peer == peer) { + it->second->status->Update(StatusCode::ERR_BAD_STATE, msg); + it = pendingOutbound.erase(it); + } else { + ++it; + } + } + waitingOps.erase(peer); + inboundWrites.erase(peer); + earlyWrites.erase(peer); +} + +void TcpTransport::HandleCtrlReadable(Connection* c) { + while (true) { + uint8_t tmp[16384]; + ssize_t n = ::recv(c->fd, tmp, sizeof(tmp), 0); + if (n < 0) { + if (IsWouldBlock(errno)) break; + MORI_IO_ERROR("TCP: recv(ctrl) fd {} failed: {}", c->fd, strerror(errno)); + ClosePeerByFd(c->fd); + return; + } + if (n == 0) { + ClosePeerByFd(c->fd); + return; + } + c->inbuf.insert(c->inbuf.end(), tmp, tmp + n); + } + + while (true) { + tcp::CtrlHeaderView hv; + if (!tcp::TryParseCtrlHeader(c->inbuf.data(), c->inbuf.size(), &hv)) { + if (c->inbuf.size() >= tcp::kCtrlHeaderSize) { + MORI_IO_ERROR("TCP: bad ctrl header on fd {}, closing", c->fd); + ClosePeerByFd(c->fd); + } + break; + } + if (c->inbuf.size() < tcp::kCtrlHeaderSize + hv.bodyLen) break; + + const uint8_t* body = c->inbuf.data() + tcp::kCtrlHeaderSize; + HandleCtrlFrame(c, hv.type, body, hv.bodyLen); + c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + tcp::kCtrlHeaderSize + hv.bodyLen); + + if (c->helloReceived && c->ch == tcp::Channel::DATA) return; + } +} + +void TcpTransport::HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const uint8_t* body, + size_t len) { + if (type == tcp::CtrlMsgType::HELLO) { + HandleHello(c, body, len); + return; + } + if (!c->helloReceived) { + MORI_IO_WARN("TCP: received ctrl message before HELLO, dropping"); + return; + } + switch (type) { + case tcp::CtrlMsgType::WRITE_REQ: + HandleWriteReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::READ_REQ: + HandleReadReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::BATCH_WRITE_REQ: + HandleBatchWriteReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::BATCH_READ_REQ: + HandleBatchReadReq(c->peerKey, body, len); + break; + case tcp::CtrlMsgType::COMPLETION: + HandleCompletion(c->peerKey, body, len); + break; + default: + MORI_IO_WARN("TCP: unknown ctrl msg type {}", static_cast(type)); + } +} + +void TcpTransport::HandleHello(Connection* c, const uint8_t* body, size_t len) { + if (len < 1 + 4) { + MORI_IO_WARN("TCP: bad HELLO len {}", len); + ClosePeerByFd(c->fd); + return; + } + size_t off = 0; + const uint8_t chRaw = body[off++]; + uint32_t keyLen = 0; + if (!tcp::ReadU32BE(body, len, &off, &keyLen)) { + ClosePeerByFd(c->fd); + return; + } + if (off + keyLen > len) { + ClosePeerByFd(c->fd); + return; + } + EngineKey peerKey(reinterpret_cast(body + off), keyLen); + off += keyLen; + + c->peerKey = peerKey; + c->ch = + (chRaw == static_cast(tcp::Channel::DATA)) ? tcp::Channel::DATA : tcp::Channel::CTRL; + c->helloReceived = true; + MORI_IO_TRACE("TCP: recv HELLO fd={} peer={} ch={} outgoing={}", c->fd, c->peerKey.c_str(), + static_cast(c->ch), c->isOutgoing); + + if (!c->helloSent) { + QueueHello(c->fd); + UpdateWriteInterest(c->fd); + } + if (c->ch == tcp::Channel::CTRL) ConfigureCtrlSocket(c->fd, config); + AssignConnToPeer(c); + MaybeDispatchQueuedOps(peerKey); +} + +std::optional TcpTransport::LookupLocalMem(MemoryUniqueId id) { + std::lock_guard lock(memMu); + auto it = localMems.find(id); + if (it == localMems.end()) return std::nullopt; + return it->second; +} + +void TcpTransport::RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, + const std::string& msg) { + std::lock_guard lock(inboundMu); + inboundStatus[peer][id] = InboundStatusEntry{code, msg}; +} + +void TcpTransport::SendCompletionAndRecord(const EngineKey& peer, TransferUniqueId opId, + StatusCode code, const std::string& msg) { + Connection* ctrl = PeerCtrl(peer); + if (ctrl != nullptr) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(code), msg)); + UpdateWriteInterest(ctrl->fd); + } + RecordInboundStatus(peer, opId, code, msg); +} + +Connection* TcpTransport::PeerCtrl(const EngineKey& peer) { + auto it = peers.find(peer); + if (it == peers.end() || !it->second.CtrlUp()) return nullptr; + return conns[it->second.ctrlFd].get(); +} + +void TcpTransport::FinalizeInboundWriteSetup(const EngineKey& peer, TransferUniqueId opId, + InboundWriteState& ws) { + if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { + const uint64_t total = SumLens(ws.dstSegs); + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; + } + inboundWrites[peer][opId] = ws; + SetupInboundWriteWorkerTarget(peer, opId, ws); + TryConsumeEarlyWriteLanes(peer, opId); +} + +void TcpTransport::SetupInboundWriteWorkerTarget(const EngineKey& peer, TransferUniqueId opId, + const InboundWriteState& ws) { + WorkerRecvTarget target; + target.lanesTotal = ws.lanesTotal; + target.totalLen = SumLens(ws.dstSegs); + target.discard = ws.discard; + if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { + target.toGpu = true; + target.pinned = ws.pinned; + } else if (!ws.discard) { + target.toGpu = false; + target.cpuBase = reinterpret_cast(ws.dst.data); + target.segs = ws.dstSegs; + } + RegisterRecvTargetWithWorkers(peer, opId, target); +} + +void TcpTransport::HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { + LinearReqView req; + if (!ParseLinearReq(body, len, &req)) { + MORI_IO_WARN("TCP: malformed WRITE_REQ"); + return; + } + HandleWriteReqSegments(peer, req.opId, req.memId, {{req.remoteOff, req.size}}, req.lanesTotal); +} + +void TcpTransport::HandleWriteReqSegments(const EngineKey& peer, TransferUniqueId opId, + MemoryUniqueId memId, std::vector segs, + uint8_t lanesTotal) { + auto memOpt = LookupLocalMem(memId); + InboundWriteState ws; + ws.peer = peer; + ws.id = opId; + ws.lanesTotal = lanesTotal; + ws.discard = true; + if (memOpt.has_value()) { + ws.dst = memOpt.value(); + if (SegmentsInRange(segs, ws.dst.size)) { + ws.discard = false; + ws.dstSegs = std::move(segs); + } + } + FinalizeInboundWriteSetup(peer, opId, ws); +} + +void TcpTransport::HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { + BatchReqView req; + if (!ParseBatchReq(body, len, &req)) { + MORI_IO_WARN("TCP: malformed BATCH_WRITE_REQ"); + return; + } + HandleWriteReqSegments(peer, req.opId, req.memId, std::move(req.segs), req.lanesTotal); +} + +void TcpTransport::HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { + LinearReqView req; + if (!ParseLinearReq(body, len, &req)) { + MORI_IO_WARN("TCP: malformed READ_REQ"); + return; + } + HandleReadReqSegments(peer, req.opId, req.memId, {{req.remoteOff, req.size}}, req.lanesTotal, + false); +} + +void TcpTransport::HandleReadReqSegments(const EngineKey& peer, TransferUniqueId opId, + MemoryUniqueId memId, std::vector segs, + uint8_t lanesTotal, bool batchReq) { + const StatusCode badRangeCode = + batchReq ? StatusCode::ERR_INVALID_ARGS : StatusCode::ERR_NOT_FOUND; + const char* badRangeMsg = + batchReq ? "TCP: batch read out of range" : "TCP: remote mem not found/out of range"; + const char* notFoundMsg = + batchReq ? "TCP: remote mem not found" : "TCP: remote mem not found/out of range"; + auto memOpt = LookupLocalMem(memId); + if (!memOpt.has_value()) { + SendCompletionAndRecord(peer, opId, StatusCode::ERR_NOT_FOUND, notFoundMsg); + return; + } + MemoryDesc src = memOpt.value(); + if (!SegmentsInRange(segs, src.size)) { + SendCompletionAndRecord(peer, opId, badRangeCode, badRangeMsg); + return; + } + QueueDataSendForRead(peer, opId, src, segs, lanesTotal); +} + +void TcpTransport::HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { + BatchReqView req; + if (!ParseBatchReq(body, len, &req)) { + MORI_IO_WARN("TCP: malformed BATCH_READ_REQ"); + return; + } + HandleReadReqSegments(peer, req.opId, req.memId, std::move(req.segs), req.lanesTotal, true); +} + +void TcpTransport::HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { + CompletionView msg; + if (!ParseCompletion(body, len, &msg)) { + MORI_IO_WARN("TCP: malformed COMPLETION"); + return; + } + + auto it = pendingOutbound.find(msg.opId); + if (it == pendingOutbound.end()) return; + OutboundOpState& st = *it->second; + st.completionReceived = true; + st.completionCode = static_cast(msg.statusCode); + st.completionMsg = std::move(msg.msg); + + if (st.completionCode != StatusCode::SUCCESS) { + RemoveRecvTargetFromWorkers(st.peer, msg.opId); + st.status->Update(st.completionCode, st.completionMsg); + pendingOutbound.erase(it); + return; + } + MaybeCompleteOutbound(st); +} + +void TcpTransport::MaybeCompleteOutbound(OutboundOpState& st) { + if (!st.completionReceived) return; + if (st.isRead) { + const uint8_t allMask = LanesAllMask(st.lanesTotal); + if (st.lanesDoneMask != allMask) return; + if (st.rxBytes != st.expectedRxBytes) return; + if (st.gpuCopyPending) return; + } + RemoveRecvTargetFromWorkers(st.peer, st.id); + st.status->Update(StatusCode::SUCCESS, ""); + pendingOutbound.erase(st.id); +} + +void TcpTransport::MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId) { + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt == inboundWrites.end()) return; + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt == iwPeerIt->second.end()) return; + + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + const uint8_t allMask = LanesAllMask(ws.lanesTotal); + if ((ws.lanesDoneMask & allMask) != allMask) return; + + RemoveRecvTargetFromWorkers(peer, opId); + + if (ws.discard) { + SendCompletionAndRecord(peer, opId, StatusCode::ERR_INVALID_ARGS, "TCP: write discarded"); + } else if (ws.dst.loc == MemoryLocationType::GPU) { + if (!ws.pinned) { + SendCompletionAndRecord(peer, opId, StatusCode::ERR_BAD_STATE, + "TCP: missing pinned staging (write)"); + } else { + auto pinnedRef = ws.pinned; + bool ok = ScheduleGpuHtoD(ws.dst.deviceId, ws.dst, ws.dstSegs, pinnedRef, + [this, peer, opId, pinnedRef]() { + SendCompletionAndRecord(peer, opId, StatusCode::SUCCESS, ""); + }); + if (!ok) { + SendCompletionAndRecord(peer, opId, StatusCode::ERR_BAD_STATE, + "TCP: failed HIP stream/event"); + } + } + } else { + SendCompletionAndRecord(peer, opId, StatusCode::SUCCESS, ""); + } + + iwPeerIt->second.erase(wsIt); + if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt != earlyWrites.end()) { + ewPeerIt->second.erase(opId); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + } +} + +void TcpTransport::TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId) { + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt == inboundWrites.end()) return; + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt == iwPeerIt->second.end()) return; + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + + auto ewPeerIt = earlyWrites.find(peer); + if (ewPeerIt == earlyWrites.end()) return; + auto ewIt = ewPeerIt->second.find(opId); + if (ewIt == ewPeerIt->second.end()) return; + + EarlyWriteState& early = ewIt->second; + const uint64_t total = SumLens(ws.dstSegs); + uint8_t* dstBase = reinterpret_cast(ws.dst.data); + + for (auto it = early.lanes.begin(); it != early.lanes.end();) { + const uint8_t lane = it->first; + EarlyWriteLaneState& laneState = it->second; + if (!laneState.complete) { + ++it; + continue; + } + + if (lane >= ws.lanesTotal) ws.discard = true; + const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); + if (span.len != laneState.payloadLen) ws.discard = true; + + if (!ws.discard && laneState.pinned) { + uint8_t* src = reinterpret_cast(laneState.pinned->ptr); + if (ws.dst.loc == MemoryLocationType::GPU) { + if (!ws.pinned) { + ws.pinned = staging.Acquire(static_cast(total)); + if (!ws.pinned) ws.discard = true; + } + if (!ws.discard && ws.pinned) { + std::memcpy(reinterpret_cast(ws.pinned->ptr) + span.off, src, + static_cast(span.len)); + } + } else { + const auto segs = SliceSegments(ws.dstSegs, span.off, span.len); + uint64_t copied = 0; + for (const auto& s : segs) { + std::memcpy(dstBase + s.off, src + copied, static_cast(s.len)); + copied += s.len; + } + } + } + + if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + it = early.lanes.erase(it); + } + + if (early.lanes.empty()) { + ewPeerIt->second.erase(ewIt); + if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + } + MaybeFinalizeInboundWrite(peer, opId); +} + +void TcpTransport::ProcessEventsFrom(DataConnectionWorker* worker) { + std::deque events; + worker->DrainEvents(events); + for (auto& ev : events) { + switch (ev.type) { + case WorkerEventType::RECV_DONE: + HandleWorkerRecvDone(ev); + break; + case WorkerEventType::EARLY_DATA: + HandleWorkerEarlyData(ev); + break; + case WorkerEventType::SEND_CALLBACK: + if (ev.callback) ev.callback(); + break; + case WorkerEventType::CONN_ERROR: + MORI_IO_WARN("TCP: worker error for peer {}: {}", ev.peerKey, ev.errorMsg); + ClosePeerByKey(ev.peerKey, ev.errorMsg); + break; + } + } +} + +void TcpTransport::ProcessWorkerEvents() { + std::vector notifyFds; + notifyFds.reserve(workerNotifyMap.size()); + for (const auto& kv : workerNotifyMap) notifyFds.push_back(kv.first); + for (int notifyFd : notifyFds) { + auto it = workerNotifyMap.find(notifyFd); + if (it == workerNotifyMap.end()) continue; + ProcessEventsFrom(it->second); + } +} + +void TcpTransport::HandleWorkerRecvDone(const WorkerEvent& ev) { + const EngineKey& peer = ev.peerKey; + const TransferUniqueId opId = ev.opId; + const uint8_t lane = ev.lane; + const uint64_t laneLen = ev.laneLen; + + auto iwPeerIt = inboundWrites.find(peer); + if (iwPeerIt != inboundWrites.end()) { + auto wsIt = iwPeerIt->second.find(opId); + if (wsIt != iwPeerIt->second.end()) { + InboundWriteState& ws = wsIt->second; + ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); + if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + MaybeFinalizeInboundWrite(peer, opId); + return; + } + } + + auto obIt = pendingOutbound.find(opId); + if (obIt == pendingOutbound.end()) return; + OutboundOpState& st = *obIt->second; + st.lanesTotal = ClampLanesTotal(st.lanesTotal); + const uint8_t bit = static_cast(1U << lane); + if ((st.lanesDoneMask & bit) == 0) { + st.lanesDoneMask |= bit; + st.rxBytes += laneLen; + } + + if (st.local.loc == MemoryLocationType::GPU) { + const uint8_t allMask = LanesAllMask(st.lanesTotal); + if ((st.lanesDoneMask & allMask) != allMask) { + MaybeCompleteOutbound(st); + return; + } + if (st.gpuCopyPending) return; + if (!st.pinned) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging (read)"); + RemoveRecvTargetFromWorkers(st.peer, opId); + pendingOutbound.erase(obIt); + return; + } + + st.gpuCopyPending = true; + auto pinnedRef = st.pinned; + bool ok = ScheduleGpuHtoD(st.local.deviceId, st.local, st.localSegs, pinnedRef, + [this, opId, pinnedRef]() { + auto it2 = pendingOutbound.find(opId); + if (it2 == pendingOutbound.end()) return; + it2->second->gpuCopyPending = false; + MaybeCompleteOutbound(*it2->second); + }); + if (!ok) { + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed HIP stream/event (read)"); + RemoveRecvTargetFromWorkers(st.peer, opId); + pendingOutbound.erase(obIt); + } + return; + } + + MaybeCompleteOutbound(st); +} + +void TcpTransport::HandleWorkerEarlyData(const WorkerEvent& ev) { + const EngineKey& peer = ev.peerKey; + const TransferUniqueId opId = ev.opId; + const uint8_t lane = ev.lane; + + auto& perPeer = earlyWrites[peer]; + EarlyWriteState& early = perPeer[opId]; + if (early.lanes.find(lane) != early.lanes.end()) { + MORI_IO_WARN("TCP: duplicate early data for op {} lane {} from peer {}", opId, (uint32_t)lane, + peer); + return; + } + early.lanes.emplace(lane, EarlyWriteLaneState{ev.laneLen, ev.earlyBuf, true}); + TryConsumeEarlyWriteLanes(peer, opId); +} + +void TcpTransport::ScanTimeouts() { + if (config.opTimeoutMs <= 0) return; + const auto now = Clock::now(); + const auto timeout = std::chrono::milliseconds(config.opTimeoutMs); + for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + if ((now - it->second->startTs) > timeout) { + RemoveRecvTargetFromWorkers(it->second->peer, it->first); + it->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: op timeout"); + it = pendingOutbound.erase(it); + } else { + ++it; + } + } +} + +void TcpTransport::IoLoop() { + constexpr int kMaxEvents = 128; + epoll_event events[kMaxEvents]; + + while (running.load()) { + PollGpuTasks(); + ProcessWorkerEvents(); + ScanTimeouts(); + + const bool hasActive = + !gpuTasks.empty() || !pendingOutbound.empty() || !workerNotifyMap.empty(); + const int timeoutMs = hasActive ? 0 : 2; + int nfds = epoll_wait(epfd, events, kMaxEvents, timeoutMs); + if (nfds < 0) { + if (errno == EINTR) continue; + MORI_IO_ERROR("TCP: epoll_wait failed: {}", strerror(errno)); + break; + } + + for (int i = 0; i < nfds; ++i) { + int fd = events[i].data.fd; + uint32_t evMask = events[i].events; + + if (fd == listenFd) { + AcceptNew(); + continue; + } + if (fd == wakeFd) { + DrainWakeFd(); + continue; + } + + auto wnit = workerNotifyMap.find(fd); + if (wnit != workerNotifyMap.end()) { + ProcessEventsFrom(wnit->second); + continue; + } + + Connection* c = nullptr; + auto it = conns.find(fd); + if (it != conns.end()) c = it->second.get(); + if (!c) continue; + + if (evMask & (EPOLLERR | EPOLLHUP)) { + ClosePeerByFd(fd); + continue; + } + + if (evMask & EPOLLIN) { + HandleCtrlReadable(c); + auto it2 = conns.find(fd); + if (it2 == conns.end()) continue; + c = it2->second.get(); + if (!c) continue; + } + if (evMask & EPOLLOUT) HandleConnWritable(c); + } + } +} + +} // namespace io +} // namespace mori diff --git a/src/io/tcp/transport.hpp b/src/io/tcp/transport.hpp new file mode 100644 index 000000000..74ea99a44 --- /dev/null +++ b/src/io/tcp/transport.hpp @@ -0,0 +1,210 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include + +#include +#include + +#include "mori/application/utils/check.hpp" +#include "mori/io/engine.hpp" +#include "src/io/tcp/data_worker.hpp" +#include "src/io/tcp/tcp_types.hpp" +#include "src/io/xgmi/hip_resource_pool.hpp" + +namespace mori { +namespace io { + +class TcpTransport { + public: + TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg); + ~TcpTransport(); + + TcpTransport(const TcpTransport&) = delete; + TcpTransport& operator=(const TcpTransport&) = delete; + + void Start(); + void Shutdown(); + + std::optional GetListenPort() const; + + void RegisterRemoteEngine(const EngineDesc& desc); + void DeregisterRemoteEngine(const EngineDesc& desc); + void RegisterMemory(const MemoryDesc& desc); + void DeregisterMemory(const MemoryDesc& desc); + + bool PopInboundTransferStatus(const EngineKey& remote, TransferUniqueId id, + TransferStatus* status); + + void SubmitReadWrite(const MemoryDesc& local, size_t localOffset, const MemoryDesc& remote, + size_t remoteOffset, size_t size, TransferStatus* status, + TransferUniqueId id, bool isRead); + + void SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& localOffsets, + const MemoryDesc& remote, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead); + + private: + void EnqueueOp(std::unique_ptr op); + + void AddEpoll(int fd, bool wantRead, bool wantWrite); + void ModEpoll(int fd, bool wantRead, bool wantWrite); + void DelEpoll(int fd); + void CloseConnInternal(Connection* c); + + bool PreferOutgoingFor(const EngineKey& peerKey) const; + void AssignConnToPeer(Connection* c); + void MaybeDispatchQueuedOps(const EngineKey& peerKey); + void EnsurePeerChannels(const EngineKey& peerKey); + void ConnectChannel(const EngineKey& peerKey, tcp::Channel ch); + void QueueHello(int fd); + void AcceptNew(); + void DrainWakeFd(); + bool IsPeerReady(const EngineKey& peerKey); + + void RegisterRecvTargetWithWorkers(const EngineKey& peerKey, TransferUniqueId opId, + const WorkerRecvTarget& target); + void RemoveRecvTargetFromWorkers(const EngineKey& peerKey, TransferUniqueId opId); + + void DispatchOp(std::unique_ptr op); + void QueueSend(int fd, std::vector bytes, std::function onDone = nullptr); + + void QueueSegmentSend(DataConnectionWorker* worker, uint64_t wireOpId, uint8_t* base, + const std::vector& segs, uint64_t totalLen, + std::function onDone = nullptr); + void QueueStripedCpuSend(const std::vector& workers, uint64_t opId, + uint8_t lanesTotal, uint8_t* base, uint64_t baseOff, uint64_t total, + std::function onLaneDone = nullptr); + void QueueGpuSend(const std::vector& workers, uint64_t opId, + uint8_t lanesTotal, const MemoryDesc& src, const std::vector& srcSegs, + std::function onLaneDone = nullptr); + + void QueueDataSendForWrite(const std::vector& workerList, + OutboundOpState& st); + void QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, + const std::vector& srcSegs, uint8_t lanesTotal); + void QueueDataSendCommon(const std::vector& workerList, + const MemoryDesc& src, const std::vector& srcSegs, + uint64_t opId, uint8_t lanesTotal, + std::function onLaneDone = nullptr); + + bool ScheduleGpuDtoH(int deviceId, const MemoryDesc& src, const std::vector& srcSegs, + std::shared_ptr pinned, std::function onComplete); + bool ScheduleGpuHtoD(int deviceId, const MemoryDesc& dst, const std::vector& dstSegs, + std::shared_ptr pinned, std::function onComplete); + + void PollGpuTasks(); + void UpdateWriteInterest(int fd); + void HandleConnWritable(Connection* c); + void FlushSend(Connection* c); + + void CloseAndRemoveFd(int fd); + EngineKey FindPeerByFd(int fd); + void ClosePeerByFd(int fd); + void ClosePeerByKey(const EngineKey& peer, const std::string& reason); + void FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg); + + void HandleCtrlReadable(Connection* c); + void HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const uint8_t* body, size_t len); + void HandleHello(Connection* c, const uint8_t* body, size_t len); + + std::optional LookupLocalMem(MemoryUniqueId id); + void RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, + const std::string& msg); + void SendCompletionAndRecord(const EngineKey& peer, TransferUniqueId opId, StatusCode code, + const std::string& msg); + Connection* PeerCtrl(const EngineKey& peer); + + void FinalizeInboundWriteSetup(const EngineKey& peer, TransferUniqueId opId, + InboundWriteState& ws); + void SetupInboundWriteWorkerTarget(const EngineKey& peer, TransferUniqueId opId, + const InboundWriteState& ws); + void HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len); + void HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len); + void HandleWriteReqSegments(const EngineKey& peer, TransferUniqueId opId, MemoryUniqueId memId, + std::vector segs, uint8_t lanesTotal); + void HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len); + void HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len); + void HandleReadReqSegments(const EngineKey& peer, TransferUniqueId opId, MemoryUniqueId memId, + std::vector segs, uint8_t lanesTotal, bool batchReq); + void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len); + + void MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId); + void TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId); + + void MaybeCompleteOutbound(OutboundOpState& st); + void ProcessEventsFrom(DataConnectionWorker* worker); + void ProcessWorkerEvents(); + void HandleWorkerRecvDone(const WorkerEvent& ev); + void HandleWorkerEarlyData(const WorkerEvent& ev); + void ScanTimeouts(); + + void IoLoop(); + + private: + EngineKey myEngKey; + IOEngineConfig engConfig; + TcpBackendConfig config; + + int epfd{-1}; + int listenFd{-1}; + int wakeFd{-1}; + uint16_t listenPort{0}; + + std::atomic running{false}; + std::thread ioThread; + + std::mutex submitMu; + std::deque> submitQ; + + std::mutex remoteMu; + std::unordered_map remoteEngines; + + std::mutex memMu; + std::unordered_map localMems; + + std::mutex inboundMu; + std::unordered_map> + inboundStatus; + + std::unordered_map> conns; + std::unordered_map peers; + std::unordered_map>> waitingOps; + std::unordered_map> pendingOutbound; + std::unordered_map> + inboundWrites; + std::unordered_map> earlyWrites; + + std::unordered_map> dataWorkers; + std::unordered_map workerNotifyMap; + + PinnedStagingPool staging; + StreamPool streamPool{8}; + EventPool eventPool{64}; + std::deque gpuTasks; +}; + +} // namespace io +} // namespace mori From 568e5097efe8c01ce9099309ffe1aa37e1c7c5a3 Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Sat, 28 Feb 2026 00:47:33 -0600 Subject: [PATCH 11/12] rm unused file --- include/mori/io/backend.hpp | 22 +- src/io/tcp/backend_impl.cpp | 28 +- src/io/tcp/backend_impl.hpp | 29 +- src/io/tcp/data_worker.hpp | 426 ------ src/io/tcp/protocol.hpp | 273 ---- src/io/tcp/tcp_types.hpp | 421 ------ src/io/tcp/transport.cpp | 1926 +++++++++++++++------------- src/io/tcp/transport.hpp | 692 ++++++++-- src/pybind/mori.cpp | 13 +- tests/python/io/benchmark.py | 8 - tests/python/io/test_engine.py | 97 ++ tests/python/io/test_engine_tcp.py | 122 -- 12 files changed, 1811 insertions(+), 2246 deletions(-) delete mode 100644 src/io/tcp/data_worker.hpp delete mode 100644 src/io/tcp/protocol.hpp delete mode 100644 src/io/tcp/tcp_types.hpp delete mode 100644 tests/python/io/test_engine_tcp.py diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index 99cf72711..a347ae4c0 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -83,12 +83,11 @@ inline std::ostream& operator<<(std::ostream& os, const XgmiBackendConfig& c) { struct TcpBackendConfig : public BackendConfig { TcpBackendConfig() : BackendConfig(BackendType::TCP) {} - TcpBackendConfig(int numIoThreads_, int sockSndbufBytes_, int sockRcvbufBytes_, int opTimeoutMs_, + TcpBackendConfig(int sockSndbufBytes_, int sockRcvbufBytes_, int opTimeoutMs_, bool enableKeepalive_, int keepaliveIdleSec_, int keepaliveIntvlSec_, - int keepaliveCnt_, bool enableCtrlNodelay_, bool enableDataNodelay_, - int numDataConns_, int stripingThresholdBytes_) + int keepaliveCnt_, bool enableCtrlNodelay_, int numDataConns_, + int stripingThresholdBytes_) : BackendConfig(BackendType::TCP), - numIoThreads(numIoThreads_), sockSndbufBytes(sockSndbufBytes_), sockRcvbufBytes(sockRcvbufBytes_), opTimeoutMs(opTimeoutMs_), @@ -97,12 +96,9 @@ struct TcpBackendConfig : public BackendConfig { keepaliveIntvlSec(keepaliveIntvlSec_), keepaliveCnt(keepaliveCnt_), enableCtrlNodelay(enableCtrlNodelay_), - enableDataNodelay(enableDataNodelay_), numDataConns(numDataConns_), stripingThresholdBytes(stripingThresholdBytes_) {} - int numIoThreads{1}; - int sockSndbufBytes{32 * 1024 * 1024}; int sockRcvbufBytes{32 * 1024 * 1024}; @@ -114,7 +110,6 @@ struct TcpBackendConfig : public BackendConfig { int keepaliveCnt{3}; bool enableCtrlNodelay{true}; - bool enableDataNodelay{true}; // Number of parallel DATA TCP connections per peer (iperf-like multi-stream striping). // Effective only when peer has >= numDataConns established and transfer is contiguous. @@ -124,12 +119,11 @@ struct TcpBackendConfig : public BackendConfig { }; inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) { - return os << "numIoThreads[" << c.numIoThreads << "] sockSndbufBytes[" << c.sockSndbufBytes - << "] sockRcvbufBytes[" << c.sockRcvbufBytes << "] opTimeoutMs[" << c.opTimeoutMs - << "] enableKeepalive[" << c.enableKeepalive << "] keepaliveIdleSec[" - << c.keepaliveIdleSec << "] keepaliveIntvlSec[" << c.keepaliveIntvlSec - << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" << c.enableCtrlNodelay - << "] enableDataNodelay[" << c.enableDataNodelay << "] numDataConns[" << c.numDataConns + return os << "sockSndbufBytes[" << c.sockSndbufBytes << "] sockRcvbufBytes[" << c.sockRcvbufBytes + << "] opTimeoutMs[" << c.opTimeoutMs << "] enableKeepalive[" << c.enableKeepalive + << "] keepaliveIdleSec[" << c.keepaliveIdleSec << "] keepaliveIntvlSec[" + << c.keepaliveIntvlSec << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay[" + << c.enableCtrlNodelay << "] numDataConns[" << c.numDataConns << "] stripingThresholdBytes[" << c.stripingThresholdBytes << "]"; } diff --git a/src/io/tcp/backend_impl.cpp b/src/io/tcp/backend_impl.cpp index 309527cac..3632c3dfb 100644 --- a/src/io/tcp/backend_impl.cpp +++ b/src/io/tcp/backend_impl.cpp @@ -19,7 +19,7 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. -// Copyright © Advanced Micro Devices, Inc. All rights reserved. + #include "src/io/tcp/backend_impl.hpp" #include "src/io/tcp/transport.hpp" @@ -51,7 +51,7 @@ TcpBackend::TcpBackend(EngineKey k, const IOEngineConfig& engCfg, const TcpBacke : myEngKey(std::move(k)), config(cfg) { transport = std::make_unique(myEngKey, engCfg, cfg); transport->Start(); - MORI_IO_INFO("TcpBackend created key={}", myEngKey.c_str()); + MORI_IO_INFO("TcpBackend created key={}", myEngKey); } TcpBackend::~TcpBackend() { transport->Shutdown(); } @@ -61,35 +61,27 @@ std::optional TcpBackend::GetListenPort() const { return transport->Ge void TcpBackend::RegisterRemoteEngine(const EngineDesc& desc) { transport->RegisterRemoteEngine(desc); } - void TcpBackend::DeregisterRemoteEngine(const EngineDesc& desc) { transport->DeregisterRemoteEngine(desc); } - void TcpBackend::RegisterMemory(MemoryDesc& desc) { transport->RegisterMemory(desc); } - void TcpBackend::DeregisterMemory(const MemoryDesc& desc) { transport->DeregisterMemory(desc); } -void TcpBackend::ReadWrite(const MemoryDesc& localDest, size_t localOffset, - const MemoryDesc& remoteSrc, size_t remoteOffset, size_t size, - TransferStatus* status, TransferUniqueId id, bool isRead) { +void TcpBackend::ReadWrite(const MemoryDesc& ld, size_t lo, const MemoryDesc& rs, size_t ro, + size_t sz, TransferStatus* st, TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; - transport->SubmitReadWrite(localDest, localOffset, remoteSrc, remoteOffset, size, status, id, - isRead); + transport->SubmitReadWrite(ld, lo, rs, ro, sz, st, id, isRead); } -void TcpBackend::BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, - const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, - const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, - bool isRead) { +void TcpBackend::BatchReadWrite(const MemoryDesc& ld, const SizeVec& lo, const MemoryDesc& rs, + const SizeVec& ro, const SizeVec& sz, TransferStatus* st, + TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; - transport->SubmitBatchReadWrite(localDest, localOffsets, remoteSrc, remoteOffsets, sizes, status, - id, isRead); + transport->SubmitBatchReadWrite(ld, lo, rs, ro, sz, st, id, isRead); } BackendSession* TcpBackend::CreateSession(const MemoryDesc& local, const MemoryDesc& remote) { - auto* sess = new TcpBackendSession(config, local, remote, transport.get()); - return sess; + return new TcpBackendSession(config, local, remote, transport.get()); } bool TcpBackend::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, diff --git a/src/io/tcp/backend_impl.hpp b/src/io/tcp/backend_impl.hpp index 153fe6169..1924d6a07 100644 --- a/src/io/tcp/backend_impl.hpp +++ b/src/io/tcp/backend_impl.hpp @@ -19,6 +19,7 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. + #pragma once #include @@ -34,22 +35,18 @@ namespace io { class TcpTransport; -/* ---------------------------------------------------------------------------------------------- */ -/* TcpBackendSession */ -/* ---------------------------------------------------------------------------------------------- */ class TcpBackendSession : public BackendSession { public: TcpBackendSession() = default; - TcpBackendSession(const TcpBackendConfig& config, const MemoryDesc& local, const MemoryDesc& remote, - TcpTransport* transport); + TcpBackendSession(const TcpBackendConfig& config, const MemoryDesc& local, + const MemoryDesc& remote, TcpTransport* transport); ~TcpBackendSession() override = default; void ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, bool isRead) override; - - void BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, - TransferStatus* status, TransferUniqueId id, bool isRead) override; - + void BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) override; bool Alive() const override; private: @@ -59,9 +56,6 @@ class TcpBackendSession : public BackendSession { TcpTransport* transport{nullptr}; }; -/* ---------------------------------------------------------------------------------------------- */ -/* TcpBackend */ -/* ---------------------------------------------------------------------------------------------- */ class TcpBackend : public Backend { public: TcpBackend(EngineKey, const IOEngineConfig&, const TcpBackendConfig&); @@ -71,21 +65,20 @@ class TcpBackend : public Backend { void RegisterRemoteEngine(const EngineDesc&) override; void DeregisterRemoteEngine(const EngineDesc&) override; - void RegisterMemory(MemoryDesc& desc) override; void DeregisterMemory(const MemoryDesc& desc) override; void ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc, size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, bool isRead) override; - void BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets, - const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, const SizeVec& sizes, - TransferStatus* status, TransferUniqueId id, bool isRead) override; + const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets, + const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, + bool isRead) override; BackendSession* CreateSession(const MemoryDesc& local, const MemoryDesc& remote) override; - - bool PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, TransferStatus* status) override; + bool PopInboundTransferStatus(EngineKey remote, TransferUniqueId id, + TransferStatus* status) override; private: EngineKey myEngKey; diff --git a/src/io/tcp/data_worker.hpp b/src/io/tcp/data_worker.hpp deleted file mode 100644 index 58c4999d8..000000000 --- a/src/io/tcp/data_worker.hpp +++ /dev/null @@ -1,426 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include - -#include - -#include "src/io/tcp/tcp_types.hpp" - -namespace mori { -namespace io { - -class DataConnectionWorker { - public: - DataConnectionWorker(int fd, EngineKey peer, PinnedStagingPool* staging) - : fd_(fd), peerKey_(std::move(peer)), staging_(staging) { - notifyFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); - wakeFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); - } - - ~DataConnectionWorker() { - Stop(); - if (notifyFd_ >= 0) { - close(notifyFd_); - notifyFd_ = -1; - } - if (wakeFd_ >= 0) { - close(wakeFd_); - wakeFd_ = -1; - } - } - - DataConnectionWorker(const DataConnectionWorker&) = delete; - DataConnectionWorker& operator=(const DataConnectionWorker&) = delete; - - void Start() { - if (running_.load()) return; - running_.store(true); - thread_ = std::thread(&DataConnectionWorker::Run, this); - } - - void Stop() { - bool was = running_.exchange(false); - if (!was) return; - WakeWorker(); - if (thread_.joinable()) thread_.join(); - } - - int NotifyFd() const { return notifyFd_; } - int Fd() const { return fd_; } - - void SubmitSend(SendItem item) { - { - std::lock_guard lk(sendMu_); - sendQ_.push_back(std::move(item)); - } - WakeWorker(); - } - - void RegisterRecvTarget(TransferUniqueId opId, const WorkerRecvTarget& target) { - std::lock_guard lk(targetMu_); - recvTargets_[opId] = target; - } - - void RemoveRecvTarget(TransferUniqueId opId) { - std::lock_guard lk(targetMu_); - recvTargets_.erase(opId); - } - - void DrainEvents(std::deque& out) { - uint64_t v; - while (::read(notifyFd_, &v, sizeof(v)) > 0) { - } - std::lock_guard lk(eventMu_); - while (!eventQ_.empty()) { - out.push_back(std::move(eventQ_.front())); - eventQ_.pop_front(); - } - } - - private: - void WakeWorker() { - uint64_t one = 1; - ::write(wakeFd_, &one, sizeof(one)); - } - - void NotifyMain() { - uint64_t one = 1; - ::write(notifyFd_, &one, sizeof(one)); - } - - void PostEvent(WorkerEvent ev) { - { - std::lock_guard lk(eventMu_); - eventQ_.push_back(std::move(ev)); - } - NotifyMain(); - } - - void PostRecvDone(TransferUniqueId opId, uint8_t lane, uint64_t laneLen, bool discarded = false) { - WorkerEvent ev; - ev.type = WorkerEventType::RECV_DONE; - ev.peerKey = peerKey_; - ev.opId = opId; - ev.lane = lane; - ev.laneLen = laneLen; - ev.discarded = discarded; - PostEvent(std::move(ev)); - } - - void PostEarlyData(TransferUniqueId opId, uint8_t lane, uint64_t laneLen, - std::shared_ptr buf) { - WorkerEvent ev; - ev.type = WorkerEventType::EARLY_DATA; - ev.peerKey = peerKey_; - ev.opId = opId; - ev.lane = lane; - ev.laneLen = laneLen; - ev.earlyBuf = std::move(buf); - PostEvent(std::move(ev)); - } - - void PostCallback(std::function cb) { - WorkerEvent ev; - ev.type = WorkerEventType::SEND_CALLBACK; - ev.callback = std::move(cb); - PostEvent(std::move(ev)); - } - - void PostError(const std::string& msg) { - WorkerEvent ev; - ev.type = WorkerEventType::CONN_ERROR; - ev.peerKey = peerKey_; - ev.errorMsg = msg; - PostEvent(std::move(ev)); - } - - void Run() { - MORI_IO_TRACE("TCP: DataWorker fd={} peer={} started", fd_, peerKey_); - struct pollfd pfds[2]; - pfds[0].fd = fd_; - pfds[1].fd = wakeFd_; - pfds[1].events = POLLIN; - - while (running_.load()) { - bool hasSend; - { - std::lock_guard lk(sendMu_); - hasSend = !sendQ_.empty(); - } - - pfds[0].events = POLLIN | (hasSend ? POLLOUT : 0); - pfds[0].revents = 0; - pfds[1].revents = 0; - - int n = ::poll(pfds, 2, hasSend ? 0 : 1); - if (n < 0) { - if (errno == EINTR) continue; - PostError(std::string("poll failed: ") + strerror(errno)); - break; - } - - if (pfds[1].revents & POLLIN) { - uint64_t v; - while (::read(wakeFd_, &v, sizeof(v)) > 0) { - } - } - - if (pfds[0].revents & (POLLERR | POLLHUP | POLLNVAL)) { - PostError("data connection error/hangup"); - break; - } - - if (pfds[0].revents & POLLOUT) { - if (!ProcessSend()) break; - } - - if (pfds[0].revents & POLLIN) { - if (!ProcessRecv()) break; - } - } - MORI_IO_TRACE("TCP: DataWorker fd={} peer={} exiting", fd_, peerKey_); - } - - bool ProcessSend() { - std::deque batch; - { - std::lock_guard lk(sendMu_); - batch.swap(sendQ_); - } - - for (auto& item : batch) { - while (!item.Done()) { - constexpr size_t kMaxIov = 64; - iovec iov[kMaxIov]; - size_t cnt = 0; - for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { - iov[cnt] = item.iov[i]; - if (i == item.idx && item.off > 0) { - iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; - iov[cnt].iov_len -= item.off; - } - cnt++; - } - msghdr msg{}; - msg.msg_iov = iov; - msg.msg_iovlen = cnt; - ssize_t n = ::sendmsg(fd_, &msg, MSG_NOSIGNAL | item.flags); - if (n < 0) { - if (IsWouldBlock(errno)) { - goto requeue_batch; - } - PostError(std::string("sendmsg failed: ") + strerror(errno)); - return false; - } - if (n == 0) { - goto requeue_batch; - } - item.Advance(static_cast(n)); - } - if (item.onDone) { - PostCallback(std::move(item.onDone)); - } - } - return true; - - requeue_batch: { - std::lock_guard lk(sendMu_); - for (auto rit = batch.rbegin(); rit != batch.rend(); ++rit) { - if (!rit->Done()) { - sendQ_.push_front(std::move(*rit)); - } - } - } - return true; - } - - bool ProcessRecv() { - while (true) { - while (hdrGot_ < tcp::kDataHeaderSize) { - ssize_t n = ::recv(fd_, hdrBuf_ + hdrGot_, tcp::kDataHeaderSize - hdrGot_, 0); - if (n < 0) { - if (IsWouldBlock(errno)) return true; - PostError(std::string("recv header failed: ") + strerror(errno)); - return false; - } - if (n == 0) { - PostError("data connection closed by peer"); - return false; - } - hdrGot_ += static_cast(n); - } - hdrGot_ = 0; - - tcp::DataHeaderView hv; - if (!tcp::TryParseDataHeader(hdrBuf_, tcp::kDataHeaderSize, &hv)) { - PostError("bad data header"); - return false; - } - - const uint8_t lane = static_cast(hv.opId & kLaneMask); - const TransferUniqueId userOpId = static_cast(ToUserOpId(hv.opId)); - const uint64_t payloadLen = hv.payloadLen; - - WorkerRecvTarget target; - bool hasTarget = false; - { - std::lock_guard lk(targetMu_); - auto it = recvTargets_.find(userOpId); - if (it != recvTargets_.end()) { - target = it->second; - hasTarget = true; - } - } - - if (hasTarget && !target.discard) { - const LaneSpan span = ComputeLaneSpan(target.totalLen, target.lanesTotal, lane); - if (span.len != payloadLen) { - MORI_IO_WARN("TCP: worker recv op {} lane {} len mismatch expected={} got={}", userOpId, - (uint32_t)lane, span.len, payloadLen); - if (!DiscardPayload(payloadLen)) return false; - PostRecvDone(userOpId, lane, payloadLen, true); - } else if (target.toGpu) { - uint8_t* dst = reinterpret_cast(target.pinned->ptr) + span.off; - if (!RecvExact(dst, payloadLen)) return false; - PostRecvDone(userOpId, lane, payloadLen); - } else { - auto segs = SliceSegments(target.segs, span.off, span.len); - if (!RecvIntoSegments(reinterpret_cast(target.cpuBase), segs, payloadLen)) - return false; - PostRecvDone(userOpId, lane, payloadLen); - } - } else if (hasTarget && target.discard) { - if (!DiscardPayload(payloadLen)) return false; - PostRecvDone(userOpId, lane, payloadLen, true); - } else { - if (payloadLen == 0) { - PostEarlyData(userOpId, lane, 0, nullptr); - } else { - auto buf = staging_->Acquire(static_cast(payloadLen)); - if (!buf) { - if (!DiscardPayload(payloadLen)) return false; - PostRecvDone(userOpId, lane, payloadLen, true); - } else { - if (!RecvExact(reinterpret_cast(buf->ptr), payloadLen)) return false; - PostEarlyData(userOpId, lane, payloadLen, std::move(buf)); - } - } - } - } - return true; - } - - bool RecvExact(uint8_t* dst, uint64_t len) { - uint64_t got = 0; - while (got < len) { - const size_t want = static_cast(std::min(len - got, 16ULL * 1024 * 1024)); - ssize_t n = ::recv(fd_, dst + got, want, 0); - if (n < 0) { - if (IsWouldBlock(errno)) continue; - PostError(std::string("recv payload failed: ") + strerror(errno)); - return false; - } - if (n == 0) { - PostError("data connection closed during recv"); - return false; - } - got += static_cast(n); - } - return true; - } - - bool RecvIntoSegments(uint8_t* base, const std::vector& segs, uint64_t totalLen) { - uint64_t remaining = totalLen; - size_t segIdx = 0; - uint64_t segOff = 0; - while (remaining > 0 && segIdx < segs.size()) { - const Segment& seg = segs[segIdx]; - const uint64_t segRemain = seg.len - segOff; - const size_t want = static_cast( - std::min(remaining, std::min(segRemain, 16ULL * 1024 * 1024))); - uint8_t* dst = base + seg.off + segOff; - ssize_t n = ::recv(fd_, dst, want, 0); - if (n < 0) { - if (IsWouldBlock(errno)) continue; - PostError(std::string("recv seg failed: ") + strerror(errno)); - return false; - } - if (n == 0) { - PostError("data connection closed during seg recv"); - return false; - } - remaining -= static_cast(n); - segOff += static_cast(n); - if (segOff >= seg.len) { - segIdx++; - segOff = 0; - } - } - return (remaining == 0); - } - - bool DiscardPayload(uint64_t len) { - uint8_t tmp[65536]; - uint64_t remaining = len; - while (remaining > 0) { - const size_t want = static_cast(std::min(remaining, sizeof(tmp))); - ssize_t n = ::recv(fd_, tmp, want, 0); - if (n < 0) { - if (IsWouldBlock(errno)) continue; - PostError(std::string("recv discard failed: ") + strerror(errno)); - return false; - } - if (n == 0) { - PostError("data connection closed during discard"); - return false; - } - remaining -= static_cast(n); - } - return true; - } - - int fd_; - EngineKey peerKey_; - PinnedStagingPool* staging_; - std::atomic running_{false}; - std::thread thread_; - int notifyFd_{-1}; - int wakeFd_{-1}; - - std::mutex sendMu_; - std::deque sendQ_; - - std::mutex targetMu_; - std::unordered_map recvTargets_; - - std::mutex eventMu_; - std::deque eventQ_; - - uint8_t hdrBuf_[tcp::kDataHeaderSize]{}; - size_t hdrGot_{0}; -}; - -} // namespace io -} // namespace mori diff --git a/src/io/tcp/protocol.hpp b/src/io/tcp/protocol.hpp deleted file mode 100644 index aebbaae55..000000000 --- a/src/io/tcp/protocol.hpp +++ /dev/null @@ -1,273 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include - -#include -#include -#include -#include -#include - -namespace mori { -namespace io { -namespace tcp { - -constexpr uint32_t kCtrlMagic = 0x4D544330; // "MTC0" -constexpr uint32_t kDataMagic = 0x4D544430; // "MTD0" -constexpr uint16_t kProtoVersion = 2; - -enum class Channel : uint8_t { CTRL = 1, DATA = 2 }; - -enum class CtrlMsgType : uint8_t { - HELLO = 1, - WRITE_REQ = 2, - READ_REQ = 3, - BATCH_WRITE_REQ = 4, - BATCH_READ_REQ = 5, - COMPLETION = 6, -}; - -// Fixed framing: -// CtrlHeader (12B) + body -// magic(4) ver(2) type(1) reserved(1) body_len(4) -// DataHeader (24B) + payload -// magic(4) ver(2) flags(2) op_id(8) payload_len(8) -constexpr size_t kCtrlHeaderSize = 12; -constexpr size_t kDataHeaderSize = 24; - -struct CtrlHeaderView { - CtrlMsgType type{CtrlMsgType::HELLO}; - uint32_t bodyLen{0}; -}; - -struct DataHeaderView { - uint16_t flags{0}; - uint64_t opId{0}; - uint64_t payloadLen{0}; -}; - -inline uint64_t HostToBe64(uint64_t v) { return htobe64(v); } -inline uint64_t BeToHost64(uint64_t v) { return be64toh(v); } - -inline void AppendU8(std::vector& out, uint8_t v) { out.push_back(v); } - -inline void AppendU16BE(std::vector& out, uint16_t v) { - uint16_t be = htons(v); - uint8_t* p = reinterpret_cast(&be); - out.insert(out.end(), p, p + sizeof(be)); -} - -inline void AppendU32BE(std::vector& out, uint32_t v) { - uint32_t be = htonl(v); - uint8_t* p = reinterpret_cast(&be); - out.insert(out.end(), p, p + sizeof(be)); -} - -inline void AppendU64BE(std::vector& out, uint64_t v) { - uint64_t be = HostToBe64(v); - uint8_t* p = reinterpret_cast(&be); - out.insert(out.end(), p, p + sizeof(be)); -} - -inline bool ReadU16BE(const uint8_t* p, size_t len, size_t* off, uint16_t* out) { - if (*off + sizeof(uint16_t) > len) return false; - uint16_t be; - std::memcpy(&be, p + *off, sizeof(be)); - *out = ntohs(be); - *off += sizeof(be); - return true; -} - -inline bool ReadU32BE(const uint8_t* p, size_t len, size_t* off, uint32_t* out) { - if (*off + sizeof(uint32_t) > len) return false; - uint32_t be; - std::memcpy(&be, p + *off, sizeof(be)); - *out = ntohl(be); - *off += sizeof(be); - return true; -} - -inline bool ReadU64BE(const uint8_t* p, size_t len, size_t* off, uint64_t* out) { - if (*off + sizeof(uint64_t) > len) return false; - uint64_t be; - std::memcpy(&be, p + *off, sizeof(be)); - *out = BeToHost64(be); - *off += sizeof(be); - return true; -} - -inline bool TryParseCtrlHeader(const uint8_t* buf, size_t len, CtrlHeaderView* out) { - if (len < kCtrlHeaderSize) return false; - size_t off = 0; - - uint32_t magic = 0; - uint16_t ver = 0; - uint8_t type = 0; - uint8_t reserved = 0; - uint32_t bodyLen = 0; - - if (!ReadU32BE(buf, len, &off, &magic)) return false; - if (!ReadU16BE(buf, len, &off, &ver)) return false; - if (off + 2 > len) return false; - type = buf[off++]; - reserved = buf[off++]; - (void)reserved; - if (!ReadU32BE(buf, len, &off, &bodyLen)) return false; - - if (magic != kCtrlMagic || ver != kProtoVersion) return false; - out->type = static_cast(type); - out->bodyLen = bodyLen; - return true; -} - -inline bool TryParseDataHeader(const uint8_t* buf, size_t len, DataHeaderView* out) { - if (len < kDataHeaderSize) return false; - size_t off = 0; - - uint32_t magic = 0; - uint16_t ver = 0; - uint16_t flags = 0; - uint64_t opId = 0; - uint64_t payloadLen = 0; - - if (!ReadU32BE(buf, len, &off, &magic)) return false; - if (!ReadU16BE(buf, len, &off, &ver)) return false; - if (!ReadU16BE(buf, len, &off, &flags)) return false; - if (!ReadU64BE(buf, len, &off, &opId)) return false; - if (!ReadU64BE(buf, len, &off, &payloadLen)) return false; - - if (magic != kDataMagic || ver != kProtoVersion) return false; - out->flags = flags; - out->opId = opId; - out->payloadLen = payloadLen; - return true; -} - -inline std::vector BuildCtrlFrame(CtrlMsgType type, const std::vector& body) { - std::vector out; - out.reserve(kCtrlHeaderSize + body.size()); - AppendU32BE(out, kCtrlMagic); - AppendU16BE(out, kProtoVersion); - AppendU8(out, static_cast(type)); - AppendU8(out, 0 /*reserved*/); - AppendU32BE(out, static_cast(body.size())); - out.insert(out.end(), body.begin(), body.end()); - return out; -} - -inline std::vector BuildHello(Channel ch, const std::string& engineKey) { - std::vector body; - body.reserve(1 + 4 + engineKey.size()); - AppendU8(body, static_cast(ch)); - AppendU32BE(body, static_cast(engineKey.size())); - body.insert(body.end(), engineKey.begin(), engineKey.end()); - return BuildCtrlFrame(CtrlMsgType::HELLO, body); -} - -inline std::vector BuildCompletion(uint64_t opId, uint32_t statusCode, - const std::string& msg) { - std::vector body; - body.reserve(8 + 4 + 4 + msg.size()); - AppendU64BE(body, opId); - AppendU32BE(body, statusCode); - AppendU32BE(body, static_cast(msg.size())); - body.insert(body.end(), msg.begin(), msg.end()); - return BuildCtrlFrame(CtrlMsgType::COMPLETION, body); -} - -inline std::vector BuildWriteReq(uint64_t opId, uint32_t remoteMemId, uint64_t remoteOff, - uint64_t size, uint8_t lanesTotal = 1) { - std::vector body; - body.reserve(8 + 4 + 8 + 8 + 1); - AppendU64BE(body, opId); - AppendU32BE(body, remoteMemId); - AppendU64BE(body, remoteOff); - AppendU64BE(body, size); - AppendU8(body, lanesTotal); - return BuildCtrlFrame(CtrlMsgType::WRITE_REQ, body); -} - -inline std::vector BuildReadReq(uint64_t opId, uint32_t srcMemId, uint64_t srcOff, - uint64_t size, uint8_t lanesTotal = 1) { - std::vector body; - body.reserve(8 + 4 + 8 + 8 + 1); - AppendU64BE(body, opId); - AppendU32BE(body, srcMemId); - AppendU64BE(body, srcOff); - AppendU64BE(body, size); - AppendU8(body, lanesTotal); - return BuildCtrlFrame(CtrlMsgType::READ_REQ, body); -} - -inline std::vector BuildBatchWriteReq(uint64_t opId, uint32_t remoteMemId, - const std::vector& remoteOffs, - const std::vector& sizes, - uint8_t lanesTotal = 1) { - std::vector body; - const uint32_t n = static_cast(sizes.size()); - body.reserve(8 + 4 + 4 + n * (8 + 8) + 1); - AppendU64BE(body, opId); - AppendU32BE(body, remoteMemId); - AppendU32BE(body, n); - for (uint32_t i = 0; i < n; ++i) { - AppendU64BE(body, remoteOffs[i]); - AppendU64BE(body, sizes[i]); - } - AppendU8(body, lanesTotal); - return BuildCtrlFrame(CtrlMsgType::BATCH_WRITE_REQ, body); -} - -inline std::vector BuildBatchReadReq(uint64_t opId, uint32_t srcMemId, - const std::vector& srcOffs, - const std::vector& sizes, - uint8_t lanesTotal = 1) { - std::vector body; - const uint32_t n = static_cast(sizes.size()); - body.reserve(8 + 4 + 4 + n * (8 + 8) + 1); - AppendU64BE(body, opId); - AppendU32BE(body, srcMemId); - AppendU32BE(body, n); - for (uint32_t i = 0; i < n; ++i) { - AppendU64BE(body, srcOffs[i]); - AppendU64BE(body, sizes[i]); - } - AppendU8(body, lanesTotal); - return BuildCtrlFrame(CtrlMsgType::BATCH_READ_REQ, body); -} - -inline std::vector BuildDataHeader(uint64_t opId, uint64_t payloadLen, uint16_t flags) { - std::vector out; - out.reserve(kDataHeaderSize); - AppendU32BE(out, kDataMagic); - AppendU16BE(out, kProtoVersion); - AppendU16BE(out, flags); - AppendU64BE(out, opId); - AppendU64BE(out, payloadLen); - return out; -} - -} // namespace tcp -} // namespace io -} // namespace mori diff --git a/src/io/tcp/tcp_types.hpp b/src/io/tcp/tcp_types.hpp deleted file mode 100644 index 49f06f5e0..000000000 --- a/src/io/tcp/tcp_types.hpp +++ /dev/null @@ -1,421 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "mori/io/backend.hpp" -#include "mori/io/common.hpp" -#include "mori/io/logging.hpp" -#include "src/io/tcp/protocol.hpp" - -namespace mori { -namespace io { - -class DataConnectionWorker; - -using Clock = std::chrono::steady_clock; - -inline bool IsWouldBlock(int err) { return (err == EAGAIN) || (err == EWOULDBLOCK); } - -inline int SetNonBlocking(int fd) { - int flags = fcntl(fd, F_GETFL, 0); - if (flags < 0) return -1; - if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) return -1; - return 0; -} - -inline int SetBlocking(int fd) { - int flags = fcntl(fd, F_GETFL, 0); - if (flags < 0) return -1; - if (fcntl(fd, F_SETFL, flags & ~O_NONBLOCK) < 0) return -1; - return 0; -} - -inline void SetSockOptOrLog(int fd, int level, int optname, const void* optval, socklen_t optlen, - const char* name) { - if (setsockopt(fd, level, optname, optval, optlen) != 0) { - MORI_IO_WARN("TCP: setsockopt {} failed: {}", name, strerror(errno)); - } -} - -inline void ConfigureSocketCommon(int fd, const TcpBackendConfig& cfg) { - if (cfg.enableKeepalive) { - int on = 1; - SetSockOptOrLog(fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on), "SO_KEEPALIVE"); - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPIDLE, &cfg.keepaliveIdleSec, - sizeof(cfg.keepaliveIdleSec), "TCP_KEEPIDLE"); - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPINTVL, &cfg.keepaliveIntvlSec, - sizeof(cfg.keepaliveIntvlSec), "TCP_KEEPINTVL"); - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_KEEPCNT, &cfg.keepaliveCnt, sizeof(cfg.keepaliveCnt), - "TCP_KEEPCNT"); - } -} - -inline void ConfigureCtrlSocket(int fd, const TcpBackendConfig& cfg) { - ConfigureSocketCommon(fd, cfg); - if (cfg.enableCtrlNodelay) { - int on = 1; - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(ctrl)"); - } -} - -inline void ConfigureDataSocket(int fd, const TcpBackendConfig& cfg) { - ConfigureSocketCommon(fd, cfg); - { - int on = 1; - SetSockOptOrLog(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(data)"); - } - if (cfg.sockSndbufBytes > 0) { - SetSockOptOrLog(fd, SOL_SOCKET, SO_SNDBUF, &cfg.sockSndbufBytes, sizeof(cfg.sockSndbufBytes), - "SO_SNDBUF"); - } - if (cfg.sockRcvbufBytes > 0) { - SetSockOptOrLog(fd, SOL_SOCKET, SO_RCVBUF, &cfg.sockRcvbufBytes, sizeof(cfg.sockRcvbufBytes), - "SO_RCVBUF"); - } -} - -inline std::optional ParseIpv4(const std::string& host, uint16_t port) { - sockaddr_in addr{}; - addr.sin_family = AF_INET; - addr.sin_port = htons(port); - if (inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) return std::nullopt; - return addr; -} - -inline uint16_t GetBoundPort(int fd) { - sockaddr_in addr{}; - socklen_t len = sizeof(addr); - if (getsockname(fd, reinterpret_cast(&addr), &len) != 0) return 0; - return ntohs(addr.sin_port); -} - -inline size_t RoundUpPow2(size_t v) { - if (v <= 1) return 1; - size_t p = 1; - while (p < v) p <<= 1; - return p; -} - -struct PinnedBuf { - void* ptr{nullptr}; - size_t cap{0}; -}; - -class PinnedStagingPool { - public: - PinnedStagingPool() = default; - ~PinnedStagingPool() { Clear(); } - - PinnedStagingPool(const PinnedStagingPool&) = delete; - PinnedStagingPool& operator=(const PinnedStagingPool&) = delete; - - std::shared_ptr Acquire(size_t size) { - const size_t cap = RoundUpPow2(size); - { - std::lock_guard lock(mu); - auto it = free.find(cap); - if (it != free.end() && !it->second.empty()) { - void* p = it->second.back(); - it->second.pop_back(); - return std::shared_ptr(new PinnedBuf{p, cap}, - [this](PinnedBuf* b) { this->Release(b); }); - } - } - void* p = nullptr; - hipError_t err = hipHostMalloc(&p, cap, hipHostMallocDefault); - if (err != hipSuccess) { - MORI_IO_ERROR("TCP: hipHostMalloc({}) failed: {}", cap, hipGetErrorString(err)); - return nullptr; - } - return std::shared_ptr(new PinnedBuf{p, cap}, - [this](PinnedBuf* b) { this->Release(b); }); - } - - void Clear() { - std::lock_guard lock(mu); - for (auto& kv : free) { - for (void* p : kv.second) { - hipError_t err = hipHostFree(p); - if (err != hipSuccess) { - MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); - } - } - kv.second.clear(); - } - free.clear(); - } - - private: - void Release(PinnedBuf* b) { - if (b == nullptr) return; - const size_t cap = b->cap; - void* p = b->ptr; - delete b; - constexpr size_t kMaxCachedPerClass = 8; - std::lock_guard lock(mu); - auto& vec = free[cap]; - if (vec.size() < kMaxCachedPerClass) { - vec.push_back(p); - } else { - hipError_t err = hipHostFree(p); - if (err != hipSuccess) { - MORI_IO_WARN("TCP: hipHostFree failed: {}", hipGetErrorString(err)); - } - } - } - - std::mutex mu; - std::unordered_map> free; -}; - -struct Segment { - uint64_t off{0}; - uint64_t len{0}; -}; - -constexpr uint8_t kLaneBits = 3; -constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; - -inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { - return (userOpId << kLaneBits) | lane; -} -inline uint64_t ToUserOpId(uint64_t wireOpId) { return wireOpId >> kLaneBits; } - -struct LaneSpan { - uint64_t off{0}; - uint64_t len{0}; -}; - -inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanesTotal, uint8_t lane) { - if (lanesTotal <= 1) return LaneSpan{0, total}; - const uint64_t base = total / lanesTotal; - const uint64_t rem = total % lanesTotal; - return LaneSpan{static_cast(lane) * base + std::min(lane, rem), - base + (lane < rem ? 1 : 0)}; -} - -inline uint8_t LanesAllMask(uint8_t lanesTotal) { - if (lanesTotal >= (1U << kLaneBits)) return 0xFF; - return static_cast((1U << lanesTotal) - 1U); -} - -inline uint8_t ClampLanesTotal(uint8_t lanesTotal) { - if (lanesTotal == 0) return 1; - return std::min(lanesTotal, static_cast(1U << kLaneBits)); -} - -inline std::vector SliceSegments(const std::vector& segs, uint64_t start, - uint64_t len) { - std::vector out; - if (len == 0) return out; - uint64_t skip = start; - uint64_t remaining = len; - for (const auto& s : segs) { - if (remaining == 0) break; - if (skip >= s.len) { - skip -= s.len; - continue; - } - const uint64_t take = std::min(s.len - skip, remaining); - out.push_back({s.off + skip, take}); - remaining -= take; - skip = 0; - } - return out; -} - -inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, - uint64_t* outLen) { - if (!outOff || !outLen || segs.empty()) return false; - uint64_t off = segs[0].off; - uint64_t end = off + segs[0].len; - for (size_t i = 1; i < segs.size(); ++i) { - if (segs[i].off != end) return false; - end += segs[i].len; - } - *outOff = off; - *outLen = end - off; - return true; -} - -inline uint64_t SumLens(const std::vector& segs) { - uint64_t total = 0; - for (const auto& s : segs) total += s.len; - return total; -} - -struct SendItem { - std::vector header; - std::vector iov; - size_t idx{0}; - size_t off{0}; - int flags{0}; - std::shared_ptr keepalive; - std::function onDone; - - bool Done() const { return idx >= iov.size(); } - - void Advance(size_t n) { - while (n > 0 && idx < iov.size()) { - size_t avail = iov[idx].iov_len - off; - if (n < avail) { - off += n; - n = 0; - break; - } - n -= avail; - idx++; - off = 0; - } - } -}; - -struct Connection { - int fd{-1}; - bool isOutgoing{false}; - bool connecting{false}; - bool helloSent{false}; - bool helloReceived{false}; - tcp::Channel ch{tcp::Channel::CTRL}; - EngineKey peerKey{}; - std::vector inbuf; - std::deque sendq; -}; - -struct PeerLinks { - int ctrlFd{-1}; - std::vector dataFds; - std::vector workers; - int ctrlPending{0}; - int dataPending{0}; - size_t rr{0}; - bool CtrlUp() const { return ctrlFd >= 0; } - bool DataUp() const { return !dataFds.empty(); } -}; - -struct InboundStatusEntry { - StatusCode code{StatusCode::INIT}; - std::string msg; -}; - -struct OutboundOpState { - EngineKey peer; - TransferUniqueId id{0}; - bool isRead{false}; - TransferStatus* status{nullptr}; - MemoryDesc local{}; - std::vector localSegs; - MemoryDesc remote{}; - std::vector remoteSegs; - uint64_t expectedRxBytes{0}; - uint64_t rxBytes{0}; - bool completionReceived{false}; - uint8_t lanesTotal{1}; - uint8_t lanesDoneMask{0}; - StatusCode completionCode{StatusCode::SUCCESS}; - std::string completionMsg; - bool gpuCopyPending{false}; - std::shared_ptr pinned; - Clock::time_point startTs{Clock::now()}; -}; - -struct InboundWriteState { - EngineKey peer; - TransferUniqueId id{0}; - MemoryDesc dst{}; - std::vector dstSegs; - bool discard{false}; - uint8_t lanesTotal{1}; - uint8_t lanesDoneMask{0}; - std::shared_ptr pinned; -}; - -struct EarlyWriteLaneState { - uint64_t payloadLen{0}; - std::shared_ptr pinned; - bool complete{false}; -}; - -struct EarlyWriteState { - std::unordered_map lanes; -}; - -struct WorkerRecvTarget { - uint8_t lanesTotal{1}; - uint64_t totalLen{0}; - bool discard{false}; - bool toGpu{false}; - void* cpuBase{nullptr}; - std::vector segs; - std::shared_ptr pinned; -}; - -enum class WorkerEventType : uint8_t { - RECV_DONE = 0, - EARLY_DATA = 1, - SEND_CALLBACK = 2, - CONN_ERROR = 3, -}; - -struct WorkerEvent { - WorkerEventType type{WorkerEventType::RECV_DONE}; - EngineKey peerKey; - TransferUniqueId opId{0}; - uint8_t lane{0}; - uint64_t laneLen{0}; - bool discarded{false}; - std::shared_ptr earlyBuf; - std::function callback; - std::string errorMsg; -}; - -struct GpuTask { - int deviceId{-1}; - hipEvent_t ev{nullptr}; - std::function onReady; -}; - -} // namespace io -} // namespace mori diff --git a/src/io/tcp/transport.cpp b/src/io/tcp/transport.cpp index 6f7ba61ec..d6d4841ec 100644 --- a/src/io/tcp/transport.cpp +++ b/src/io/tcp/transport.cpp @@ -19,6 +19,7 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. + #include "src/io/tcp/transport.hpp" #include @@ -26,58 +27,449 @@ namespace mori { namespace io { -namespace { +// =========================================================================== +// PinnedStagingPool +// =========================================================================== +std::shared_ptr PinnedStagingPool::Acquire(size_t size) { + const size_t cap = RoundUp(size); + { + std::lock_guard lk(mu_); + auto it = free_.find(cap); + if (it != free_.end() && !it->second.empty()) { + void* p = it->second.back(); + it->second.pop_back(); + return std::shared_ptr(new PinnedBuf{p, cap}, + [this](PinnedBuf* b) { Release(b); }); + } + } + void* p = nullptr; + if (hipHostMalloc(&p, cap, hipHostMallocDefault) != hipSuccess) { + MORI_IO_ERROR("TCP: hipHostMalloc({}) failed", cap); + return nullptr; + } + return std::shared_ptr(new PinnedBuf{p, cap}, [this](PinnedBuf* b) { Release(b); }); +} -struct LinearReqView { - uint64_t opId{0}; - uint32_t memId{0}; - uint64_t remoteOff{0}; - uint64_t size{0}; - uint8_t lanesTotal{1}; -}; +void PinnedStagingPool::Clear() { + std::lock_guard lk(mu_); + for (auto& kv : free_) + for (void* p : kv.second) hipHostFree(p); + free_.clear(); +} -bool ParseLinearReq(const uint8_t* body, size_t len, LinearReqView* out) { - if (out == nullptr) return false; - size_t off = 0; - if (!tcp::ReadU64BE(body, len, &off, &out->opId) || - !tcp::ReadU32BE(body, len, &off, &out->memId) || - !tcp::ReadU64BE(body, len, &off, &out->remoteOff) || - !tcp::ReadU64BE(body, len, &off, &out->size)) { - return false; +void PinnedStagingPool::Release(PinnedBuf* b) { + if (!b) return; + constexpr size_t kMaxCached = 8; + size_t cap = b->cap; + void* p = b->ptr; + delete b; + std::lock_guard lk(mu_); + auto& vec = free_[cap]; + if (vec.size() < kMaxCached) + vec.push_back(p); + else + hipHostFree(p); +} + +// =========================================================================== +// SendItem +// =========================================================================== +void SendItem::Advance(size_t n) { + while (n > 0 && idx < iov.size()) { + size_t avail = iov[idx].iov_len - off; + if (n < avail) { + off += n; + return; + } + n -= avail; + idx++; + off = 0; } - if (off < len) out->lanesTotal = body[off]; - out->lanesTotal = ClampLanesTotal(out->lanesTotal); +} + +// =========================================================================== +// DataConnectionWorker +// =========================================================================== +DataConnectionWorker::DataConnectionWorker(int fd, EngineKey peer, PinnedStagingPool* staging) + : fd_(fd), peerKey_(std::move(peer)), staging_(staging) { + notifyFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + wakeFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); +} + +DataConnectionWorker::~DataConnectionWorker() { + Stop(); + if (notifyFd_ >= 0) close(notifyFd_); + if (wakeFd_ >= 0) close(wakeFd_); +} + +void DataConnectionWorker::Start() { + if (running_.load()) return; + running_.store(true); + thread_ = std::thread(&DataConnectionWorker::Run, this); +} + +void DataConnectionWorker::Stop() { + if (!running_.exchange(false)) return; + WakeWorker(); + if (thread_.joinable()) thread_.join(); +} + +void DataConnectionWorker::SubmitSend(SendItem item) { + { + std::lock_guard lk(sendMu_); + sendQ_.push_back(std::move(item)); + } + WakeWorker(); +} + +void DataConnectionWorker::RegisterRecvTarget(TransferUniqueId opId, + const WorkerRecvTarget& target) { + std::lock_guard lk(targetMu_); + recvTargets_[opId] = target; +} + +void DataConnectionWorker::RemoveRecvTarget(TransferUniqueId opId) { + std::lock_guard lk(targetMu_); + recvTargets_.erase(opId); +} + +void DataConnectionWorker::DrainEvents(std::deque& out) { + uint64_t v; + while (::read(notifyFd_, &v, sizeof(v)) > 0) { + } + std::lock_guard lk(eventMu_); + while (!eventQ_.empty()) { + out.push_back(std::move(eventQ_.front())); + eventQ_.pop_front(); + } +} + +void DataConnectionWorker::WakeWorker() { + uint64_t one = 1; + ::write(wakeFd_, &one, sizeof(one)); +} + +void DataConnectionWorker::NotifyMain() { + uint64_t one = 1; + ::write(notifyFd_, &one, sizeof(one)); +} + +void DataConnectionWorker::PostEvent(WorkerEvent ev) { + { + std::lock_guard lk(eventMu_); + eventQ_.push_back(std::move(ev)); + } + NotifyMain(); +} + +void DataConnectionWorker::Run() { + MORI_IO_TRACE("TCP: DataWorker fd={} peer={} started", fd_, peerKey_); + pollfd pfds[2]; + pfds[0].fd = fd_; + pfds[1].fd = wakeFd_; + pfds[1].events = POLLIN; + + while (running_.load()) { + bool hasSend; + { + std::lock_guard lk(sendMu_); + hasSend = !sendQ_.empty(); + } + + pfds[0].events = POLLIN | (hasSend ? POLLOUT : 0); + pfds[0].revents = pfds[1].revents = 0; + + int n = ::poll(pfds, 2, hasSend ? 0 : 1); + if (n < 0) { + if (errno == EINTR) continue; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("poll failed: ") + strerror(errno)}); + break; + } + + if (pfds[1].revents & POLLIN) { + uint64_t v; + while (::read(wakeFd_, &v, sizeof(v)) > 0) { + } + } + if (pfds[0].revents & (POLLERR | POLLHUP | POLLNVAL)) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "data connection error/hangup"}); + break; + } + if ((pfds[0].revents & POLLOUT) && !ProcessSend()) break; + if ((pfds[0].revents & POLLIN) && !ProcessRecv()) break; + } + MORI_IO_TRACE("TCP: DataWorker fd={} peer={} exiting", fd_, peerKey_); +} + +bool DataConnectionWorker::ProcessSend() { + std::deque batch; + { + std::lock_guard lk(sendMu_); + batch.swap(sendQ_); + } + + for (auto& item : batch) { + while (!item.Done()) { + constexpr size_t kMaxIov = 64; + iovec iov[kMaxIov]; + size_t cnt = 0; + for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { + iov[cnt] = item.iov[i]; + if (i == item.idx && item.off > 0) { + iov[cnt].iov_base = static_cast(iov[cnt].iov_base) + item.off; + iov[cnt].iov_len -= item.off; + } + cnt++; + } + msghdr msg{}; + msg.msg_iov = iov; + msg.msg_iovlen = cnt; + ssize_t n = ::sendmsg(fd_, &msg, MSG_NOSIGNAL | item.flags); + if (n < 0) { + if (IsWouldBlock(errno)) goto requeue; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("sendmsg failed: ") + strerror(errno)}); + return false; + } + if (n == 0) goto requeue; + item.Advance(static_cast(n)); + } + if (item.onDone) { + WorkerEvent ev; + ev.type = WorkerEventType::SEND_CALLBACK; + ev.callback = std::move(item.onDone); + PostEvent(std::move(ev)); + } + } + return true; + +requeue: { + std::lock_guard lk(sendMu_); + for (auto rit = batch.rbegin(); rit != batch.rend(); ++rit) + if (!rit->Done()) sendQ_.push_front(std::move(*rit)); +} return true; } -struct BatchReqView { +bool DataConnectionWorker::ProcessRecv() { + while (true) { + // Read data header + while (hdrGot_ < tcp::kDataHeaderSize) { + ssize_t n = ::recv(fd_, hdrBuf_ + hdrGot_, tcp::kDataHeaderSize - hdrGot_, 0); + if (n < 0) { + if (IsWouldBlock(errno)) return true; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("recv header failed: ") + strerror(errno)}); + return false; + } + if (n == 0) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "data connection closed by peer"}); + return false; + } + hdrGot_ += static_cast(n); + } + hdrGot_ = 0; + + tcp::DataHeaderView hv; + if (!tcp::TryParseDataHeader(hdrBuf_, tcp::kDataHeaderSize, &hv)) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "bad data header"}); + return false; + } + + const uint8_t lane = static_cast(hv.opId & kLaneMask); + const TransferUniqueId userOpId = static_cast(ToUserOpId(hv.opId)); + const uint64_t payloadLen = hv.payloadLen; + + // Look up recv target + WorkerRecvTarget target; + bool hasTarget = false; + { + std::lock_guard lk(targetMu_); + auto it = recvTargets_.find(userOpId); + if (it != recvTargets_.end()) { + target = it->second; + hasTarget = true; + } + } + + auto postRecvDone = [&](bool discarded = false) { + WorkerEvent ev; + ev.type = WorkerEventType::RECV_DONE; + ev.peerKey = peerKey_; + ev.opId = userOpId; + ev.lane = lane; + ev.laneLen = payloadLen; + ev.discarded = discarded; + PostEvent(std::move(ev)); + }; + + if (hasTarget && !target.discard) { + const LaneSpan span = ComputeLaneSpan(target.totalLen, target.lanesTotal, lane); + if (span.len != payloadLen) { + MORI_IO_WARN("TCP: worker recv op {} lane {} len mismatch expected={} got={}", userOpId, + (uint32_t)lane, span.len, payloadLen); + if (!DiscardPayload(payloadLen)) return false; + postRecvDone(true); + } else if (target.toGpu) { + if (!RecvExact(reinterpret_cast(target.pinned->ptr) + span.off, payloadLen)) + return false; + postRecvDone(); + } else { + if (!RecvIntoSegments(reinterpret_cast(target.cpuBase), + SliceSegments(target.segs, span.off, span.len), payloadLen)) + return false; + postRecvDone(); + } + } else if (hasTarget) { + if (!DiscardPayload(payloadLen)) return false; + postRecvDone(true); + } else { + // Early data: no target registered yet + if (payloadLen == 0) { + WorkerEvent ev; + ev.type = WorkerEventType::EARLY_DATA; + ev.peerKey = peerKey_; + ev.opId = userOpId; + ev.lane = lane; + ev.laneLen = 0; + PostEvent(std::move(ev)); + } else { + auto buf = staging_->Acquire(static_cast(payloadLen)); + if (!buf) { + if (!DiscardPayload(payloadLen)) return false; + postRecvDone(true); + } else { + if (!RecvExact(reinterpret_cast(buf->ptr), payloadLen)) return false; + WorkerEvent ev; + ev.type = WorkerEventType::EARLY_DATA; + ev.peerKey = peerKey_; + ev.opId = userOpId; + ev.lane = lane; + ev.laneLen = payloadLen; + ev.earlyBuf = std::move(buf); + PostEvent(std::move(ev)); + } + } + } + } + return true; +} + +bool DataConnectionWorker::RecvExact(uint8_t* dst, uint64_t len) { + uint64_t got = 0; + while (got < len) { + ssize_t n = + ::recv(fd_, dst + got, static_cast(std::min(len - got, 16ULL << 20)), 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("recv payload failed: ") + strerror(errno)}); + return false; + } + if (n == 0) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "data connection closed during recv"}); + return false; + } + got += static_cast(n); + } + return true; +} + +bool DataConnectionWorker::RecvIntoSegments(uint8_t* base, const std::vector& segs, + uint64_t totalLen) { + uint64_t remaining = totalLen; + size_t segIdx = 0; + uint64_t segOff = 0; + while (remaining > 0 && segIdx < segs.size()) { + const Segment& seg = segs[segIdx]; + size_t want = static_cast( + std::min(remaining, std::min(seg.len - segOff, 16ULL << 20))); + ssize_t n = ::recv(fd_, base + seg.off + segOff, want, 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("recv seg failed: ") + strerror(errno)}); + return false; + } + if (n == 0) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "data connection closed during seg recv"}); + return false; + } + remaining -= static_cast(n); + segOff += static_cast(n); + if (segOff >= seg.len) { + segIdx++; + segOff = 0; + } + } + return (remaining == 0); +} + +bool DataConnectionWorker::DiscardPayload(uint64_t len) { + uint8_t tmp[65536]; + uint64_t remaining = len; + while (remaining > 0) { + ssize_t n = + ::recv(fd_, tmp, static_cast(std::min(remaining, sizeof(tmp))), 0); + if (n < 0) { + if (IsWouldBlock(errno)) continue; + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + std::string("recv discard failed: ") + strerror(errno)}); + return false; + } + if (n == 0) { + PostEvent({WorkerEventType::CONN_ERROR, peerKey_, 0, 0, 0, false, nullptr, nullptr, + "data connection closed during discard"}); + return false; + } + remaining -= static_cast(n); + } + return true; +} + +// =========================================================================== +// Request parsing helpers (anonymous namespace) +// =========================================================================== +namespace { + +struct RequestView { uint64_t opId{0}; uint32_t memId{0}; std::vector segs; uint8_t lanesTotal{1}; }; -bool ParseBatchReq(const uint8_t* body, size_t len, BatchReqView* out) { - if (out == nullptr) return false; - size_t off = 0; - uint32_t n = 0; - if (!tcp::ReadU64BE(body, len, &off, &out->opId) || - !tcp::ReadU32BE(body, len, &off, &out->memId) || !tcp::ReadU32BE(body, len, &off, &n)) { - return false; - } - - out->segs.clear(); - out->segs.reserve(n); - for (uint32_t i = 0; i < n; ++i) { - uint64_t remoteOff = 0; - uint64_t size = 0; - if (!tcp::ReadU64BE(body, len, &off, &remoteOff) || !tcp::ReadU64BE(body, len, &off, &size)) { - return false; +// Parse either a linear or batch request into a uniform RequestView +bool ParseRequest(tcp::CtrlMsgType type, const uint8_t* body, size_t len, RequestView* out) { + tcp::WireReader r{body, len}; + if (!r.u64(&out->opId) || !r.u32(&out->memId)) return false; + + bool isBatch = + (type == tcp::CtrlMsgType::BATCH_WRITE_REQ || type == tcp::CtrlMsgType::BATCH_READ_REQ); + if (isBatch) { + uint32_t n = 0; + if (!r.u32(&n)) return false; + out->segs.reserve(n); + for (uint32_t i = 0; i < n; ++i) { + uint64_t off, sz; + if (!r.u64(&off) || !r.u64(&sz)) return false; + if (sz > 0) out->segs.push_back({off, sz}); } - if (size > 0) out->segs.push_back({remoteOff, size}); + } else { + uint64_t off, sz; + if (!r.u64(&off) || !r.u64(&sz)) return false; + out->segs.push_back({off, sz}); + } + if (r.off < r.len) { + uint8_t lt; + if (r.u8(<)) out->lanesTotal = lt; } - - if (off < len) out->lanesTotal = body[off]; out->lanesTotal = ClampLanesTotal(out->lanesTotal); return true; } @@ -89,133 +481,114 @@ struct CompletionView { }; bool ParseCompletion(const uint8_t* body, size_t len, CompletionView* out) { - if (out == nullptr) return false; - size_t off = 0; + tcp::WireReader r{body, len}; uint32_t msgLen = 0; - if (!tcp::ReadU64BE(body, len, &off, &out->opId) || - !tcp::ReadU32BE(body, len, &off, &out->statusCode) || - !tcp::ReadU32BE(body, len, &off, &msgLen)) { - return false; - } - if (off + msgLen > len) return false; - out->msg.assign(reinterpret_cast(body + off), msgLen); - return true; -} - -bool SegmentsInRange(const std::vector& segs, uint64_t memSize) { - for (const auto& seg : segs) { - if (seg.off + seg.len > memSize) return false; - } + if (!r.u64(&out->opId) || !r.u32(&out->statusCode) || !r.u32(&msgLen)) return false; + if (r.off + msgLen > r.len) return false; + out->msg.assign(reinterpret_cast(body + r.off), msgLen); return true; } } // namespace +// =========================================================================== +// TcpTransport +// =========================================================================== TcpTransport::TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg) - : myEngKey(std::move(myKey)), engConfig(engCfg), config(cfg) {} + : myEngKey_(std::move(myKey)), engConfig_(engCfg), config_(cfg) {} TcpTransport::~TcpTransport() { Shutdown(); } void TcpTransport::Start() { - if (running.load()) return; + if (running_.load()) return; - epfd = epoll_create1(EPOLL_CLOEXEC); - assert(epfd >= 0); + epfd_ = epoll_create1(EPOLL_CLOEXEC); + assert(epfd_ >= 0); - listenFd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); - assert(listenFd >= 0); + listenFd_ = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); + assert(listenFd_ >= 0); int one = 1; - SetSockOptOrLog(listenFd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one), "SO_REUSEADDR"); + SetSockOpt(listenFd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one), "SO_REUSEADDR"); - auto addrOpt = - ParseIpv4(engConfig.host.empty() ? std::string("0.0.0.0") : engConfig.host, engConfig.port); + auto addrOpt = ParseIpv4(engConfig_.host.empty() ? "0.0.0.0" : engConfig_.host, engConfig_.port); assert(addrOpt.has_value()); - sockaddr_in addr = addrOpt.value(); - if (bind(listenFd, reinterpret_cast(&addr), sizeof(addr)) != 0) { - MORI_IO_ERROR("TCP: bind {}:{} failed: {}", engConfig.host, engConfig.port, strerror(errno)); + sockaddr_in addr = *addrOpt; + if (bind(listenFd_, reinterpret_cast(&addr), sizeof(addr)) != 0) { + MORI_IO_ERROR("TCP: bind {}:{} failed: {}", engConfig_.host, engConfig_.port, strerror(errno)); assert(false && "bind failed"); } - if (listen(listenFd, 256) != 0) { + if (listen(listenFd_, 256) != 0) { MORI_IO_ERROR("TCP: listen failed: {}", strerror(errno)); assert(false && "listen failed"); } - listenPort = GetBoundPort(listenFd); - MORI_IO_INFO("TCP: listen on {}:{} (port={})", engConfig.host, engConfig.port, listenPort); + listenPort_ = GetBoundPort(listenFd_); + MORI_IO_INFO("TCP: listen on {}:{} (port={})", engConfig_.host, engConfig_.port, listenPort_); - wakeFd = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); - assert(wakeFd >= 0); + wakeFd_ = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC); + assert(wakeFd_ >= 0); - AddEpoll(listenFd, true, false); - AddEpoll(wakeFd, true, false); + AddEpoll(listenFd_, true, false); + AddEpoll(wakeFd_, true, false); - running.store(true); - ioThread = std::thread([this] { this->IoLoop(); }); + running_.store(true); + ioThread_ = std::thread([this] { IoLoop(); }); } void TcpTransport::Shutdown() { - bool wasRunning = running.exchange(false); - if (!wasRunning) return; - - if (wakeFd >= 0) { + if (!running_.exchange(false)) return; + if (wakeFd_ >= 0) { uint64_t one = 1; - ::write(wakeFd, &one, sizeof(one)); + ::write(wakeFd_, &one, sizeof(one)); } - if (ioThread.joinable()) ioThread.join(); + if (ioThread_.joinable()) ioThread_.join(); - for (auto& kv : dataWorkers) kv.second->Stop(); - dataWorkers.clear(); - workerNotifyMap.clear(); + for (auto& kv : dataWorkers_) kv.second->Stop(); + dataWorkers_.clear(); + workerNotifyMap_.clear(); - for (auto& kv : conns) CloseConnInternal(kv.second.get()); - conns.clear(); - peers.clear(); + for (auto& kv : conns_) CloseConnInternal(kv.second.get()); + conns_.clear(); + peers_.clear(); - if (listenFd >= 0) { - close(listenFd); - listenFd = -1; - } - if (wakeFd >= 0) { - close(wakeFd); - wakeFd = -1; - } - if (epfd >= 0) { - close(epfd); - epfd = -1; - } + auto closeFd = [](int& fd) { + if (fd >= 0) { + close(fd); + fd = -1; + } + }; + closeFd(listenFd_); + closeFd(wakeFd_); + closeFd(epfd_); } std::optional TcpTransport::GetListenPort() const { - if (listenPort == 0) return std::nullopt; - return listenPort; + return listenPort_ ? std::optional(listenPort_) : std::nullopt; } void TcpTransport::RegisterRemoteEngine(const EngineDesc& desc) { - std::lock_guard lock(remoteMu); - remoteEngines[desc.key] = desc; + std::lock_guard lk(remoteMu_); + remoteEngines_[desc.key] = desc; } - void TcpTransport::DeregisterRemoteEngine(const EngineDesc& desc) { - std::lock_guard lock(remoteMu); - remoteEngines.erase(desc.key); + std::lock_guard lk(remoteMu_); + remoteEngines_.erase(desc.key); } - void TcpTransport::RegisterMemory(const MemoryDesc& desc) { - std::lock_guard lock(memMu); - localMems[desc.id] = desc; + std::lock_guard lk(memMu_); + localMems_[desc.id] = desc; } - void TcpTransport::DeregisterMemory(const MemoryDesc& desc) { - std::lock_guard lock(memMu); - localMems.erase(desc.id); + std::lock_guard lk(memMu_); + localMems_.erase(desc.id); } bool TcpTransport::PopInboundTransferStatus(const EngineKey& remote, TransferUniqueId id, TransferStatus* status) { - std::lock_guard lock(inboundMu); - auto it = inboundStatus.find(remote); - if (it == inboundStatus.end()) return false; + std::lock_guard lk(inboundMu_); + auto it = inboundStatus_.find(remote); + if (it == inboundStatus_.end()) return false; auto it2 = it->second.find(id); if (it2 == it->second.end()) return false; status->Update(it2->second.code, it2->second.msg); @@ -223,15 +596,18 @@ bool TcpTransport::PopInboundTransferStatus(const EngineKey& remote, TransferUni return true; } +// --------------------------------------------------------------------------- +// Submission +// --------------------------------------------------------------------------- void TcpTransport::SubmitReadWrite(const MemoryDesc& local, size_t localOffset, const MemoryDesc& remote, size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, bool isRead) { - if (status == nullptr) return; + if (!status) return; if (size == 0) { status->SetCode(StatusCode::SUCCESS); return; } - if ((localOffset + size) > local.size || (remoteOffset + size) > remote.size) { + if (localOffset + size > local.size || remoteOffset + size > remote.size) { status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: offset+size out of range"); return; } @@ -243,11 +619,9 @@ void TcpTransport::SubmitReadWrite(const MemoryDesc& local, size_t localOffset, op->status = status; op->local = local; op->remote = remote; - op->localSegs = {{static_cast(localOffset), static_cast(size)}}; - op->remoteSegs = {{static_cast(remoteOffset), static_cast(size)}}; - op->expectedRxBytes = isRead ? static_cast(size) : 0; - op->startTs = Clock::now(); - + op->localSegs = {{uint64_t(localOffset), uint64_t(size)}}; + op->remoteSegs = {{uint64_t(remoteOffset), uint64_t(size)}}; + op->expectedRxBytes = isRead ? uint64_t(size) : 0; status->SetCode(StatusCode::IN_PROGRESS); EnqueueOp(std::move(op)); } @@ -256,7 +630,7 @@ void TcpTransport::SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& const MemoryDesc& remote, const SizeVec& remoteOffsets, const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, bool isRead) { - if (status == nullptr) return; + if (!status) return; const size_t n = sizes.size(); if (n == 0) { status->SetCode(StatusCode::SUCCESS); @@ -267,45 +641,43 @@ void TcpTransport::SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& return; } - std::vector localSegs, remoteSegs; - localSegs.reserve(n); - remoteSegs.reserve(n); + std::vector lSegs, rSegs; + lSegs.reserve(n); + rSegs.reserve(n); uint64_t total = 0; for (size_t i = 0; i < n; ++i) { - const size_t lo = localOffsets[i], ro = remoteOffsets[i], sz = sizes[i]; - if (sz == 0) continue; - if ((lo + sz) > local.size || (ro + sz) > remote.size) { + if (sizes[i] == 0) continue; + if (localOffsets[i] + sizes[i] > local.size || remoteOffsets[i] + sizes[i] > remote.size) { status->Update(StatusCode::ERR_INVALID_ARGS, "TCP: batch offset+size out of range"); return; } - localSegs.push_back({static_cast(lo), static_cast(sz)}); - remoteSegs.push_back({static_cast(ro), static_cast(sz)}); - total += static_cast(sz); - } - - if (localSegs.size() > 1) { - std::vector newLocal, newRemote; - newLocal.reserve(localSegs.size()); - newRemote.reserve(remoteSegs.size()); - Segment curL = localSegs[0], curR = remoteSegs[0]; - for (size_t i = 1; i < localSegs.size(); ++i) { - const Segment& l = localSegs[i]; - const Segment& r = remoteSegs[i]; - if ((curL.off + curL.len == l.off) && (curR.off + curR.len == r.off) && - (curL.len == curR.len) && (l.len == r.len)) { - curL.len += l.len; - curR.len += r.len; + lSegs.push_back({uint64_t(localOffsets[i]), uint64_t(sizes[i])}); + rSegs.push_back({uint64_t(remoteOffsets[i]), uint64_t(sizes[i])}); + total += sizes[i]; + } + + // Merge adjacent contiguous segments + if (lSegs.size() > 1) { + std::vector ml, mr; + ml.reserve(lSegs.size()); + mr.reserve(rSegs.size()); + Segment cl = lSegs[0], cr = rSegs[0]; + for (size_t i = 1; i < lSegs.size(); ++i) { + if (cl.off + cl.len == lSegs[i].off && cr.off + cr.len == rSegs[i].off && cl.len == cr.len && + lSegs[i].len == rSegs[i].len) { + cl.len += lSegs[i].len; + cr.len += rSegs[i].len; } else { - newLocal.push_back(curL); - newRemote.push_back(curR); - curL = l; - curR = r; + ml.push_back(cl); + mr.push_back(cr); + cl = lSegs[i]; + cr = rSegs[i]; } } - newLocal.push_back(curL); - newRemote.push_back(curR); - localSegs = std::move(newLocal); - remoteSegs = std::move(newRemote); + ml.push_back(cl); + mr.push_back(cr); + lSegs = std::move(ml); + rSegs = std::move(mr); } auto op = std::make_unique(); @@ -315,205 +687,190 @@ void TcpTransport::SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& op->status = status; op->local = local; op->remote = remote; - op->localSegs = std::move(localSegs); - op->remoteSegs = std::move(remoteSegs); + op->localSegs = std::move(lSegs); + op->remoteSegs = std::move(rSegs); op->expectedRxBytes = isRead ? total : 0; - op->startTs = Clock::now(); - status->SetCode(StatusCode::IN_PROGRESS); EnqueueOp(std::move(op)); } void TcpTransport::EnqueueOp(std::unique_ptr op) { { - std::lock_guard lock(submitMu); - submitQ.push_back(std::move(op)); + std::lock_guard lk(submitMu_); + submitQ_.push_back(std::move(op)); } uint64_t one = 1; - ::write(wakeFd, &one, sizeof(one)); + ::write(wakeFd_, &one, sizeof(one)); } -void TcpTransport::AddEpoll(int fd, bool wantRead, bool wantWrite) { +// --------------------------------------------------------------------------- +// Epoll helpers +// --------------------------------------------------------------------------- +void TcpTransport::AddEpoll(int fd, bool rd, bool wr) { epoll_event ev{}; ev.data.fd = fd; - ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); - SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev)); + ev.events = EPOLLET | (rd ? EPOLLIN : 0) | (wr ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd_, EPOLL_CTL_ADD, fd, &ev)); } -void TcpTransport::ModEpoll(int fd, bool wantRead, bool wantWrite) { +void TcpTransport::ModEpoll(int fd, bool rd, bool wr) { epoll_event ev{}; ev.data.fd = fd; - ev.events = EPOLLET | (wantRead ? EPOLLIN : 0) | (wantWrite ? EPOLLOUT : 0); - SYSCALL_RETURN_ZERO(epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev)); + ev.events = EPOLLET | (rd ? EPOLLIN : 0) | (wr ? EPOLLOUT : 0); + SYSCALL_RETURN_ZERO(epoll_ctl(epfd_, EPOLL_CTL_MOD, fd, &ev)); } -void TcpTransport::DelEpoll(int fd) { epoll_ctl(epfd, EPOLL_CTL_DEL, fd, nullptr); } +void TcpTransport::DelEpoll(int fd) { epoll_ctl(epfd_, EPOLL_CTL_DEL, fd, nullptr); } void TcpTransport::CloseConnInternal(Connection* c) { - if (c == nullptr) return; - if (c->fd >= 0) { - DelEpoll(c->fd); - shutdown(c->fd, SHUT_RDWR); - close(c->fd); - c->fd = -1; - } + if (!c || c->fd < 0) return; + DelEpoll(c->fd); + shutdown(c->fd, SHUT_RDWR); + close(c->fd); + c->fd = -1; } -bool TcpTransport::PreferOutgoingFor(const EngineKey& peerKey) const { return myEngKey < peerKey; } - +// --------------------------------------------------------------------------- +// Connection management +// --------------------------------------------------------------------------- void TcpTransport::AssignConnToPeer(Connection* c) { assert(c && c->helloReceived); - PeerLinks& link = peers[c->peerKey]; - const bool preferOutgoing = PreferOutgoingFor(c->peerKey); - const bool wasOutgoing = c->isOutgoing; - const tcp::Channel ch = c->ch; + PeerLinks& link = peers_[c->peerKey]; + const bool preferOut = myEngKey_ < c->peerKey; - if (wasOutgoing) { - if (ch == tcp::Channel::CTRL) { + if (c->isOutgoing) { + if (c->ch == tcp::Channel::CTRL) { if (link.ctrlPending > 0) link.ctrlPending--; } else { if (link.dataPending > 0) link.dataPending--; } } - auto replace_if_needed = [&](int& slotFd) { - if (slotFd < 0) { - slotFd = c->fd; + if (c->ch == tcp::Channel::CTRL) { + // Replace ctrl connection if needed + if (link.ctrlFd < 0) { + link.ctrlFd = c->fd; return; } - const int existingFd = slotFd; - const int newFd = c->fd; - Connection* existing = conns[existingFd].get(); - if (!existing) { - slotFd = newFd; + int existFd = link.ctrlFd; + auto eIt = conns_.find(existFd); + if (eIt == conns_.end()) { + link.ctrlFd = c->fd; return; } - const bool keepNew = (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); + bool keepNew = (preferOut && c->isOutgoing) || (!preferOut && !c->isOutgoing); if (keepNew) { - MORI_IO_WARN("TCP: peer {} channel {} replacing fd {} with fd {}", c->peerKey, - static_cast(c->ch), existing->fd, c->fd); - CloseConnInternal(existing); - conns.erase(existingFd); - slotFd = newFd; + MORI_IO_WARN("TCP: peer {} CTRL replacing fd {} with {}", c->peerKey, existFd, c->fd); + CloseConnInternal(eIt->second.get()); + conns_.erase(existFd); + link.ctrlFd = c->fd; } else { - MORI_IO_WARN("TCP: peer {} channel {} dropping duplicate fd {}", c->peerKey, - static_cast(c->ch), c->fd); + MORI_IO_WARN("TCP: peer {} CTRL dropping duplicate fd {}", c->peerKey, c->fd); + int fd = c->fd; CloseConnInternal(c); - conns.erase(newFd); + conns_.erase(fd); } - }; - - if (ch == tcp::Channel::CTRL) { - replace_if_needed(link.ctrlFd); return; } - const bool keepPreferred = - (preferOutgoing && c->isOutgoing) || (!preferOutgoing && !c->isOutgoing); - if (!keepPreferred) { - MORI_IO_TRACE("TCP: peer {} dropping non-preferred DATA fd {} outgoing={}", c->peerKey, c->fd, - c->isOutgoing); - const int fd = c->fd; + // DATA channel + bool keepPref = (preferOut && c->isOutgoing) || (!preferOut && !c->isOutgoing); + if (!keepPref) { + MORI_IO_TRACE("TCP: peer {} dropping non-preferred DATA fd {}", c->peerKey, c->fd); + int fd = c->fd; CloseConnInternal(c); - conns.erase(fd); + conns_.erase(fd); return; } - - const size_t want = static_cast(std::max(1, config.numDataConns)); + size_t want = static_cast(std::max(1, config_.numDataConns)); if (link.dataFds.size() >= want) { - MORI_IO_TRACE("TCP: peer {} dropping extra DATA fd {} (have {} want {})", c->peerKey, c->fd, - link.dataFds.size(), want); - const int fd = c->fd; + MORI_IO_TRACE("TCP: peer {} dropping extra DATA fd {}", c->peerKey, c->fd); + int fd = c->fd; CloseConnInternal(c); - conns.erase(fd); + conns_.erase(fd); return; } - const int dataFd = c->fd; + int dataFd = c->fd; link.dataFds.push_back(dataFd); - MORI_IO_TRACE("TCP: peer {} DATA conn up {}/{}", c->peerKey.c_str(), link.dataFds.size(), want); + MORI_IO_TRACE("TCP: peer {} DATA conn up {}/{}", c->peerKey, link.dataFds.size(), want); DelEpoll(dataFd); SetNonBlocking(dataFd); - ConfigureDataSocket(dataFd, config); + ConfigureDataSocket(dataFd, config_); - auto worker = std::make_unique(dataFd, c->peerKey, &staging); + auto worker = std::make_unique(dataFd, c->peerKey, &staging_); worker->Start(); AddEpoll(worker->NotifyFd(), true, false); - workerNotifyMap[worker->NotifyFd()] = worker.get(); + workerNotifyMap_[worker->NotifyFd()] = worker.get(); link.workers.push_back(worker.get()); - dataWorkers[dataFd] = std::move(worker); + dataWorkers_[dataFd] = std::move(worker); } -void TcpTransport::MaybeDispatchQueuedOps(const EngineKey& peerKey) { - auto it = peers.find(peerKey); - if (it == peers.end()) return; - if (!it->second.CtrlUp() || !it->second.DataUp()) return; - Connection* ctrl = conns[it->second.ctrlFd].get(); - if (!ctrl || !ctrl->helloReceived) return; - if (it->second.workers.empty()) return; - - auto qit = waitingOps.find(peerKey); - if (qit == waitingOps.end()) return; +void TcpTransport::MaybeDispatchQueuedOps(const EngineKey& peer) { + auto it = peers_.find(peer); + if (it == peers_.end() || !it->second.CtrlUp() || !it->second.DataUp()) return; + Connection* ctrl = conns_[it->second.ctrlFd].get(); + if (!ctrl || !ctrl->helloReceived || it->second.workers.empty()) return; + auto qit = waitingOps_.find(peer); + if (qit == waitingOps_.end()) return; auto ops = std::move(qit->second); - waitingOps.erase(qit); - MORI_IO_TRACE("TCP: peer {} ready, dispatch {} queued ops", peerKey.c_str(), ops.size()); + waitingOps_.erase(qit); + MORI_IO_TRACE("TCP: peer {} ready, dispatch {} queued ops", peer, ops.size()); for (auto& op : ops) DispatchOp(std::move(op)); } -void TcpTransport::EnsurePeerChannels(const EngineKey& peerKey) { - PeerLinks& link = peers[peerKey]; - if (!link.CtrlUp() && link.ctrlPending == 0) ConnectChannel(peerKey, tcp::Channel::CTRL); - const int want = std::max(1, config.numDataConns); - while (static_cast(link.dataFds.size()) + link.dataPending < want) - ConnectChannel(peerKey, tcp::Channel::DATA); +void TcpTransport::EnsurePeerChannels(const EngineKey& peer) { + PeerLinks& link = peers_[peer]; + if (!link.CtrlUp() && link.ctrlPending == 0) ConnectChannel(peer, tcp::Channel::CTRL); + int want = std::max(1, config_.numDataConns); + while (int(link.dataFds.size()) + link.dataPending < want) + ConnectChannel(peer, tcp::Channel::DATA); } -void TcpTransport::ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { +void TcpTransport::ConnectChannel(const EngineKey& peer, tcp::Channel ch) { EngineDesc desc; { - std::lock_guard lock(remoteMu); - auto it = remoteEngines.find(peerKey); - if (it == remoteEngines.end()) { - MORI_IO_ERROR("TCP: remote engine {} not registered", peerKey.c_str()); + std::lock_guard lk(remoteMu_); + auto it = remoteEngines_.find(peer); + if (it == remoteEngines_.end()) { + MORI_IO_ERROR("TCP: remote engine {} not registered", peer); return; } desc = it->second; } - auto peerAddrOpt = ParseIpv4(desc.host, static_cast(desc.port)); - if (!peerAddrOpt.has_value()) { + auto peerAddr = ParseIpv4(desc.host, static_cast(desc.port)); + if (!peerAddr) { MORI_IO_ERROR("TCP: invalid remote host {}:{}", desc.host, desc.port); return; } - sockaddr_in peerAddr = peerAddrOpt.value(); int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0); if (fd < 0) { MORI_IO_ERROR("TCP: socket() failed: {}", strerror(errno)); return; } - MORI_IO_TRACE("TCP: connect start peer={} ch={} fd={}", peerKey.c_str(), static_cast(ch), - fd); + MORI_IO_TRACE("TCP: connect start peer={} ch={} fd={}", peer, int(ch), fd); - if (!engConfig.host.empty()) { - auto localAddrOpt = ParseIpv4(engConfig.host, 0); - if (localAddrOpt.has_value()) { - sockaddr_in localAddr = localAddrOpt.value(); - if (bind(fd, reinterpret_cast(&localAddr), sizeof(localAddr)) != 0) { - MORI_IO_WARN("TCP: bind(local) {} failed: {}", engConfig.host, strerror(errno)); - } + if (!engConfig_.host.empty()) { + auto la = ParseIpv4(engConfig_.host, 0); + if (la) { + sockaddr_in localAddr = *la; + if (bind(fd, reinterpret_cast(&localAddr), sizeof(localAddr)) != 0) + MORI_IO_WARN("TCP: bind(local) {} failed: {}", engConfig_.host, strerror(errno)); } } - int rc = connect(fd, reinterpret_cast(&peerAddr), sizeof(peerAddr)); + sockaddr_in pa = *peerAddr; + int rc = connect(fd, reinterpret_cast(&pa), sizeof(pa)); bool connecting = false; if (rc != 0) { - if (errno == EINPROGRESS) { + if (errno == EINPROGRESS) connecting = true; - } else { - MORI_IO_ERROR("TCP: connect to {}:{} failed: {}", desc.host, desc.port, strerror(errno)); + else { + MORI_IO_ERROR("TCP: connect failed: {}", strerror(errno)); close(fd); return; } @@ -523,22 +880,18 @@ void TcpTransport::ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { conn->fd = fd; conn->isOutgoing = true; conn->connecting = connecting; - conn->peerKey = peerKey; + conn->peerKey = peer; conn->ch = ch; conn->inbuf.reserve(4096); + if (ch == tcp::Channel::CTRL) ConfigureCtrlSocket(fd, config_); + AddEpoll(fd, true, connecting || !conn->sendq.empty()); + conns_[fd] = std::move(conn); - if (ch == tcp::Channel::CTRL) ConfigureCtrlSocket(fd, config); - - const bool wantWrite = connecting || !conn->sendq.empty(); - AddEpoll(fd, true, wantWrite); - conns[fd] = std::move(conn); - - PeerLinks& link = peers[peerKey]; + PeerLinks& link = peers_[peer]; if (ch == tcp::Channel::CTRL) link.ctrlPending++; else link.dataPending++; - if (!connecting) { QueueHello(fd); ModEpoll(fd, true, true); @@ -546,18 +899,15 @@ void TcpTransport::ConnectChannel(const EngineKey& peerKey, tcp::Channel ch) { } void TcpTransport::QueueHello(int fd) { - Connection* c = conns[fd].get(); + auto it = conns_.find(fd); + if (it == conns_.end()) return; + Connection* c = it->second.get(); if (!c || c->helloSent) return; c->helloSent = true; - auto hello = tcp::BuildHello(c->ch, myEngKey); - MORI_IO_TRACE("TCP: queue HELLO fd={} ch={} peer={}", fd, static_cast(c->ch), - c->peerKey.c_str()); - + MORI_IO_TRACE("TCP: queue HELLO fd={} ch={}", fd, int(c->ch)); SendItem item; - item.header = std::move(hello); - item.iov.resize(1); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); + item.header = tcp::BuildHello(c->ch, myEngKey_); + item.iov = {{item.header.data(), item.header.size()}}; c->sendq.push_back(std::move(item)); } @@ -566,408 +916,306 @@ void TcpTransport::AcceptNew() { sockaddr_in peer{}; socklen_t len = sizeof(peer); int fd = - accept4(listenFd, reinterpret_cast(&peer), &len, SOCK_NONBLOCK | SOCK_CLOEXEC); + accept4(listenFd_, reinterpret_cast(&peer), &len, SOCK_NONBLOCK | SOCK_CLOEXEC); if (fd < 0) { if (IsWouldBlock(errno)) break; MORI_IO_WARN("TCP: accept failed: {}", strerror(errno)); break; } MORI_IO_TRACE("TCP: accept fd={}", fd); - auto conn = std::make_unique(); conn->fd = fd; - conn->isOutgoing = false; conn->inbuf.reserve(4096); AddEpoll(fd, true, false); - conns[fd] = std::move(conn); + conns_[fd] = std::move(conn); } } void TcpTransport::DrainWakeFd() { - uint64_t v = 0; - while (true) { - ssize_t n = ::read(wakeFd, &v, sizeof(v)); - if (n <= 0) break; + uint64_t v; + while (::read(wakeFd_, &v, sizeof(v)) > 0) { } - std::deque> ops; { - std::lock_guard lock(submitMu); - ops.swap(submitQ); + std::lock_guard lk(submitMu_); + ops.swap(submitQ_); } - for (auto& op : ops) { EnsurePeerChannels(op->peer); - if (!IsPeerReady(op->peer)) { - waitingOps[op->peer].push_back(std::move(op)); - continue; - } - DispatchOp(std::move(op)); + if (IsPeerReady(op->peer)) + DispatchOp(std::move(op)); + else + waitingOps_[op->peer].push_back(std::move(op)); } } -bool TcpTransport::IsPeerReady(const EngineKey& peerKey) { - auto it = peers.find(peerKey); - if (it == peers.end()) return false; - if (!it->second.CtrlUp() || !it->second.DataUp()) return false; - Connection* ctrl = conns[it->second.ctrlFd].get(); - if (!ctrl || !ctrl->helloReceived) return false; +bool TcpTransport::IsPeerReady(const EngineKey& peer) { + auto it = peers_.find(peer); + if (it == peers_.end() || !it->second.CtrlUp() || !it->second.DataUp()) return false; + auto cit = conns_.find(it->second.ctrlFd); + if (cit == conns_.end() || !cit->second->helloReceived) return false; return !it->second.workers.empty(); } -void TcpTransport::RegisterRecvTargetWithWorkers(const EngineKey& peerKey, TransferUniqueId opId, +void TcpTransport::RegisterRecvTargetWithWorkers(const EngineKey& peer, TransferUniqueId opId, const WorkerRecvTarget& target) { - auto pit = peers.find(peerKey); - if (pit == peers.end()) return; + auto pit = peers_.find(peer); + if (pit == peers_.end()) return; for (auto* w : pit->second.workers) w->RegisterRecvTarget(opId, target); } -void TcpTransport::RemoveRecvTargetFromWorkers(const EngineKey& peerKey, TransferUniqueId opId) { - auto pit = peers.find(peerKey); - if (pit == peers.end()) return; +void TcpTransport::RemoveRecvTargetFromWorkers(const EngineKey& peer, TransferUniqueId opId) { + auto pit = peers_.find(peer); + if (pit == peers_.end()) return; for (auto* w : pit->second.workers) w->RemoveRecvTarget(opId); } +// --------------------------------------------------------------------------- +// DispatchOp - initiate an outbound operation +// --------------------------------------------------------------------------- void TcpTransport::DispatchOp(std::unique_ptr op) { if (!op) { MORI_IO_ERROR("TCP: DispatchOp got null op"); return; } const EngineKey peerKey = op->peer; - auto it = peers.find(peerKey); - if (it == peers.end() || !it->second.CtrlUp() || !it->second.DataUp()) { + auto pit = peers_.find(peerKey); + if (pit == peers_.end() || !pit->second.CtrlUp() || !pit->second.DataUp()) { op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer not connected"); return; } - Connection* ctrl = conns[it->second.ctrlFd].get(); + Connection* ctrl = conns_[pit->second.ctrlFd].get(); if (!ctrl) { - op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: peer ctrl connection missing"); + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: ctrl missing"); return; } - auto& workerList = it->second.workers; + auto& workerList = pit->second.workers; if (workerList.empty()) { - op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: no data workers for peer"); + op->status->Update(StatusCode::ERR_BAD_STATE, "TCP: no data workers"); return; } const TransferUniqueId opId = op->id; - auto [itIns, inserted] = pendingOutbound.emplace(opId, std::move(op)); + auto [itIns, inserted] = pendingOutbound_.emplace(opId, std::move(op)); if (!inserted) { - MORI_IO_ERROR("TCP: duplicate outbound op id={} for peer={}", opId, peerKey.c_str()); + MORI_IO_ERROR("TCP: duplicate op id={}", opId); itIns->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: duplicate op id"); - pendingOutbound.erase(itIns); + pendingOutbound_.erase(itIns); return; } OutboundOpState* st = itIns->second.get(); - if (!st) { - MORI_IO_ERROR("TCP: failed to store outbound op id={}", opId); - pendingOutbound.erase(itIns); - return; - } + // Decide lane count for striping const uint64_t totalBytes = SumLens(st->localSegs); - int wantLanes = std::max(1, config.numDataConns); - wantLanes = std::min(wantLanes, (1U << kLaneBits)); + int wantLanes = std::min(std::max(1, config_.numDataConns), 1U << kLaneBits); uint8_t lanesTotal = 1; - const bool canStripe = (wantLanes > 1) && (config.stripingThresholdBytes > 0) && - (totalBytes >= static_cast(config.stripingThresholdBytes)) && - (st->localSegs.size() == 1) && (st->remoteSegs.size() == 1) && - (workerList.size() >= 2); - if (canStripe) { - lanesTotal = - static_cast(std::min(static_cast(wantLanes), workerList.size())); + if (wantLanes > 1 && config_.stripingThresholdBytes > 0 && + totalBytes >= uint64_t(config_.stripingThresholdBytes) && st->localSegs.size() == 1 && + st->remoteSegs.size() == 1 && workerList.size() >= 2) { + lanesTotal = uint8_t(std::min(wantLanes, workerList.size())); } st->lanesTotal = lanesTotal; + // Allocate pinned staging for GPU reads if (st->isRead && st->local.loc == MemoryLocationType::GPU) { - st->pinned = staging.Acquire(static_cast(totalBytes)); + st->pinned = staging_.Acquire(static_cast(totalBytes)); if (!st->pinned) { - st->status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed to allocate pinned staging"); - pendingOutbound.erase(opId); + st->status->Update(StatusCode::ERR_BAD_STATE, "TCP: staging alloc failed"); + pendingOutbound_.erase(opId); return; } } + // Set up recv targets for reads if (st->isRead) { WorkerRecvTarget target; target.lanesTotal = lanesTotal; target.totalLen = totalBytes; - target.discard = false; if (st->local.loc == MemoryLocationType::GPU) { target.toGpu = true; target.pinned = st->pinned; } else { - target.toGpu = false; target.cpuBase = reinterpret_cast(st->local.data); target.segs = st->localSegs; } RegisterRecvTargetWithWorkers(peerKey, opId, target); } + // Build and send ctrl frame std::vector ctrlFrame; if (st->localSegs.size() == 1) { - if (st->isRead) - ctrlFrame = tcp::BuildReadReq(st->id, st->remote.id, st->remoteSegs[0].off, + auto type = st->isRead ? tcp::CtrlMsgType::READ_REQ : tcp::CtrlMsgType::WRITE_REQ; + ctrlFrame = tcp::BuildLinearReq(type, st->id, st->remote.id, st->remoteSegs[0].off, st->remoteSegs[0].len, lanesTotal); - else - ctrlFrame = tcp::BuildWriteReq(st->id, st->remote.id, st->remoteSegs[0].off, - st->remoteSegs[0].len, lanesTotal); } else { + auto type = st->isRead ? tcp::CtrlMsgType::BATCH_READ_REQ : tcp::CtrlMsgType::BATCH_WRITE_REQ; std::vector roffs, szs; roffs.reserve(st->remoteSegs.size()); szs.reserve(st->remoteSegs.size()); - for (const auto& s : st->remoteSegs) { + for (auto& s : st->remoteSegs) { roffs.push_back(s.off); szs.push_back(s.len); } - if (st->isRead) - ctrlFrame = tcp::BuildBatchReadReq(st->id, st->remote.id, roffs, szs, lanesTotal); - else - ctrlFrame = tcp::BuildBatchWriteReq(st->id, st->remote.id, roffs, szs, lanesTotal); + ctrlFrame = tcp::BuildBatchReq(type, st->id, st->remote.id, roffs, szs, lanesTotal); } - QueueSend(ctrl->fd, std::move(ctrlFrame)); - if (!st->isRead) QueueDataSendForWrite(workerList, *st); + + if (!st->isRead) QueueDataSend(workerList, st->local, st->localSegs, st->id, lanesTotal); UpdateWriteInterest(ctrl->fd); } void TcpTransport::QueueSend(int fd, std::vector bytes, std::function onDone) { - Connection* c = conns[fd].get(); - if (!c) return; + auto it = conns_.find(fd); + if (it == conns_.end()) return; SendItem item; item.header = std::move(bytes); - item.iov.resize(1); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.onDone = std::move(onDone); - c->sendq.push_back(std::move(item)); -} - -void TcpTransport::QueueSegmentSend(DataConnectionWorker* worker, uint64_t wireOpId, uint8_t* base, - const std::vector& segs, uint64_t totalLen, - std::function onDone) { - SendItem item; - item.header = tcp::BuildDataHeader(wireOpId, totalLen, 0); - item.iov.reserve(1 + segs.size()); - item.iov.push_back({item.header.data(), item.header.size()}); - for (const auto& s : segs) item.iov.push_back({base + s.off, static_cast(s.len)}); + item.iov = {{item.header.data(), item.header.size()}}; item.onDone = std::move(onDone); - worker->SubmitSend(std::move(item)); + it->second->sendq.push_back(std::move(item)); } -void TcpTransport::QueueStripedCpuSend(const std::vector& workers, - uint64_t opId, uint8_t lanesTotal, uint8_t* base, - uint64_t baseOff, uint64_t total, - std::function onLaneDone) { - for (uint8_t lane = 0; lane < lanesTotal; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); - auto* worker = workers[lane % workers.size()]; - SendItem item; - item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = base + baseOff + span.off; - item.iov[1].iov_len = static_cast(span.len); - item.onDone = onLaneDone; - worker->SubmitSend(std::move(item)); +// --------------------------------------------------------------------------- +// Data send (unified for write-send and read-response) +// --------------------------------------------------------------------------- +void TcpTransport::QueueDataSend(const std::vector& workers, + const MemoryDesc& src, const std::vector& srcSegs, + uint64_t opId, uint8_t lanesTotal, + std::function onLaneDone) { + if (workers.empty()) return; + const uint64_t total = SumLens(srcSegs); + lanesTotal = ClampLanesTotal(lanesTotal); + lanesTotal = std::min(lanesTotal, static_cast(workers.size())); + if (lanesTotal > 1 && srcSegs.size() != 1) { + MORI_IO_WARN("TCP: striping requires 1 segment, fallback to 1 lane"); + lanesTotal = 1; } -} -bool TcpTransport::ScheduleGpuDtoH(int deviceId, const MemoryDesc& src, - const std::vector& srcSegs, - std::shared_ptr pinned, - std::function onComplete) { - const uint64_t total = SumLens(srcSegs); - hipStream_t stream = streamPool.GetNextStream(deviceId); - hipEvent_t ev = eventPool.GetEvent(deviceId); - if (stream == nullptr || ev == nullptr) { - MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU DtoH"); - if (ev) eventPool.PutEvent(ev, deviceId); - return false; + if (src.loc == MemoryLocationType::GPU) { + // GPU path: DtoH copy, then send from pinned buffer + auto pinned = staging_.Acquire(static_cast(total)); + if (!pinned) { + MORI_IO_ERROR("TCP: staging alloc failed for GPU send"); + return; + } + + auto workersCopy = + std::vector(workers.begin(), workers.begin() + lanesTotal); + auto sendCb = [workersCopy, pinned, opId, lanesTotal, total, + onLaneDone = std::move(onLaneDone)]() { + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); + item.iov = {{item.header.data(), item.header.size()}, + {static_cast(pinned->ptr) + span.off, size_t(span.len)}}; + item.keepalive = pinned; + item.onDone = onLaneDone; + workersCopy[lane % workersCopy.size()]->SubmitSend(std::move(item)); + } + }; + ScheduleGpuCopy(src.deviceId, false, src, srcSegs, pinned, std::move(sendCb)); + return; } - HIP_RUNTIME_CHECK(hipSetDevice(deviceId)); - uint8_t* dst = reinterpret_cast(pinned->ptr); - uint64_t spanOff = 0, spanLen = 0; - if (IsSingleContiguousSpan(srcSegs, &spanOff, &spanLen) && spanLen == total) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst, gpuPtr, static_cast(total), stream)); + // CPU path + uint8_t* base = reinterpret_cast(src.data); + if (lanesTotal == 1) { + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, 0), total, 0); + item.iov.reserve(1 + srcSegs.size()); + item.iov.push_back({item.header.data(), item.header.size()}); + for (auto& s : srcSegs) item.iov.push_back({base + s.off, size_t(s.len)}); + item.onDone = std::move(onLaneDone); + workers[0]->SubmitSend(std::move(item)); } else { - uint64_t off = 0; - for (const auto& s : srcSegs) { - hipDeviceptr_t gpuPtr = reinterpret_cast(src.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(dst + off, gpuPtr, static_cast(s.len), stream)); - off += s.len; + // Multi-lane striping (contiguous single-segment) + for (uint8_t lane = 0; lane < lanesTotal; ++lane) { + LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); + SendItem item; + item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); + item.iov = {{item.header.data(), item.header.size()}, + {base + srcSegs[0].off + span.off, size_t(span.len)}}; + item.onDone = onLaneDone; + workers[lane % workers.size()]->SubmitSend(std::move(item)); } } - HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - gpuTasks.push_back({deviceId, ev, std::move(onComplete)}); - return true; } -bool TcpTransport::ScheduleGpuHtoD(int deviceId, const MemoryDesc& dst, - const std::vector& dstSegs, +// --------------------------------------------------------------------------- +// GPU copy (unified DtoH / HtoD) +// --------------------------------------------------------------------------- +bool TcpTransport::ScheduleGpuCopy(int deviceId, bool toDevice, const MemoryDesc& mem, + const std::vector& segs, std::shared_ptr pinned, std::function onComplete) { - const uint64_t total = SumLens(dstSegs); - hipStream_t stream = streamPool.GetNextStream(deviceId); - hipEvent_t ev = eventPool.GetEvent(deviceId); - if (stream == nullptr || ev == nullptr) { - MORI_IO_ERROR("TCP: failed to get HIP stream/event for GPU HtoD"); - if (ev) eventPool.PutEvent(ev, deviceId); + const uint64_t total = SumLens(segs); + hipStream_t stream = streamPool_.GetNextStream(deviceId); + hipEvent_t ev = eventPool_.GetEvent(deviceId); + if (!stream || !ev) { + MORI_IO_ERROR("TCP: failed to get HIP stream/event"); + if (ev) eventPool_.PutEvent(ev, deviceId); return false; } HIP_RUNTIME_CHECK(hipSetDevice(deviceId)); - uint8_t* src = reinterpret_cast(pinned->ptr); + uint8_t* hostPtr = reinterpret_cast(pinned->ptr); + uint64_t spanOff = 0, spanLen = 0; - if (IsSingleContiguousSpan(dstSegs, &spanOff, &spanLen) && spanLen == total) { - void* gpuPtr = reinterpret_cast(dst.data + spanOff); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src, static_cast(total), stream)); + if (IsSingleContiguousSpan(segs, &spanOff, &spanLen) && spanLen == total) { + if (toDevice) { + void* gpu = reinterpret_cast(mem.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpu, hostPtr, size_t(total), stream)); + } else { + hipDeviceptr_t gpu = reinterpret_cast(mem.data + spanOff); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(hostPtr, gpu, size_t(total), stream)); + } } else { uint64_t off = 0; - for (const auto& s : dstSegs) { - void* gpuPtr = reinterpret_cast(dst.data + s.off); - HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpuPtr, src + off, static_cast(s.len), stream)); + for (auto& s : segs) { + if (toDevice) { + void* gpu = reinterpret_cast(mem.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyHtoDAsync(gpu, hostPtr + off, size_t(s.len), stream)); + } else { + hipDeviceptr_t gpu = reinterpret_cast(mem.data + s.off); + HIP_RUNTIME_CHECK(hipMemcpyDtoHAsync(hostPtr + off, gpu, size_t(s.len), stream)); + } off += s.len; } } HIP_RUNTIME_CHECK(hipEventRecord(ev, stream)); - gpuTasks.push_back({deviceId, ev, std::move(onComplete)}); + gpuTasks_.push_back({deviceId, ev, std::move(onComplete)}); return true; } -void TcpTransport::QueueGpuSend(const std::vector& workers, uint64_t opId, - uint8_t lanesTotal, const MemoryDesc& src, - const std::vector& srcSegs, - std::function onLaneDone) { - const uint64_t total = SumLens(srcSegs); - auto pinned = staging.Acquire(static_cast(total)); - if (!pinned) { - MORI_IO_ERROR("TCP: failed to allocate pinned staging for GPU send"); - return; - } - - auto sendCallback = [workers, pinned, opId, lanesTotal, total, - onLaneDone = std::move(onLaneDone)]() { - for (uint8_t lane = 0; lane < lanesTotal; ++lane) { - const LaneSpan span = ComputeLaneSpan(total, lanesTotal, lane); - auto* worker = workers[lane % workers.size()]; - SendItem item; - item.header = tcp::BuildDataHeader(ToWireOpId(opId, lane), span.len, 0); - item.iov.resize(2); - item.iov[0].iov_base = item.header.data(); - item.iov[0].iov_len = item.header.size(); - item.iov[1].iov_base = static_cast(pinned->ptr) + static_cast(span.off); - item.iov[1].iov_len = static_cast(span.len); - item.keepalive = pinned; - item.onDone = onLaneDone; - worker->SubmitSend(std::move(item)); - } - }; - - ScheduleGpuDtoH(src.deviceId, src, srcSegs, pinned, std::move(sendCallback)); -} - -void TcpTransport::QueueDataSendForWrite(const std::vector& workerList, - OutboundOpState& st) { - if (workerList.empty()) return; - uint8_t lanesTotal = std::max(1, st.lanesTotal); - - if (lanesTotal > 1 && st.localSegs.size() != 1) { - MORI_IO_WARN("TCP: striping requested but localSegs.size={}, fallback to 1 lane", - st.localSegs.size()); - lanesTotal = 1; - st.lanesTotal = 1; - } - - QueueDataSendCommon(workerList, st.local, st.localSegs, st.id, lanesTotal); -} - -void TcpTransport::QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, - const std::vector& srcSegs, uint8_t lanesTotal) { - auto pit = peers.find(peer); - if (pit == peers.end()) return; - auto& workerList = pit->second.workers; - if (workerList.empty()) return; - - struct DoneState { - EngineKey peer; - uint64_t opId{0}; - std::atomic remaining{0}; - }; - auto done = std::make_shared(); - done->peer = peer; - done->opId = opId; - uint8_t useLanes = std::min( - ClampLanesTotal(lanesTotal), static_cast(std::max(1, workerList.size()))); - if (useLanes > 1 && srcSegs.size() != 1) useLanes = 1; - done->remaining.store(useLanes); - auto laneDone = [this, done]() mutable { - if (done->remaining.fetch_sub(1) > 1) return; - SendCompletionAndRecord(done->peer, done->opId, StatusCode::SUCCESS, ""); - }; - - QueueDataSendCommon(workerList, src, srcSegs, opId, useLanes, std::move(laneDone)); -} - -void TcpTransport::QueueDataSendCommon(const std::vector& workerList, - const MemoryDesc& src, const std::vector& srcSegs, - uint64_t opId, uint8_t lanesTotal, - std::function onLaneDone) { - if (workerList.empty()) return; - const uint64_t total = SumLens(srcSegs); - lanesTotal = ClampLanesTotal(lanesTotal); - lanesTotal = std::min(lanesTotal, static_cast(workerList.size())); - if (lanesTotal > 1 && srcSegs.size() != 1) { - MORI_IO_WARN("TCP: striping requested for {} segments, fallback to 1 lane", srcSegs.size()); - lanesTotal = 1; - } - - if (src.loc == MemoryLocationType::GPU) { - std::vector workers(workerList.begin(), workerList.begin() + lanesTotal); - QueueGpuSend(workers, opId, lanesTotal, src, srcSegs, std::move(onLaneDone)); - return; - } - - uint8_t* base = reinterpret_cast(src.data); - if (lanesTotal == 1) { - QueueSegmentSend(workerList[0], ToWireOpId(opId, 0), base, srcSegs, total, - std::move(onLaneDone)); - } else { - QueueStripedCpuSend(workerList, opId, lanesTotal, base, srcSegs[0].off, total, - std::move(onLaneDone)); - } -} - void TcpTransport::PollGpuTasks() { - for (auto it = gpuTasks.begin(); it != gpuTasks.end();) { + for (auto it = gpuTasks_.begin(); it != gpuTasks_.end();) { hipError_t st = hipEventQuery(it->ev); if (st == hipSuccess) { - eventPool.PutEvent(it->ev, it->deviceId); + eventPool_.PutEvent(it->ev, it->deviceId); if (it->onReady) it->onReady(); - it = gpuTasks.erase(it); + it = gpuTasks_.erase(it); } else if (st == hipErrorNotReady) { ++it; } else { MORI_IO_ERROR("TCP: hipEventQuery failed: {}", hipGetErrorString(st)); - eventPool.PutEvent(it->ev, it->deviceId); - it = gpuTasks.erase(it); + eventPool_.PutEvent(it->ev, it->deviceId); + it = gpuTasks_.erase(it); } } } +// --------------------------------------------------------------------------- +// Ctrl-connection I/O +// --------------------------------------------------------------------------- void TcpTransport::UpdateWriteInterest(int fd) { - auto it = conns.find(fd); - if (it == conns.end()) return; + auto it = conns_.find(fd); + if (it == conns_.end()) return; Connection* c = it->second.get(); if (!c || c->fd < 0) return; - if (!c->connecting && !c->sendq.empty()) { FlushSend(c); - it = conns.find(fd); - if (it == conns.end()) return; + it = conns_.find(fd); + if (it == conns_.end()) return; c = it->second.get(); if (!c || c->fd < 0) return; } @@ -979,7 +1227,7 @@ void TcpTransport::HandleConnWritable(Connection* c) { int err = 0; socklen_t len = sizeof(err); if (getsockopt(c->fd, SOL_SOCKET, SO_ERROR, &err, &len) != 0 || err != 0) { - MORI_IO_ERROR("TCP: connect failed fd {}: {}", c->fd, strerror(err == 0 ? errno : err)); + MORI_IO_ERROR("TCP: connect failed fd {}: {}", c->fd, strerror(err ? err : errno)); ClosePeerByFd(c->fd); return; } @@ -999,7 +1247,6 @@ void TcpTransport::FlushSend(Connection* c) { if (cb) cb(); continue; } - iovec iov[kMaxIov]; size_t cnt = 0; for (size_t i = item.idx; i < item.iov.size() && cnt < kMaxIov; ++i) { @@ -1013,84 +1260,83 @@ void TcpTransport::FlushSend(Connection* c) { msghdr msg{}; msg.msg_iov = iov; msg.msg_iovlen = cnt; - ssize_t n = sendmsg(c->fd, &msg, MSG_NOSIGNAL | item.flags); + ssize_t n = ::sendmsg(c->fd, &msg, MSG_NOSIGNAL); if (n < 0) { - if (IsWouldBlock(errno)) break; - MORI_IO_ERROR("TCP: sendmsg fd {} failed: {}", c->fd, strerror(errno)); + if (IsWouldBlock(errno)) return; + MORI_IO_ERROR("TCP: sendmsg ctrl fd {} failed: {}", c->fd, strerror(errno)); ClosePeerByFd(c->fd); - break; + return; } - if (n == 0) break; + if (n == 0) return; item.Advance(static_cast(n)); } } +// --------------------------------------------------------------------------- +// Peer lifecycle +// --------------------------------------------------------------------------- void TcpTransport::CloseAndRemoveFd(int fd) { - if (fd < 0) return; - auto wit = dataWorkers.find(fd); - if (wit != dataWorkers.end()) { + auto wit = dataWorkers_.find(fd); + if (wit != dataWorkers_.end()) { wit->second->Stop(); - auto nit = workerNotifyMap.find(wit->second->NotifyFd()); - if (nit != workerNotifyMap.end()) { - DelEpoll(nit->first); - workerNotifyMap.erase(nit); - } - dataWorkers.erase(wit); + int nfd = wit->second->NotifyFd(); + DelEpoll(nfd); + workerNotifyMap_.erase(nfd); + dataWorkers_.erase(wit); } - auto it = conns.find(fd); - if (it != conns.end()) { - CloseConnInternal(it->second.get()); - conns.erase(it); + auto cit = conns_.find(fd); + if (cit != conns_.end()) { + CloseConnInternal(cit->second.get()); + conns_.erase(cit); } } EngineKey TcpTransport::FindPeerByFd(int fd) { - for (auto& kv : peers) { - if (kv.second.ctrlFd == fd) return kv.first; - for (int dfd : kv.second.dataFds) { - if (dfd == fd) return kv.first; - } + for (auto& [key, link] : peers_) { + if (link.ctrlFd == fd) return key; + for (int dfd : link.dataFds) + if (dfd == fd) return key; } return {}; } void TcpTransport::ClosePeerByFd(int fd) { - MORI_IO_TRACE("TCP: close fd={}", fd); EngineKey peer = FindPeerByFd(fd); - if (!peer.empty()) { + if (!peer.empty()) ClosePeerByKey(peer, "TCP: connection lost"); - } else { + else CloseAndRemoveFd(fd); - } } void TcpTransport::ClosePeerByKey(const EngineKey& peer, const std::string& reason) { - auto pit = peers.find(peer); - if (pit == peers.end()) return; + auto pit = peers_.find(peer); + if (pit == peers_.end()) return; auto link = pit->second; CloseAndRemoveFd(link.ctrlFd); for (int dfd : link.dataFds) CloseAndRemoveFd(dfd); - peers.erase(peer); + peers_.erase(peer); FailPendingOpsForPeer(peer, reason); } void TcpTransport::FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg) { - for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + for (auto it = pendingOutbound_.begin(); it != pendingOutbound_.end();) { if (it->second->peer == peer) { it->second->status->Update(StatusCode::ERR_BAD_STATE, msg); - it = pendingOutbound.erase(it); - } else { + it = pendingOutbound_.erase(it); + } else ++it; - } } - waitingOps.erase(peer); - inboundWrites.erase(peer); - earlyWrites.erase(peer); + waitingOps_.erase(peer); + inboundWrites_.erase(peer); + earlyWrites_.erase(peer); } +// --------------------------------------------------------------------------- +// Ctrl message handling +// --------------------------------------------------------------------------- void TcpTransport::HandleCtrlReadable(Connection* c) { while (true) { - uint8_t tmp[16384]; + uint8_t tmp[65536]; ssize_t n = ::recv(c->fd, tmp, sizeof(tmp), 0); if (n < 0) { if (IsWouldBlock(errno)) break; @@ -1109,7 +1355,7 @@ void TcpTransport::HandleCtrlReadable(Connection* c) { tcp::CtrlHeaderView hv; if (!tcp::TryParseCtrlHeader(c->inbuf.data(), c->inbuf.size(), &hv)) { if (c->inbuf.size() >= tcp::kCtrlHeaderSize) { - MORI_IO_ERROR("TCP: bad ctrl header on fd {}, closing", c->fd); + MORI_IO_ERROR("TCP: bad ctrl header fd {}", c->fd); ClosePeerByFd(c->fd); } break; @@ -1119,7 +1365,6 @@ void TcpTransport::HandleCtrlReadable(Connection* c) { const uint8_t* body = c->inbuf.data() + tcp::kCtrlHeaderSize; HandleCtrlFrame(c, hv.type, body, hv.bodyLen); c->inbuf.erase(c->inbuf.begin(), c->inbuf.begin() + tcp::kCtrlHeaderSize + hv.bodyLen); - if (c->helloReceived && c->ch == tcp::Channel::DATA) return; } } @@ -1131,109 +1376,186 @@ void TcpTransport::HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const u return; } if (!c->helloReceived) { - MORI_IO_WARN("TCP: received ctrl message before HELLO, dropping"); + MORI_IO_WARN("TCP: ctrl message before HELLO, dropping"); return; } + switch (type) { case tcp::CtrlMsgType::WRITE_REQ: - HandleWriteReq(c->peerKey, body, len); - break; case tcp::CtrlMsgType::READ_REQ: - HandleReadReq(c->peerKey, body, len); - break; case tcp::CtrlMsgType::BATCH_WRITE_REQ: - HandleBatchWriteReq(c->peerKey, body, len); - break; case tcp::CtrlMsgType::BATCH_READ_REQ: - HandleBatchReadReq(c->peerKey, body, len); + HandleRequest(c->peerKey, type, body, len); break; case tcp::CtrlMsgType::COMPLETION: HandleCompletion(c->peerKey, body, len); break; default: - MORI_IO_WARN("TCP: unknown ctrl msg type {}", static_cast(type)); + MORI_IO_WARN("TCP: unknown ctrl msg type {}", uint32_t(type)); } } void TcpTransport::HandleHello(Connection* c, const uint8_t* body, size_t len) { - if (len < 1 + 4) { + if (len < 5) { MORI_IO_WARN("TCP: bad HELLO len {}", len); ClosePeerByFd(c->fd); return; } - size_t off = 0; - const uint8_t chRaw = body[off++]; - uint32_t keyLen = 0; - if (!tcp::ReadU32BE(body, len, &off, &keyLen)) { - ClosePeerByFd(c->fd); - return; - } - if (off + keyLen > len) { + tcp::WireReader r{body, len}; + uint8_t chRaw; + uint32_t keyLen; + if (!r.u8(&chRaw) || !r.u32(&keyLen) || r.off + keyLen > len) { ClosePeerByFd(c->fd); return; } - EngineKey peerKey(reinterpret_cast(body + off), keyLen); - off += keyLen; - c->peerKey = peerKey; - c->ch = - (chRaw == static_cast(tcp::Channel::DATA)) ? tcp::Channel::DATA : tcp::Channel::CTRL; + c->peerKey.assign(reinterpret_cast(body + r.off), keyLen); + c->ch = (chRaw == uint8_t(tcp::Channel::DATA)) ? tcp::Channel::DATA : tcp::Channel::CTRL; c->helloReceived = true; - MORI_IO_TRACE("TCP: recv HELLO fd={} peer={} ch={} outgoing={}", c->fd, c->peerKey.c_str(), - static_cast(c->ch), c->isOutgoing); + MORI_IO_TRACE("TCP: recv HELLO fd={} peer={} ch={} out={}", c->fd, c->peerKey, int(c->ch), + c->isOutgoing); if (!c->helloSent) { QueueHello(c->fd); UpdateWriteInterest(c->fd); } - if (c->ch == tcp::Channel::CTRL) ConfigureCtrlSocket(c->fd, config); + if (c->ch == tcp::Channel::CTRL) ConfigureCtrlSocket(c->fd, config_); AssignConnToPeer(c); - MaybeDispatchQueuedOps(peerKey); + MaybeDispatchQueuedOps(c->peerKey); } +// Unified handler for WRITE_REQ, READ_REQ, BATCH_WRITE_REQ, BATCH_READ_REQ +void TcpTransport::HandleRequest(const EngineKey& peer, tcp::CtrlMsgType type, const uint8_t* body, + size_t len) { + RequestView req; + if (!ParseRequest(type, body, len, &req)) { + MORI_IO_WARN("TCP: malformed request type={}", uint8_t(type)); + return; + } + + bool isWrite = (type == tcp::CtrlMsgType::WRITE_REQ || type == tcp::CtrlMsgType::BATCH_WRITE_REQ); + bool isBatch = + (type == tcp::CtrlMsgType::BATCH_WRITE_REQ || type == tcp::CtrlMsgType::BATCH_READ_REQ); + + auto memOpt = LookupLocalMem(req.memId); + + if (isWrite) { + // Inbound write: remote is sending data to us + InboundWriteState ws; + ws.peer = peer; + ws.id = req.opId; + ws.lanesTotal = req.lanesTotal; + ws.discard = true; + if (memOpt && SegmentsInRange(req.segs, memOpt->size)) { + ws.dst = *memOpt; + ws.dstSegs = std::move(req.segs); + ws.discard = false; + } + FinalizeInboundWriteSetup(peer, req.opId, ws); + } else { + // Inbound read: remote wants data from us + if (!memOpt) { + SendCompletionAndRecord( + peer, req.opId, StatusCode::ERR_NOT_FOUND, + isBatch ? "TCP: remote mem not found" : "TCP: remote mem not found/out of range"); + return; + } + if (!SegmentsInRange(req.segs, memOpt->size)) { + auto code = isBatch ? StatusCode::ERR_INVALID_ARGS : StatusCode::ERR_NOT_FOUND; + SendCompletionAndRecord( + peer, req.opId, code, + isBatch ? "TCP: batch read out of range" : "TCP: remote mem not found/out of range"); + return; + } + + // Send data back to requester + auto pit = peers_.find(peer); + if (pit == peers_.end() || pit->second.workers.empty()) return; + auto& workerList = pit->second.workers; + uint8_t useLanes = std::min(ClampLanesTotal(req.lanesTotal), + uint8_t(std::max(1, workerList.size()))); + if (useLanes > 1 && req.segs.size() != 1) useLanes = 1; + + struct DoneState { + EngineKey peer; + uint64_t opId; + std::atomic remaining{0}; + }; + auto done = std::make_shared(); + done->peer = peer; + done->opId = req.opId; + done->remaining.store(useLanes); + auto laneDone = [this, done]() { + if (done->remaining.fetch_sub(1) > 1) return; + SendCompletionAndRecord(done->peer, done->opId, StatusCode::SUCCESS, ""); + }; + QueueDataSend(workerList, *memOpt, req.segs, req.opId, useLanes, std::move(laneDone)); + } +} + +void TcpTransport::HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { + CompletionView msg; + if (!ParseCompletion(body, len, &msg)) { + MORI_IO_WARN("TCP: malformed COMPLETION"); + return; + } + + auto it = pendingOutbound_.find(msg.opId); + if (it == pendingOutbound_.end()) return; + OutboundOpState& st = *it->second; + st.completionReceived = true; + st.completionCode = static_cast(msg.statusCode); + st.completionMsg = std::move(msg.msg); + if (st.completionCode != StatusCode::SUCCESS) { + RemoveRecvTargetFromWorkers(st.peer, msg.opId); + st.status->Update(st.completionCode, st.completionMsg); + pendingOutbound_.erase(it); + return; + } + MaybeCompleteOutbound(st); +} + +// --------------------------------------------------------------------------- +// Inbound / Outbound state machines +// --------------------------------------------------------------------------- std::optional TcpTransport::LookupLocalMem(MemoryUniqueId id) { - std::lock_guard lock(memMu); - auto it = localMems.find(id); - if (it == localMems.end()) return std::nullopt; - return it->second; + std::lock_guard lk(memMu_); + auto it = localMems_.find(id); + return (it != localMems_.end()) ? std::optional(it->second) : std::nullopt; } void TcpTransport::RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, const std::string& msg) { - std::lock_guard lock(inboundMu); - inboundStatus[peer][id] = InboundStatusEntry{code, msg}; + std::lock_guard lk(inboundMu_); + inboundStatus_[peer][id] = {code, msg}; } void TcpTransport::SendCompletionAndRecord(const EngineKey& peer, TransferUniqueId opId, StatusCode code, const std::string& msg) { Connection* ctrl = PeerCtrl(peer); - if (ctrl != nullptr) { - QueueSend(ctrl->fd, tcp::BuildCompletion(opId, static_cast(code), msg)); + if (ctrl) { + QueueSend(ctrl->fd, tcp::BuildCompletion(opId, uint32_t(code), msg)); UpdateWriteInterest(ctrl->fd); } RecordInboundStatus(peer, opId, code, msg); } Connection* TcpTransport::PeerCtrl(const EngineKey& peer) { - auto it = peers.find(peer); - if (it == peers.end() || !it->second.CtrlUp()) return nullptr; - return conns[it->second.ctrlFd].get(); + auto it = peers_.find(peer); + if (it == peers_.end() || !it->second.CtrlUp()) return nullptr; + auto cit = conns_.find(it->second.ctrlFd); + return (cit != conns_.end()) ? cit->second.get() : nullptr; } void TcpTransport::FinalizeInboundWriteSetup(const EngineKey& peer, TransferUniqueId opId, InboundWriteState& ws) { if (!ws.discard && ws.dst.loc == MemoryLocationType::GPU) { - const uint64_t total = SumLens(ws.dstSegs); - ws.pinned = staging.Acquire(static_cast(total)); + ws.pinned = staging_.Acquire(static_cast(SumLens(ws.dstSegs))); if (!ws.pinned) ws.discard = true; } - inboundWrites[peer][opId] = ws; - SetupInboundWriteWorkerTarget(peer, opId, ws); - TryConsumeEarlyWriteLanes(peer, opId); -} + inboundWrites_[peer][opId] = ws; -void TcpTransport::SetupInboundWriteWorkerTarget(const EngineKey& peer, TransferUniqueId opId, - const InboundWriteState& ws) { + // Set up worker recv targets WorkerRecvTarget target; target.lanesTotal = ws.lanesTotal; target.totalLen = SumLens(ws.dstSegs); @@ -1242,137 +1564,23 @@ void TcpTransport::SetupInboundWriteWorkerTarget(const EngineKey& peer, Transfer target.toGpu = true; target.pinned = ws.pinned; } else if (!ws.discard) { - target.toGpu = false; target.cpuBase = reinterpret_cast(ws.dst.data); target.segs = ws.dstSegs; } RegisterRecvTargetWithWorkers(peer, opId, target); -} -void TcpTransport::HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { - LinearReqView req; - if (!ParseLinearReq(body, len, &req)) { - MORI_IO_WARN("TCP: malformed WRITE_REQ"); - return; - } - HandleWriteReqSegments(peer, req.opId, req.memId, {{req.remoteOff, req.size}}, req.lanesTotal); -} - -void TcpTransport::HandleWriteReqSegments(const EngineKey& peer, TransferUniqueId opId, - MemoryUniqueId memId, std::vector segs, - uint8_t lanesTotal) { - auto memOpt = LookupLocalMem(memId); - InboundWriteState ws; - ws.peer = peer; - ws.id = opId; - ws.lanesTotal = lanesTotal; - ws.discard = true; - if (memOpt.has_value()) { - ws.dst = memOpt.value(); - if (SegmentsInRange(segs, ws.dst.size)) { - ws.discard = false; - ws.dstSegs = std::move(segs); - } - } - FinalizeInboundWriteSetup(peer, opId, ws); -} - -void TcpTransport::HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len) { - BatchReqView req; - if (!ParseBatchReq(body, len, &req)) { - MORI_IO_WARN("TCP: malformed BATCH_WRITE_REQ"); - return; - } - HandleWriteReqSegments(peer, req.opId, req.memId, std::move(req.segs), req.lanesTotal); -} - -void TcpTransport::HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { - LinearReqView req; - if (!ParseLinearReq(body, len, &req)) { - MORI_IO_WARN("TCP: malformed READ_REQ"); - return; - } - HandleReadReqSegments(peer, req.opId, req.memId, {{req.remoteOff, req.size}}, req.lanesTotal, - false); -} - -void TcpTransport::HandleReadReqSegments(const EngineKey& peer, TransferUniqueId opId, - MemoryUniqueId memId, std::vector segs, - uint8_t lanesTotal, bool batchReq) { - const StatusCode badRangeCode = - batchReq ? StatusCode::ERR_INVALID_ARGS : StatusCode::ERR_NOT_FOUND; - const char* badRangeMsg = - batchReq ? "TCP: batch read out of range" : "TCP: remote mem not found/out of range"; - const char* notFoundMsg = - batchReq ? "TCP: remote mem not found" : "TCP: remote mem not found/out of range"; - auto memOpt = LookupLocalMem(memId); - if (!memOpt.has_value()) { - SendCompletionAndRecord(peer, opId, StatusCode::ERR_NOT_FOUND, notFoundMsg); - return; - } - MemoryDesc src = memOpt.value(); - if (!SegmentsInRange(segs, src.size)) { - SendCompletionAndRecord(peer, opId, badRangeCode, badRangeMsg); - return; - } - QueueDataSendForRead(peer, opId, src, segs, lanesTotal); -} - -void TcpTransport::HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len) { - BatchReqView req; - if (!ParseBatchReq(body, len, &req)) { - MORI_IO_WARN("TCP: malformed BATCH_READ_REQ"); - return; - } - HandleReadReqSegments(peer, req.opId, req.memId, std::move(req.segs), req.lanesTotal, true); -} - -void TcpTransport::HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len) { - CompletionView msg; - if (!ParseCompletion(body, len, &msg)) { - MORI_IO_WARN("TCP: malformed COMPLETION"); - return; - } - - auto it = pendingOutbound.find(msg.opId); - if (it == pendingOutbound.end()) return; - OutboundOpState& st = *it->second; - st.completionReceived = true; - st.completionCode = static_cast(msg.statusCode); - st.completionMsg = std::move(msg.msg); - - if (st.completionCode != StatusCode::SUCCESS) { - RemoveRecvTargetFromWorkers(st.peer, msg.opId); - st.status->Update(st.completionCode, st.completionMsg); - pendingOutbound.erase(it); - return; - } - MaybeCompleteOutbound(st); -} - -void TcpTransport::MaybeCompleteOutbound(OutboundOpState& st) { - if (!st.completionReceived) return; - if (st.isRead) { - const uint8_t allMask = LanesAllMask(st.lanesTotal); - if (st.lanesDoneMask != allMask) return; - if (st.rxBytes != st.expectedRxBytes) return; - if (st.gpuCopyPending) return; - } - RemoveRecvTargetFromWorkers(st.peer, st.id); - st.status->Update(StatusCode::SUCCESS, ""); - pendingOutbound.erase(st.id); + TryConsumeEarlyWriteLanes(peer, opId); } void TcpTransport::MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId) { - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt == inboundWrites.end()) return; - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt == iwPeerIt->second.end()) return; + auto iwIt = inboundWrites_.find(peer); + if (iwIt == inboundWrites_.end()) return; + auto wsIt = iwIt->second.find(opId); + if (wsIt == iwIt->second.end()) return; InboundWriteState& ws = wsIt->second; ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - const uint8_t allMask = LanesAllMask(ws.lanesTotal); - if ((ws.lanesDoneMask & allMask) != allMask) return; + if ((ws.lanesDoneMask & LanesAllMask(ws.lanesTotal)) != LanesAllMask(ws.lanesTotal)) return; RemoveRecvTargetFromWorkers(peer, opId); @@ -1381,92 +1589,102 @@ void TcpTransport::MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniq } else if (ws.dst.loc == MemoryLocationType::GPU) { if (!ws.pinned) { SendCompletionAndRecord(peer, opId, StatusCode::ERR_BAD_STATE, - "TCP: missing pinned staging (write)"); + "TCP: missing staging (write)"); } else { auto pinnedRef = ws.pinned; - bool ok = ScheduleGpuHtoD(ws.dst.deviceId, ws.dst, ws.dstSegs, pinnedRef, + bool ok = ScheduleGpuCopy(ws.dst.deviceId, true, ws.dst, ws.dstSegs, pinnedRef, [this, peer, opId, pinnedRef]() { SendCompletionAndRecord(peer, opId, StatusCode::SUCCESS, ""); }); - if (!ok) { - SendCompletionAndRecord(peer, opId, StatusCode::ERR_BAD_STATE, - "TCP: failed HIP stream/event"); - } + if (!ok) + SendCompletionAndRecord(peer, opId, StatusCode::ERR_BAD_STATE, "TCP: HIP copy failed"); } } else { SendCompletionAndRecord(peer, opId, StatusCode::SUCCESS, ""); } - iwPeerIt->second.erase(wsIt); - if (iwPeerIt->second.empty()) inboundWrites.erase(iwPeerIt); - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt != earlyWrites.end()) { - ewPeerIt->second.erase(opId); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + iwIt->second.erase(wsIt); + if (iwIt->second.empty()) inboundWrites_.erase(iwIt); + auto ewIt = earlyWrites_.find(peer); + if (ewIt != earlyWrites_.end()) { + ewIt->second.erase(opId); + if (ewIt->second.empty()) earlyWrites_.erase(ewIt); } } void TcpTransport::TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId) { - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt == inboundWrites.end()) return; - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt == iwPeerIt->second.end()) return; + auto iwIt = inboundWrites_.find(peer); + if (iwIt == inboundWrites_.end()) return; + auto wsIt = iwIt->second.find(opId); + if (wsIt == iwIt->second.end()) return; InboundWriteState& ws = wsIt->second; ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - auto ewPeerIt = earlyWrites.find(peer); - if (ewPeerIt == earlyWrites.end()) return; - auto ewIt = ewPeerIt->second.find(opId); - if (ewIt == ewPeerIt->second.end()) return; + auto ewIt = earlyWrites_.find(peer); + if (ewIt == earlyWrites_.end()) return; + auto elIt = ewIt->second.find(opId); + if (elIt == ewIt->second.end()) return; - EarlyWriteState& early = ewIt->second; + EarlyWriteState& early = elIt->second; const uint64_t total = SumLens(ws.dstSegs); uint8_t* dstBase = reinterpret_cast(ws.dst.data); for (auto it = early.lanes.begin(); it != early.lanes.end();) { - const uint8_t lane = it->first; - EarlyWriteLaneState& laneState = it->second; - if (!laneState.complete) { + uint8_t lane = it->first; + EarlyWriteLaneState& ls = it->second; + if (!ls.complete) { ++it; continue; } if (lane >= ws.lanesTotal) ws.discard = true; - const LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); - if (span.len != laneState.payloadLen) ws.discard = true; + LaneSpan span = ComputeLaneSpan(total, ws.lanesTotal, lane); + if (span.len != ls.payloadLen) ws.discard = true; - if (!ws.discard && laneState.pinned) { - uint8_t* src = reinterpret_cast(laneState.pinned->ptr); + if (!ws.discard && ls.pinned) { + uint8_t* src = reinterpret_cast(ls.pinned->ptr); if (ws.dst.loc == MemoryLocationType::GPU) { if (!ws.pinned) { - ws.pinned = staging.Acquire(static_cast(total)); + ws.pinned = staging_.Acquire(size_t(total)); if (!ws.pinned) ws.discard = true; } - if (!ws.discard && ws.pinned) { - std::memcpy(reinterpret_cast(ws.pinned->ptr) + span.off, src, - static_cast(span.len)); - } + if (!ws.discard && ws.pinned) + std::memcpy(reinterpret_cast(ws.pinned->ptr) + span.off, src, size_t(span.len)); } else { - const auto segs = SliceSegments(ws.dstSegs, span.off, span.len); + auto segs = SliceSegments(ws.dstSegs, span.off, span.len); uint64_t copied = 0; - for (const auto& s : segs) { - std::memcpy(dstBase + s.off, src + copied, static_cast(s.len)); + for (auto& s : segs) { + std::memcpy(dstBase + s.off, src + copied, size_t(s.len)); copied += s.len; } } } - - if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + if (lane < 8) ws.lanesDoneMask |= uint8_t(1U << lane); it = early.lanes.erase(it); } if (early.lanes.empty()) { - ewPeerIt->second.erase(ewIt); - if (ewPeerIt->second.empty()) earlyWrites.erase(ewPeerIt); + ewIt->second.erase(elIt); + if (ewIt->second.empty()) earlyWrites_.erase(ewIt); } MaybeFinalizeInboundWrite(peer, opId); } +void TcpTransport::MaybeCompleteOutbound(OutboundOpState& st) { + if (!st.completionReceived) return; + if (st.isRead) { + uint8_t allMask = LanesAllMask(st.lanesTotal); + if (st.lanesDoneMask != allMask || st.rxBytes != st.expectedRxBytes || st.gpuCopyPending) + return; + } + RemoveRecvTargetFromWorkers(st.peer, st.id); + st.status->Update(StatusCode::SUCCESS, ""); + pendingOutbound_.erase(st.id); +} + +// --------------------------------------------------------------------------- +// Worker event processing +// --------------------------------------------------------------------------- void TcpTransport::ProcessEventsFrom(DataConnectionWorker* worker) { std::deque events; worker->DrainEvents(events); @@ -1482,7 +1700,7 @@ void TcpTransport::ProcessEventsFrom(DataConnectionWorker* worker) { if (ev.callback) ev.callback(); break; case WorkerEventType::CONN_ERROR: - MORI_IO_WARN("TCP: worker error for peer {}: {}", ev.peerKey, ev.errorMsg); + MORI_IO_WARN("TCP: worker error peer {}: {}", ev.peerKey, ev.errorMsg); ClosePeerByKey(ev.peerKey, ev.errorMsg); break; } @@ -1490,165 +1708,157 @@ void TcpTransport::ProcessEventsFrom(DataConnectionWorker* worker) { } void TcpTransport::ProcessWorkerEvents() { - std::vector notifyFds; - notifyFds.reserve(workerNotifyMap.size()); - for (const auto& kv : workerNotifyMap) notifyFds.push_back(kv.first); - for (int notifyFd : notifyFds) { - auto it = workerNotifyMap.find(notifyFd); - if (it == workerNotifyMap.end()) continue; - ProcessEventsFrom(it->second); + // Snapshot notify fds to avoid invalidation during iteration + std::vector fds; + fds.reserve(workerNotifyMap_.size()); + for (auto& kv : workerNotifyMap_) fds.push_back(kv.first); + for (int nfd : fds) { + auto it = workerNotifyMap_.find(nfd); + if (it != workerNotifyMap_.end()) ProcessEventsFrom(it->second); } } void TcpTransport::HandleWorkerRecvDone(const WorkerEvent& ev) { const EngineKey& peer = ev.peerKey; const TransferUniqueId opId = ev.opId; - const uint8_t lane = ev.lane; - const uint64_t laneLen = ev.laneLen; - - auto iwPeerIt = inboundWrites.find(peer); - if (iwPeerIt != inboundWrites.end()) { - auto wsIt = iwPeerIt->second.find(opId); - if (wsIt != iwPeerIt->second.end()) { - InboundWriteState& ws = wsIt->second; - ws.lanesTotal = ClampLanesTotal(ws.lanesTotal); - if (lane < 8) ws.lanesDoneMask |= static_cast(1U << lane); + + // Check inbound writes first + auto iwIt = inboundWrites_.find(peer); + if (iwIt != inboundWrites_.end()) { + auto wsIt = iwIt->second.find(opId); + if (wsIt != iwIt->second.end()) { + wsIt->second.lanesTotal = ClampLanesTotal(wsIt->second.lanesTotal); + if (ev.lane < 8) wsIt->second.lanesDoneMask |= uint8_t(1U << ev.lane); MaybeFinalizeInboundWrite(peer, opId); return; } } - auto obIt = pendingOutbound.find(opId); - if (obIt == pendingOutbound.end()) return; + // Check outbound reads + auto obIt = pendingOutbound_.find(opId); + if (obIt == pendingOutbound_.end()) return; OutboundOpState& st = *obIt->second; st.lanesTotal = ClampLanesTotal(st.lanesTotal); - const uint8_t bit = static_cast(1U << lane); - if ((st.lanesDoneMask & bit) == 0) { + uint8_t bit = uint8_t(1U << ev.lane); + if (!(st.lanesDoneMask & bit)) { st.lanesDoneMask |= bit; - st.rxBytes += laneLen; + st.rxBytes += ev.laneLen; } if (st.local.loc == MemoryLocationType::GPU) { - const uint8_t allMask = LanesAllMask(st.lanesTotal); + uint8_t allMask = LanesAllMask(st.lanesTotal); if ((st.lanesDoneMask & allMask) != allMask) { MaybeCompleteOutbound(st); return; } if (st.gpuCopyPending) return; if (!st.pinned) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: missing pinned staging (read)"); + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: missing staging (read)"); RemoveRecvTargetFromWorkers(st.peer, opId); - pendingOutbound.erase(obIt); + pendingOutbound_.erase(obIt); return; } - st.gpuCopyPending = true; auto pinnedRef = st.pinned; - bool ok = ScheduleGpuHtoD(st.local.deviceId, st.local, st.localSegs, pinnedRef, + bool ok = ScheduleGpuCopy(st.local.deviceId, true, st.local, st.localSegs, pinnedRef, [this, opId, pinnedRef]() { - auto it2 = pendingOutbound.find(opId); - if (it2 == pendingOutbound.end()) return; + auto it2 = pendingOutbound_.find(opId); + if (it2 == pendingOutbound_.end()) return; it2->second->gpuCopyPending = false; MaybeCompleteOutbound(*it2->second); }); if (!ok) { - st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: failed HIP stream/event (read)"); + st.status->Update(StatusCode::ERR_BAD_STATE, "TCP: HIP copy failed (read)"); RemoveRecvTargetFromWorkers(st.peer, opId); - pendingOutbound.erase(obIt); + pendingOutbound_.erase(obIt); } return; } - MaybeCompleteOutbound(st); } void TcpTransport::HandleWorkerEarlyData(const WorkerEvent& ev) { - const EngineKey& peer = ev.peerKey; - const TransferUniqueId opId = ev.opId; - const uint8_t lane = ev.lane; - - auto& perPeer = earlyWrites[peer]; - EarlyWriteState& early = perPeer[opId]; - if (early.lanes.find(lane) != early.lanes.end()) { - MORI_IO_WARN("TCP: duplicate early data for op {} lane {} from peer {}", opId, (uint32_t)lane, - peer); + auto& perPeer = earlyWrites_[ev.peerKey]; + auto& early = perPeer[ev.opId]; + if (early.lanes.count(ev.lane)) { + MORI_IO_WARN("TCP: duplicate early data op {} lane {} peer {}", ev.opId, uint32_t(ev.lane), + ev.peerKey); return; } - early.lanes.emplace(lane, EarlyWriteLaneState{ev.laneLen, ev.earlyBuf, true}); - TryConsumeEarlyWriteLanes(peer, opId); + early.lanes.emplace(ev.lane, EarlyWriteLaneState{ev.laneLen, ev.earlyBuf, true}); + TryConsumeEarlyWriteLanes(ev.peerKey, ev.opId); } void TcpTransport::ScanTimeouts() { - if (config.opTimeoutMs <= 0) return; - const auto now = Clock::now(); - const auto timeout = std::chrono::milliseconds(config.opTimeoutMs); - for (auto it = pendingOutbound.begin(); it != pendingOutbound.end();) { + if (config_.opTimeoutMs <= 0) return; + auto now = Clock::now(); + auto timeout = std::chrono::milliseconds(config_.opTimeoutMs); + for (auto it = pendingOutbound_.begin(); it != pendingOutbound_.end();) { if ((now - it->second->startTs) > timeout) { RemoveRecvTargetFromWorkers(it->second->peer, it->first); it->second->status->Update(StatusCode::ERR_BAD_STATE, "TCP: op timeout"); - it = pendingOutbound.erase(it); - } else { + it = pendingOutbound_.erase(it); + } else ++it; - } } } +// --------------------------------------------------------------------------- +// Main I/O loop +// --------------------------------------------------------------------------- void TcpTransport::IoLoop() { constexpr int kMaxEvents = 128; epoll_event events[kMaxEvents]; - while (running.load()) { + while (running_.load()) { PollGpuTasks(); ProcessWorkerEvents(); ScanTimeouts(); - const bool hasActive = - !gpuTasks.empty() || !pendingOutbound.empty() || !workerNotifyMap.empty(); - const int timeoutMs = hasActive ? 0 : 2; - int nfds = epoll_wait(epfd, events, kMaxEvents, timeoutMs); + bool hasActive = !gpuTasks_.empty() || !pendingOutbound_.empty() || !workerNotifyMap_.empty(); + int nfds = epoll_wait(epfd_, events, kMaxEvents, hasActive ? 0 : 2); if (nfds < 0) { if (errno == EINTR) continue; - MORI_IO_ERROR("TCP: epoll_wait failed: {}", strerror(errno)); + MORI_IO_ERROR("TCP: epoll_wait: {}", strerror(errno)); break; } for (int i = 0; i < nfds; ++i) { int fd = events[i].data.fd; - uint32_t evMask = events[i].events; + uint32_t ev = events[i].events; - if (fd == listenFd) { + if (fd == listenFd_) { AcceptNew(); continue; } - if (fd == wakeFd) { + if (fd == wakeFd_) { DrainWakeFd(); continue; } - auto wnit = workerNotifyMap.find(fd); - if (wnit != workerNotifyMap.end()) { + auto wnit = workerNotifyMap_.find(fd); + if (wnit != workerNotifyMap_.end()) { ProcessEventsFrom(wnit->second); continue; } - Connection* c = nullptr; - auto it = conns.find(fd); - if (it != conns.end()) c = it->second.get(); + auto cit = conns_.find(fd); + if (cit == conns_.end()) continue; + Connection* c = cit->second.get(); if (!c) continue; - if (evMask & (EPOLLERR | EPOLLHUP)) { + if (ev & (EPOLLERR | EPOLLHUP)) { ClosePeerByFd(fd); continue; } - - if (evMask & EPOLLIN) { + if (ev & EPOLLIN) { HandleCtrlReadable(c); - auto it2 = conns.find(fd); - if (it2 == conns.end()) continue; - c = it2->second.get(); + cit = conns_.find(fd); + if (cit == conns_.end()) continue; + c = cit->second.get(); if (!c) continue; } - if (evMask & EPOLLOUT) HandleConnWritable(c); + if (ev & EPOLLOUT) HandleConnWritable(c); } } } diff --git a/src/io/tcp/transport.hpp b/src/io/tcp/transport.hpp index 74ea99a44..c56e0e2e7 100644 --- a/src/io/tcp/transport.hpp +++ b/src/io/tcp/transport.hpp @@ -19,34 +19,584 @@ // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. + #pragma once +#include +#include +#include +#include +#include +#include +#include +#include #include #include +#include +#include +#include +#include +#include #include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include +#include #include "mori/application/utils/check.hpp" +#include "mori/io/backend.hpp" +#include "mori/io/common.hpp" #include "mori/io/engine.hpp" -#include "src/io/tcp/data_worker.hpp" -#include "src/io/tcp/tcp_types.hpp" +#include "mori/io/logging.hpp" #include "src/io/xgmi/hip_resource_pool.hpp" namespace mori { namespace io { +// --------------------------------------------------------------------------- +// Socket utilities +// --------------------------------------------------------------------------- +inline bool IsWouldBlock(int err) { return err == EAGAIN || err == EWOULDBLOCK; } + +inline int SetNonBlocking(int fd) { + int f = fcntl(fd, F_GETFL, 0); + return (f < 0 || fcntl(fd, F_SETFL, f | O_NONBLOCK) < 0) ? -1 : 0; +} + +inline void SetSockOpt(int fd, int level, int opt, const void* val, socklen_t len, const char* nm) { + if (setsockopt(fd, level, opt, val, len) != 0) + MORI_IO_WARN("TCP: setsockopt {} failed: {}", nm, strerror(errno)); +} + +inline void ConfigureSocketCommon(int fd, const TcpBackendConfig& cfg) { + if (!cfg.enableKeepalive) return; + int on = 1; + SetSockOpt(fd, SOL_SOCKET, SO_KEEPALIVE, &on, sizeof(on), "SO_KEEPALIVE"); + SetSockOpt(fd, IPPROTO_TCP, TCP_KEEPIDLE, &cfg.keepaliveIdleSec, sizeof(cfg.keepaliveIdleSec), + "TCP_KEEPIDLE"); + SetSockOpt(fd, IPPROTO_TCP, TCP_KEEPINTVL, &cfg.keepaliveIntvlSec, sizeof(cfg.keepaliveIntvlSec), + "TCP_KEEPINTVL"); + SetSockOpt(fd, IPPROTO_TCP, TCP_KEEPCNT, &cfg.keepaliveCnt, sizeof(cfg.keepaliveCnt), + "TCP_KEEPCNT"); +} + +inline void ConfigureCtrlSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + if (cfg.enableCtrlNodelay) { + int on = 1; + SetSockOpt(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(ctrl)"); + } +} + +inline void ConfigureDataSocket(int fd, const TcpBackendConfig& cfg) { + ConfigureSocketCommon(fd, cfg); + int on = 1; + SetSockOpt(fd, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on), "TCP_NODELAY(data)"); + if (cfg.sockSndbufBytes > 0) + SetSockOpt(fd, SOL_SOCKET, SO_SNDBUF, &cfg.sockSndbufBytes, sizeof(cfg.sockSndbufBytes), + "SO_SNDBUF"); + if (cfg.sockRcvbufBytes > 0) + SetSockOpt(fd, SOL_SOCKET, SO_RCVBUF, &cfg.sockRcvbufBytes, sizeof(cfg.sockRcvbufBytes), + "SO_RCVBUF"); +} + +inline std::optional ParseIpv4(const std::string& host, uint16_t port) { + sockaddr_in a{}; + a.sin_family = AF_INET; + a.sin_port = htons(port); + if (inet_pton(AF_INET, host.c_str(), &a.sin_addr) != 1) return std::nullopt; + return a; +} + +inline uint16_t GetBoundPort(int fd) { + sockaddr_in a{}; + socklen_t l = sizeof(a); + return (getsockname(fd, reinterpret_cast(&a), &l) == 0) ? ntohs(a.sin_port) : 0; +} + +// --------------------------------------------------------------------------- +// Wire protocol constants and helpers +// --------------------------------------------------------------------------- +namespace tcp { + +constexpr uint32_t kCtrlMagic = 0x4D544330; // "MTC0" +constexpr uint32_t kDataMagic = 0x4D544430; // "MTD0" +constexpr uint16_t kProtoVersion = 2; + +enum class Channel : uint8_t { CTRL = 1, DATA = 2 }; + +enum class CtrlMsgType : uint8_t { + HELLO = 1, + WRITE_REQ = 2, + READ_REQ = 3, + BATCH_WRITE_REQ = 4, + BATCH_READ_REQ = 5, + COMPLETION = 6, +}; + +constexpr size_t kCtrlHeaderSize = 12; +constexpr size_t kDataHeaderSize = 24; + +struct CtrlHeaderView { + CtrlMsgType type{CtrlMsgType::HELLO}; + uint32_t bodyLen{0}; +}; +struct DataHeaderView { + uint16_t flags{0}; + uint64_t opId{0}; + uint64_t payloadLen{0}; +}; + +// Compact wire writer (appends big-endian values) +struct WireWriter { + std::vector buf; + void reserve(size_t n) { buf.reserve(n); } + void u8(uint8_t v) { buf.push_back(v); } + void u16(uint16_t v) { + v = htons(v); + auto* p = reinterpret_cast(&v); + buf.insert(buf.end(), p, p + 2); + } + void u32(uint32_t v) { + v = htonl(v); + auto* p = reinterpret_cast(&v); + buf.insert(buf.end(), p, p + 4); + } + void u64(uint64_t v) { + v = htobe64(v); + auto* p = reinterpret_cast(&v); + buf.insert(buf.end(), p, p + 8); + } + void bytes(const void* d, size_t n) { + auto* p = static_cast(d); + buf.insert(buf.end(), p, p + n); + } +}; + +// Compact wire reader (reads big-endian values from buffer) +struct WireReader { + const uint8_t* data; + size_t len; + size_t off{0}; + bool u8(uint8_t* o) { + if (off + 1 > len) return false; + *o = data[off++]; + return true; + } + bool u16(uint16_t* o) { + if (off + 2 > len) return false; + uint16_t v; + memcpy(&v, data + off, 2); + *o = ntohs(v); + off += 2; + return true; + } + bool u32(uint32_t* o) { + if (off + 4 > len) return false; + uint32_t v; + memcpy(&v, data + off, 4); + *o = ntohl(v); + off += 4; + return true; + } + bool u64(uint64_t* o) { + if (off + 8 > len) return false; + uint64_t v; + memcpy(&v, data + off, 8); + *o = be64toh(v); + off += 8; + return true; + } +}; + +inline bool TryParseCtrlHeader(const uint8_t* buf, size_t len, CtrlHeaderView* h) { + if (len < kCtrlHeaderSize) return false; + WireReader r{buf, len}; + uint32_t magic; + uint16_t ver; + uint8_t type, reserved; + if (!r.u32(&magic) || !r.u16(&ver) || !r.u8(&type) || !r.u8(&reserved) || !r.u32(&h->bodyLen)) + return false; + if (magic != kCtrlMagic || ver != kProtoVersion) return false; + h->type = static_cast(type); + return true; +} + +inline bool TryParseDataHeader(const uint8_t* buf, size_t len, DataHeaderView* h) { + if (len < kDataHeaderSize) return false; + WireReader r{buf, len}; + uint32_t magic; + uint16_t ver; + if (!r.u32(&magic) || !r.u16(&ver) || !r.u16(&h->flags) || !r.u64(&h->opId) || + !r.u64(&h->payloadLen)) + return false; + return (magic == kDataMagic && ver == kProtoVersion); +} + +// Build a ctrl frame: header(12B) + body +inline std::vector BuildCtrlFrame(CtrlMsgType type, + const std::function& writeBody) { + WireWriter body; + writeBody(body); + WireWriter frame; + frame.reserve(kCtrlHeaderSize + body.buf.size()); + frame.u32(kCtrlMagic); + frame.u16(kProtoVersion); + frame.u8(static_cast(type)); + frame.u8(0); + frame.u32(static_cast(body.buf.size())); + frame.bytes(body.buf.data(), body.buf.size()); + return std::move(frame.buf); +} + +inline std::vector BuildHello(Channel ch, const EngineKey& key) { + return BuildCtrlFrame(CtrlMsgType::HELLO, [&](WireWriter& w) { + w.u8(static_cast(ch)); + w.u32(static_cast(key.size())); + w.bytes(key.data(), key.size()); + }); +} + +// Unified request builder for WRITE_REQ, READ_REQ (single segment) +inline std::vector BuildLinearReq(CtrlMsgType type, uint64_t opId, uint32_t memId, + uint64_t off, uint64_t size, uint8_t lanes) { + return BuildCtrlFrame(type, [&](WireWriter& w) { + w.u64(opId); + w.u32(memId); + w.u64(off); + w.u64(size); + w.u8(lanes); + }); +} + +// Unified request builder for BATCH_WRITE_REQ, BATCH_READ_REQ +inline std::vector BuildBatchReq(CtrlMsgType type, uint64_t opId, uint32_t memId, + const std::vector& offs, + const std::vector& sizes, uint8_t lanes) { + return BuildCtrlFrame(type, [&](WireWriter& w) { + w.u64(opId); + w.u32(memId); + w.u32(static_cast(offs.size())); + for (size_t i = 0; i < offs.size(); ++i) { + w.u64(offs[i]); + w.u64(sizes[i]); + } + w.u8(lanes); + }); +} + +inline std::vector BuildCompletion(uint64_t opId, uint32_t code, const std::string& msg) { + return BuildCtrlFrame(CtrlMsgType::COMPLETION, [&](WireWriter& w) { + w.u64(opId); + w.u32(code); + w.u32(static_cast(msg.size())); + w.bytes(msg.data(), msg.size()); + }); +} + +inline std::vector BuildDataHeader(uint64_t opId, uint64_t payloadLen, uint16_t flags) { + WireWriter w; + w.reserve(kDataHeaderSize); + w.u32(kDataMagic); + w.u16(kProtoVersion); + w.u16(flags); + w.u64(opId); + w.u64(payloadLen); + return std::move(w.buf); +} + +} // namespace tcp + +// --------------------------------------------------------------------------- +// Segment and lane helpers +// --------------------------------------------------------------------------- +struct Segment { + uint64_t off{0}; + uint64_t len{0}; +}; + +constexpr uint8_t kLaneBits = 3; +constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; + +inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { + return (userOpId << kLaneBits) | lane; +} +inline uint64_t ToUserOpId(uint64_t wireOpId) { return wireOpId >> kLaneBits; } + +struct LaneSpan { + uint64_t off{0}; + uint64_t len{0}; +}; + +inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanes, uint8_t lane) { + if (lanes <= 1) return {0, total}; + uint64_t base = total / lanes, rem = total % lanes; + return {uint64_t(lane) * base + std::min(lane, rem), base + (lane < rem ? 1 : 0)}; +} + +inline uint8_t LanesAllMask(uint8_t n) { + return (n >= (1U << kLaneBits)) ? 0xFF : uint8_t((1U << n) - 1); +} +inline uint8_t ClampLanesTotal(uint8_t n) { + return n == 0 ? 1 : std::min(n, 1U << kLaneBits); +} + +inline uint64_t SumLens(const std::vector& segs) { + uint64_t t = 0; + for (auto& s : segs) t += s.len; + return t; +} + +inline std::vector SliceSegments(const std::vector& segs, uint64_t start, + uint64_t len) { + std::vector out; + if (len == 0) return out; + uint64_t skip = start, remaining = len; + for (auto& s : segs) { + if (remaining == 0) break; + if (skip >= s.len) { + skip -= s.len; + continue; + } + uint64_t take = std::min(s.len - skip, remaining); + out.push_back({s.off + skip, take}); + remaining -= take; + skip = 0; + } + return out; +} + +inline bool IsSingleContiguousSpan(const std::vector& segs, uint64_t* outOff, + uint64_t* outLen) { + if (segs.empty()) return false; + uint64_t off = segs[0].off, end = off + segs[0].len; + for (size_t i = 1; i < segs.size(); ++i) { + if (segs[i].off != end) return false; + end += segs[i].len; + } + *outOff = off; + *outLen = end - off; + return true; +} + +inline bool SegmentsInRange(const std::vector& segs, uint64_t memSize) { + for (auto& s : segs) + if (s.off + s.len > memSize) return false; + return true; +} + +// --------------------------------------------------------------------------- +// Pinned staging pool (HIP host memory) +// --------------------------------------------------------------------------- +struct PinnedBuf { + void* ptr{nullptr}; + size_t cap{0}; +}; + +class PinnedStagingPool { + public: + PinnedStagingPool() = default; + ~PinnedStagingPool() { Clear(); } + PinnedStagingPool(const PinnedStagingPool&) = delete; + PinnedStagingPool& operator=(const PinnedStagingPool&) = delete; + + std::shared_ptr Acquire(size_t size); + void Clear(); + + private: + void Release(PinnedBuf* b); + static size_t RoundUp(size_t v) { + size_t p = 1; + while (p < v) p <<= 1; + return p; + } + + std::mutex mu_; + std::unordered_map> free_; +}; + +// --------------------------------------------------------------------------- +// Send / Connection / Peer state +// --------------------------------------------------------------------------- +using Clock = std::chrono::steady_clock; + +struct SendItem { + std::vector header; + std::vector iov; + size_t idx{0}, off{0}; + int flags{0}; + std::shared_ptr keepalive; + std::function onDone; + bool Done() const { return idx >= iov.size(); } + void Advance(size_t n); +}; + +struct Connection { + int fd{-1}; + bool isOutgoing{false}, connecting{false}, helloSent{false}, helloReceived{false}; + tcp::Channel ch{tcp::Channel::CTRL}; + EngineKey peerKey; + std::vector inbuf; + std::deque sendq; +}; + +class DataConnectionWorker; + +struct PeerLinks { + int ctrlFd{-1}; + std::vector dataFds; + std::vector workers; + int ctrlPending{0}, dataPending{0}; + bool CtrlUp() const { return ctrlFd >= 0; } + bool DataUp() const { return !dataFds.empty(); } +}; + +struct InboundStatusEntry { + StatusCode code{StatusCode::INIT}; + std::string msg; +}; + +struct OutboundOpState { + EngineKey peer; + TransferUniqueId id{0}; + bool isRead{false}; + TransferStatus* status{nullptr}; + MemoryDesc local{}, remote{}; + std::vector localSegs, remoteSegs; + uint64_t expectedRxBytes{0}, rxBytes{0}; + bool completionReceived{false}, gpuCopyPending{false}; + uint8_t lanesTotal{1}, lanesDoneMask{0}; + StatusCode completionCode{StatusCode::SUCCESS}; + std::string completionMsg; + std::shared_ptr pinned; + Clock::time_point startTs{Clock::now()}; +}; + +struct InboundWriteState { + EngineKey peer; + TransferUniqueId id{0}; + MemoryDesc dst{}; + std::vector dstSegs; + bool discard{false}; + uint8_t lanesTotal{1}, lanesDoneMask{0}; + std::shared_ptr pinned; +}; + +struct EarlyWriteLaneState { + uint64_t payloadLen{0}; + std::shared_ptr pinned; + bool complete{false}; +}; +struct EarlyWriteState { + std::unordered_map lanes; +}; + +// Worker ←→ IO thread communication +struct WorkerRecvTarget { + uint8_t lanesTotal{1}; + uint64_t totalLen{0}; + bool discard{false}, toGpu{false}; + void* cpuBase{nullptr}; + std::vector segs; + std::shared_ptr pinned; +}; + +enum class WorkerEventType : uint8_t { + RECV_DONE = 0, + EARLY_DATA = 1, + SEND_CALLBACK = 2, + CONN_ERROR = 3 +}; + +struct WorkerEvent { + WorkerEventType type{WorkerEventType::RECV_DONE}; + EngineKey peerKey; + TransferUniqueId opId{0}; + uint8_t lane{0}; + uint64_t laneLen{0}; + bool discarded{false}; + std::shared_ptr earlyBuf; + std::function callback; + std::string errorMsg; +}; + +struct GpuTask { + int deviceId{-1}; + hipEvent_t ev{nullptr}; + std::function onReady; +}; + +// --------------------------------------------------------------------------- +// DataConnectionWorker — runs one thread per data connection +// --------------------------------------------------------------------------- +class DataConnectionWorker { + public: + DataConnectionWorker(int fd, EngineKey peer, PinnedStagingPool* staging); + ~DataConnectionWorker(); + DataConnectionWorker(const DataConnectionWorker&) = delete; + DataConnectionWorker& operator=(const DataConnectionWorker&) = delete; + + void Start(); + void Stop(); + int NotifyFd() const { return notifyFd_; } + int Fd() const { return fd_; } + + void SubmitSend(SendItem item); + void RegisterRecvTarget(TransferUniqueId opId, const WorkerRecvTarget& target); + void RemoveRecvTarget(TransferUniqueId opId); + void DrainEvents(std::deque& out); + + private: + void WakeWorker(); + void NotifyMain(); + void PostEvent(WorkerEvent ev); + void Run(); + bool ProcessSend(); + bool ProcessRecv(); + bool RecvExact(uint8_t* dst, uint64_t len); + bool RecvIntoSegments(uint8_t* base, const std::vector& segs, uint64_t totalLen); + bool DiscardPayload(uint64_t len); + + int fd_; + EngineKey peerKey_; + PinnedStagingPool* staging_; + std::atomic running_{false}; + std::thread thread_; + int notifyFd_{-1}, wakeFd_{-1}; + + std::mutex sendMu_; + std::deque sendQ_; + + std::mutex targetMu_; + std::unordered_map recvTargets_; + + std::mutex eventMu_; + std::deque eventQ_; + + uint8_t hdrBuf_[tcp::kDataHeaderSize]{}; + size_t hdrGot_{0}; +}; + +// --------------------------------------------------------------------------- +// TcpTransport — main transport layer +// --------------------------------------------------------------------------- class TcpTransport { public: TcpTransport(EngineKey myKey, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg); ~TcpTransport(); - TcpTransport(const TcpTransport&) = delete; TcpTransport& operator=(const TcpTransport&) = delete; void Start(); void Shutdown(); - std::optional GetListenPort() const; void RegisterRemoteEngine(const EngineDesc& desc); @@ -60,76 +610,69 @@ class TcpTransport { void SubmitReadWrite(const MemoryDesc& local, size_t localOffset, const MemoryDesc& remote, size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id, bool isRead); - void SubmitBatchReadWrite(const MemoryDesc& local, const SizeVec& localOffsets, const MemoryDesc& remote, const SizeVec& remoteOffsets, const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, bool isRead); private: + // Operation submission void EnqueueOp(std::unique_ptr op); - void AddEpoll(int fd, bool wantRead, bool wantWrite); - void ModEpoll(int fd, bool wantRead, bool wantWrite); + // Epoll helpers + void AddEpoll(int fd, bool rd, bool wr); + void ModEpoll(int fd, bool rd, bool wr); void DelEpoll(int fd); void CloseConnInternal(Connection* c); - bool PreferOutgoingFor(const EngineKey& peerKey) const; + // Connection management void AssignConnToPeer(Connection* c); - void MaybeDispatchQueuedOps(const EngineKey& peerKey); - void EnsurePeerChannels(const EngineKey& peerKey); - void ConnectChannel(const EngineKey& peerKey, tcp::Channel ch); + void MaybeDispatchQueuedOps(const EngineKey& peer); + void EnsurePeerChannels(const EngineKey& peer); + void ConnectChannel(const EngineKey& peer, tcp::Channel ch); void QueueHello(int fd); void AcceptNew(); void DrainWakeFd(); - bool IsPeerReady(const EngineKey& peerKey); + bool IsPeerReady(const EngineKey& peer); - void RegisterRecvTargetWithWorkers(const EngineKey& peerKey, TransferUniqueId opId, + // Worker coordination + void RegisterRecvTargetWithWorkers(const EngineKey& peer, TransferUniqueId opId, const WorkerRecvTarget& target); - void RemoveRecvTargetFromWorkers(const EngineKey& peerKey, TransferUniqueId opId); + void RemoveRecvTargetFromWorkers(const EngineKey& peer, TransferUniqueId opId); + // Data transfer void DispatchOp(std::unique_ptr op); void QueueSend(int fd, std::vector bytes, std::function onDone = nullptr); - - void QueueSegmentSend(DataConnectionWorker* worker, uint64_t wireOpId, uint8_t* base, - const std::vector& segs, uint64_t totalLen, - std::function onDone = nullptr); - void QueueStripedCpuSend(const std::vector& workers, uint64_t opId, - uint8_t lanesTotal, uint8_t* base, uint64_t baseOff, uint64_t total, - std::function onLaneDone = nullptr); - void QueueGpuSend(const std::vector& workers, uint64_t opId, - uint8_t lanesTotal, const MemoryDesc& src, const std::vector& srcSegs, - std::function onLaneDone = nullptr); - - void QueueDataSendForWrite(const std::vector& workerList, - OutboundOpState& st); - void QueueDataSendForRead(const EngineKey& peer, uint64_t opId, const MemoryDesc& src, - const std::vector& srcSegs, uint8_t lanesTotal); - void QueueDataSendCommon(const std::vector& workerList, - const MemoryDesc& src, const std::vector& srcSegs, - uint64_t opId, uint8_t lanesTotal, - std::function onLaneDone = nullptr); - - bool ScheduleGpuDtoH(int deviceId, const MemoryDesc& src, const std::vector& srcSegs, - std::shared_ptr pinned, std::function onComplete); - bool ScheduleGpuHtoD(int deviceId, const MemoryDesc& dst, const std::vector& dstSegs, - std::shared_ptr pinned, std::function onComplete); - + void QueueDataSend(const std::vector& workers, const MemoryDesc& src, + const std::vector& srcSegs, uint64_t opId, uint8_t lanesTotal, + std::function onLaneDone = nullptr); + + // GPU memory transfers + bool ScheduleGpuCopy(int deviceId, bool toDevice, const MemoryDesc& mem, + const std::vector& segs, std::shared_ptr pinned, + std::function onComplete); void PollGpuTasks(); + + // Ctrl-connection I/O void UpdateWriteInterest(int fd); void HandleConnWritable(Connection* c); void FlushSend(Connection* c); + // Peer lifecycle void CloseAndRemoveFd(int fd); EngineKey FindPeerByFd(int fd); void ClosePeerByFd(int fd); void ClosePeerByKey(const EngineKey& peer, const std::string& reason); void FailPendingOpsForPeer(const EngineKey& peer, const std::string& msg); + // Ctrl message handling void HandleCtrlReadable(Connection* c); void HandleCtrlFrame(Connection* c, tcp::CtrlMsgType type, const uint8_t* body, size_t len); void HandleHello(Connection* c, const uint8_t* body, size_t len); + void HandleRequest(const EngineKey& peer, tcp::CtrlMsgType type, const uint8_t* body, size_t len); + void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len); + // Inbound / outbound state machines std::optional LookupLocalMem(MemoryUniqueId id); void RecordInboundStatus(const EngineKey& peer, TransferUniqueId id, StatusCode code, const std::string& msg); @@ -139,22 +682,11 @@ class TcpTransport { void FinalizeInboundWriteSetup(const EngineKey& peer, TransferUniqueId opId, InboundWriteState& ws); - void SetupInboundWriteWorkerTarget(const EngineKey& peer, TransferUniqueId opId, - const InboundWriteState& ws); - void HandleWriteReq(const EngineKey& peer, const uint8_t* body, size_t len); - void HandleBatchWriteReq(const EngineKey& peer, const uint8_t* body, size_t len); - void HandleWriteReqSegments(const EngineKey& peer, TransferUniqueId opId, MemoryUniqueId memId, - std::vector segs, uint8_t lanesTotal); - void HandleReadReq(const EngineKey& peer, const uint8_t* body, size_t len); - void HandleBatchReadReq(const EngineKey& peer, const uint8_t* body, size_t len); - void HandleReadReqSegments(const EngineKey& peer, TransferUniqueId opId, MemoryUniqueId memId, - std::vector segs, uint8_t lanesTotal, bool batchReq); - void HandleCompletion(const EngineKey& peer, const uint8_t* body, size_t len); - void MaybeFinalizeInboundWrite(const EngineKey& peer, TransferUniqueId opId); void TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniqueId opId); - void MaybeCompleteOutbound(OutboundOpState& st); + + // Worker event processing void ProcessEventsFrom(DataConnectionWorker* worker); void ProcessWorkerEvents(); void HandleWorkerRecvDone(const WorkerEvent& ev); @@ -164,46 +696,44 @@ class TcpTransport { void IoLoop(); private: - EngineKey myEngKey; - IOEngineConfig engConfig; - TcpBackendConfig config; + EngineKey myEngKey_; + IOEngineConfig engConfig_; + TcpBackendConfig config_; - int epfd{-1}; - int listenFd{-1}; - int wakeFd{-1}; - uint16_t listenPort{0}; + int epfd_{-1}, listenFd_{-1}, wakeFd_{-1}; + uint16_t listenPort_{0}; - std::atomic running{false}; - std::thread ioThread; + std::atomic running_{false}; + std::thread ioThread_; - std::mutex submitMu; - std::deque> submitQ; + std::mutex submitMu_; + std::deque> submitQ_; - std::mutex remoteMu; - std::unordered_map remoteEngines; + std::mutex remoteMu_; + std::unordered_map remoteEngines_; - std::mutex memMu; - std::unordered_map localMems; + std::mutex memMu_; + std::unordered_map localMems_; - std::mutex inboundMu; + std::mutex inboundMu_; std::unordered_map> - inboundStatus; + inboundStatus_; - std::unordered_map> conns; - std::unordered_map peers; - std::unordered_map>> waitingOps; - std::unordered_map> pendingOutbound; + std::unordered_map> conns_; + std::unordered_map peers_; + std::unordered_map>> waitingOps_; + std::unordered_map> pendingOutbound_; std::unordered_map> - inboundWrites; - std::unordered_map> earlyWrites; + inboundWrites_; + std::unordered_map> earlyWrites_; - std::unordered_map> dataWorkers; - std::unordered_map workerNotifyMap; + std::unordered_map> dataWorkers_; + std::unordered_map workerNotifyMap_; - PinnedStagingPool staging; - StreamPool streamPool{8}; - EventPool eventPool{64}; - std::deque gpuTasks; + PinnedStagingPool staging_; + StreamPool streamPool_{8}; + EventPool eventPool_{64}; + std::deque gpuTasks_; }; } // namespace io diff --git a/src/pybind/mori.cpp b/src/pybind/mori.cpp index 28b310917..72a40662a 100644 --- a/src/pybind/mori.cpp +++ b/src/pybind/mori.cpp @@ -759,14 +759,13 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("num_events", &mori::io::XgmiBackendConfig::numEvents); py::class_(m, "TcpBackendConfig") - .def(py::init(), - py::arg("num_io_threads") = 1, py::arg("sock_sndbuf_bytes") = 4 * 1024 * 1024, + .def(py::init(), + py::arg("sock_sndbuf_bytes") = 4 * 1024 * 1024, py::arg("sock_rcvbuf_bytes") = 4 * 1024 * 1024, py::arg("op_timeout_ms") = 30 * 1000, py::arg("enable_keepalive") = true, py::arg("keepalive_idle_sec") = 30, py::arg("keepalive_intvl_sec") = 10, py::arg("keepalive_cnt") = 3, - py::arg("enable_ctrl_nodelay") = true, py::arg("enable_data_nodelay") = false, - py::arg("num_data_conns") = 4, py::arg("striping_threshold_bytes") = 256 * 1024) - .def_readwrite("num_io_threads", &mori::io::TcpBackendConfig::numIoThreads) + py::arg("enable_ctrl_nodelay") = true, py::arg("num_data_conns") = 4, + py::arg("striping_threshold_bytes") = 256 * 1024) .def_readwrite("sock_sndbuf_bytes", &mori::io::TcpBackendConfig::sockSndbufBytes) .def_readwrite("sock_rcvbuf_bytes", &mori::io::TcpBackendConfig::sockRcvbufBytes) .def_readwrite("op_timeout_ms", &mori::io::TcpBackendConfig::opTimeoutMs) @@ -775,9 +774,9 @@ void RegisterMoriIo(pybind11::module_& m) { .def_readwrite("keepalive_intvl_sec", &mori::io::TcpBackendConfig::keepaliveIntvlSec) .def_readwrite("keepalive_cnt", &mori::io::TcpBackendConfig::keepaliveCnt) .def_readwrite("enable_ctrl_nodelay", &mori::io::TcpBackendConfig::enableCtrlNodelay) - .def_readwrite("enable_data_nodelay", &mori::io::TcpBackendConfig::enableDataNodelay) .def_readwrite("num_data_conns", &mori::io::TcpBackendConfig::numDataConns) - .def_readwrite("striping_threshold_bytes", &mori::io::TcpBackendConfig::stripingThresholdBytes); + .def_readwrite("striping_threshold_bytes", + &mori::io::TcpBackendConfig::stripingThresholdBytes); py::class_(m, "IOEngineConfig") .def(py::init(), py::arg("host") = "", py::arg("port") = 0) diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 8a879a690..59b8cd457 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -90,14 +90,6 @@ def parse_args(): default="read", help="Type of ops, choices [read, write], default to 'read'", ) - # Backward-compatible alias for existing scripts. - parser.add_argument( - "--op", - dest="op_type", - type=str, - choices=["read", "write"], - help="Alias for --op-type", - ) parser.add_argument( "--buffer-size", type=int, diff --git a/tests/python/io/test_engine.py b/tests/python/io/test_engine.py index 819c0743a..aebfac321 100644 --- a/tests/python/io/test_engine.py +++ b/tests/python/io/test_engine.py @@ -34,6 +34,7 @@ StatusCode, MemoryLocationType, RdmaBackendConfig, + TcpBackendConfig, XgmiBackendConfig, set_log_level, ) @@ -713,3 +714,99 @@ def test_xgmi_cross_engine_transfer(): status.Wait() assert status.Succeeded() assert torch.equal(src_tensor.cpu(), dst_tensor.cpu()) + + +# --------------------------------------------------------------------------- # +# TCP backend tests # +# --------------------------------------------------------------------------- # + + +def _wait_inbound_status(engine, remote_engine_key, remote_transfer_uid, timeout_s=5.0): + deadline = time.time() + timeout_s + while time.time() < deadline: + st = engine.pop_inbound_transfer_status(remote_engine_key, remote_transfer_uid) + if st is not None: + return st + time.sleep(0.001) + raise RuntimeError("Timed out waiting for inbound status") + + +def _create_tcp_engine_pair(name_prefix, port_a=0, port_b=0): + cfg_a = IOEngineConfig(host="127.0.0.1", port=port_a) + cfg_b = IOEngineConfig(host="127.0.0.1", port=port_b) + + a = IOEngine(key=f"{name_prefix}_a", config=cfg_a) + b = IOEngine(key=f"{name_prefix}_b", config=cfg_b) + + a.create_backend(BackendType.TCP, TcpBackendConfig()) + b.create_backend(BackendType.TCP, TcpBackendConfig()) + + a_desc = a.get_engine_desc() + b_desc = b.get_engine_desc() + a.register_remote_engine(b_desc) + b.register_remote_engine(a_desc) + return a, b, a_desc, b_desc + + +def test_tcp_engine_desc_port_zero_auto_bind(): + set_log_level("error") + engine = IOEngine( + key="engine_tcp_port0", config=IOEngineConfig(host="127.0.0.1", port=0) + ) + engine.create_backend(BackendType.TCP, TcpBackendConfig()) + desc = engine.get_engine_desc() + assert desc.port > 0 + + +def test_tcp_cpu_write_read_and_batch(): + set_log_level("error") + a, b, a_desc, b_desc = _create_tcp_engine_pair( + "tcp_cpu", get_free_port(), get_free_port() + ) + + # Allocate CPU tensors and register memory. + src = torch.arange(0, 1024 * 4, dtype=torch.uint8) + dst = torch.zeros_like(src) + src_md = a.register_torch_tensor(src) + dst_md = b.register_torch_tensor(dst) + + # MemoryDesc serialization should work for TCP too. + packed = dst_md.pack() + dst_md_remote = MemoryDesc.unpack(packed) + assert dst_md == dst_md_remote + + # Single write + uid = a.allocate_transfer_uid() + st = a.write(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + assert torch.equal(src, dst) + + # Single read (b -> a) + dst.zero_() + uid = a.allocate_transfer_uid() + st = a.read(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + assert torch.equal(src, dst) + + # Batch write via session + sess = a.create_session(src_md, dst_md) + assert sess is not None + offsets = [0, 256, 512, 768] + sizes = [128, 128, 128, 128] + + dst.zero_() + uid = sess.allocate_transfer_uid() + st = sess.batch_write(offsets, offsets, sizes, uid) + st.Wait() + assert st.Succeeded(), st.Message() + bst = _wait_inbound_status(b, a_desc.key, uid) + assert bst.Succeeded(), bst.Message() + + for off, sz in zip(offsets, sizes): + assert torch.equal(src[off : off + sz], dst[off : off + sz]) diff --git a/tests/python/io/test_engine_tcp.py b/tests/python/io/test_engine_tcp.py deleted file mode 100644 index 5bfa1c10f..000000000 --- a/tests/python/io/test_engine_tcp.py +++ /dev/null @@ -1,122 +0,0 @@ -# Copyright © Advanced Micro Devices, Inc. All rights reserved. -# -# MIT License -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -import time - -import torch -from tests.python.utils import get_free_port - -from mori.io import ( - BackendType, - IOEngine, - IOEngineConfig, - MemoryDesc, - TcpBackendConfig, - set_log_level, -) - - -def _wait_inbound_status(engine, remote_engine_key, remote_transfer_uid, timeout_s=5.0): - deadline = time.time() + timeout_s - while time.time() < deadline: - st = engine.pop_inbound_transfer_status(remote_engine_key, remote_transfer_uid) - if st is not None: - return st - time.sleep(0.001) - raise RuntimeError("Timed out waiting for inbound status") - - -def _create_tcp_engine_pair(name_prefix, port_a=0, port_b=0): - cfg_a = IOEngineConfig(host="127.0.0.1", port=port_a) - cfg_b = IOEngineConfig(host="127.0.0.1", port=port_b) - - a = IOEngine(key=f"{name_prefix}_a", config=cfg_a) - b = IOEngine(key=f"{name_prefix}_b", config=cfg_b) - - a.create_backend(BackendType.TCP, TcpBackendConfig()) - b.create_backend(BackendType.TCP, TcpBackendConfig()) - - a_desc = a.get_engine_desc() - b_desc = b.get_engine_desc() - a.register_remote_engine(b_desc) - b.register_remote_engine(a_desc) - return a, b, a_desc, b_desc - - -def test_tcp_engine_desc_port_zero_auto_bind(): - set_log_level("error") - engine = IOEngine(key="engine_tcp_port0", config=IOEngineConfig(host="127.0.0.1", port=0)) - engine.create_backend(BackendType.TCP, TcpBackendConfig()) - desc = engine.get_engine_desc() - assert desc.port > 0 - - -def test_tcp_cpu_write_read_and_batch(): - set_log_level("error") - a, b, a_desc, b_desc = _create_tcp_engine_pair("tcp_cpu", get_free_port(), get_free_port()) - - # Allocate CPU tensors and register memory. - src = torch.arange(0, 1024 * 4, dtype=torch.uint8) - dst = torch.zeros_like(src) - src_md = a.register_torch_tensor(src) - dst_md = b.register_torch_tensor(dst) - - # MemoryDesc serialization should work for TCP too. - packed = dst_md.pack() - dst_md_remote = MemoryDesc.unpack(packed) - assert dst_md == dst_md_remote - - # Single write - uid = a.allocate_transfer_uid() - st = a.write(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) - st.Wait() - assert st.Succeeded(), st.Message() - bst = _wait_inbound_status(b, a_desc.key, uid) - assert bst.Succeeded(), bst.Message() - assert torch.equal(src, dst) - - # Single read (b -> a) - dst.zero_() - uid = a.allocate_transfer_uid() - st = a.read(src_md, 0, dst_md, 0, src.numel() * src.element_size(), uid) - st.Wait() - assert st.Succeeded(), st.Message() - bst = _wait_inbound_status(b, a_desc.key, uid) - assert bst.Succeeded(), bst.Message() - assert torch.equal(src, dst) - - # Batch write via session - sess = a.create_session(src_md, dst_md) - assert sess is not None - offsets = [0, 256, 512, 768] - sizes = [128, 128, 128, 128] - - dst.zero_() - uid = sess.allocate_transfer_uid() - st = sess.batch_write(offsets, offsets, sizes, uid) - st.Wait() - assert st.Succeeded(), st.Message() - bst = _wait_inbound_status(b, a_desc.key, uid) - assert bst.Succeeded(), bst.Message() - - for off, sz in zip(offsets, sizes): - assert torch.equal(src[off : off + sz], dst[off : off + sz]) - From e46d0742d32bdd94ad055322f4384a0bfdec80aa Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Sat, 28 Feb 2026 01:57:35 -0600 Subject: [PATCH 12/12] io/tcp: use more conn --- include/mori/io/backend.hpp | 2 +- src/io/engine.cpp | 1 + src/io/tcp/transport.cpp | 4 ++-- src/io/tcp/transport.hpp | 12 +++++++----- src/pybind/mori.cpp | 2 +- 5 files changed, 12 insertions(+), 9 deletions(-) diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index a347ae4c0..798149c5e 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -113,7 +113,7 @@ struct TcpBackendConfig : public BackendConfig { // Number of parallel DATA TCP connections per peer (iperf-like multi-stream striping). // Effective only when peer has >= numDataConns established and transfer is contiguous. - int numDataConns{4}; + int numDataConns{8}; // Stripe large contiguous transfers; keep small transfers on a single stream for latency. int stripingThresholdBytes{64 * 1024}; }; diff --git a/src/io/engine.cpp b/src/io/engine.cpp index db5bfa003..93338c525 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -153,6 +153,7 @@ void IOEngine::CreateBackend(BackendType type, const BackendConfig& beConfig) { } backends.insert({type, std::move(backend)}); + InvalidateRouteCache(); } else { assert(false && "not implemented"); } diff --git a/src/io/tcp/transport.cpp b/src/io/tcp/transport.cpp index d6d4841ec..73eb12208 100644 --- a/src/io/tcp/transport.cpp +++ b/src/io/tcp/transport.cpp @@ -1673,7 +1673,7 @@ void TcpTransport::TryConsumeEarlyWriteLanes(const EngineKey& peer, TransferUniq void TcpTransport::MaybeCompleteOutbound(OutboundOpState& st) { if (!st.completionReceived) return; if (st.isRead) { - uint8_t allMask = LanesAllMask(st.lanesTotal); + uint16_t allMask = LanesAllMask(st.lanesTotal); if (st.lanesDoneMask != allMask || st.rxBytes != st.expectedRxBytes || st.gpuCopyPending) return; } @@ -1746,7 +1746,7 @@ void TcpTransport::HandleWorkerRecvDone(const WorkerEvent& ev) { } if (st.local.loc == MemoryLocationType::GPU) { - uint8_t allMask = LanesAllMask(st.lanesTotal); + uint16_t allMask = LanesAllMask(st.lanesTotal); if ((st.lanesDoneMask & allMask) != allMask) { MaybeCompleteOutbound(st); return; diff --git a/src/io/tcp/transport.hpp b/src/io/tcp/transport.hpp index c56e0e2e7..e46fa1d52 100644 --- a/src/io/tcp/transport.hpp +++ b/src/io/tcp/transport.hpp @@ -323,7 +323,7 @@ struct Segment { uint64_t len{0}; }; -constexpr uint8_t kLaneBits = 3; +constexpr uint8_t kLaneBits = 4; constexpr uint64_t kLaneMask = (1ULL << kLaneBits) - 1ULL; inline uint64_t ToWireOpId(uint64_t userOpId, uint8_t lane) { @@ -342,8 +342,8 @@ inline LaneSpan ComputeLaneSpan(uint64_t total, uint8_t lanes, uint8_t lane) { return {uint64_t(lane) * base + std::min(lane, rem), base + (lane < rem ? 1 : 0)}; } -inline uint8_t LanesAllMask(uint8_t n) { - return (n >= (1U << kLaneBits)) ? 0xFF : uint8_t((1U << n) - 1); +inline uint16_t LanesAllMask(uint8_t n) { + return (n >= (1U << kLaneBits)) ? 0xFFFF : uint16_t((1U << n) - 1); } inline uint8_t ClampLanesTotal(uint8_t n) { return n == 0 ? 1 : std::min(n, 1U << kLaneBits); @@ -473,7 +473,8 @@ struct OutboundOpState { std::vector localSegs, remoteSegs; uint64_t expectedRxBytes{0}, rxBytes{0}; bool completionReceived{false}, gpuCopyPending{false}; - uint8_t lanesTotal{1}, lanesDoneMask{0}; + uint8_t lanesTotal{1}; + uint16_t lanesDoneMask{0}; StatusCode completionCode{StatusCode::SUCCESS}; std::string completionMsg; std::shared_ptr pinned; @@ -486,7 +487,8 @@ struct InboundWriteState { MemoryDesc dst{}; std::vector dstSegs; bool discard{false}; - uint8_t lanesTotal{1}, lanesDoneMask{0}; + uint8_t lanesTotal{1}; + uint16_t lanesDoneMask{0}; std::shared_ptr pinned; }; diff --git a/src/pybind/mori.cpp b/src/pybind/mori.cpp index 72a40662a..34f6ef2be 100644 --- a/src/pybind/mori.cpp +++ b/src/pybind/mori.cpp @@ -764,7 +764,7 @@ void RegisterMoriIo(pybind11::module_& m) { py::arg("sock_rcvbuf_bytes") = 4 * 1024 * 1024, py::arg("op_timeout_ms") = 30 * 1000, py::arg("enable_keepalive") = true, py::arg("keepalive_idle_sec") = 30, py::arg("keepalive_intvl_sec") = 10, py::arg("keepalive_cnt") = 3, - py::arg("enable_ctrl_nodelay") = true, py::arg("num_data_conns") = 4, + py::arg("enable_ctrl_nodelay") = true, py::arg("num_data_conns") = 8, py::arg("striping_threshold_bytes") = 256 * 1024) .def_readwrite("sock_sndbuf_bytes", &mori::io::TcpBackendConfig::sockSndbufBytes) .def_readwrite("sock_rcvbuf_bytes", &mori::io::TcpBackendConfig::sockRcvbufBytes)