engine.cc

#include "engine.h"
#include "compression.h"
#include "endpoint_wrapper.h"
#include "util/debug.h"
#include "util/gpu_rt.h"
#include "util/pause.h"
#include "util/util.h"
#include <arpa/inet.h>
#include <cc/cc_state.h>
#include <netinet/in.h>
#include <array>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <future>
#include <iostream>
#include <optional>
#include <sstream>
#include <thread>
#include <Python.h>
#include <sys/socket.h>
#include <unistd.h>

thread_local bool inside_python = false;

std::size_t VectorUint8Hash::operator()(std::vector<uint8_t> const& vec) const {
  std::size_t hash = vec.size();
  for (uint8_t byte : vec) {
    hash = hash * 31 + static_cast<std::size_t>(byte);
  }
  return hash;
}

ShmMsg::ShmMsg() : type(ShmMsgType::COMPLETION), completion(0) {
  std::memset(src_bdf, 0, sizeof(src_bdf));
}

// TaskBatch (declaration order: after ShmMsg, before Endpoint)
TaskBatch::TaskBatch() : num_iovs(0) {}

TaskBatch::TaskBatch(TaskBatch&& other) noexcept
    : num_iovs(other.num_iovs),
      data_ptr(std::move(other.data_ptr)),
      size_ptr(std::move(other.size_ptr)),
      mr_id_ptr(std::move(other.mr_id_ptr)),
      slot_item_ptr(std::move(other.slot_item_ptr)) {}

TaskBatch& TaskBatch::operator=(TaskBatch&& other) noexcept {
  if (this != &other) {
    num_iovs = other.num_iovs;
    data_ptr = std::move(other.data_ptr);
    size_ptr = std::move(other.size_ptr);
    mr_id_ptr = std::move(other.mr_id_ptr);
    slot_item_ptr = std::move(other.slot_item_ptr);
  }
  return *this;
}

void** TaskBatch::data_v() const {
  if (!data_ptr) return nullptr;
  return data_ptr->data();
}
size_t* TaskBatch::size_v() const {
  if (!size_ptr) return nullptr;
  return size_ptr->data();
}
uint64_t* TaskBatch::mr_id_v() const {
  if (!mr_id_ptr) return nullptr;
  return mr_id_ptr->data();
}
FifoItem* TaskBatch::slot_item_v() const {
  if (!slot_item_ptr) return nullptr;
  return slot_item_ptr->data();
}

// UnifiedTask (declaration order: after TaskBatch, before Endpoint)
UnifiedTask::UnifiedTask()
    : type(TaskType::WRITE_NET),
      data(nullptr),
      size(0),
      conn_id(0),
      mr_id(0),
      status_ptr(nullptr),
      specific() {}

UnifiedTask::~UnifiedTask() {
  if (is_batch_task()) {
    specific.batch.task_batch.~TaskBatch();
  }
}

FifoItem& UnifiedTask::slot_item() { return specific.net.slot_item; }

FifoItem const& UnifiedTask::slot_item() const {
  return specific.net.slot_item;
}

IpcTransferInfo& UnifiedTask::ipc_info() { return specific.ipc.ipc_info; }

IpcTransferInfo const& UnifiedTask::ipc_info() const {
  return specific.ipc.ipc_info;
}

TaskBatch& UnifiedTask::task_batch() { return specific.batch.task_batch; }

TaskBatch const& UnifiedTask::task_batch() const {
  return specific.batch.task_batch;
}

bool UnifiedTask::is_batch_task() const {
  return type == TaskType::WRITEV || type == TaskType::READV;
}

UnifiedTask::SpecificData::SpecificData() : base{} {}

UnifiedTask::SpecificData::~SpecificData() {}

static inline void check_python_signals() {
  PyGILState_STATE gstate = PyGILState_Ensure();
  if (PyErr_CheckSignals() != 0) {
    UCCL_LOG(FATAL) << "Python signal caught, exiting...";
  }
  PyGILState_Release(gstate);
}

static inline bool is_retryable_post_failure(int rc) {
  if (is_cxi_transport()) return rc == UCCL_POST_TRANSIENT;
  return rc == -1;
}

static inline bool raw_one_sided_batch_eligible(
    std::vector<size_t> const& size_v, size_t num_iovs) {
  for (size_t i = 0; i < num_iovs; ++i) {
    if (ChunkSplitStrategy::get_message_chunk_count(size_v[i]) != 1)
      return false;
  }
  return true;
}

static inline size_t max_iov_bytes(std::vector<size_t> const& size_v,
                                   size_t num_iovs) {
  size_t max_bytes = 0;
  for (size_t i = 0; i < num_iovs; ++i) {
    max_bytes = std::max(max_bytes, size_v[i]);
  }
  return max_bytes;
}

size_t Endpoint::max_one_sided_inflight_ops() {
  size_t limit = kMaxInflightOps;
  if (is_cxi_transport()) limit = 32;

  char const* env = std::getenv("UCCL_P2P_MAX_INFLIGHT_OPS");
  if (env && env[0] != '\0') {
    char* end = nullptr;
    unsigned long requested = std::strtoul(env, &end, 10);
    if (end != env && requested > 0) {
      limit = std::min<size_t>(requested, kMaxInflightOps);
    }
  }
  return std::max<size_t>(1, limit);
}

// ShmChannel helper function
static inline std::string shm_ring_name(std::string const& from_bdf,
                                        std::string const& to_bdf) {
  return "/uccl_gpu_jring_" + uccl::sanitize_bdf(from_bdf) + "_" +
         uccl::sanitize_bdf(to_bdf);
}

static inline void shm_ring_send(jring_t* ring, ShmMsg const& msg) {
  // jring bulk copy requires 16B aligned buffer
  alignas(16) ShmMsg tmp = msg;

  while (jring_mp_enqueue_bulk(ring, (void*)&tmp, 1, nullptr) != 1) {
    machnet_pause();
  }
}

static inline void shm_ring_recv(jring_t* ring, ShmMsg& msg) {
  alignas(16) ShmMsg tmp;

  while (jring_sc_dequeue_bulk(ring, (void*)&tmp, 1, nullptr) != 1) {
    machnet_pause();
  }
  msg = tmp;
}

uccl::UCCLLogLevel Endpoint::parse_log_level_from_env() {
  char const* env = std::getenv("UCCL_P2P_LOG_LEVEL");
  if (!env) {
    return uccl::WARN;
  }

  if (!strcasecmp(env, "INFO")) return uccl::INFO;
  if (!strcasecmp(env, "WARN") || !strcasecmp(env, "WARNING"))
    return uccl::WARN;
  if (!strcasecmp(env, "ERROR")) return uccl::ERROR;
  if (!strcasecmp(env, "FATAL")) return uccl::FATAL;

  char* end = nullptr;
  long val = std::strtol(env, &end, 10);
  if (end != env && val >= 0 && val <= 3) {
    return static_cast<uccl::UCCLLogLevel>(val);
  }

  return uccl::WARN;
}

// -----------------------------------------------------------------------------
// The main P2P Endpoint class implementation
// -----------------------------------------------------------------------------

Endpoint::Endpoint(uint32_t const gpu_idx) : passive_accept_(false) {
  // gpu_idx is always a local CUDA device ordinal.
  int ngpus = 0;
  GPU_RT_CHECK(gpuGetDeviceCount(&ngpus));
  UCCL_CHECK(static_cast<int>(gpu_idx) < ngpus)
      << "GPU index " << gpu_idx << " out of range (visible devices: " << ngpus
      << ")";

  local_gpu_idx_ = gpu_idx;

  // Get PCI Bus ID as cross-process identity
  char bdf_buf[64];
  GPU_RT_CHECK(gpuDeviceGetPCIBusId(bdf_buf, sizeof(bdf_buf), gpu_idx));
  gpu_bus_id_ = uccl::normalize_pci_bus_id(bdf_buf);

  std::cout << "Creating Engine with GPU index: " << gpu_idx
            << " (bus_id: " << gpu_bus_id_ << ")" << std::endl;
  int n_streams = std::max(1, (int)kNumGpuRtStreams);

  ipc_streams_.resize(ngpus);
  for (int i = 0; i < ngpus; ++i) {
    GPU_RT_CHECK(gpuSetDevice(i));
    ipc_streams_[i].resize(n_streams);
    for (int j = 0; j < n_streams; ++j) {
      GPU_RT_CHECK(
          gpuStreamCreateWithFlags(&ipc_streams_[i][j], gpuStreamNonBlocking));
    }
  }
  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));

  uccl::ucclLogger.setLogLevel(Endpoint::parse_log_level_from_env());

  if (is_nccl_transport()) {
    ep_ = std::make_shared<NCCLEndpoint>(local_gpu_idx_, 0);
    numa_node_ = get_numa_node_from_iface();
  } else if (is_cxi_transport()) {
    ep_ = std::make_shared<CxiEndpoint>(local_gpu_idx_, 0);
    numa_node_ = get_numa_node_from_iface();
  } else {
    ep_ = std::shared_ptr<RDMAEndpoint>(new RDMAEndpoint(local_gpu_idx_, 0));
  }

  std::cout << "Engine initialized for GPU " << local_gpu_idx_ << std::endl;
  engine_initialized_ = true;

  send_unified_task_ring_ =
      uccl::create_ring(sizeof(UnifiedTask*), kTaskRingSize);
  recv_unified_task_ring_ =
      uccl::create_ring(sizeof(UnifiedTask*), kTaskRingSize);

  if (!is_nccl_transport() && !is_cxi_transport()) {
    numa_node_ = RdmaDeviceManager::instance().get_numa_node(
        RdmaDeviceManager::instance().get_best_dev_idx(local_gpu_idx_)[0]);
  }

  send_proxy_thread_ = std::thread(&Endpoint::send_proxy_thread_func, this);
  recv_proxy_thread_ = std::thread(&Endpoint::recv_proxy_thread_func, this);

  ipc_inflight_ring_ = uccl::create_ring(sizeof(IpcInflightOp*), kTaskRingSize);
  ipc_poller_thread_ = std::thread(&Endpoint::ipc_poller_thread_func, this);

  // Initialize ShmChannel for local connections (use BDF for ring names)
  size_t elem_sz = sizeof(ShmMsg);
  size_t elem_cnt = ShmRingDefaultElemCnt;
  auto all_bdfs = uccl::enumerate_all_gpu_bdfs();
  for (auto const& sender_bdf : all_bdfs) {
    auto name = shm_ring_name(sender_bdf, gpu_bus_id_);
    ShmRingHandle handle;
    handle.shm_name = name;
    bool is_creator = false;
    handle.ring =
        uccl::create_shared_ring(name.c_str(), elem_sz, elem_cnt, handle.shm_fd,
                                 handle.shm_size, &is_creator);
    inbox_rings_[sender_bdf] = handle;
    inbox_creators_[sender_bdf] = is_creator;
  }

  std::cout << "Endpoint initialized successfully" << std::endl;
}

Endpoint::Endpoint() : local_gpu_idx_(INVALID_GPU), passive_accept_(false) {
  std::cout << "Creating Engine" << std::endl;
  int n_streams = std::max(1, (int)kNumGpuRtStreams);

  int ngpus = 0;
  GPU_RT_CHECK(gpuGetDeviceCount(&ngpus));

  // Save the current device before the loop so we can restore it afterwards.
  // The loop calls gpuSetDevice(i) for each GPU, leaving the active device as
  // GPU ngpus-1, which would corrupt gpu_bus_id_ and the advertised BDF.
  int cur_dev = 0;
  GPU_RT_CHECK(gpuGetDevice(&cur_dev));

  ipc_streams_.resize(ngpus);
  for (int i = 0; i < ngpus; ++i) {
    GPU_RT_CHECK(gpuSetDevice(i));
    ipc_streams_[i].resize(n_streams);
    for (int j = 0; j < n_streams; ++j) {
      GPU_RT_CHECK(
          gpuStreamCreateWithFlags(&ipc_streams_[i][j], gpuStreamNonBlocking));
    }
  }
  GPU_RT_CHECK(gpuSetDevice(cur_dev));

  // Get the current GPU's BDF for metadata (needed before lazy engine init).
  char bdf_buf[64];
  GPU_RT_CHECK(gpuDeviceGetPCIBusId(bdf_buf, sizeof(bdf_buf), cur_dev));
  gpu_bus_id_ = uccl::normalize_pci_bus_id(bdf_buf);

  if (is_nccl_transport()) {
    ep_ = std::make_shared<NCCLEndpoint>(local_gpu_idx_, 0);
  } else if (is_cxi_transport()) {
    ep_ = std::make_shared<CxiEndpoint>(INVALID_GPU, 0);
  } else {
    ep_ = std::shared_ptr<RDMAEndpoint>(new RDMAEndpoint(INVALID_GPU, 0));
  }

  std::cout << "Endpoint initialized successfully" << std::endl;
}

Endpoint::~Endpoint() {
  std::cout << "Destroying Engine..." << std::endl;

  stop_.store(true, std::memory_order_release);

  if (passive_accept_) {
    passive_accept_stop_.store(true, std::memory_order_release);
    if (passive_accept_thread_.joinable()) {
      passive_accept_thread_.join();
    }
    if (passive_accept_local_thread_.joinable()) {
      passive_accept_local_thread_.join();
    }
  }

  if (send_proxy_thread_.joinable()) {
    send_proxy_thread_.join();
  }
  if (recv_proxy_thread_.joinable()) {
    recv_proxy_thread_.join();
  }
  if (ipc_poller_thread_.joinable()) {
    ipc_poller_thread_.join();
  }

  if (send_unified_task_ring_ != nullptr) {
    free(send_unified_task_ring_);
  }
  if (recv_unified_task_ring_ != nullptr) {
    free(recv_unified_task_ring_);
  }
  if (ipc_inflight_ring_ != nullptr) {
    free(ipc_inflight_ring_);
  }

  {
    std::shared_lock<std::shared_mutex> lock(conn_mu_);
    for (auto& [conn_id, conn] : conn_id_to_conn_) {
      auto& shm_channel_ = conn->remote_inbox_;

      if (conn->shm_attached_) {
        uccl::detach_shared_ring(shm_channel_.ring, shm_channel_.shm_fd,
                                 shm_channel_.shm_size);
      }
      delete conn;
    }
  }

  for (auto& [bdf, handle] : inbox_rings_) {
    auto it = inbox_creators_.find(bdf);
    if (it != inbox_creators_.end() && it->second)
      uccl::destroy_shared_ring(handle.shm_name.c_str(), handle.ring,
                                handle.shm_fd, handle.shm_size);
    else
      uccl::detach_shared_ring(handle.ring, handle.shm_fd, handle.shm_size);
  }

  {
    std::shared_lock<std::shared_mutex> lock(mr_mu_);
    for (auto& [mr_id, mr] : mr_id_to_mr_) {
      delete mr->mhandle_;
      delete mr;
    }
  }

  std::cout << "Engine destroyed" << std::endl;
}

void Endpoint::stop_accept() { uccl_stop_accept(ep_); }

bool Endpoint::connect(std::string ip_addr, int remote_gpu_idx, int remote_port,
                       uint64_t& conn_id) {
  std::cout << "Attempting to connect to " << ip_addr << ":" << remote_gpu_idx
            << " via port " << remote_port << std::endl;
  // Create a new connection ID
  conn_id = next_conn_id_.fetch_add(1);

  assert(local_gpu_idx_ != INVALID_GPU);

  std::future<ConnID> uccl_conn_id_future = std::async(
      std::launch::async, [this, remote_gpu_idx, &ip_addr, remote_port]() {
        return uccl_connect(ep_, remote_gpu_idx, ip_addr, remote_port);
      });

  // Check for Python signals (eg, ctrl+c) while waiting for connection
  while (uccl_conn_id_future.wait_for(std::chrono::seconds(0)) !=
         std::future_status::ready) {
    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;
    std::this_thread::sleep_for(std::chrono::seconds(1));
  }
  ConnID uccl_conn_id = uccl_conn_id_future.get();

  // Store the connection ID.
  {
    std::unique_lock<std::shared_mutex> lock(conn_mu_);
    conn_id_to_conn_[conn_id] = new Conn{
        conn_id, uccl_conn_id, ip_addr, static_cast<uint16_t>(remote_port), ""};
  }
  return true;
}

std::vector<uint8_t> Endpoint::get_metadata() {
  std::string ip_str = uccl::get_oob_ip();
  uint16_t port = get_p2p_listen_port(ep_);

  bool is_ipv6 = ip_str.find(':') != std::string::npos;
  size_t ip_len = is_ipv6 ? 16 : 4;

  // 2 bytes port + 1 byte BDF length + N bytes BDF string
  uint8_t bdf_len = static_cast<uint8_t>(gpu_bus_id_.size());
  size_t total_len = ip_len + 2 + 1 + bdf_len;
  std::vector<uint8_t> metadata(total_len);

  // Copy IP
  if (is_ipv6) {
    struct in6_addr ip6_bin;
    if (inet_pton(AF_INET6, ip_str.c_str(), &ip6_bin) != 1)
      throw std::runtime_error("Invalid IPv6 address: " + ip_str);
    std::memcpy(metadata.data(), &ip6_bin, 16);
  } else {
    struct in_addr ip4_bin;
    if (inet_pton(AF_INET, ip_str.c_str(), &ip4_bin) != 1)
      throw std::runtime_error("Invalid IPv4 address: " + ip_str);
    std::memcpy(metadata.data(), &ip4_bin, 4);
  }

  // Copy port in network byte order
  uint16_t net_port = htons(port);
  std::memcpy(metadata.data() + ip_len, &net_port, 2);

  // Copy BDF string (1 byte length + string data)
  metadata[ip_len + 2] = bdf_len;
  std::memcpy(metadata.data() + ip_len + 3, gpu_bus_id_.data(), bdf_len);

  return metadata;
}
std::tuple<std::string, uint16_t, std::string> Endpoint::parse_metadata(
    std::vector<uint8_t> const& metadata) {
  // Minimum: 4 (IPv4) + 2 (port) + 1 (bdf_len) = 7
  if (metadata.size() < 7) {
    throw std::runtime_error("Metadata too short: " +
                             std::to_string(metadata.size()));
  }

  std::string ip_str;
  size_t ip_len;

  // Determine IP version from metadata length and BDF field
  // Try IPv4 first: 4 bytes IP + 2 bytes port + 1 byte bdf_len + N bytes bdf
  uint8_t bdf_len_v4 = metadata[6];
  if (metadata.size() == static_cast<size_t>(4 + 2 + 1 + bdf_len_v4)) {
    // IPv4
    ip_len = 4;
    char buf[INET_ADDRSTRLEN];
    if (inet_ntop(AF_INET, metadata.data(), buf, sizeof(buf)) == nullptr) {
      throw std::runtime_error("Failed to parse IPv4 address from metadata");
    }
    ip_str = buf;
  } else {
    // IPv6: 16 bytes IP + 2 bytes port + 1 byte bdf_len + N bytes bdf
    ip_len = 16;
    if (metadata.size() < 19) {
      throw std::runtime_error("Metadata too short for IPv6: " +
                               std::to_string(metadata.size()));
    }
    uint8_t bdf_len_v6 = metadata[18];
    if (metadata.size() != static_cast<size_t>(19 + bdf_len_v6)) {
      throw std::runtime_error("Metadata size mismatch for IPv6: got " +
                               std::to_string(metadata.size()) + ", expected " +
                               std::to_string(19 + bdf_len_v6));
    }
    char buf[INET6_ADDRSTRLEN];
    if (inet_ntop(AF_INET6, metadata.data(), buf, sizeof(buf)) == nullptr) {
      throw std::runtime_error("Failed to parse IPv6 address from metadata");
    }
    ip_str = buf;
  }

  uint16_t net_port;
  std::memcpy(&net_port, metadata.data() + ip_len, 2);
  uint16_t port = ntohs(net_port);

  uint8_t bdf_len = metadata[ip_len + 2];
  std::string gpu_bdf;
  if (bdf_len > 0) {
    gpu_bdf = std::string(
        reinterpret_cast<char const*>(metadata.data() + ip_len + 3), bdf_len);
  }

  return std::make_tuple(ip_str, port, gpu_bdf);
}

bool Endpoint::accept(std::string& ip_addr, int& remote_gpu_idx,
                      uint64_t& conn_id) {
  std::cout << "Waiting to accept incoming connection..." << std::endl;

  // For demo purposes, simulate accepted connection
  conn_id = next_conn_id_.fetch_add(1);

  // Wait until engine is intialized to get the correct local_gpu_idx_
  while (!engine_initialized_) {
    std::this_thread::sleep_for(std::chrono::milliseconds(1));
  }
  std::future<ConnID> uccl_conn_id_future =
      std::async(std::launch::async, [this, &ip_addr, &remote_gpu_idx]() {
        return uccl_accept(ep_, ip_addr, &remote_gpu_idx);
      });

  // Check for Python signals (eg, ctrl+c) while waiting for connection
  while (uccl_conn_id_future.wait_for(std::chrono::seconds(0)) !=
         std::future_status::ready) {
    if (passive_accept_ &&
        passive_accept_stop_.load(std::memory_order_acquire)) {
      std::cout << "Stop background accept..." << std::endl;
      uccl_stop_accept(ep_);
    }
    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;
    std::this_thread::sleep_for(std::chrono::seconds(1));
  }
  ConnID uccl_conn_id = uccl_conn_id_future.get();

  // Return if Conn ID is invalid
  if (uccl_conn_id.sock_fd < 0 || uccl_conn_id.peer_id == UINT64_MAX) {
    return false;
  }

  // Store the connection ID.
  {
    std::unique_lock<std::shared_mutex> lock(conn_mu_);
    conn_id_to_conn_[conn_id] = new Conn{conn_id, uccl_conn_id, ip_addr, 0, ""};
  }

  return true;
}

bool Endpoint::reg(void const* data, size_t size, uint64_t& mr_id,
                   FloatType float_type) {
  mr_id = next_mr_id_.fetch_add(1);

  if (!engine_initialized_) {
    int idx = uccl::get_dev_idx((void*)data);
    if (idx != -1) {
      local_gpu_idx_ = idx;
    } else {
      local_gpu_idx_ = 0;
    }
    // Get PCI Bus ID for cross-process identity
    char bdf_buf[64];
    GPU_RT_CHECK(
        gpuDeviceGetPCIBusId(bdf_buf, sizeof(bdf_buf), local_gpu_idx_));
    gpu_bus_id_ = uccl::normalize_pci_bus_id(bdf_buf);
    initialize_engine();
    engine_initialized_ = true;
  }

  P2PMhandle* mhandle = new P2PMhandle();
  mhandle->compress_ctx = make_compress_ctx(float_type);
  if (!uccl_regmr(ep_, const_cast<void*>(data), size, mhandle)) {
    delete mhandle;
    return false;
  }
  {
    std::unique_lock<std::shared_mutex> lock(mr_mu_);
    mr_id_to_mr_[mr_id] = new MR{mr_id, mhandle};
  }

  return true;
}

