distributed_coordinator.h

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <string>
#include <vector>
#include <functional>
#include <map>
#include <atomic>
 
#ifdef HAS_MPI
#include <mpi.h>
#endif
 
namespace tinyllama {
 
class DistributedCoordinator {
public:
    enum class NodeRole {
        MASTER,     // Coordinates requests and aggregates results
        WORKER,     // Processes inference requests
        HYBRID      // Can act as both master and worker
    };
 
    struct NodeInfo {
        int rank;
        std::string hostname;
        NodeRole role;
        bool has_gpu;
        size_t gpu_memory_mb;
        int gpu_layers;
        std::atomic<bool> is_busy{false};
        std::atomic<int> active_requests{0};
    };
 
    struct InferenceRequest {
        std::string request_id;
        std::string prompt;
        int max_tokens;
        float temperature;
        int top_k;
        float top_p;
        std::string system_prompt;
        int target_node = -1;  // -1 for auto-selection
    };
 
    struct InferenceResult {
        std::string request_id;
        std::string generated_text;
        int tokens_generated;
        float inference_time_ms;
        int processing_node;
        bool success;
        std::string error_message;
    };
 
    DistributedCoordinator(NodeRole role, 
                          const std::string& model_path,
                          const std::string& tokenizer_path,
                          int gpu_layers = -1);
 
    ~DistributedCoordinator();
 
    bool initialize();
 
    void shutdown();
 
    bool submit_request(const InferenceRequest& request,
                       std::function<void(const InferenceResult&)> callback);
 
    void run_worker();
 
    void run_master();
 
    std::vector<NodeInfo> get_cluster_info() const;
 
    int get_rank() const { return rank_; }
 
    int get_size() const { return size_; }
 
    bool is_mpi_available() const;
 
    struct LoadStats {
        int total_requests;
        int completed_requests;
        int failed_requests;
        float average_response_time_ms;
        std::map<int, int> requests_per_node;
    };
    LoadStats get_load_stats() const;
 
private:
    NodeRole role_;
    std::string model_path_;
    std::string tokenizer_path_;
    int gpu_layers_;
    
    int rank_ = -1;
    int size_ = -1;
    bool mpi_initialized_ = false;
    
    std::vector<NodeInfo> cluster_nodes_;
    std::map<std::string, std::function<void(const InferenceResult&)>> pending_callbacks_;
    
    // Load balancing
    std::atomic<int> total_requests_{0};
    std::atomic<int> completed_requests_{0};
    std::atomic<int> failed_requests_{0};
    
    // MPI message tags
    static constexpr int TAG_REQUEST = 1;
    static constexpr int TAG_RESULT = 2;
    static constexpr int TAG_HEARTBEAT = 3;
    static constexpr int TAG_SHUTDOWN = 4;
    static constexpr int TAG_NODE_INFO = 5;
 
    void discover_cluster();
 
    int select_optimal_node(const InferenceRequest& request);
 
    bool send_request_to_worker(int worker_rank, const InferenceRequest& request);
 
    bool receive_and_process_request();
 
    bool send_result_to_master(const InferenceResult& result);
 
    bool receive_result_from_worker();
 
    std::vector<char> serialize_request(const InferenceRequest& request);
 
    InferenceRequest deserialize_request(const std::vector<char>& data);
 
    std::vector<char> serialize_result(const InferenceResult& result);
 
    InferenceResult deserialize_result(const std::vector<char>& data);
 
    InferenceResult process_inference_locally(const InferenceRequest& request);
 
    void send_heartbeat();
 
    void handle_heartbeat(int source_rank);
};
 
} // namespace tinyllama