shuffler.hpp

 
#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/communicator/metadata_payload_exchange/tag.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 = std::function<void()>;
 
    Shuffler(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        FinishedCallback&& finished_callback,
        PartitionOwner partition_owner = round_robin,
        std::unique_ptr<communicator::MetadataPayloadExchange> mpe = nullptr
    );
 
    Shuffler(
        std::shared_ptr<Communicator> comm,
        OpID op_id,
        PartID total_num_partitions,
        BufferResource* br,
        PartitionOwner partition_owner = round_robin,
        std::unique_ptr<communicator::MetadataPayloadExchange> mpe = nullptr
    )
        : Shuffler(
              comm,
              op_id,
              total_num_partitions,
              br,
              nullptr,
              partition_owner,
              std::move(mpe)
          ) {}
 
    ~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();
 
    [[nodiscard]] std::vector<PackedData> extract(PartID pid);
 
    [[nodiscard]] bool finished() const;
 
    void wait(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 Rank extract_rank(detail::ChunkID cid) {
        return safe_cast<Rank>(cid >> chunk_id_counter_bits);
    }
 
  private:
    void insert(detail::Chunk&& chunk);
 
    void insert_into_received(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};
    // Have we called `insert_finished()` on this rank.
    std::atomic<bool> locally_finished_{false};
    // Flipped to true exactly once when partitions are ready for extraction and we've
    // posted all sends we're going to
    bool can_extract_{false};
    detail::ChunksToSend to_send_;  
    detail::ReceivedChunks received_;  
 
    std::shared_ptr<Communicator> comm_;
    std::unique_ptr<communicator::MetadataPayloadExchange> mpe_;
    ProgressThread::FunctionID progress_thread_function_id_;
 
    SpillManager::SpillFunctionID spill_function_id_;
 
    std::vector<PartID> const local_partitions_;
 
    detail::FinishCounter finish_counter_;
    std::vector<detail::ChunkID> outbound_chunk_counter_;  
    std::atomic<detail::ChunkID> chunk_id_counter_{0};
 
    std::shared_ptr<Statistics> statistics_;
 
    // For notifications.
    mutable std::mutex mutex_;
    std::condition_variable cv_;
    FinishedCallback finished_callback_;  
 
    class Progress;
};
 
inline std::ostream& operator<<(std::ostream& os, Shuffler const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler