shuffler.hpp

 
#pragma once
 
#include <atomic>
#include <chrono>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/buffer/packed_data.hpp>
#include <rapidsmpf/buffer/resource.hpp>
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/nvtx.hpp>
#include <rapidsmpf/progress_thread.hpp>
#include <rapidsmpf/shuffler/chunk.hpp>
#include <rapidsmpf/shuffler/finish_counter.hpp>
#include <rapidsmpf/shuffler/postbox.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/utils.hpp>
 
 
class ShuffleInsertGroupedTest;
 
namespace rapidsmpf::shuffler {
 
class Shuffler {
    friend class ::ShuffleInsertGroupedTest;
 
  public:
    using PartitionOwner = std::function<Rank(std::shared_ptr<Communicator>, PartID)>;
 
    static Rank round_robin(std::shared_ptr<Communicator> const& comm, PartID pid) {
        return static_cast<Rank>(pid % static_cast<PartID>(comm->nranks()));
    }
 
    static std::vector<PartID> local_partitions(
        std::shared_ptr<Communicator> const& comm,
        PartID total_num_partitions,
        PartitionOwner partition_owner
    );
 
    using FinishedCallback = detail::FinishCounter::FinishedCallback;
 
    Shuffler(
        std::shared_ptr<Communicator> comm,
        std::shared_ptr<ProgressThread> progress_thread,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        FinishedCallback&& finished_callback,
        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
        PartitionOwner partition_owner = round_robin
    );
 
    Shuffler(
        std::shared_ptr<Communicator> comm,
        std::shared_ptr<ProgressThread> progress_thread,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
        PartitionOwner partition_owner = round_robin
    )
        : Shuffler(
              comm,
              progress_thread,
              op_id,
              total_num_partitions,
              br,
              nullptr,
              statistics,
              partition_owner
          ) {}
 
    ~Shuffler();
 
    Shuffler(Shuffler const&) = delete;
    Shuffler& operator=(Shuffler const&) = delete;
 
    void shutdown();
 
    void concat_insert(std::unordered_map<PartID, PackedData>&& chunks);
 
    void insert(std::unordered_map<PartID, PackedData>&& chunks);
 
    void insert_finished(PartID pid);
 
    void insert_finished(std::vector<PartID>&& pids);
 
    [[nodiscard]] std::vector<PackedData> extract(PartID pid);
 
    [[nodiscard]] bool finished() const;
 
    [[nodiscard]] bool is_finished(PartID pid) const;
 
    PartID wait_any(std::optional<std::chrono::milliseconds> timeout = {});
 
    void wait_on(PartID pid, std::optional<std::chrono::milliseconds> timeout = {});
 
    std::size_t spill(std::optional<std::size_t> amount = std::nullopt);
 
    [[nodiscard]] std::string str() const;
 
    [[nodiscard]] std::span<PartID const> local_partitions() const;
 
    static constexpr int chunk_id_counter_bits = 38;
 
    static constexpr uint64_t counter_mask = (uint64_t(1) << chunk_id_counter_bits) - 1;
 
    static constexpr uint64_t extract_counter(detail::ChunkID cid) {
        return cid & counter_mask;
    }
 
    static constexpr Rank extract_rank(detail::ChunkID cid) {
        return static_cast<Rank>(cid >> chunk_id_counter_bits);
    }
 
    static constexpr std::pair<Rank, uint64_t> extract_info(detail::ChunkID cid) {
        return std::make_pair(extract_rank(cid), extract_counter(cid));
    }
 
  private:
    void insert(detail::Chunk&& chunk);
 
    void insert_into_ready_postbox(detail::Chunk&& chunk);
 
    [[nodiscard]] detail::ChunkID get_new_cid();
 
    [[nodiscard]] detail::Chunk create_chunk(PartID pid, PackedData&& packed_data);
 
  public:
    PartID const total_num_partitions;  
    PartitionOwner const partition_owner;  
 
  private:
    BufferResource* br_;
    std::atomic<bool> active_{true};
    detail::PostBox<Rank> outgoing_postbox_;  
    detail::PostBox<PartID> ready_postbox_;  
 
    std::shared_ptr<Communicator> comm_;
    std::shared_ptr<ProgressThread> progress_thread_;
    ProgressThread::FunctionID progress_thread_function_id_;
    OpID const op_id_;
 
    SpillManager::SpillFunctionID spill_function_id_;
 
    std::vector<PartID> const local_partitions_;
 
    detail::FinishCounter finish_counter_;
    std::unordered_map<PartID, detail::ChunkID> outbound_chunk_counter_;
    mutable std::mutex outbound_chunk_counter_mutex_;
 
    // We protect ready_postbox extraction to avoid returning a chunk that is in the
    // process of being spilled by `Shuffler::spill`.
    mutable std::mutex ready_postbox_spilling_mutex_;
 
    std::atomic<detail::ChunkID> chunk_id_counter_{0};
 
    std::shared_ptr<Statistics> statistics_;
 
    class Progress;
};
 
inline std::ostream& operator<<(std::ostream& os, Shuffler const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler