worker.h

 
#pragma once
 
#include <functional>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <utility>
#include <variant>
 
#include <ucp/api/ucp.h>
 
#include <ucxx/buffer.h>
#include <ucxx/component.h>
#include <ucxx/constructors.h>
#include <ucxx/context.h>
#include <ucxx/delayed_submission.h>
#include <ucxx/future.h>
#include <ucxx/inflight_requests.h>
#include <ucxx/notifier.h>
#include <ucxx/typedefs.h>
#include <ucxx/worker_progress_thread.h>
 
namespace ucxx {
 
namespace experimental {
class WorkerBuilder;
}  // namespace experimental
 
class Address;
class Buffer;
class Endpoint;
class Listener;
class RequestAm;
 
namespace internal {
class AmData;
}  // namespace internal
 
class Worker : public Component {
 private:
  ucp_worker_h _handle{nullptr};        
  int _epollFileDescriptor{-1};         
  int _workerFileDescriptor{-1};        
  std::mutex _inflightRequestsMutex{};  
  std::unique_ptr<InflightRequests> _inflightRequests{
    std::make_unique<InflightRequests>()};  
  std::mutex
    _inflightRequestsToCancelMutex{};  
  std::unique_ptr<InflightRequests> _inflightRequestsToCancel{
    std::make_unique<InflightRequests>()};  
  WorkerProgressThread _progressThread{};   
  std::thread::id _progressThreadId{};      
  std::function<void(void*)> _progressThreadStartCallback{
    nullptr};  
  void* _progressThreadStartCallbackArg{
    nullptr};  
  std::shared_ptr<DelayedSubmissionCollection> _delayedSubmissionCollection{
    nullptr};  
 
  friend std::shared_ptr<RequestAm> createRequestAm(
    std::shared_ptr<Endpoint> endpoint,
    const std::variant<data::AmSend, data::AmReceive> requestData,
    const bool enablePythonFuture,
    RequestCallbackUserFunction callbackFunction,
    RequestCallbackUserData callbackData);
 
 protected:
  bool _enableFuture{
    false};  
  bool _enableRequestAttributes{
    false};  
  std::mutex _futuresPoolMutex{};  
  std::queue<std::shared_ptr<Future>>
    _futuresPool{};  
  std::shared_ptr<Notifier> _notifier{nullptr};  
  std::shared_ptr<internal::AmData>
    _amData;  
  BufferType _cudaBufferType{BufferType::Invalid};  
 
 private:
  void drainWorkerTagRecv();
 
  [[nodiscard]] std::shared_ptr<RequestAm> getAmRecv(
    ucp_ep_h ep, std::function<std::shared_ptr<RequestAm>()> createAmRecvRequestFunction);
 
  void stopProgressThreadNoWarn();
 
  [[nodiscard]] std::shared_ptr<Request> registerInflightRequest(std::shared_ptr<Request> request);
 
  bool progressPending();
 
  void setCudaBufferType(BufferType bufferType);
 
 protected:
  explicit Worker(std::shared_ptr<Context> context,
                  const bool enableDelayedSubmission = false,
                  const bool enableFuture            = false);
 
 public:
  Worker()                         = delete;
  Worker(const Worker&)            = delete;
  Worker& operator=(Worker const&) = delete;
  Worker(Worker&& o)               = delete;
  Worker& operator=(Worker&& o)    = delete;
 
  friend std::shared_ptr<Worker> createWorker(std::shared_ptr<Context> context,
                                              const bool enableDelayedSubmission,
                                              const bool enableFuture);
 
  friend class experimental::WorkerBuilder;
 
  virtual ~Worker();
 
  [[nodiscard]] ucp_worker_h getHandle();
 
  [[nodiscard]] std::string getInfo();
 
  void initBlockingProgressMode();
 
  int getEpollFileDescriptor();
 
  bool arm();
 
  bool progressWorkerEvent(const int epollTimeout = -1);
 
  void signal();
 
  bool waitProgress();
 
  bool progressOnce();
 
