finish_counter.hpp

 
#pragma once
 
#include <array>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <mutex>
#include <optional>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/shuffler/chunk.hpp>
#include <rapidsmpf/shuffler/postbox.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf::shuffler {
 
namespace detail {
 
class FinishCounter {
  public:
    using FinishedCallback = std::function<void(PartID)>;
 
    FinishCounter(
        Rank nranks,
        std::vector<PartID> const& local_partitions,
        FinishedCallback&& finished_callback = nullptr
    );
 
    ~FinishCounter() = default;
 
    void move_goalpost(PartID pid, ChunkID nchunks);
 
    void add_finished_chunk(PartID pid);
 
    [[nodiscard]] bool all_finished() const;
 
    PartID wait_any(std::optional<std::chrono::milliseconds> timeout = {});
 
    void wait_on(PartID pid, std::optional<std::chrono::milliseconds> timeout = {});
 
    [[nodiscard]] std::string str() const;
 
  private:
    Rank const nranks_;
    PartID
        n_unfinished_partitions_;  
 
    struct PartitionInfo {
        Rank rank_count{0};  
        ChunkID chunk_goal{0};  
        ChunkID finished_chunk_count{
            0
        };  
 
        constexpr PartitionInfo() = default;
 
        constexpr void move_goalpost(ChunkID nchunks, Rank nranks) {
            RAPIDSMPF_EXPECTS(nchunks != 0, "the goalpost was moved by 0 chunks");
            RAPIDSMPF_EXPECTS(
                ++rank_count <= nranks, "the goalpost was moved more than one per rank"
            );
            chunk_goal += nchunks;
        }
 
        constexpr void add_finished_chunk(Rank nranks) {
            finished_chunk_count++;
            // only throw if rank_count == nranks
            RAPIDSMPF_EXPECTS(
                (rank_count < nranks) || (finished_chunk_count <= chunk_goal),
                "finished chunk exceeds the goal"
            );
        }
 
        // The partition is finished if the goalpost has been set by all ranks
        // and the number of finished chunks has reached the goal.
        [[nodiscard]] constexpr bool is_finished(Rank nranks) const {
            return rank_count == nranks && finished_chunk_count == chunk_goal;
        }
 
        [[nodiscard]] constexpr ChunkID data_chunk_goal() const {
            // there will always be a control message from each rank indicating how many
            // chunks it's sending. Chunk goal contains this control message for each
            // rank. Therefore, to get the data chunk goal, we need to subtract the number
            // of ranks that have reported their chunk count from the chunk goal.
            return chunk_goal - static_cast<ChunkID>(rank_count);
        }
    };
 
    // The goalpost of each partition. The goal is a rank counter to track how many ranks
    // has reported their goal, and a chunk counter that specifies the goal. It is only
    // when all ranks has reported their goal that the goalpost is final.
    std::unordered_map<PartID, PartitionInfo> goalposts_;
 
    mutable std::mutex mutex_;  // TODO: use a shared_mutex lock?
    mutable std::condition_variable wait_cv_;
 
    FinishedCallback finished_callback_ =
        nullptr;  
};
 
}  // namespace detail
 
inline std::ostream& operator<<(std::ostream& os, detail::FinishCounter const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler