262 lines
7.2 KiB
C++
262 lines
7.2 KiB
C++
// Copyright 2024 KVCache.AI
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#ifndef RDMA_CONTEXT_H
|
|
#define RDMA_CONTEXT_H
|
|
|
|
#include <gflags/gflags.h>
|
|
#include <glog/logging.h>
|
|
#include <infiniband/verbs.h>
|
|
|
|
#ifdef USE_MLX5DV
|
|
#include <infiniband/mlx5dv.h>
|
|
#endif
|
|
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <cstdint>
|
|
#include <list>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <unordered_map>
|
|
|
|
#include "common.h"
|
|
#include "rdma_gid_probe.h"
|
|
#include "rdma_transport.h"
|
|
#include "transport/transport.h"
|
|
|
|
namespace mooncake {
|
|
|
|
class RdmaEndPoint;
|
|
class RdmaTransport;
|
|
class RdmaContextTestPeer;
|
|
class WorkerPool;
|
|
class EndpointStore;
|
|
|
|
// Enum to represent the network state of the GID found
|
|
enum class GidNetworkState {
|
|
GID_WITH_NETWORK = 0, // Found a GID with network device (best choice)
|
|
GID_WITHOUT_NETWORK = 1, // Found a GID without network device
|
|
GID_NOT_FOUND = 2 // No suitable GID found
|
|
};
|
|
|
|
struct GidSelectionSnapshot {
|
|
std::string gid;
|
|
int gid_index = -1;
|
|
};
|
|
|
|
struct RdmaCq {
|
|
RdmaCq() : native(nullptr), outstanding(0) {}
|
|
ibv_cq *native;
|
|
volatile int outstanding;
|
|
};
|
|
|
|
struct MemoryRegionMeta {
|
|
// mr->addr is not set to starting address for iova based mr. Therefore we
|
|
// track it ourselves.
|
|
void *addr;
|
|
struct ibv_mr *mr;
|
|
};
|
|
|
|
// RdmaContext represents the set of resources controlled by each local NIC,
|
|
// including Memory Region, CQ, EndPoint (QPs), etc.
|
|
class RdmaContext {
|
|
public:
|
|
friend class RdmaContextTestPeer;
|
|
|
|
RdmaContext(RdmaTransport &engine, const std::string &device_name);
|
|
|
|
~RdmaContext();
|
|
|
|
int construct(size_t num_cq_list = 1, size_t num_comp_channels = 1,
|
|
uint8_t port = 1, int gid_index = -1, size_t max_cqe = 4096,
|
|
int max_endpoints = 256);
|
|
|
|
private:
|
|
int deconstruct();
|
|
|
|
public:
|
|
// Memory Region Management
|
|
int registerMemoryRegion(void *addr, size_t length, int access);
|
|
|
|
int unregisterMemoryRegion(void *addr);
|
|
|
|
int preTouchMemory(void *addr, size_t length);
|
|
|
|
uint32_t rkey(void *addr);
|
|
|
|
uint32_t lkey(void *addr);
|
|
|
|
private:
|
|
int registerMemoryRegionInternal(void *addr, size_t length, int access,
|
|
MemoryRegionMeta &mrMeta);
|
|
|
|
using MemoryRegionMap = std::map<uintptr_t, MemoryRegionMeta>;
|
|
|
|
MemoryRegionMap::iterator findMemoryRegionContaining(uintptr_t addr);
|
|
|
|
MemoryRegionMap::const_iterator findMemoryRegionContaining(
|
|
uintptr_t addr) const;
|
|
|
|
public:
|
|
bool active() const { return active_; }
|
|
|
|
void set_active(bool flag) { active_ = flag; }
|
|
|
|
public:
|
|
// EndPoint Management
|
|
std::shared_ptr<RdmaEndPoint> endpoint(const std::string &peer_nic_path);
|
|
|
|
std::shared_ptr<RdmaEndPoint> getEndpointByPtr(
|
|
const RdmaEndPoint *endpoint_ptr);
|
|
|
|
int deleteEndpoint(const std::string &peer_nic_path);
|
|
|
|
// Drain the endpoint store's waiting list. Safe to call on any thread;
|
|
// intended to be invoked periodically from monitorWorker so reclaim is
|
|
// not gated on new endpoint insertions (which can stall under failure
|
|
// load while evictions/deletions continue). See issue #1845.
|
|
void reclaimEndpoints();
|
|
|
|
// Number of endpoints awaiting reclaim. For tests and operator
|
|
// observability.
|
|
size_t waitingListSize() const;
|
|
|
|
// Test-only: push a pre-constructed endpoint into the store's
|
|
// waiting_list_ so the reclaim path can be exercised without standing up
|
|
// a real RDMA QP.
|
|
void testOnlyInsertWaiting(std::shared_ptr<RdmaEndPoint> ep);
|
|
|
|
int disconnectAllEndpoints();
|
|
|
|
// Get the total number of QPs across all endpoints in this context
|
|
size_t getTotalQPNumber() const;
|
|
|
|
public:
|
|
// Device name, such as `mlx5_3`
|
|
std::string deviceName() const { return device_name_; }
|
|
|
|
// NIC Path, such as `192.168.3.76@mlx5_3`
|
|
std::string nicPath() const;
|
|
|
|
public:
|
|
uint16_t lid() const { return lid_; }
|
|
|
|
std::string gid() const;
|
|
|
|
GidSelectionSnapshot gidSelection() const;
|
|
|
|
int gidIndex() const;
|
|
|
|
bool autoGidSelectionEnabled() const { return auto_gid_selection_enabled_; }
|
|
|
|
bool reprobeAutoGid(
|
|
const GidSelectionSnapshot &expected_selection,
|
|
const std::vector<AutoGidSelectionIdentity> &tried_selections = {},
|
|
std::string *previous_gid = nullptr, std::string *next_gid = nullptr);
|
|
|
|
ibv_context *context() const { return context_; }
|
|
|
|
RdmaTransport &engine() const { return engine_; }
|
|
|
|
ibv_pd *pd() const { return pd_; }
|
|
|
|
uint8_t portNum() const { return port_; }
|
|
|
|
uint8_t numLagPorts() const { return num_lag_ports_; }
|
|
|
|
int activeSpeed() const { return active_speed_; }
|
|
|
|
ibv_mtu activeMTU() const { return active_mtu_; }
|
|
|
|
ibv_comp_channel *compChannel();
|
|
|
|
int compVector();
|
|
|
|
int eventFd() const { return event_fd_; }
|
|
|
|
ibv_cq *cq();
|
|
|
|
volatile int *cqOutstandingCount(int cq_index) {
|
|
return &cq_list_[cq_index].outstanding;
|
|
}
|
|
|
|
int cqCount() const { return cq_list_.size(); }
|
|
|
|
int poll(int num_entries, ibv_wc *wc, int cq_index = 0);
|
|
|
|
int socketId();
|
|
|
|
private:
|
|
int openRdmaDevice(const std::string &device_name, uint8_t port,
|
|
int gid_index);
|
|
|
|
int joinNonblockingPollList(int event_fd, int data_fd);
|
|
|
|
GidNetworkState findBestGidIndex(const std::string &device_name,
|
|
struct ibv_context *context,
|
|
ibv_port_attr &port_attr, uint8_t port,
|
|
int &gid_index);
|
|
|
|
public:
|
|
int submitPostSend(const std::vector<Transport::Slice *> &slice_list);
|
|
|
|
private:
|
|
const std::string device_name_;
|
|
RdmaTransport &engine_;
|
|
|
|
ibv_context *context_ = nullptr;
|
|
ibv_pd *pd_ = nullptr;
|
|
uint64_t max_mr_size;
|
|
int event_fd_ = -1;
|
|
|
|
size_t num_comp_channel_ = 0;
|
|
ibv_comp_channel **comp_channel_ = nullptr;
|
|
|
|
uint8_t port_ = 0;
|
|
uint16_t lid_ = 0;
|
|
int gid_index_ = -1;
|
|
int active_speed_ = -1;
|
|
ibv_mtu active_mtu_;
|
|
uint8_t num_lag_ports_ = 0; // 0/1 = not in LAG; ≥2 = LAG active
|
|
ibv_gid gid_;
|
|
mutable std::mutex gid_lock_;
|
|
mutable std::mutex gid_reprobe_lock_;
|
|
bool auto_gid_selection_enabled_ = false;
|
|
|
|
RWSpinlock memory_regions_lock_;
|
|
MemoryRegionMap memory_region_map_;
|
|
std::vector<RdmaCq> cq_list_;
|
|
|
|
std::shared_ptr<EndpointStore> endpoint_store_;
|
|
|
|
std::vector<std::thread> background_thread_;
|
|
std::atomic<bool> threads_running_;
|
|
|
|
std::atomic<int> next_comp_channel_index_;
|
|
std::atomic<int> next_comp_vector_index_;
|
|
std::atomic<int> next_cq_list_index_;
|
|
|
|
std::shared_ptr<WorkerPool> worker_pool_;
|
|
|
|
volatile bool active_;
|
|
};
|
|
|
|
} // namespace mooncake
|
|
|
|
#endif // RDMA_CONTEXT_H
|