bool Endpoint::regv(std::vector<void const*> const& data_v,
                    std::vector<size_t> const& size_v,
                    std::vector<uint64_t>& mr_id_v) {
  if (data_v.size() != size_v.size())
    throw std::invalid_argument(
        "[Endpoint::regv] data_v/size_v length mismatch");

  size_t const n = data_v.size();

  // Early return if empty
  if (n == 0) {
    mr_id_v.clear();
    return true;
  }

  if (!engine_initialized_) {
    int idx = uccl::get_dev_idx((void*)data_v[0]);
    if (idx != -1) {
      local_gpu_idx_ = idx;
    } else {
      local_gpu_idx_ = 0;
    }
    char bdf_buf[64];
    GPU_RT_CHECK(
        gpuDeviceGetPCIBusId(bdf_buf, sizeof(bdf_buf), local_gpu_idx_));
    gpu_bus_id_ = uccl::normalize_pci_bus_id(bdf_buf);
    initialize_engine();
    engine_initialized_ = true;
  }

  mr_id_v.resize(n);

  {
    std::unique_lock<std::shared_mutex> lock(mr_mu_);
    mr_id_to_mr_.reserve(mr_id_to_mr_.size() + n);
  }

  for (size_t i = 0; i < n; ++i) {
    uint64_t id = next_mr_id_.fetch_add(1);
    P2PMhandle* mhandle = new P2PMhandle();

    if (!uccl_regmr(ep_, const_cast<void*>(data_v[i]), size_v[i], mhandle)) {
      std::cerr << "[Endpoint::regv] registration failed at i=" << i << '\n';
      return false;
    }

    {
      std::unique_lock<std::shared_mutex> lock(mr_mu_);
      mr_id_to_mr_[id] = new MR{id, mhandle};
    }
    mr_id_v[i] = id;
  }
  return true;
}

bool Endpoint::dereg(uint64_t mr_id) {
  MR* mr = nullptr;
  {
    std::unique_lock<std::shared_mutex> lock(mr_mu_);
    auto it = mr_id_to_mr_.find(mr_id);
    if (it == mr_id_to_mr_.end()) {
      std::cerr << "[dereg] Error: Invalid mr_id " << mr_id << std::endl;
      return false;
    }
    mr = it->second;
    mr_id_to_mr_.erase(mr_id);
  }
  mr->mhandle_->compress_ctx.reset();
  uccl_deregmr(ep_, mr->mhandle_);
  delete mr->mhandle_;
  delete mr;
  return true;
}

bool Endpoint::read(uint64_t conn_id, uint64_t mr_id, void* dst, size_t size,
                    FifoItem const& slot_item) {
  UCCL_DCHECK(size <= 0xffffffff) << "size must be < 4 GB";
  auto* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[read] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  P2PMhandle* mhandle = get_mhandle(mr_id);
  if (unlikely(mhandle == nullptr)) {
    std::cerr << "[read] Error: Invalid mr_id " << mr_id << std::endl;
    return false;
  }

  UcclRequest ureq = {};
  FifoItem curr_slot_item = slot_item;
  curr_slot_item.size = size;

  int rc;
  do {
    rc = uccl_read_async(ep_, conn, mhandle, dst, size, curr_slot_item, &ureq);
  } while (is_retryable_post_failure(rc));
  if (is_cxi_transport() && rc < 0) {
    UCCL_LOG(ERROR) << "read failed to post: rc=" << rc;
    return false;
  }

  bool done = false;
  while (!done) {
    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;
    uccl_drive_recv(ep_);
    if (uccl_poll_ureq_once(ep_, &ureq)) {
      done = true;
    }
  }

  return true;
}