  bool progress();
 
  void registerDelayedSubmission(std::shared_ptr<Request> request,
                                 DelayedSubmissionCallbackType callback);
 
  [[nodiscard]] bool registerGenericPre(DelayedSubmissionCallbackType callback,
                                        uint64_t period = 0);
 
  [[nodiscard]] bool registerGenericPost(DelayedSubmissionCallbackType callback,
                                         uint64_t period = 0);
 
  [[nodiscard]] bool isDelayedRequestSubmissionEnabled() const;
 
  [[nodiscard]] bool isFutureEnabled() const;
 
  [[nodiscard]] bool isRequestAttributesEnabled() const noexcept;
 
  [[nodiscard]] BufferType getCudaBufferType() const;
 
  virtual void populateFuturesPool();
 
  virtual void clearFuturesPool();
 
  [[nodiscard]] virtual std::shared_ptr<Future> getFuture();
 
  [[nodiscard]] virtual RequestNotifierWaitState waitRequestNotifier(uint64_t periodNs);
 
  virtual void runRequestNotifier();
 
  virtual void stopRequestNotifierThread();
 
  void setProgressThreadStartCallback(std::function<void(void*)> callback, void* callbackArg);
 
  void startProgressThread(const bool pollingMode = false, const int epollTimeout = 1);
 
  void stopProgressThread();
 
  [[nodiscard]] bool isProgressThreadRunning();
 
  [[nodiscard]] std::thread::id getProgressThreadId();
 
  size_t cancelInflightRequests(uint64_t period = 0, uint64_t maxAttempts = 1);
 
  void scheduleRequestCancel(TrackedRequests trackedRequests);
 
  void removeInflightRequest(std::shared_ptr<Request> request);
 
  [[nodiscard]] std::shared_ptr<TagProbeInfo> tagProbe(const Tag tag,
                                                       const TagMask tagMask = TagMaskFull,
                                                       const bool remove     = false) const;
 
  [[nodiscard]] std::shared_ptr<Request> tagRecv(
    void* buffer,
    size_t length,
    Tag tag,
    TagMask tagMask,
    const bool enableFuture                      = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Request> tagRecvWithHandle(
    void* buffer,
    std::shared_ptr<TagProbeInfo> probeInfo,
    const bool enableFuture                      = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  [[nodiscard]] std::shared_ptr<Address> getAddress();
 
  [[nodiscard]] std::shared_ptr<Endpoint> createEndpointFromHostname(
    std::string ipAddress, uint16_t port, bool endpointErrorHandling = true);
 
  [[nodiscard]] std::shared_ptr<Endpoint> createEndpointFromWorkerAddress(
    std::shared_ptr<Address> address, bool endpointErrorHandling = true);
 
  [[nodiscard]] std::shared_ptr<Listener> createListener(uint16_t port,
                                                         ucp_listener_conn_callback_t callback,
                                                         void* callbackArgs);
 
  void registerAmAllocator(ucs_memory_type_t memoryType, AmAllocatorType allocator);
 
  void registerAmReceiverCallback(AmReceiverCallbackInfo info, AmReceiverCallbackType callback);
 
  [[nodiscard]] bool amProbe(const ucp_ep_h endpointHandle) const;
 
  [[nodiscard]] std::shared_ptr<Request> flush(
    const bool enablePythonFuture                = false,
    RequestCallbackUserFunction callbackFunction = nullptr,
    RequestCallbackUserData callbackData         = nullptr);
 
  struct Attributes {
    ucs_thread_mode_t threadMode{UCS_THREAD_MODE_MULTI};
    size_t maxAmHeader{0};
    std::string name{};
    size_t maxDebugString{0};
  };
 
  [[nodiscard]] Attributes queryAttributes() const;
};
 
std::shared_ptr<Worker> createWorker(std::shared_ptr<Context> context,
                                     const bool enableDelayedSubmission,
                                     const bool enableFuture);
 
}  // namespace ucxx
 
// Include experimental features
#include <ucxx/experimental/worker_builder.h>