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);
};
 
struct ErrorCallbackData {
  Endpoint*
    _endpoint;  
  std::unique_ptr<std::mutex>
    _mutex;  
  ucs_status_t status;                                 
  std::shared_ptr<InflightRequests> inflightRequests;  
  EndpointCloseCallbackUserFunction closeCallback;     
  EndpointCloseCallbackUserData closeCallbackArg;      
  std::shared_ptr<Worker> worker;  
};
 
class Endpoint : public Component {
 private:
  ucp_ep_h _handle{nullptr};          
  ucp_ep_h _originalHandle{nullptr};  
  bool _endpointErrorHandling{true};  
  std::unique_ptr<ErrorCallbackData> _callbackData{
    nullptr};  
  std::shared_ptr<InflightRequests> _inflightRequests{
    std::make_shared<InflightRequests>()};  
 
  Endpoint(std::shared_ptr<Component> workerOrListener,
           ucp_ep_params_t* params,
           bool endpointErrorHandling);
 
  std::shared_ptr<Request> registerInflightRequest(std::shared_ptr<Request> request);
 
 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);
 
  ucp_ep_h getHandle();
 
  bool isAlive() const;
 
  void raiseOnError();
 
  void removeInflightRequest(const Request* const request);
 
  size_t cancelInflightRequests();
 
  size_t getCancelingSize() const;
 
  size_t cancelInflightRequestsBlocking(uint64_t period = 0, uint64_t maxAttempts = 1);
 
  void setCloseCallback(EndpointCloseCallbackUserFunction closeCallback,
                        EndpointCloseCallbackUserData closeCallbackArg);
 
  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);
 
  std::shared_ptr<Request> amRecv(const bool enablePythonFuture                = false,
                                  RequestCallbackUserFunction callbackFunction = nullptr,
                                  RequestCallbackUserData callbackData         = nullptr);
 
  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);
 
  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);
 
  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);
 
  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);
 
  std::shared_ptr<Request> streamSend(void* buffer, size_t length, const bool enablePythonFuture);
 
  std::shared_ptr<Request> streamRecv(void* buffer, size_t length, const bool enablePythonFuture);
 
  std::shared_ptr<Request> tagSend(void* buffer,
                                   size_t length,
                                   Tag tag,
                                   const bool enablePythonFuture                = false,
                                   RequestCallbackUserFunction callbackFunction = nullptr,
                                   RequestCallbackUserData callbackData         = nullptr);
 
  std::shared_ptr<Request> tagRecv(void* buffer,
                                   size_t length,
                                   Tag tag,
                                   TagMask tagMask,
                                   const bool enablePythonFuture                = false,
                                   RequestCallbackUserFunction callbackFunction = nullptr,
                                   RequestCallbackUserData callbackData         = nullptr);
 
  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);
 
  std::shared_ptr<Request> tagMultiRecv(const Tag tag,
                                        const TagMask tagMask,
                                        const bool enablePythonFuture);
 
  std::shared_ptr<Request> flush(const bool enablePythonFuture                = false,
                                 RequestCallbackUserFunction callbackFunction = nullptr,
                                 RequestCallbackUserData callbackData         = nullptr);
 
  std::shared_ptr<Worker> getWorker();
 
  static void errorCallback(void* arg, ucp_ep_h ep, ucs_status_t status);
 
  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