ucxx.hpp

 
#pragma once
 
#include <cstdlib>
#include <memory>
#include <utility>
 
#include <ucxx/api.h>
 
#include <rmm/device_buffer.hpp>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/progress_thread.hpp>
 
namespace rapidsmpf {
 
namespace ucxx {
 
 
using HostPortPair =
    std::pair<std::string, std::uint16_t>;  
using RemoteAddress = std::variant<
    HostPortPair,
    std::shared_ptr<::ucxx::Address>>;  
 
enum class ProgressMode : std::uint8_t {
    Blocking = 0,
    Polling,
    ThreadBlocking,
    ThreadPolling
};
 
class ListenerAddress {
  public:
    RemoteAddress address;  
    Rank rank{};  
};
 
class SharedResources;
 
class InitializedRank {
  public:
    InitializedRank(std::shared_ptr<SharedResources> shared_resources);
 
    std::shared_ptr<SharedResources> shared_resources_{
        nullptr
    };  
};
 
std::unique_ptr<rapidsmpf::ucxx::InitializedRank> init(
    std::shared_ptr<::ucxx::Worker> worker,
    Rank nranks,
    std::optional<RemoteAddress> remote_address,
    config::Options options
);
 
class UCXX final : public Communicator {
  public:
    class Future : public Communicator::Future {
        friend class UCXX;
 
      public:
        Future(std::shared_ptr<::ucxx::Request> req, std::unique_ptr<Buffer> data_buffer)
            : req_{std::move(req)}, data_buffer_{std::move(data_buffer)} {}
 
        Future(
            std::shared_ptr<::ucxx::Request> req,
            std::unique_ptr<std::vector<std::uint8_t>> synced_host_data
        )
            : req_{std::move(req)}, synced_host_data_{std::move(synced_host_data)} {}
 
        ~Future() noexcept override = default;
 
      private:
        std::shared_ptr<::ucxx::Request> req_;
        std::unique_ptr<Buffer> data_buffer_;
        // Dedicated storage for host data that is valid at the time of construction.
        std::unique_ptr<std::vector<std::uint8_t>> synced_host_data_;
    };
 
    UCXX(
        std::unique_ptr<InitializedRank> ucxx_initialized_rank,
        config::Options options,
        std::shared_ptr<ProgressThread> progress_thread
    );
 
    ~UCXX() noexcept override;
 
    [[nodiscard]] Rank rank() const override;
 
    [[nodiscard]] Rank nranks() const override;
 
    [[nodiscard]] std::unique_ptr<Communicator::Future> send(
        std::unique_ptr<std::vector<std::uint8_t>> msg, Rank rank, Tag tag
    ) override;
 
    // clang-format off
    // clang-format on
    [[nodiscard]] std::unique_ptr<Communicator::Future> send(
        std::unique_ptr<Buffer> msg, Rank rank, Tag tag
    ) override;
 
    [[nodiscard]] std::unique_ptr<Communicator::Future> recv(
        Rank rank, Tag tag, std::unique_ptr<Buffer> recv_buffer
    ) override;
 
    // clang-format off
    // clang-format on
    [[nodiscard]] std::unique_ptr<Communicator::Future> recv_sync_host_data(
        Rank rank, Tag tag, std::unique_ptr<std::vector<std::uint8_t>> synced_buffer
    ) override;
 
    [[nodiscard]] std::pair<std::unique_ptr<std::vector<std::uint8_t>>, Rank> recv_any(
        Tag tag
    ) override;
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> recv_from(
        Rank src, Tag tag
    ) override;
 
    std::pair<
        std::vector<std::unique_ptr<Communicator::Future>>,
        std::vector<std::size_t>>
    test_some(std::vector<std::unique_ptr<Communicator::Future>>& future_vector) override;
 
    // clang-format off
    // clang-format on
    std::vector<std::size_t> test_some(
        std::unordered_map<std::size_t, std::unique_ptr<Communicator::Future>> const&
            future_map
    ) override;
 
    bool test(std::unique_ptr<Communicator::Future>& future) override;
    std::vector<std::unique_ptr<Buffer>> wait_all(
        std::vector<std::unique_ptr<Communicator::Future>>&& futures
    ) override;
 
    [[nodiscard]] std::unique_ptr<Buffer> wait(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::unique_ptr<Buffer> release_data(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> release_sync_host_data(
        std::unique_ptr<Communicator::Future> future
    ) override;
 
    [[nodiscard]] std::shared_ptr<Communicator::Logger> const& logger() override {
        return logger_;
    }
 
    [[nodiscard]] std::shared_ptr<ProgressThread> const&
    progress_thread() const override {
        return progress_thread_;
    }
 
    [[nodiscard]] std::string str() const override;
 
    void barrier();
 
    ListenerAddress listener_address();
 
    std::shared_ptr<UCXX> split();
 
  private:
    std::shared_ptr<SharedResources> shared_resources_;
    config::Options options_;
    std::shared_ptr<Logger> logger_;
    std::shared_ptr<ProgressThread> progress_thread_;
 
    std::shared_ptr<::ucxx::Endpoint> get_endpoint(Rank rank);
    void progress_worker();
};
 
}  // namespace ucxx
 
}  // namespace rapidsmpf