bool Endpoint::read_async(uint64_t conn_id, uint64_t mr_id, void* dst,
                          size_t size, FifoItem const& slot_item,
                          uint64_t* transfer_id) {
  if (size <= kDirectAsyncNetThreshold) {
    auto* conn = get_conn(conn_id);
    if (unlikely(conn == nullptr)) {
      std::cerr << "[read_async] Error: Invalid conn_id " << conn_id
                << std::endl;
      return false;
    }
    P2PMhandle* mhandle = get_mhandle(mr_id);
    if (unlikely(mhandle == nullptr)) {
      std::cerr << "[read_async] Error: Invalid mr_id " << mr_id << std::endl;
      return false;
    }

    UcclRequest ureq = {};
    FifoItem curr_slot_item = slot_item;
    curr_slot_item.size = size;
    int rc;
    do {
      rc =
          uccl_read_async(ep_, conn, mhandle, dst, size, curr_slot_item, &ureq);
    } while (is_retryable_post_failure(rc));
    if (is_cxi_transport() && rc < 0) {
      UCCL_LOG(ERROR) << "read_async failed to post: rc=" << rc;
      return false;
    }

    auto* status = new TransferStatus();
    status->poll_net_ureq = true;
    status->ureq = ureq;
    *transfer_id = reinterpret_cast<uint64_t>(status);
    return true;
  }

  auto task_ptr =
      create_net_task(conn_id, mr_id, TaskType::READ_NET, dst, size, slot_item);
  if (unlikely(task_ptr == nullptr)) {
    return false;
  }

  auto* status = new TransferStatus();
  status->task_ptr = task_ptr;
  task_ptr->status_ptr = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  UnifiedTask* task_raw = task_ptr.get();

  while (jring_mp_enqueue_bulk(recv_unified_task_ring_, &task_raw, 1,
                               nullptr) != 1) {
  }
  recv_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::readv(uint64_t conn_id, std::vector<uint64_t> const& mr_id_v,
                     std::vector<void*> const& dst_v,
                     std::vector<size_t> const& size_v,
                     std::vector<FifoItem> const& slot_item_v,
                     size_t num_iovs) {
  auto* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[readv] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  // Pre-resolve all mhandles once to avoid per-iov mutex + map lookup.
  std::vector<P2PMhandle*> mhandle_v(num_iovs);
  for (size_t i = 0; i < num_iovs; ++i) {
    mhandle_v[i] = get_mhandle(mr_id_v[i]);
    if (unlikely(mhandle_v[i] == nullptr)) {
      std::cerr << "[readv] Error: Invalid mr_id " << mr_id_v[i] << std::endl;
      return false;
    }
  }

  // Enable doorbell-batched posting for the duration of this call.
  struct BatchGuard {
    BatchGuard() { g_uccl_batch_post = true; }
    ~BatchGuard() { g_uccl_batch_post = false; }
  } batch_guard;

  // Resolve the send group once so per-slot completion checks avoid the
  // per-call shared_lock + map.find().
  SendConnection* send_group =
      uccl_resolve_send_group(ep_, conn->uccl_conn_id_.peer_id);

  if (send_group != nullptr && raw_one_sided_batch_eligible(size_v, num_iovs) &&
      send_group->can_use_raw_one_sided_batch(
          SendType::Read, max_iov_bytes(size_v, num_iovs))) {
    return run_raw_one_sided_pipeline(send_group, SendType::Read, mhandle_v,
                                      dst_v, size_v, slot_item_v, num_iovs);
  }

  UcclRequest ureq[kMaxInflightOps] = {};
  bool active[kMaxInflightOps] = {false};
  size_t const max_inflight_ops = max_one_sided_inflight_ops();

  size_t next_iov = 0;
  size_t num_completed = 0;
  size_t num_inflight = 0;

  while (num_completed < num_iovs) {
    while (next_iov < num_iovs && num_inflight < max_inflight_ops) {
      size_t slot = 0;
      while (slot < kMaxInflightOps && active[slot]) {
        slot++;
      }
      UCCL_DCHECK(slot < kMaxInflightOps);

      auto rc =
          send_group != nullptr
              ? uccl_read_async_on_group(send_group, conn, mhandle_v[next_iov],
                                         dst_v[next_iov], size_v[next_iov],
                                         slot_item_v[next_iov], &ureq[slot])
              : uccl_read_async(ep_, conn, mhandle_v[next_iov], dst_v[next_iov],
                                size_v[next_iov], slot_item_v[next_iov],
                                &ureq[slot]);
      if (is_retryable_post_failure(rc)) {
        if (is_cxi_transport() && num_inflight == 0) {
          uccl_drive_send(ep_);
          uccl_drive_recv(ep_);
          std::this_thread::yield();
          continue;
        }
        break;
      }
      if (is_cxi_transport() && rc < 0) {
        UCCL_LOG(ERROR) << "readv failed to post iov " << next_iov
                        << ": rc=" << rc;
        return false;
      }
      active[slot] = true;
      next_iov++;
      num_inflight++;
    }

    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;

    // Flush batched WRs in one doorbell per channel, then drive send once.
    // (readv issues RDMA reads via the send path, so drive send here.)
    if (num_inflight > 0) {
      uccl_flush_send(ep_);
      uccl_drive_send(ep_);
      uccl_drive_recv(ep_);
    }

    if (send_group != nullptr) {
      for (size_t slot = 0; slot < kMaxInflightOps; slot++) {
        if (!active[slot]) continue;
        if (uccl_check_wr_fast(send_group, ureq[slot].engine_idx)) {
          active[slot] = false;
          num_completed++;
          num_inflight--;
        }
      }
    } else {
      for (size_t slot = 0; slot < kMaxInflightOps; slot++) {
        if (!active[slot]) continue;
        if (uccl_check_ureq_once(ep_, &ureq[slot])) {
          active[slot] = false;
          num_completed++;
          num_inflight--;
        }
      }
    }
  }

  return true;
}

bool Endpoint::readv_async(uint64_t conn_id, std::vector<uint64_t> mr_id_v,
                           std::vector<void*> dst_v, std::vector<size_t> size_v,
                           std::vector<FifoItem> slot_item_v, size_t num_iovs,
                           uint64_t* transfer_id) {
  // Use move semantics to reduce memory copies
  auto task_ptr =
      create_readv_task(conn_id, std::move(dst_v), std::move(size_v),
                        std::move(mr_id_v), std::move(slot_item_v));
  if (unlikely(task_ptr == nullptr)) {
    return false;
  }

  auto* status = new TransferStatus();
  status->task_ptr = task_ptr;
  task_ptr->status_ptr = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  UnifiedTask* task_raw = task_ptr.get();

  while (jring_mp_enqueue_bulk(recv_unified_task_ring_, &task_raw, 1,
                               nullptr) != 1) {
  }
  recv_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::run_raw_one_sided_pipeline(
    SendConnection* send_group, SendType send_type,
    std::vector<P2PMhandle*> const& mhandle_v,
    std::vector<void*> const& local_addr_v, std::vector<size_t> const& size_v,
    std::vector<FifoItem> const& slot_item_v, size_t num_iovs) {
  struct PendingRawWait {
    SendConnection::RawBatchWait wait;
    bool active = false;
  };

  std::array<PendingRawWait, kMaxInflightOps> pending_waits{};
  size_t const max_inflight_ops = max_one_sided_inflight_ops();
  size_t next_iov = 0;
  size_t num_completed = 0;
  size_t num_inflight = 0;

  auto post_raw_window = [&]() -> bool {
    if (next_iov >= num_iovs || num_inflight >= max_inflight_ops) return true;

    size_t const batch_iovs =
        std::min<size_t>(max_inflight_ops - num_inflight, num_iovs - next_iov);
    std::array<SendConnection::OneSidedBatchOp, kMaxInflightOps> ops{};
    int64_t wr_ids[kMaxInflightOps] = {};
    std::array<SendConnection::RawBatchWait, kMaxInflightOps> waits{};
    size_t num_waits = 0;

    for (size_t i = 0; i < batch_iovs; ++i) {
      size_t const iov = next_iov + i;
      ops[i].local_addr = local_addr_v[iov];
      ops[i].size = size_v[iov];
      ops[i].local_mr_array = &mhandle_v[iov]->mr_array;
      ops[i].slot_item = &slot_item_v[iov];
    }

    if (!send_group->post_write_or_read_batch(send_type, ops.data(), batch_iovs,
                                              wr_ids, waits.data(),
                                              &num_waits)) {
      return false;
    }
    if (unlikely(num_waits == 0)) return false;

    for (size_t i = 0; i < num_waits; ++i) {
      bool inserted = false;
      for (auto& pending : pending_waits) {
        if (pending.active) continue;
        pending.wait = waits[i];
        pending.active = true;
        inserted = true;
        break;
      }
      if (unlikely(!inserted)) return false;
    }

    next_iov += batch_iovs;
    num_inflight += batch_iovs;
    return true;
  };

  while (num_completed < num_iovs) {
    while (next_iov < num_iovs && num_inflight < max_inflight_ops) {
      if (!post_raw_window()) return false;
    }

    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;
    uccl_drive_send(ep_);

    for (auto& pending : pending_waits) {
      if (!pending.active) continue;
      if (!uccl_check_wr_fast(send_group, pending.wait.wr_id)) continue;

      pending.active = false;
      num_completed += pending.wait.iov_count;
      UCCL_DCHECK(num_inflight >= pending.wait.iov_count);
      num_inflight -= pending.wait.iov_count;
      pending.wait = {};
    }
  }

  return true;
}

bool Endpoint::write(uint64_t conn_id, uint64_t mr_id, void* src, size_t size,
                     FifoItem const& slot_item) {
  UCCL_DCHECK(size <= 0xffffffff) << "size must be < 4 GB";
  auto* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[write] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }
  P2PMhandle* mhandle = get_mhandle(mr_id);
  if (unlikely(mhandle == nullptr)) {
    std::cerr << "[write] Error: Invalid mr_id " << mr_id << std::endl;
    return false;
  }
  SendConnection* send_group =
      uccl_resolve_send_group(ep_, conn->uccl_conn_id_.peer_id);
  UcclRequest ureq = {};
  FifoItem curr_slot_item = slot_item;
  curr_slot_item.size = size;

  int rc;
  do {
    rc = send_group != nullptr
             ? uccl_write_async_on_group(send_group, conn, mhandle, src, size,
                                         curr_slot_item, &ureq)
             : uccl_write_async(ep_, conn, mhandle, src, size, curr_slot_item,
                                &ureq);
  } while (is_retryable_post_failure(rc));
  if (is_cxi_transport() && rc < 0) {
    UCCL_LOG(ERROR) << "write failed to post: rc=" << rc;
    return false;
  }

  bool done = false;
  while (!done) {
    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;
    if (uccl_poll_ureq_once(ep_, &ureq)) {
      done = true;
    }
  }

  return true;
}

bool Endpoint::write_async(uint64_t conn_id, uint64_t mr_id, void* src,
                           size_t size, FifoItem const& slot_item,
                           uint64_t* transfer_id) {
  if (size <= kDirectAsyncNetThreshold) {
    auto* conn = get_conn(conn_id);
    if (unlikely(conn == nullptr)) {
      std::cerr << "[write_async] Error: Invalid conn_id " << conn_id
                << std::endl;
      return false;
    }
    P2PMhandle* mhandle = get_mhandle(mr_id);
    if (unlikely(mhandle == nullptr)) {
      std::cerr << "[write_async] Error: Invalid mr_id " << mr_id << std::endl;
      return false;
    }

    UcclRequest ureq = {};
    FifoItem curr_slot_item = slot_item;
    curr_slot_item.size = size;
    int rc;
    do {
      rc = uccl_write_async(ep_, conn, mhandle, src, size, curr_slot_item,
                            &ureq);
    } while (is_retryable_post_failure(rc));
    if (is_cxi_transport() && rc < 0) {
      UCCL_LOG(ERROR) << "write_async failed to post: rc=" << rc;
      return false;
    }

    auto* status = new TransferStatus();
    status->poll_net_ureq = true;
    status->ureq = ureq;
    *transfer_id = reinterpret_cast<uint64_t>(status);
    return true;
  }

  auto task_ptr = create_net_task(conn_id, mr_id, TaskType::WRITE_NET, src,
                                  size, slot_item);
  if (unlikely(task_ptr == nullptr)) {
    return false;
  }

  auto* status = new TransferStatus();
  status->task_ptr = task_ptr;
  task_ptr->status_ptr = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  UnifiedTask* task_raw = task_ptr.get();

  while (jring_mp_enqueue_bulk(send_unified_task_ring_, &task_raw, 1,
                               nullptr) != 1) {
  }
  send_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::writev(uint64_t conn_id, std::vector<uint64_t> const& mr_id_v,
                      std::vector<void*> const& src_v,
                      std::vector<size_t> const& size_v,
                      std::vector<FifoItem> const& slot_item_v,
                      size_t num_iovs) {
  auto* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[writev] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  // Pre-resolve all mhandles once to avoid per-iov mutex + map lookup.
  std::vector<P2PMhandle*> mhandle_v(num_iovs);
  for (size_t i = 0; i < num_iovs; ++i) {
    mhandle_v[i] = get_mhandle(mr_id_v[i]);
    if (unlikely(mhandle_v[i] == nullptr)) {
      std::cerr << "[writev] Error: Invalid mr_id " << mr_id_v[i] << std::endl;
      return false;
    }
  }

  // Enable doorbell-batched posting for the duration of this call. All
  // submit_request() calls from this thread accumulate WRs per channel; we
  // flush once per outer pass before polling.
  struct BatchGuard {
    BatchGuard() { g_uccl_batch_post = true; }
    ~BatchGuard() { g_uccl_batch_post = false; }
  } batch_guard;

  // Resolve the send group once so per-slot completion checks avoid the
  // per-call shared_lock + map.find().
  SendConnection* send_group =
      uccl_resolve_send_group(ep_, conn->uccl_conn_id_.peer_id);

  if (send_group != nullptr && raw_one_sided_batch_eligible(size_v, num_iovs) &&
      send_group->can_use_raw_one_sided_batch(
          SendType::Write, max_iov_bytes(size_v, num_iovs))) {
    return run_raw_one_sided_pipeline(send_group, SendType::Write, mhandle_v,
                                      src_v, size_v, slot_item_v, num_iovs);
  }

  UcclRequest ureq[kMaxInflightOps] = {};
  bool active[kMaxInflightOps] = {false};
  size_t const max_inflight_ops = max_one_sided_inflight_ops();

  size_t next_iov = 0;
  size_t num_completed = 0;
  size_t num_inflight = 0;

  while (num_completed < num_iovs) {
    while (next_iov < num_iovs && num_inflight < max_inflight_ops) {
      size_t slot = 0;
      while (slot < kMaxInflightOps && active[slot]) {
        slot++;
      }
      UCCL_DCHECK(slot < kMaxInflightOps);

      auto rc =
          send_group != nullptr
              ? uccl_write_async_on_group(send_group, conn, mhandle_v[next_iov],
                                          src_v[next_iov], size_v[next_iov],
                                          slot_item_v[next_iov], &ureq[slot])
              : uccl_write_async(ep_, conn, mhandle_v[next_iov],
                                 src_v[next_iov], size_v[next_iov],
                                 slot_item_v[next_iov], &ureq[slot]);
      if (is_retryable_post_failure(rc)) {
        if (is_cxi_transport() && num_inflight == 0) {
          uccl_drive_send(ep_);
          uccl_drive_recv(ep_);
          std::this_thread::yield();
          continue;
        }
        break;
      }
      if (is_cxi_transport() && rc < 0) {
        UCCL_LOG(ERROR) << "writev failed to post iov " << next_iov
                        << ": rc=" << rc;
        return false;
      }
      active[slot] = true;
      next_iov++;
      num_inflight++;
    }

    auto _ = inside_python ? (check_python_signals(), nullptr) : nullptr;

    // Flush all batched WRs in one doorbell per channel, then drive the send
    // poller once for this pass.
    if (num_inflight > 0) {
      uccl_flush_send(ep_);
      uccl_drive_send(ep_);
    }

    if (send_group != nullptr) {
      // Fast path: avoid per-slot mutex + map lookup.
      for (size_t slot = 0; slot < kMaxInflightOps; slot++) {
        if (!active[slot]) continue;
        if (uccl_check_wr_fast(send_group, ureq[slot].engine_idx)) {
          active[slot] = false;
          num_completed++;
          num_inflight--;
        }
      }
    } else {
      for (size_t slot = 0; slot < kMaxInflightOps; slot++) {
        if (!active[slot]) continue;
        if (uccl_check_ureq_once(ep_, &ureq[slot])) {
          active[slot] = false;
          num_completed++;
          num_inflight--;
        }
      }
    }
  }

  return true;
}

bool Endpoint::writev_async(uint64_t conn_id, std::vector<uint64_t> mr_id_v,
                            std::vector<void*> src_v,
                            std::vector<size_t> size_v,
                            std::vector<FifoItem> slot_item_v, size_t num_iovs,
                            uint64_t* transfer_id) {
  // Use move semantics to reduce memory copies
  auto task_ptr =
      create_writev_task(conn_id, std::move(src_v), std::move(size_v),
                         std::move(mr_id_v), std::move(slot_item_v));
  if (unlikely(task_ptr == nullptr)) {
    return false;
  }

  auto* status = new TransferStatus();
  status->task_ptr = task_ptr;
  task_ptr->status_ptr = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  UnifiedTask* task_raw = task_ptr.get();

  while (jring_mp_enqueue_bulk(send_unified_task_ring_, &task_raw, 1,
                               nullptr) != 1) {
  }
  send_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::advertise(uint64_t mr_id, void* addr, size_t len,
                         char* out_buf) {
  auto mhandle = get_mhandle(mr_id);
  if (unlikely(mhandle == nullptr)) {
    std::cerr << "[advertise] Error: Invalid mr_id " << mr_id << std::endl;
    return false;
  }
  if (prepare_fifo_metadata(ep_, mhandle, addr, len, out_buf) == -1)
    return false;
  return true;
}

bool Endpoint::advertisev(std::vector<uint64_t> mr_id_v,
                          std::vector<void*> addr_v, std::vector<size_t> len_v,
                          std::vector<char*> out_buf_v, size_t num_iovs) {
  std::vector<P2PMhandle*> mhandles(num_iovs);

  for (int i = 0; i < num_iovs; i++) {
    mhandles[i] = get_mhandle(mr_id_v[i]);
    if (unlikely(mhandles[i] == nullptr)) {
      std::cerr << "[advertisev] Error: Invalid mr_id " << mr_id_v[i]
                << std::endl;
      return false;
    }
  }

  for (size_t i = 0; i < num_iovs; ++i) {
    if (prepare_fifo_metadata(ep_, mhandles[i], addr_v[i], len_v[i],
                              out_buf_v[i]) == -1) {
      return false;
    }
  }
  return true;
}

bool Endpoint::connect_local(std::string const& remote_gpu_bdf,
                             uint64_t& conn_id, bool same_process) {
  constexpr int kMaxRetry = 20;
  constexpr int kRetrySleepMs = 10;

  Conn* conn = new Conn;
  conn->remote_gpu_bdf_ = remote_gpu_bdf;

  if (same_process) {
    // Same process: use inbox_rings_ directly, no shm attach needed.
    auto it = inbox_rings_.find(remote_gpu_bdf);
    if (it != inbox_rings_.end()) {
      conn->remote_inbox_ = it->second;
    }
    conn->shm_attached_ = false;
  } else {
    // cross process: attach remote inbox via shared memory
    conn->remote_inbox_.shm_name = shm_ring_name(gpu_bus_id_, remote_gpu_bdf);

    jring_t* ring = nullptr;
    for (int attempt = 0; attempt < kMaxRetry; attempt++) {
      ring = uccl::attach_shared_ring(conn->remote_inbox_.shm_name.c_str(),
                                      conn->remote_inbox_.shm_fd,
                                      conn->remote_inbox_.shm_size);
      if (ring != nullptr) {
        break;
      }
      // retry
      std::this_thread::sleep_for(std::chrono::milliseconds(kRetrySleepMs));
    }

    if (ring == nullptr) {
      std::cout << "connect_local failed: cannot attach remote inbox ring "
                << conn->remote_inbox_.shm_name << " after " << kMaxRetry
                << " retries" << std::endl;
      delete conn;
      return false;
    }

    conn->remote_inbox_.ring = ring;
    conn->shm_attached_ = true;
  }

  // allocate conn_id
  conn_id = next_conn_id_.fetch_add(1);
  conn->conn_id_ = conn_id;

  if (!same_process) {
    // Cross-process: send CONNECT handshake via shm ring.
    ShmMsg msg;
    std::strncpy(msg.src_bdf, gpu_bus_id_.c_str(), sizeof(msg.src_bdf) - 1);
    msg.src_bdf[sizeof(msg.src_bdf) - 1] = '\0';
    msg.type = ShmMsgType::CONNECT;
    msg.completion = 0;
    shm_ring_send(conn->remote_inbox_.ring, msg);
  }

  {
    std::unique_lock lock(conn_mu_);
    ConnID dummy{nullptr, 0, 0, 0};
    conn->uccl_conn_id_ = dummy;
    conn->is_local_ = true;
    conn_id_to_conn_[conn_id] = conn;
  }

  return true;
}

bool Endpoint::accept_local(std::string& remote_gpu_bdf, uint64_t& conn_id) {
  std::cout << "Waiting for local CONNECT..." << std::endl;

  // recv CONNECT
  ShmMsg msg;
  while (true) {
    for (auto& [bdf, handle] : inbox_rings_) {
      if (bdf == gpu_bus_id_) continue;
      if (handle.ring == nullptr) continue;

      if (!jring_empty(handle.ring)) {
        shm_ring_recv(handle.ring, msg);
        goto got;
      }
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(1));
  }
got:
  UCCL_CHECK(msg.type == ShmMsgType::CONNECT);

  remote_gpu_bdf = std::string(msg.src_bdf);
  conn_id = next_conn_id_.fetch_add(1);

  Conn* conn = new Conn;
  conn->conn_id_ = conn_id;
  conn->remote_gpu_bdf_ = remote_gpu_bdf;

  if (remote_gpu_bdf == gpu_bus_id_) {
    // same GPU: no attach, bind inbox directly
    auto it = inbox_rings_.find(remote_gpu_bdf);
    if (it != inbox_rings_.end()) {
      conn->remote_inbox_ = it->second;
    }
    conn->shm_attached_ = false;
  } else {
    // cross GPU: attach client inbox
    conn->remote_inbox_.shm_name = shm_ring_name(gpu_bus_id_, remote_gpu_bdf);
    conn->remote_inbox_.ring = uccl::attach_shared_ring(
        conn->remote_inbox_.shm_name.c_str(), conn->remote_inbox_.shm_fd,
        conn->remote_inbox_.shm_size);
    conn->shm_attached_ = true;
  }

  {
    std::unique_lock lock(conn_mu_);
    ConnID dummy{nullptr, 0, 0, 0};
    conn->uccl_conn_id_ = dummy;
    conn->is_local_ = true;
    conn_id_to_conn_[conn_id] = conn;
  }

  return true;
}

bool Endpoint::write_ipc(uint64_t conn_id, void const* data, size_t size,
                         IpcTransferInfo const& info) {
  UCCL_CHECK(data != nullptr) << "write_ipc: data pointer is null!";

  // Get connection info
  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[write_ipc] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  bool is_host = (uccl::get_dev_idx(const_cast<void*>(data)) == -1);
  auto memcpy_kind = is_host ? gpuMemcpyHostToDevice : gpuMemcpyDeviceToDevice;

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  // Open the remote IPC memory handle
  void* raw_dst_ptr = nullptr;
  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));
  GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_dst_ptr, info.handle,
                                   gpuIpcMemLazyEnablePeerAccess));

  // Calculate destination pointer with offset
  void* dst_ptr = reinterpret_cast<void*>(
      reinterpret_cast<uintptr_t>(raw_dst_ptr) + info.offset);

  // Perform the memory copy using multiple streams for better performance
  std::vector<gpuStream_t>& dst_streams = ipc_streams_[local_gpu_idx_];
  int num_streams =
      std::min(dst_streams.size(),
               size < kIpcSizePerEngine ? 1 : (size_t)size / kIpcSizePerEngine);
  size_t chunk_size = size / num_streams;

  for (int i = 0; i < num_streams; ++i) {
    // Split data and dst_ptr into n_streams chunks
    void* chunk_data = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(data) + i * chunk_size);
    void* chunk_dst_ptr = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(dst_ptr) + i * chunk_size);
    auto copy_size = i == num_streams - 1 ? size - i * chunk_size : chunk_size;

    GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst_ptr, chunk_data, copy_size,
                                memcpy_kind, dst_streams[i]));
  }

  // Wait for all streams to complete
  for (auto& stream : dst_streams) {
    GPU_RT_CHECK(gpuStreamSynchronize(stream));
  }

  // Close the IPC memory handle
  GPU_RT_CHECK(gpuIpcCloseMemHandle(raw_dst_ptr));

  return true;
}

bool Endpoint::read_ipc(uint64_t conn_id, void* data, size_t size,
                        IpcTransferInfo const& info) {
  UCCL_CHECK(data != nullptr) << "read_ipc: data pointer is null!";

  // Get connection info
  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[read_ipc] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  bool is_host = (uccl::get_dev_idx(data) == -1);
  auto memcpy_kind = is_host ? gpuMemcpyDeviceToHost : gpuMemcpyDeviceToDevice;

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  // Open the remote IPC memory handle
  void* raw_src_ptr = nullptr;
  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));
  GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_src_ptr, info.handle,
                                   gpuIpcMemLazyEnablePeerAccess));

  // Calculate source pointer with offset
  void* src_ptr = reinterpret_cast<void*>(
      reinterpret_cast<uintptr_t>(raw_src_ptr) + info.offset);

  // Perform the memory copy using multiple streams for better performance
  std::vector<gpuStream_t>& src_streams = ipc_streams_[local_gpu_idx_];
  int num_streams =
      std::min(src_streams.size(),
               size < kIpcSizePerEngine ? 1 : (size_t)size / kIpcSizePerEngine);
  size_t chunk_size = size / num_streams;

  for (int i = 0; i < num_streams; ++i) {
    // Split src_ptr and data into n_streams chunks
    void* chunk_src_ptr = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(src_ptr) + i * chunk_size);
    void* chunk_data = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(data) + i * chunk_size);
    auto copy_size = i == num_streams - 1 ? size - i * chunk_size : chunk_size;

    GPU_RT_CHECK(gpuMemcpyAsync(chunk_data, chunk_src_ptr, copy_size,
                                memcpy_kind, src_streams[i]));
  }

  // Wait for all streams to complete
  for (auto& stream : src_streams) {
    GPU_RT_CHECK(gpuStreamSynchronize(stream));
  }

  // Close the IPC memory handle
  GPU_RT_CHECK(gpuIpcCloseMemHandle(raw_src_ptr));

  return true;
}

bool Endpoint::writev_ipc(uint64_t conn_id, std::vector<void const*> data_v,
                          std::vector<size_t> size_v,
                          std::vector<IpcTransferInfo> info_v,
                          size_t num_iovs) {
  UCCL_CHECK_EQ(data_v.size(), num_iovs) << "writev_ipc: data_v size mismatch";
  UCCL_CHECK_EQ(size_v.size(), num_iovs) << "writev_ipc: size_v size mismatch";
  UCCL_CHECK_EQ(info_v.size(), num_iovs) << "writev_ipc: info_v size mismatch";

  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[writev_ipc] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));
  std::vector<gpuStream_t>& streams = ipc_streams_[local_gpu_idx_];

  // Open all handles and issue all memcpys before syncing any stream.
  std::vector<void*> raw_ptrs(num_iovs, nullptr);
  for (size_t iov = 0; iov < num_iovs; ++iov) {
    UCCL_CHECK(data_v[iov] != nullptr)
        << "writev_ipc: data_v[" << iov << "] is null";
    GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_ptrs[iov], info_v[iov].handle,
                                     gpuIpcMemLazyEnablePeerAccess));
    void* dst_ptr = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(raw_ptrs[iov]) + info_v[iov].offset);

    bool is_host = (uccl::get_dev_idx(const_cast<void*>(data_v[iov])) == -1);
    auto memcpy_kind =
        is_host ? gpuMemcpyHostToDevice : gpuMemcpyDeviceToDevice;

    size_t sz = size_v[iov];
    int num_streams = std::min(
        streams.size(), sz < kIpcSizePerEngine ? 1 : sz / kIpcSizePerEngine);
    size_t chunk_size = sz / num_streams;

    for (int i = 0; i < num_streams; ++i) {
      void* chunk_src = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(data_v[iov]) + i * chunk_size);
      void* chunk_dst = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(dst_ptr) + i * chunk_size);
      auto copy_size = i == num_streams - 1 ? sz - i * chunk_size : chunk_size;
      GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst, chunk_src, copy_size, memcpy_kind,
                                  streams[i]));
    }
  }

  // Sync all streams, then close all handles.
  for (auto& stream : streams) {
    GPU_RT_CHECK(gpuStreamSynchronize(stream));
  }
  for (size_t iov = 0; iov < num_iovs; ++iov) {
    GPU_RT_CHECK(gpuIpcCloseMemHandle(raw_ptrs[iov]));
  }

  return true;
}

bool Endpoint::readv_ipc(uint64_t conn_id, std::vector<void*> data_v,
                         std::vector<size_t> size_v,
                         std::vector<IpcTransferInfo> info_v, size_t num_iovs) {
  UCCL_CHECK_EQ(data_v.size(), num_iovs) << "readv_ipc: data_v size mismatch";
  UCCL_CHECK_EQ(size_v.size(), num_iovs) << "readv_ipc: size_v size mismatch";
  UCCL_CHECK_EQ(info_v.size(), num_iovs) << "readv_ipc: info_v size mismatch";

  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[readv_ipc] Error: Invalid conn_id " << conn_id << std::endl;
    return false;
  }

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));
  std::vector<gpuStream_t>& streams = ipc_streams_[local_gpu_idx_];

  // Open all handles and issue all memcpys before syncing any stream.
  std::vector<void*> raw_ptrs(num_iovs, nullptr);
  for (size_t iov = 0; iov < num_iovs; ++iov) {
    UCCL_CHECK(data_v[iov] != nullptr)
        << "readv_ipc: data_v[" << iov << "] is null";
    GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_ptrs[iov], info_v[iov].handle,
                                     gpuIpcMemLazyEnablePeerAccess));
    void* src_ptr = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(raw_ptrs[iov]) + info_v[iov].offset);

    bool is_host = (uccl::get_dev_idx(data_v[iov]) == -1);
    auto memcpy_kind =
        is_host ? gpuMemcpyDeviceToHost : gpuMemcpyDeviceToDevice;

    size_t sz = size_v[iov];
    int num_streams = std::min(
        streams.size(), sz < kIpcSizePerEngine ? 1 : sz / kIpcSizePerEngine);
    size_t chunk_size = sz / num_streams;

    for (int i = 0; i < num_streams; ++i) {
      void* chunk_src = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(src_ptr) + i * chunk_size);
      void* chunk_dst = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(data_v[iov]) + i * chunk_size);
      auto copy_size = i == num_streams - 1 ? sz - i * chunk_size : chunk_size;
      GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst, chunk_src, copy_size, memcpy_kind,
                                  streams[i]));
    }
  }

  // Sync all streams, then close all handles.
  for (auto& stream : streams) {
    GPU_RT_CHECK(gpuStreamSynchronize(stream));
  }
  for (size_t iov = 0; iov < num_iovs; ++iov) {
    GPU_RT_CHECK(gpuIpcCloseMemHandle(raw_ptrs[iov]));
  }

  return true;
}

