shuffler.hpp

 
#pragma once
 
#include <unordered_set>
 
#include <rapidsmpf/shuffler/shuffler.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/core/node.hpp>
#include <rapidsmpf/streaming/cudf/partition.hpp>
 
namespace rapidsmpf::streaming {
 
class ShufflerAsync {
  public:
    ShufflerAsync(
        std::shared_ptr<Context> ctx,
        OpID op_id,
        shuffler::PartID total_num_partitions,
        shuffler::Shuffler::PartitionOwner partition_owner =
            shuffler::Shuffler::round_robin
    );
 
    // Prevent copying
    ShufflerAsync(ShufflerAsync const&) = delete;
    ShufflerAsync& operator=(ShufflerAsync const&) = delete;
 
    ~ShufflerAsync() noexcept;
 
    [[nodiscard]] constexpr std::shared_ptr<Context> const& ctx() const {
        return ctx_;
    }
 
    [[nodiscard]] constexpr shuffler::PartID total_num_partitions() const {
        return shuffler_.total_num_partitions;
    }
 
    [[nodiscard]] constexpr shuffler::Shuffler::PartitionOwner const&
    partition_owner() const {
        return shuffler_.partition_owner;
    }
 
    [[nodiscard]] std::span<shuffler::PartID const> local_partitions() const;
 
    void insert(std::unordered_map<shuffler::PartID, PackedData>&& chunks);
 
    [[nodiscard]] Node insert_finished(std::vector<shuffler::PartID>&& pids);
 
    [[nodiscard]] coro::task<std::optional<std::vector<PackedData>>> extract_async(
        shuffler::PartID pid
    );
 
    using ExtractResult = std::pair<shuffler::PartID, std::vector<PackedData>>;
 
    [[nodiscard]] coro::task<std::optional<ExtractResult>> extract_any_async();
 
  private:
    [[nodiscard]] Node finished_drain();
 
    std::shared_ptr<Context> ctx_;
    coro::task_container<coro::thread_pool>
        notifications_;  
    Semaphore semaphore_{0};  
    coro::latch
        latch_;  
    std::mutex mtx_;  
    shuffler::Shuffler shuffler_;
 
    std::unordered_set<shuffler::PartID> ready_pids_;
    std::unordered_set<shuffler::PartID> extracted_pids_;
};
 
namespace node {
[[nodiscard]] Node shuffler(
    std::shared_ptr<Context> ctx,
    std::shared_ptr<Channel> ch_in,
    std::shared_ptr<Channel> ch_out,
    OpID op_id,
    shuffler::PartID total_num_partitions,
    shuffler::Shuffler::PartitionOwner partition_owner = shuffler::Shuffler::round_robin
);
 
}  // namespace node
 
}  // namespace rapidsmpf::streaming