shuffler.hpp

 
#pragma once
 
#include <coro/event.hpp>
 
#include <rapidsmpf/shuffler/shuffler.hpp>
#include <rapidsmpf/streaming/core/actor.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/cudf/partition.hpp>
 
namespace rapidsmpf::streaming {
 
class ShufflerAsync {
  public:
    ShufflerAsync(
        std::shared_ptr<Context> ctx,
        std::shared_ptr<Communicator> comm,
        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]] std::shared_ptr<Communicator> const& comm() const noexcept {
        return shuffler_.comm();
    }
 
    [[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]] Actor insert_finished();
 
    [[nodiscard]] std::vector<PackedData> extract(shuffler::PartID pid);
 
  private:
    std::shared_ptr<Context> ctx_;
    coro::event
        event_{};  
    shuffler::Shuffler shuffler_;
};
 
namespace actor {
[[nodiscard]] Actor shuffler(
    std::shared_ptr<Context> ctx,
    std::shared_ptr<Communicator> comm,
    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 actor
 
}  // namespace rapidsmpf::streaming