bool Endpoint::write_ipc_async(uint64_t conn_id, void const* data, size_t size,
                               IpcTransferInfo const& info,
                               uint64_t* transfer_id) {
  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[write_ipc_async] Error: Invalid conn_id " << conn_id
              << std::endl;
    return false;
  }

  bool is_host = (uccl::get_dev_idx(const_cast<void*>(data)) == -1);
  auto memcpy_kind = is_host ? gpuMemcpyHostToDevice : gpuMemcpyDeviceToDevice;

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  int target_gpu = (info.gpu_idx >= 0) ? info.gpu_idx : local_gpu_idx_;
  GPU_RT_CHECK(gpuSetDevice(target_gpu));

  void* raw_dst_ptr = nullptr;
  void* dst_ptr = nullptr;
  if (info.direct_addr != 0) {
    // Same-process: use virtual address directly, no IPC handle needed.
    dst_ptr = reinterpret_cast<void*>(info.direct_addr);
  } else {
    GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_dst_ptr, info.handle,
                                     gpuIpcMemLazyEnablePeerAccess));
    dst_ptr = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(raw_dst_ptr) +
                                      info.offset);
  }

  std::vector<gpuStream_t>& streams = ipc_streams_[target_gpu];
  int num_streams =
      std::min(streams.size(),
               size < kIpcSizePerEngine ? 1 : (size_t)size / kIpcSizePerEngine);
  size_t chunk_size = size / num_streams;

  auto* op = new IpcInflightOp{{}, raw_dst_ptr, nullptr, target_gpu};
  op->events.resize(num_streams);
  for (int i = 0; i < num_streams; ++i) {
    void* chunk_data = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(data) + i * chunk_size);
    void* chunk_dst = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(dst_ptr) + i * chunk_size);
    auto copy_size = i == num_streams - 1 ? size - i * chunk_size : chunk_size;
    GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst, chunk_data, copy_size, memcpy_kind,
                                streams[i]));
    GPU_RT_CHECK(
        gpuEventCreateWithFlags(&op->events[i], gpuEventDisableTiming));
    GPU_RT_CHECK(gpuEventRecord(op->events[i], streams[i]));
  }

  auto* status = new TransferStatus();
  op->status = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  while (jring_mp_enqueue_bulk(ipc_inflight_ring_, &op, 1, nullptr) != 1) {
  }
  ipc_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::read_ipc_async(uint64_t conn_id, void* data, size_t size,
                              IpcTransferInfo const& info,
                              uint64_t* transfer_id) {
  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[read_ipc_async] Error: Invalid conn_id " << conn_id
              << std::endl;
    return false;
  }

  bool is_host = (uccl::get_dev_idx(data) == -1);
  auto memcpy_kind = is_host ? gpuMemcpyDeviceToHost : gpuMemcpyDeviceToDevice;

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  int target_gpu = (info.gpu_idx >= 0) ? info.gpu_idx : local_gpu_idx_;
  GPU_RT_CHECK(gpuSetDevice(target_gpu));

  void* raw_src_ptr = nullptr;
  void* src_ptr = nullptr;
  if (info.direct_addr != 0) {
    src_ptr = reinterpret_cast<void*>(info.direct_addr);
  } else {
    GPU_RT_CHECK(gpuIpcOpenMemHandle(&raw_src_ptr, info.handle,
                                     gpuIpcMemLazyEnablePeerAccess));
    src_ptr = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(raw_src_ptr) +
                                      info.offset);
  }

  std::vector<gpuStream_t>& streams = ipc_streams_[target_gpu];
  int num_streams =
      std::min(streams.size(),
               size < kIpcSizePerEngine ? 1 : (size_t)size / kIpcSizePerEngine);
  size_t chunk_size = size / num_streams;

  auto* op = new IpcInflightOp{{}, raw_src_ptr, nullptr, target_gpu};
  op->events.resize(num_streams);
  for (int i = 0; i < num_streams; ++i) {
    void* chunk_src = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(src_ptr) + i * chunk_size);
    void* chunk_data = reinterpret_cast<void*>(
        reinterpret_cast<uintptr_t>(data) + i * chunk_size);
    auto copy_size = i == num_streams - 1 ? size - i * chunk_size : chunk_size;
    GPU_RT_CHECK(gpuMemcpyAsync(chunk_data, chunk_src, copy_size, memcpy_kind,
                                streams[i]));
    GPU_RT_CHECK(
        gpuEventCreateWithFlags(&op->events[i], gpuEventDisableTiming));
    GPU_RT_CHECK(gpuEventRecord(op->events[i], streams[i]));
  }

  auto* status = new TransferStatus();
  op->status = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  while (jring_mp_enqueue_bulk(ipc_inflight_ring_, &op, 1, nullptr) != 1) {
  }
  ipc_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::writev_ipc_async(uint64_t conn_id,
                                std::vector<void const*> data_v,
                                std::vector<size_t> size_v,
                                std::vector<IpcTransferInfo> info_v,
                                size_t num_iovs, uint64_t* transfer_id) {
  UCCL_CHECK_EQ(data_v.size(), num_iovs)
      << "writev_ipc_async: data_v size mismatch";
  UCCL_CHECK_EQ(size_v.size(), num_iovs)
      << "writev_ipc_async: size_v size mismatch";
  UCCL_CHECK_EQ(info_v.size(), num_iovs)
      << "writev_ipc_async: info_v size mismatch";

  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[writev_ipc_async] Error: Invalid conn_id " << conn_id
              << std::endl;
    return false;
  }

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  int target_gpu = (num_iovs > 0 && info_v[0].gpu_idx >= 0) ? info_v[0].gpu_idx
                                                            : local_gpu_idx_;
  GPU_RT_CHECK(gpuSetDevice(target_gpu));
  std::vector<gpuStream_t>& streams = ipc_streams_[target_gpu];

  // Use raw_ptr=nullptr to signal vectorized op to the poller thread.
  auto* op = new IpcInflightOp{{}, nullptr, nullptr, -1};
  op->raw_ptrs_v.resize(num_iovs);
  op->gpu_idxs_v.assign(num_iovs, target_gpu);

  for (size_t iov = 0; iov < num_iovs; ++iov) {
    void* dst_ptr = nullptr;
    if (info_v[iov].direct_addr != 0) {
      dst_ptr = reinterpret_cast<void*>(info_v[iov].direct_addr);
    } else {
      GPU_RT_CHECK(gpuIpcOpenMemHandle(&op->raw_ptrs_v[iov], info_v[iov].handle,
                                       gpuIpcMemLazyEnablePeerAccess));
      dst_ptr = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(op->raw_ptrs_v[iov]) +
          info_v[iov].offset);
    }

    bool is_host = (uccl::get_dev_idx(const_cast<void*>(data_v[iov])) == -1);
    auto memcpy_kind =
        is_host ? gpuMemcpyHostToDevice : gpuMemcpyDeviceToDevice;

    size_t sz = size_v[iov];
    int num_streams = std::min(
        streams.size(), sz < kIpcSizePerEngine ? 1 : sz / kIpcSizePerEngine);
    size_t chunk_size = sz / num_streams;

    for (int i = 0; i < num_streams; ++i) {
      void* chunk_src = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(data_v[iov]) + i * chunk_size);
      void* chunk_dst = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(dst_ptr) + i * chunk_size);
      auto copy_size = i == num_streams - 1 ? sz - i * chunk_size : chunk_size;
      GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst, chunk_src, copy_size, memcpy_kind,
                                  streams[i]));
      gpuEvent_t ev;
      GPU_RT_CHECK(gpuEventCreateWithFlags(&ev, gpuEventDisableTiming));
      GPU_RT_CHECK(gpuEventRecord(ev, streams[i]));
      op->events.push_back(ev);
    }
  }

  auto* status = new TransferStatus();
  op->status = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  while (jring_mp_enqueue_bulk(ipc_inflight_ring_, &op, 1, nullptr) != 1) {
  }
  ipc_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::readv_ipc_async(uint64_t conn_id, std::vector<void*> data_v,
                               std::vector<size_t> size_v,
                               std::vector<IpcTransferInfo> info_v,
                               size_t num_iovs, uint64_t* transfer_id) {
  UCCL_CHECK_EQ(data_v.size(), num_iovs)
      << "readv_ipc_async: data_v size mismatch";
  UCCL_CHECK_EQ(size_v.size(), num_iovs)
      << "readv_ipc_async: size_v size mismatch";
  UCCL_CHECK_EQ(info_v.size(), num_iovs)
      << "readv_ipc_async: info_v size mismatch";

  Conn* conn = get_conn(conn_id);
  if (unlikely(conn == nullptr)) {
    std::cerr << "[readv_ipc_async] Error: Invalid conn_id " << conn_id
              << std::endl;
    return false;
  }

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  int target_gpu = (num_iovs > 0 && info_v[0].gpu_idx >= 0) ? info_v[0].gpu_idx
                                                            : local_gpu_idx_;
  GPU_RT_CHECK(gpuSetDevice(target_gpu));
  std::vector<gpuStream_t>& streams = ipc_streams_[target_gpu];

  // Use raw_ptr=nullptr to signal vectorized op to the poller thread.
  auto* op = new IpcInflightOp{{}, nullptr, nullptr, -1};
  op->raw_ptrs_v.resize(num_iovs);
  op->gpu_idxs_v.assign(num_iovs, target_gpu);

  for (size_t iov = 0; iov < num_iovs; ++iov) {
    void* src_ptr = nullptr;
    if (info_v[iov].direct_addr != 0) {
      src_ptr = reinterpret_cast<void*>(info_v[iov].direct_addr);
    } else {
      GPU_RT_CHECK(gpuIpcOpenMemHandle(&op->raw_ptrs_v[iov], info_v[iov].handle,
                                       gpuIpcMemLazyEnablePeerAccess));
      src_ptr = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(op->raw_ptrs_v[iov]) +
          info_v[iov].offset);
    }

    bool is_host = (uccl::get_dev_idx(data_v[iov]) == -1);
    auto memcpy_kind =
        is_host ? gpuMemcpyDeviceToHost : gpuMemcpyDeviceToDevice;

    size_t sz = size_v[iov];
    int num_streams = std::min(
        streams.size(), sz < kIpcSizePerEngine ? 1 : sz / kIpcSizePerEngine);
    size_t chunk_size = sz / num_streams;

    for (int i = 0; i < num_streams; ++i) {
      void* chunk_src = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(src_ptr) + i * chunk_size);
      void* chunk_dst = reinterpret_cast<void*>(
          reinterpret_cast<uintptr_t>(data_v[iov]) + i * chunk_size);
      auto copy_size = i == num_streams - 1 ? sz - i * chunk_size : chunk_size;
      GPU_RT_CHECK(gpuMemcpyAsync(chunk_dst, chunk_src, copy_size, memcpy_kind,
                                  streams[i]));
      gpuEvent_t ev;
      GPU_RT_CHECK(gpuEventCreateWithFlags(&ev, gpuEventDisableTiming));
      GPU_RT_CHECK(gpuEventRecord(ev, streams[i]));
      op->events.push_back(ev);
    }
  }

  auto* status = new TransferStatus();
  op->status = status;
  *transfer_id = reinterpret_cast<uint64_t>(status);

  while (jring_mp_enqueue_bulk(ipc_inflight_ring_, &op, 1, nullptr) != 1) {
  }
  ipc_proxy_adaptive_sleeper_.maybe_wake_proxy_thread();

  return true;
}

