endpoint.h

 
#pragma once
 
#include <netdb.h>
 
#include <memory>
#include <string>
#include <vector>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/address.h>
#include <ucxx/component.h>
#include <ucxx/exception.h>
#include <ucxx/inflight_requests.h>
#include <ucxx/listener.h>
#include <ucxx/request.h>
#include <ucxx/typedefs.h>
#include <ucxx/utils/sockaddr.h>
#include <ucxx/worker.h>
 
namespace ucxx {
 
struct EpParamsDeleter {
  void operator()(ucp_ep_params_t* ptr);
};
 
class Endpoint : public Component {
 private:
  ucp_ep_h _handle{nullptr};          
  ucp_ep_h _originalHandle{nullptr};  
  bool _endpointErrorHandling{true};  
  std::unique_ptr<InflightRequests> _inflightRequests{
    std::make_unique<InflightRequests>()};  
  std::mutex _mutex{};  
  ucs_status_t _status{UCS_INPROGRESS};  
  std::atomic<bool> _closing{false};     
  EndpointCloseCallbackUserFunction _closeCallback{nullptr};  
  EndpointCloseCallbackUserData _closeCallbackArg{
    nullptr};  
 
  Endpoint(std::shared_ptr<Component> workerOrListener, bool endpointErrorHandling);
 
  void create(ucp_ep_params_t* params);
 
  [[nodiscard]] std::shared_ptr<Request> registerInflightRequest(std::shared_ptr<Request> request);
 
  friend void endpointErrorCallback(void* arg, ucp_ep_h ep, ucs_status_t status);
 
 public:
  Endpoint()                           = delete;
  Endpoint(const Endpoint&)            = delete;
  Endpoint& operator=(Endpoint const&) = delete;
  Endpoint(Endpoint&& o)               = delete;
  Endpoint& operator=(Endpoint&& o)    = delete;
 
  ~Endpoint();
 
  friend std::shared_ptr<Endpoint> createEndpointFromHostname(std::shared_ptr<Worker> worker,
                                                              std::string ipAddress,
                                                              uint16_t port,
                                                              bool endpointErrorHandling);
 
  friend std::shared_ptr<Endpoint> createEndpointFromConnRequest(std::shared_ptr<Listener> listener,
                                                                 ucp_conn_request_h connRequest,
                                                                 bool endpointErrorHandling);
 
  friend std::shared_ptr<Endpoint> createEndpointFromWorkerAddress(std::shared_ptr<Worker> worker,
                                                                   std::shared_ptr<Address> address,
                                                                   bool endpointErrorHandling);
 
  [[nodiscard]] ucp_ep_h getHandle();
 
  [[nodiscard]] bool isAlive() const;
 
  void raiseOnError();
 
  void removeInflightRequest(const Request* const request);
 
  size_t cancelInflightRequests();
 
  [[nodiscard]] size_t getCancelingSize() const;
 
  size_t cancelInflightRequestsBlocking(uint64_t period = 0, uint64_t maxAttempts = 1);
 
  void setCloseCallback(EndpointCloseCallbackUserFunction closeCallback,
                        EndpointCloseCallbackUserData closeCallbackArg);
 
  [[nodiscard]] std::shared_ptr<Request> amSend(
    void* buffer,
    const size_t length,
    const ucs_memory_type_t memoryType,
    const std::optional<AmReceiverCallbackInfo> receiverCallbackInfo = std::nullopt,
    const bool enablePythonFuture                                    = false,
    RequestCallbackUserFunction callbackFunction                     = nullptr,
    RequestCallbackUserData callbackData                             = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> amRecv(
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> memPut(
    void* buffer,
    size_t length,
    uint64_t remote_addr,
    ucp_rkey_h rkey,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> memPut(
    void* buffer,
    size_t length,
    std::shared_ptr<ucxx::RemoteKey> remoteKey,
    uint64_t remoteAddrOffset                    = 0,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> memGet(
    void* buffer,
    size_t length,
    uint64_t remoteAddr,
    ucp_rkey_h rkey,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> memGet(
    void* buffer,
    size_t length,
    std::shared_ptr<ucxx::RemoteKey> remoteKey,
    uint64_t remoteAddrOffset                    = 0,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> streamSend(void* buffer,
                                                    size_t length,
                                                    const bool enablePythonFuture);
 
  [[nodiscard]] std::shared_ptr<Request> streamRecv(void* buffer,
                                                    size_t length,
                                                    const bool enablePythonFuture);
 
  [[nodiscard]] std::shared_ptr<Request> tagSend(
    void* buffer,
    size_t length,
    Tag tag,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> tagRecv(
    void* buffer,
    size_t length,
    Tag tag,
    TagMask tagMask,
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> tagMultiSend(const std::vector<void*>& buffer,
                                                      const std::vector<size_t>& size,
                                                      const std::vector<int>& isCUDA,
                                                      const Tag tag,
                                                      const bool enablePythonFuture);
 
  [[nodiscard]] std::shared_ptr<Request> tagMultiRecv(const Tag tag,
                                                      const TagMask tagMask,
                                                      const bool enablePythonFuture);
 
  [[nodiscard]] std::shared_ptr<Request> flush(
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Worker> getWorker();
 
  [[nodiscard]] std::shared_ptr<Request> close(
    const bool enablePythonFuture                      = false,
    EndpointCloseCallbackUserFunction callbackFunction = nullptr,
    EndpointCloseCallbackUserData callbackData         = nullptr);
 
  void closeBlocking(uint64_t period = 0, uint64_t maxAttempts = 1);
};
 
}  // namespace ucxx