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/communicator/communicator.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/packed_data.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/misc.hpp>
 
namespace rapidsmpf::shuffler {
 
class Shuffler {
  public:
    using PartitionOwner =
        std::function<Rank(std::shared_ptr<Communicator> const&, PartID, PartID)>;
 
    static Rank round_robin(
        std::shared_ptr<Communicator> const& comm,
        PartID pid,
        [[maybe_unused]] PartID total_num_partitions
    ) {
        return safe_cast<Rank>(pid % safe_cast<PartID>(comm->nranks()));
    }
 
    static Rank contiguous(
        std::shared_ptr<Communicator> const& comm, PartID pid, PartID total_num_partitions
    ) {
        return safe_cast<Rank>(
            (pid * safe_cast<PartID>(comm->nranks())) / total_num_partitions
        );
    }
 
    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,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        FinishedCallback&& finished_callback,
        PartitionOwner partition_owner = round_robin
    );
 
    Shuffler(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        PartitionOwner partition_owner = round_robin
    )
        : Shuffler(comm, op_id, total_num_partitions, br, nullptr, partition_owner) {}
 
    ~Shuffler();
 
    [[nodiscard]] std::shared_ptr<Communicator> const& comm() const noexcept {
        return comm_;
    }
 
    Shuffler(Shuffler const&) = delete;
    Shuffler& operator=(Shuffler const&) = delete;
 
    void shutdown();
 
    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 std::uint64_t counter_mask =
        (std::uint64_t{1} << chunk_id_counter_bits) - 1;
 
    static constexpr std::uint64_t extract_counter(detail::ChunkID cid) {
        return cid & counter_mask;
    }
 
    static constexpr Rank extract_rank(detail::ChunkID cid) {
        return safe_cast<Rank>(cid >> chunk_id_counter_bits);
    }
 
    static constexpr std::pair<Rank, std::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_;
    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