bool Endpoint::advertise_ipc(uint64_t conn_id, void* addr, size_t len,
                             char* out_buf) {
  UCCL_CHECK(addr != nullptr) << "advertise_ipc: addr pointer is null!";
  UCCL_CHECK(out_buf != nullptr) << "advertise_ipc: out_buf pointer is null!";

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));

  // Generate IPC memory handle for the address
  IpcTransferInfo transfer_info = {};  // Initialize to zero
  transfer_info.size = len;
  transfer_info.operation = 1;  // response

  // Calculate aligned address and offset
  auto addr_aligned = reinterpret_cast<uintptr_t>(addr) & ~(kIpcAlignment - 1);
  auto addr_offset = reinterpret_cast<uintptr_t>(addr) - addr_aligned;
  transfer_info.offset = addr_offset;

  GPU_RT_CHECK(gpuIpcGetMemHandle(&transfer_info.handle,
                                  reinterpret_cast<void*>(addr_aligned)));

  // Copy the transfer info to output buffer
  std::memcpy(out_buf, &transfer_info, sizeof(transfer_info));

  return true;
}

bool Endpoint::advertisev_ipc(uint64_t conn_id, std::vector<void*> addr_v,
                              std::vector<size_t> len_v,
                              std::vector<char*> out_buf_v, size_t num_iovs) {
  UCCL_CHECK_EQ(addr_v.size(), num_iovs) << "addr_v size mismatch";
  UCCL_CHECK_EQ(len_v.size(), num_iovs) << "len_v size mismatch";
  UCCL_CHECK_EQ(out_buf_v.size(), num_iovs) << "out_buf_v size mismatch";

  int orig_device;
  GPU_RT_CHECK(gpuGetDevice(&orig_device));
  auto dev_reset =
      uccl::finally([&]() { GPU_RT_CHECK(gpuSetDevice(orig_device)); });

  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));

  for (size_t i = 0; i < num_iovs; ++i) {
    UCCL_CHECK(addr_v[i] != nullptr)
        << "advertisev_ipc: addr_v[" << i << "] is null!";
    UCCL_CHECK(out_buf_v[i] != nullptr)
        << "advertisev_ipc: out_buf_v[" << i << "] is null!";

    // Generate IPC memory handle for each address
    IpcTransferInfo transfer_info = {};  // Initialize to zero
    transfer_info.size = len_v[i];
    transfer_info.operation = 1;  // response

    // Calculate aligned address and offset
    auto addr_aligned =
        reinterpret_cast<uintptr_t>(addr_v[i]) & ~(kIpcAlignment - 1);
    auto addr_offset = reinterpret_cast<uintptr_t>(addr_v[i]) - addr_aligned;
    transfer_info.offset = addr_offset;

    GPU_RT_CHECK(gpuIpcGetMemHandle(&transfer_info.handle,
                                    reinterpret_cast<void*>(addr_aligned)));

    // Copy the transfer info to output buffer
    std::memcpy(out_buf_v[i], &transfer_info, sizeof(transfer_info));
  }

  return true;
}

bool Endpoint::add_remote_endpoint(std::vector<uint8_t> const& metadata,
                                   uint64_t& conn_id) {
  {
    // Check if we have connected to the remote endpoint before
    std::shared_lock<std::shared_mutex> lock(conn_mu_);
    auto it = remote_endpoint_to_conn_id_.find(metadata);
    if (it != remote_endpoint_to_conn_id_.end()) {
      conn_id = it->second;
      return true;
    }
  }
  auto [remote_ip, remote_port, remote_gpu_bdf] = parse_metadata(metadata);
  std::string local_ip = uccl::get_oob_ip();
  bool is_local = (remote_ip == local_ip);
  bool success;
  if (is_local && !remote_gpu_bdf.empty()) {
    success = connect_local(remote_gpu_bdf, conn_id);
  } else {
    success = connect(remote_ip, 0, remote_port, conn_id);
  }
  if (success) {
    std::unique_lock<std::shared_mutex> lock(conn_mu_);
    auto conn_it = conn_id_to_conn_.find(conn_id);
    if (conn_it != conn_id_to_conn_.end()) {
      conn_it->second->ip_addr_ = remote_ip;
      conn_it->second->remote_port_ = remote_port;
      conn_it->second->remote_gpu_bdf_ = remote_gpu_bdf;
    }
    remote_endpoint_to_conn_id_[metadata] = conn_id;
    return true;
  }
  return false;
}

bool Endpoint::remove_remote_endpoint(uint64_t conn_id) {
  std::unique_lock<std::shared_mutex> lock(conn_mu_);

  auto it = conn_id_to_conn_.find(conn_id);
  if (it == conn_id_to_conn_.end()) {
    std::cerr << "[remove_remote_endpoint] Error: Invalid conn_id " << conn_id
              << std::endl;
    return false;
  }

  Conn* conn = it->second;
  uint64_t loopback_conn_id = conn->rdma_loopback_conn_id_;

  // Detach shared memory if this was a local connection
  if (conn->shm_attached_) {
    auto& shm = conn->remote_inbox_;
    uccl::detach_shared_ring(shm.ring, shm.shm_fd, shm.shm_size);
  }

  // Remove from the metadata memoization map
  for (auto mit = remote_endpoint_to_conn_id_.begin();
       mit != remote_endpoint_to_conn_id_.end(); ++mit) {
    if (mit->second == conn_id) {
      remote_endpoint_to_conn_id_.erase(mit);
      break;
    }
  }

  delete conn;
  conn_id_to_conn_.erase(it);

  if (loopback_conn_id != UINT64_MAX) {
    auto loopback_it = conn_id_to_conn_.find(loopback_conn_id);
    if (loopback_it != conn_id_to_conn_.end()) {
      Conn* loopback_conn = loopback_it->second;
      if (loopback_conn->shm_attached_) {
        auto& shm = loopback_conn->remote_inbox_;
        uccl::detach_shared_ring(shm.ring, shm.shm_fd, shm.shm_size);
      }
      delete loopback_conn;
      conn_id_to_conn_.erase(loopback_it);
    }
  }

  return true;
}

bool Endpoint::start_passive_accept() {
  if (!passive_accept_) {
    passive_accept_stop_.store(false, std::memory_order_release);
    passive_accept_thread_ =
        std::thread(&Endpoint::passive_accept_thread_func, this);
    passive_accept_local_thread_ =
        std::thread(&Endpoint::passive_accept_local_thread_func, this);
    passive_accept_ = true;
  }
  return true;
}

bool Endpoint::poll_async(uint64_t transfer_id, bool* is_done) {
  auto* status = reinterpret_cast<TransferStatus*>(transfer_id);
  if (status->poll_net_ureq && !status->done.load(std::memory_order_acquire)) {
    if (uccl_poll_ureq_once(ep_, &status->ureq)) {
      status->done.store(true, std::memory_order_release);
    }
  }
  *is_done = status->done.load(std::memory_order_acquire);
  if (*is_done) {
    delete status;
  }
  return true;
}

int Endpoint::get_sock_fd(uint64_t conn_id) const {
  auto it = conn_id_to_conn_.find(conn_id);
  if (it == conn_id_to_conn_.end()) {
    return -1;
  }
  return it->second->uccl_conn_id_.sock_fd;
}

uint64_t Endpoint::conn_id_of_rank(int rank) const {
  auto it = rank2conn_.find(rank);
  return it != rank2conn_.end() ? it->second : UINT64_MAX;
}

std::shared_ptr<EpollClient> Endpoint::get_oob_client() const {
  return std::visit(
      [](auto const& s) -> std::shared_ptr<EpollClient> {
        return s->get_oob_client();
      },
      ep_);
}

std::string Endpoint::get_oob_conn_key(uint64_t conn_id) const {
  auto it = conn_id_to_conn_.find(conn_id);
  if (it == conn_id_to_conn_.end()) {
    return "";
  }
  uint64_t peer_id = it->second->uccl_conn_id_.peer_id;

  return std::visit(
      [&](auto const& s) -> std::string {
        return s->get_oob_conn_key(peer_id);
      },
      ep_);
}

int Endpoint::send_notification(uint64_t conn_id,
                                NotifyMsg const& notification) const {
  Conn* conn = get_conn(conn_id);
  if (conn == nullptr) {
    return -1;
  }
  uint64_t peer_id = conn->uccl_conn_id_.peer_id;
  if (peer_id == UINT64_MAX) {
    return -1;
  }
  if (is_nccl_transport()) {
    auto* nccl_ep = std::get_if<std::shared_ptr<NCCLEndpoint>>(&ep_);
    if (!nccl_ep || !*nccl_ep) return -1;
    return (*nccl_ep)->send_notification(peer_id, notification);
  }
  if (is_cxi_transport()) {
    auto* cxi_ep = std::get_if<std::shared_ptr<CxiEndpoint>>(&ep_);
    if (!cxi_ep || !*cxi_ep) return -1;
    return (*cxi_ep)->send_notification(peer_id, notification);
  }
  return -1;
}

MR* Endpoint::get_mr(uint64_t mr_id) const {
  std::shared_lock<std::shared_mutex> lock(mr_mu_);
  auto it = mr_id_to_mr_.find(mr_id);
  if (it == mr_id_to_mr_.end()) {
    return nullptr;
  }
  return it->second;
}

P2PMhandle* Endpoint::get_mhandle(uint64_t mr_id) const {
  auto mr = get_mr(mr_id);
  if (unlikely(mr == nullptr)) {
    return nullptr;
  }
  return mr->mhandle_;
}

Conn* Endpoint::get_conn(uint64_t conn_id) const {
  std::shared_lock<std::shared_mutex> lock(conn_mu_);
  auto it = conn_id_to_conn_.find(conn_id);
  if (it == conn_id_to_conn_.end()) {
    return nullptr;
  }
  return it->second;
}

GenericEndpoint Endpoint::get_endpoint() const { return ep_; }

void Endpoint::initialize_engine() {
  int n_streams = std::max(1, (int)kNumGpuRtStreams);
  GPU_RT_CHECK(gpuSetDevice(local_gpu_idx_));

  if (is_nccl_transport()) {
    numa_node_ = get_numa_node_from_iface();
  } else if (is_cxi_transport()) {
    numa_node_ = get_numa_node_from_iface();
  } else {
    numa_node_ = RdmaDeviceManager::instance().get_numa_node(
        RdmaDeviceManager::instance().get_best_dev_idx(local_gpu_idx_)[0]);
  }

  // Initialize rdma contexts for devices used by the GPU
  initialize_rdma_ctx_for_gpu(ep_, local_gpu_idx_);
  std::cout << "Lazy creation of engine for GPU " << local_gpu_idx_
            << std::endl;

  // Initialize task rings
  send_unified_task_ring_ =
      uccl::create_ring(sizeof(UnifiedTask*), kTaskRingSize);
  recv_unified_task_ring_ =
      uccl::create_ring(sizeof(UnifiedTask*), kTaskRingSize);

  send_proxy_thread_ = std::thread(&Endpoint::send_proxy_thread_func, this);
  recv_proxy_thread_ = std::thread(&Endpoint::recv_proxy_thread_func, this);

  ipc_inflight_ring_ = uccl::create_ring(sizeof(IpcInflightOp*), kTaskRingSize);
  ipc_poller_thread_ = std::thread(&Endpoint::ipc_poller_thread_func, this);

  // Initialize ShmChannel for local connections (use BDF for ring names)
  size_t elem_sz = sizeof(ShmMsg);
  size_t elem_cnt = ShmRingDefaultElemCnt;
  auto all_bdfs = uccl::enumerate_all_gpu_bdfs();
  for (auto const& sender_bdf : all_bdfs) {
    auto name = shm_ring_name(sender_bdf, gpu_bus_id_);
    ShmRingHandle handle;
    handle.shm_name = name;
    bool is_creator = false;
    handle.ring =
        uccl::create_shared_ring(name.c_str(), elem_sz, elem_cnt, handle.shm_fd,
                                 handle.shm_size, &is_creator);
    inbox_rings_[sender_bdf] = handle;
    inbox_creators_[sender_bdf] = is_creator;
  }
}

void Endpoint::send_proxy_thread_func() {
  uccl::pin_thread_to_numa(numa_node_);
  // Use 16-byte buffer to avoid stringop-overflow warning from jring's 16-byte
  // bulk copy
  alignas(16) char task_buffer[16];
  UnifiedTask* task;
  send_proxy_adaptive_sleeper_.update_timer();

  while (!stop_.load(std::memory_order_acquire)) {
    send_proxy_adaptive_sleeper_.maybe_sleep();

    if (jring_sc_dequeue_bulk(send_unified_task_ring_, task_buffer, 1,
                              nullptr) == 1) {
      task = *reinterpret_cast<UnifiedTask**>(task_buffer);
      switch (task->type) {
        case TaskType::WRITE_NET:
          write(task->conn_id, task->mr_id, task->data, task->size,
                task->slot_item());
          break;
        case TaskType::WRITEV: {
          TaskBatch const& batch = task->task_batch();
          writev(task->conn_id, *batch.mr_id_ptr, *batch.data_ptr,
                 *batch.size_ptr, *batch.slot_item_ptr, batch.num_iovs);
          break;
        }
        default:
          UCCL_LOG(ERROR) << "Unexpected task type in send processing: "
                          << static_cast<int>(task->type);
          break;
      }
      auto* status = task->status_ptr;
      status->task_ptr.reset();
      status->done.store(true, std::memory_order_release);
      send_proxy_adaptive_sleeper_.update_timer();
    }
  }
}

void Endpoint::recv_proxy_thread_func() {
  uccl::pin_thread_to_numa(numa_node_);
  // Use 16-byte buffer to avoid stringop-overflow warning from jring's 16-byte
  // bulk copy
  alignas(16) char task_buffer[16];
  UnifiedTask* task;
  recv_proxy_adaptive_sleeper_.update_timer();

  while (!stop_.load(std::memory_order_acquire)) {
    recv_proxy_adaptive_sleeper_.maybe_sleep();

    if (jring_sc_dequeue_bulk(recv_unified_task_ring_, task_buffer, 1,
                              nullptr) == 1) {
      task = *reinterpret_cast<UnifiedTask**>(task_buffer);
      switch (task->type) {
        case TaskType::READ_NET:
          read(task->conn_id, task->mr_id, task->data, task->size,
               task->slot_item());
          break;
        case TaskType::READV: {
          TaskBatch const& batch = task->task_batch();
          readv(task->conn_id, *batch.mr_id_ptr, *batch.data_ptr,
                *batch.size_ptr, *batch.slot_item_ptr, batch.num_iovs);
          break;
        }
        case TaskType::WRITE_NET:
        default:
          UCCL_LOG(ERROR) << "Unexpected task type in receive processing: "
                          << static_cast<int>(task->type);
          break;
      }
      auto* status = task->status_ptr;
      status->task_ptr.reset();
      status->done.store(true, std::memory_order_release);
      recv_proxy_adaptive_sleeper_.update_timer();
    } else {
      uccl_drive_recv(ep_);
    }
  }
}

void Endpoint::passive_accept_thread_func() {
  std::string ip_addr;
  int remote_gpu_idx;
  uint64_t conn_id;
  while (!stop_.load(std::memory_order_acquire)) {
    (void)accept(ip_addr, remote_gpu_idx, conn_id);
  }
}

void Endpoint::passive_accept_local_thread_func() {
  while (!passive_accept_stop_.load(std::memory_order_acquire)) {
    bool found = false;
    for (auto& [bdf, handle] : inbox_rings_) {
      if (bdf == gpu_bus_id_) continue;
      if (handle.ring == nullptr) continue;
      if (jring_empty(handle.ring)) continue;

      std::string remote_gpu_bdf;
      uint64_t conn_id;
      (void)accept_local(remote_gpu_bdf, conn_id);
      found = true;
      break;
    }
    if (!found) {
      std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
  }
}

void Endpoint::ipc_poller_thread_func() {
  uccl::pin_thread_to_numa(numa_node_);
  std::deque<IpcInflightOp*> active_ops;
  alignas(16) char buf[16];
  int cur_device = -1;
  ipc_proxy_adaptive_sleeper_.update_timer();

  while (!stop_.load(std::memory_order_acquire) || !active_ops.empty()) {
    // only try to sleep if there are no more inflight ops
    if (!active_ops.empty()) {
      ipc_proxy_adaptive_sleeper_.update_timer();
    }
    ipc_proxy_adaptive_sleeper_.maybe_sleep();

    // Drain newly submitted ops into the local active list.
    while (jring_sc_dequeue_bulk(ipc_inflight_ring_, buf, 1, nullptr) == 1) {
      active_ops.push_back(*reinterpret_cast<IpcInflightOp**>(buf));
    }

    for (auto it = active_ops.begin(); it != active_ops.end();) {
      IpcInflightOp* op = *it;
      bool all_done = true;
      for (auto& event : op->events) {
        auto result = gpuEventQuery(event);
        if (result == gpuErrorNotReady) {
          all_done = false;
          break;
        }
        GPU_RT_CHECK(result);
      }
      if (all_done) {
        for (auto& event : op->events) {
          GPU_RT_CHECK(gpuEventDestroy(event));
        }
        if (op->raw_ptr != nullptr) {
          // Scalar op: single handle to close.
          if (cur_device != op->gpu_idx) {
            GPU_RT_CHECK(gpuSetDevice(op->gpu_idx));
            cur_device = op->gpu_idx;
          }
          GPU_RT_CHECK(gpuIpcCloseMemHandle(op->raw_ptr));
        } else {
          // Vectorized op: close all handles.
          for (size_t i = 0; i < op->raw_ptrs_v.size(); ++i) {
            if (op->raw_ptrs_v[i] == nullptr) continue;  // direct_addr path
            if (cur_device != op->gpu_idxs_v[i]) {
              GPU_RT_CHECK(gpuSetDevice(op->gpu_idxs_v[i]));
              cur_device = op->gpu_idxs_v[i];
            }
            GPU_RT_CHECK(gpuIpcCloseMemHandle(op->raw_ptrs_v[i]));
          }
        }
        op->status->done.store(true, std::memory_order_release);
        delete op;
        it = active_ops.erase(it);
      } else {
        ++it;
      }
    }
  }
}

std::shared_ptr<UnifiedTask> Endpoint::create_task(uint64_t conn_id,
                                                   uint64_t mr_id,
                                                   TaskType type, void* data,
                                                   size_t size) {
  auto task = std::make_shared<UnifiedTask>();
  task->type = type;
  task->data = data;
  task->size = size;
  task->conn_id = conn_id;
  task->mr_id = mr_id;
  return task;
}

std::shared_ptr<UnifiedTask> Endpoint::create_batch_task(uint64_t conn_id,
                                                         TaskType type,
                                                         TaskBatch&& batch) {
  auto task = std::make_shared<UnifiedTask>();
  task->type = type;
  task->conn_id = conn_id;
  task->mr_id = 0;
  task->data = nullptr;
  task->size = 0;
  new (&task->specific.batch.task_batch) TaskBatch(std::move(batch));
  return task;
}

std::shared_ptr<UnifiedTask> Endpoint::create_writev_task(
    uint64_t conn_id, std::vector<void*>&& data_v, std::vector<size_t>&& size_v,
    std::vector<uint64_t>&& mr_id_v, std::vector<FifoItem>&& slot_item_v) {
  if (data_v.size() != size_v.size() || size_v.size() != mr_id_v.size() ||
      mr_id_v.size() != slot_item_v.size()) {
    return nullptr;
  }
  size_t num_iovs = data_v.size();
  auto data_ptr = std::make_shared<std::vector<void*>>(std::move(data_v));
  auto size_ptr = std::make_shared<std::vector<size_t>>(std::move(size_v));
  auto mr_id_ptr = std::make_shared<std::vector<uint64_t>>(std::move(mr_id_v));
  auto slot_item_ptr =
      std::make_shared<std::vector<FifoItem>>(std::move(slot_item_v));
  TaskBatch batch;
  batch.num_iovs = num_iovs;
  batch.data_ptr = std::move(data_ptr);
  batch.size_ptr = std::move(size_ptr);
  batch.mr_id_ptr = std::move(mr_id_ptr);
  batch.slot_item_ptr = std::move(slot_item_ptr);
  return create_batch_task(conn_id, TaskType::WRITEV, std::move(batch));
}

std::shared_ptr<UnifiedTask> Endpoint::create_readv_task(
    uint64_t conn_id, std::vector<void*>&& data_v, std::vector<size_t>&& size_v,
    std::vector<uint64_t>&& mr_id_v, std::vector<FifoItem>&& slot_item_v) {
  if (data_v.size() != size_v.size() || size_v.size() != mr_id_v.size() ||
      mr_id_v.size() != slot_item_v.size()) {
    return nullptr;
  }
  size_t num_iovs = data_v.size();
  auto data_ptr = std::make_shared<std::vector<void*>>(std::move(data_v));
  auto size_ptr = std::make_shared<std::vector<size_t>>(std::move(size_v));
  auto mr_id_ptr = std::make_shared<std::vector<uint64_t>>(std::move(mr_id_v));
  auto slot_item_ptr =
      std::make_shared<std::vector<FifoItem>>(std::move(slot_item_v));
  TaskBatch batch;
  batch.num_iovs = num_iovs;
  batch.data_ptr = std::move(data_ptr);
  batch.size_ptr = std::move(size_ptr);
  batch.mr_id_ptr = std::move(mr_id_ptr);
  batch.slot_item_ptr = std::move(slot_item_ptr);
  return create_batch_task(conn_id, TaskType::READV, std::move(batch));
}

std::shared_ptr<UnifiedTask> Endpoint::create_net_task(
    uint64_t conn_id, uint64_t mr_id, TaskType type, void* data, size_t size,
    FifoItem const& slot_item) {
  auto task = create_task(conn_id, mr_id, type, data, size);
  task->slot_item() = slot_item;
  return task;
}

std::shared_ptr<UnifiedTask> Endpoint::create_ipc_task(
    uint64_t conn_id, uint64_t mr_id, TaskType type, void* data, size_t size,
    IpcTransferInfo const& ipc_info) {
  auto task = create_task(conn_id, mr_id, type, data, size);
  task->ipc_info() = ipc_info;
  return task;
}