Distributed coordinator for multi-node LLM inference. More...

#include <distributed_coordinator.h>

Collaboration diagram for tinyllama::DistributedCoordinator:

Classes
struct	InferenceRequest

struct	InferenceResult

struct	LoadStats
	Get load balancing statistics. More...

struct	NodeInfo

Public Types
enum class	NodeRole { MASTER , WORKER , HYBRID }

Public Member Functions
	DistributedCoordinator (NodeRole role, const std::string &model_path, const std::string &tokenizer_path, int gpu_layers=-1)
	Initialize the distributed coordinator.

	~DistributedCoordinator ()

bool	initialize ()
	Initialize MPI and discover network topology.

void	shutdown ()
	Shutdown MPI and cleanup resources.

bool	submit_request (const InferenceRequest &request, std::function< void(const InferenceResult &)> callback)
	Submit an inference request (master node only)

void	run_worker ()
	Process requests as a worker node This function blocks and processes incoming requests.

void	run_master ()
	Run as master node, coordinating requests This function blocks and coordinates the distributed system.

std::vector< NodeInfo >	get_cluster_info () const
	Get information about all nodes in the cluster.

int	get_rank () const
	Get the rank of this node.

int	get_size () const
	Get the total number of nodes.

bool	is_mpi_available () const
	Check if MPI is available and initialized.

LoadStats	get_load_stats () const

Private Member Functions
void	discover_cluster ()
	Discover cluster topology and node capabilities.

int	select_optimal_node (const InferenceRequest &request)
	Select optimal node for processing request.

bool	send_request_to_worker (int worker_rank, const InferenceRequest &request)
	Send request to worker node.

bool	receive_and_process_request ()
	Receive and process request (worker node)

bool	send_result_to_master (const InferenceResult &result)
	Send result back to master.

bool	receive_result_from_worker ()
	Receive result from worker (master node)

std::vector< char >	serialize_request (const InferenceRequest &request)
	Serialize request for MPI transmission.

InferenceRequest	deserialize_request (const std::vector< char > &data)
	Deserialize request from MPI transmission.

std::vector< char >	serialize_result (const InferenceResult &result)
	Serialize result for MPI transmission.

InferenceResult	deserialize_result (const std::vector< char > &data)
	Deserialize result from MPI transmission.

InferenceResult	process_inference_locally (const InferenceRequest &request)
	Process inference request locally.

void	send_heartbeat ()
	Send heartbeat to all nodes.

void	handle_heartbeat (int source_rank)
	Handle heartbeat from other nodes.

Private Attributes
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_

std::atomic< int >	total_requests_ {0}

std::atomic< int >	completed_requests_ {0}

std::atomic< int >	failed_requests_ {0}

Static Private Attributes
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

Detailed Description

Distributed coordinator for multi-node LLM inference.

This class manages distributed inference across multiple nodes using MPI. It handles model sharding, request routing, and result aggregation.

Definition at line 22 of file distributed_coordinator.h.

Member Enumeration Documentation

◆ NodeRole

enum class tinyllama::DistributedCoordinator::NodeRole

strong

Enumerator
MASTER
WORKER
HYBRID

Definition at line 24 of file distributed_coordinator.h.

                        {
        MASTER,     // Coordinates requests and aggregates results
        WORKER,     // Processes inference requests
        HYBRID      // Can act as both master and worker
    };

Constructor & Destructor Documentation

◆ DistributedCoordinator()

tinyllama::DistributedCoordinator::DistributedCoordinator	(	NodeRole	role,
		const std::string &	model_path,
		const std::string &	tokenizer_path,
		int	gpu_layers = `-1`
	)

Initialize the distributed coordinator.

Parameters

role	The role of this node
model_path	Path to the model file
tokenizer_path	Path to the tokenizer file
gpu_layers	Number of GPU layers for this node

Definition at line 16 of file distributed_coordinator.cpp.

    : role_(role)
    , model_path_(model_path)
    , tokenizer_path_(tokenizer_path)
    , gpu_layers_(gpu_layers) {
}

◆ ~DistributedCoordinator()

tinyllama::DistributedCoordinator::~DistributedCoordinator ( )

Definition at line 26 of file distributed_coordinator.cpp.

                                                {
    shutdown();
}

References shutdown().

Member Function Documentation

◆ deserialize_request()

DistributedCoordinator::InferenceRequest tinyllama::DistributedCoordinator::deserialize_request ( const std::vector< char > & data )

private

Deserialize request from MPI transmission.

Definition at line 466 of file distributed_coordinator.cpp.

                                                                                                              {
    std::string str(data.begin(), data.end());
    std::istringstream iss(str);
    
    InferenceRequest request;
    std::string temp;
    
    if (std::getline(iss, request.request_id, '|') &&
        std::getline(iss, request.prompt, '|') &&
        std::getline(iss, temp, '|')) {
        request.max_tokens = std::stoi(temp);
    }
    if (std::getline(iss, temp, '|')) {
        request.temperature = std::stof(temp);
    }
    if (std::getline(iss, temp, '|')) {
        request.top_k = std::stoi(temp);
    }
    if (std::getline(iss, temp, '|')) {
        request.top_p = std::stof(temp);
    }
    if (std::getline(iss, request.system_prompt, '|') &&
        std::getline(iss, temp)) {
        request.target_node = std::stoi(temp);
    }
    
    return request;
}

References j, tinyllama::DistributedCoordinator::InferenceRequest::max_tokens, and str.

Referenced by receive_and_process_request().

◆ deserialize_result()

DistributedCoordinator::InferenceResult tinyllama::DistributedCoordinator::deserialize_result ( const std::vector< char > & data )

private

Deserialize result from MPI transmission.

Definition at line 509 of file distributed_coordinator.cpp.

                                                                                                            {
    std::string str(data.begin(), data.end());
    std::istringstream iss(str);
    
    InferenceResult result;
    std::string temp;
    
    if (std::getline(iss, result.request_id, '|') &&
        std::getline(iss, result.generated_text, '|') &&
        std::getline(iss, temp, '|')) {
        result.tokens_generated = std::stoi(temp);
    }
    if (std::getline(iss, temp, '|')) {
        result.inference_time_ms = std::stof(temp);
    }
    if (std::getline(iss, temp, '|')) {
        result.processing_node = std::stoi(temp);
    }
    if (std::getline(iss, temp, '|')) {
        result.success = (std::stoi(temp) == 1);
    }
    if (std::getline(iss, result.error_message)) {
        // Got error message
    }
    
    return result;
}

References j, str, and tinyllama::DistributedCoordinator::InferenceResult::tokens_generated.

Referenced by receive_result_from_worker().

◆ discover_cluster()

void tinyllama::DistributedCoordinator::discover_cluster ( )

private

Discover cluster topology and node capabilities.

Definition at line 87 of file distributed_coordinator.cpp.

                                              {
#ifdef HAS_MPI
    if (!mpi_initialized_) return;
 
    cluster_nodes_.clear();
    cluster_nodes_.resize(size_);
 
    // Get hostname
    char hostname[MPI_MAX_PROCESSOR_NAME];
    int name_len;
    MPI_Get_processor_name(hostname, &name_len);
 
    // Create node info for this node
    NodeInfo local_info;
    local_info.rank = rank_;
    local_info.hostname = std::string(hostname, name_len);
    local_info.role = role_;
    local_info.has_gpu = true; // Assume GPU available - could be detected
    local_info.gpu_memory_mb = 8192; // Default - could be detected
    local_info.gpu_layers = gpu_layers_;
 
    // Serialize local node info
    std::string serialized_info = local_info.hostname + "|" + 
                                 std::to_string(static_cast<int>(local_info.role)) + "|" +
                                 std::to_string(local_info.has_gpu ? 1 : 0) + "|" +
                                 std::to_string(local_info.gpu_memory_mb) + "|" +
                                 std::to_string(local_info.gpu_layers);
 
    // Gather all node info at rank 0
    std::vector<char> send_buffer(serialized_info.begin(), serialized_info.end());
    send_buffer.resize(256, '\0'); // Fixed size buffer
 
    std::vector<char> recv_buffer;
    if (rank_ == 0) {
        recv_buffer.resize(256 * size_);
    }
 
    MPI_Gather(send_buffer.data(), 256, MPI_CHAR, 
               recv_buffer.data(), 256, MPI_CHAR, 0, MPI_COMM_WORLD);
 
    // Parse and distribute node info
    if (rank_ == 0) {
        for (int i = 0; i < size_; ++i) {
            std::string node_data(recv_buffer.data() + i * 256);
            // Remove null terminators
            node_data = node_data.substr(0, node_data.find('\0'));
            
            std::istringstream ss(node_data);
            std::string hostname, role_str, has_gpu_str, memory_str, layers_str;
            
            if (std::getline(ss, hostname, '|') &&
                std::getline(ss, role_str, '|') &&
                std::getline(ss, has_gpu_str, '|') &&
                std::getline(ss, memory_str, '|') &&
                std::getline(ss, layers_str)) {
                
                cluster_nodes_[i].rank = i;
                cluster_nodes_[i].hostname = hostname;
                cluster_nodes_[i].role = static_cast<NodeRole>(std::stoi(role_str));
                cluster_nodes_[i].has_gpu = (std::stoi(has_gpu_str) == 1);
                cluster_nodes_[i].gpu_memory_mb = std::stoul(memory_str);
                cluster_nodes_[i].gpu_layers = std::stoi(layers_str);
            }
        }
    }
 
    // Broadcast cluster info to all nodes
    if (rank_ == 0) {
        for (int i = 1; i < size_; ++i) {
            MPI_Send(recv_buffer.data(), recv_buffer.size(), MPI_CHAR, i, TAG_NODE_INFO, MPI_COMM_WORLD);
        }
    } else {
        recv_buffer.resize(256 * size_);
        MPI_Status status;
        MPI_Recv(recv_buffer.data(), recv_buffer.size(), MPI_CHAR, 0, TAG_NODE_INFO, MPI_COMM_WORLD, &status);
        
        // Parse received cluster info
        for (int i = 0; i < size_; ++i) {
            std::string node_data(recv_buffer.data() + i * 256);
            node_data = node_data.substr(0, node_data.find('\0'));
            
            std::istringstream ss(node_data);
            std::string hostname, role_str, has_gpu_str, memory_str, layers_str;
            
            if (std::getline(ss, hostname, '|') &&
                std::getline(ss, role_str, '|') &&
                std::getline(ss, has_gpu_str, '|') &&
                std::getline(ss, memory_str, '|') &&
                std::getline(ss, layers_str)) {
                
                cluster_nodes_[i].rank = i;
                cluster_nodes_[i].hostname = hostname;
                cluster_nodes_[i].role = static_cast<NodeRole>(std::stoi(role_str));
                cluster_nodes_[i].has_gpu = (std::stoi(has_gpu_str) == 1);
                cluster_nodes_[i].gpu_memory_mb = std::stoul(memory_str);
                cluster_nodes_[i].gpu_layers = std::stoi(layers_str);
            }
        }
    }
 
    Logger::log("Discovered " + std::to_string(cluster_nodes_.size()) + " nodes in cluster", Logger::INFO);
    for (const auto& node : cluster_nodes_) {
        Logger::log("Node " + std::to_string(node.rank) + ": " + node.hostname + 
                   " (GPU: " + (node.has_gpu ? "yes" : "no") + 
                   ", Layers: " + std::to_string(node.gpu_layers) + ")", Logger::INFO);
    }
#endif
}

References cluster_nodes_, gpu_layers_, Logger::INFO, j, mpi_initialized_, tinyllama::DistributedCoordinator::NodeInfo::rank, rank_, role_, size_, and TAG_NODE_INFO.

Referenced by initialize().

◆ get_cluster_info()

std::vector< DistributedCoordinator::NodeInfo > tinyllama::DistributedCoordinator::get_cluster_info ( ) const

Get information about all nodes in the cluster.

Definition at line 586 of file distributed_coordinator.cpp.

                                                                                         {
    return cluster_nodes_;
}

References cluster_nodes_.

◆ get_load_stats()

DistributedCoordinator::LoadStats tinyllama::DistributedCoordinator::get_load_stats ( ) const

Definition at line 590 of file distributed_coordinator.cpp.

                                                                             {
    LoadStats stats;
    stats.total_requests = total_requests_.load();
    stats.completed_requests = completed_requests_.load();
    stats.failed_requests = failed_requests_.load();
    
    if (stats.completed_requests > 0) {
        stats.average_response_time_ms = 0.0f; // Would need to track this
    }
    
    for (const auto& node : cluster_nodes_) {
        stats.requests_per_node[node.rank] = node.active_requests.load();
    }
    
    return stats;
}

References cluster_nodes_, completed_requests_, failed_requests_, j, tinyllama::DistributedCoordinator::LoadStats::total_requests, and total_requests_.

◆ get_rank()

int tinyllama::DistributedCoordinator::get_rank ( ) const

inline

Get the rank of this node.

Definition at line 116 of file distributed_coordinator.h.

116{ return rank_; }

References rank_.

◆ get_size()

int tinyllama::DistributedCoordinator::get_size ( ) const

inline

Get the total number of nodes.

Definition at line 121 of file distributed_coordinator.h.

121{ return size_; }

References size_.

◆ handle_heartbeat()

void tinyllama::DistributedCoordinator::handle_heartbeat ( int source_rank )

private

Handle heartbeat from other nodes.

Definition at line 579 of file distributed_coordinator.cpp.

                                                             {
    // Update last seen time for the node
    if (source_rank >= 0 && source_rank < static_cast<int>(cluster_nodes_.size())) {
        // Could update a last_seen timestamp here
    }
}

References cluster_nodes_, and j.

◆ initialize()

bool tinyllama::DistributedCoordinator::initialize ( )

Initialize MPI and discover network topology.

Returns: true if initialization successful

Definition at line 30 of file distributed_coordinator.cpp.

                                        {
#ifdef HAS_MPI
    int provided;
    int result = MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
    if (result != MPI_SUCCESS) {
        Logger::log("Failed to initialize MPI", Logger::ERROR);
        return false;
    }
 
    if (provided < MPI_THREAD_MULTIPLE) {
        Logger::log("MPI does not support multithreading", Logger::WARNING);
    }
 
    MPI_Comm_rank(MPI_COMM_WORLD, &rank_);
    MPI_Comm_size(MPI_COMM_WORLD, &size_);
    
    mpi_initialized_ = true;
 
    Logger::log("MPI initialized - Rank: " + std::to_string(rank_) + 
                ", Size: " + std::to_string(size_), Logger::INFO);
 
    // Discover cluster topology
    discover_cluster();
 
    return true;
#else
    Logger::log("MPI support not compiled in", Logger::ERROR);
    return false;
#endif
}

References discover_cluster(), Logger::ERROR, Logger::INFO, j, mpi_initialized_, rank_, size_, and Logger::WARNING.

◆ is_mpi_available()

bool tinyllama::DistributedCoordinator::is_mpi_available ( ) const

Check if MPI is available and initialized.

Definition at line 79 of file distributed_coordinator.cpp.

                                                    {
#ifdef HAS_MPI
    return mpi_initialized_;
#else
    return false;
#endif
}

References mpi_initialized_.

◆ process_inference_locally()

DistributedCoordinator::InferenceResult tinyllama::DistributedCoordinator::process_inference_locally ( const InferenceRequest & request )

private

Process inference request locally.

Definition at line 537 of file distributed_coordinator.cpp.

                                                                                                                       {
    InferenceResult result;
    result.request_id = request.request_id;
    result.processing_node = rank_;
    
    auto start_time = std::chrono::high_resolution_clock::now();
    
    try {
        // Use the existing TinyLlama API for local inference
        // This is a simplified version - you would integrate with your actual API
        
        // For now, create a mock result
        result.generated_text = "Generated response for: " + request.prompt;
        result.tokens_generated = request.max_tokens;
        result.success = true;
        
        // TODO: Replace with actual inference call
        // std::string response = tinyllama_generate(request.prompt, request.max_tokens, 
        //                                          request.temperature, request.top_k, request.top_p);
        // result.generated_text = response;
        
    } catch (const std::exception& e) {
        result.success = false;
        result.error_message = e.what();
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
    result.inference_time_ms = static_cast<float>(duration.count());
    
    return result;
}

References j, rank_, and tinyllama::DistributedCoordinator::InferenceResult::request_id.

Referenced by receive_and_process_request(), and submit_request().

◆ receive_and_process_request()

bool tinyllama::DistributedCoordinator::receive_and_process_request ( )

private

Receive and process request (worker node)

Definition at line 346 of file distributed_coordinator.cpp.

                                                         {
#ifdef HAS_MPI
    int flag;
    MPI_Status status;
    
    // Check if there's a request available
    MPI_Iprobe(0, TAG_REQUEST, MPI_COMM_WORLD, &flag, &status);
    if (!flag) {
        return false;
    }
 
    // Receive size first
    int size;
    MPI_Recv(&size, 1, MPI_INT, 0, TAG_REQUEST, MPI_COMM_WORLD, &status);
 
    // Receive data
    std::vector<char> buffer(size);
    MPI_Recv(buffer.data(), size, MPI_CHAR, 0, TAG_REQUEST, MPI_COMM_WORLD, &status);
 
    // Deserialize request
    InferenceRequest request = deserialize_request(buffer);
 
    // Update active request count
    if (rank_ < static_cast<int>(cluster_nodes_.size())) {
        cluster_nodes_[rank_].active_requests++;
    }
 
    // Process request locally
    InferenceResult result = process_inference_locally(request);
    result.processing_node = rank_;
 
    // Send result back to master
    send_result_to_master(result);
 
    // Update active request count
    if (rank_ < static_cast<int>(cluster_nodes_.size())) {
        cluster_nodes_[rank_].active_requests--;
    }
 
    return true;
#else
    return false;
#endif
}

References cluster_nodes_, deserialize_request(), j, process_inference_locally(), rank_, send_result_to_master(), and TAG_REQUEST.

Referenced by run_worker().

◆ receive_result_from_worker()

bool tinyllama::DistributedCoordinator::receive_result_from_worker ( )

private

Receive result from worker (master node)

Definition at line 391 of file distributed_coordinator.cpp.

                                                        {
#ifdef HAS_MPI
    int flag;
    MPI_Status status;
    
    // Check if there's a result available from any worker
    MPI_Iprobe(MPI_ANY_SOURCE, TAG_RESULT, MPI_COMM_WORLD, &flag, &status);
    if (!flag) {
        return false;
    }
 
    // Receive size first
    int size;
    MPI_Recv(&size, 1, MPI_INT, status.MPI_SOURCE, TAG_RESULT, MPI_COMM_WORLD, &status);
 
    // Receive data
    std::vector<char> buffer(size);
    MPI_Recv(buffer.data(), size, MPI_CHAR, status.MPI_SOURCE, TAG_RESULT, MPI_COMM_WORLD, &status);
 
    // Deserialize result
    InferenceResult result = deserialize_result(buffer);
 
    // Find and call the callback
    auto callback_it = pending_callbacks_.find(result.request_id);
    if (callback_it != pending_callbacks_.end()) {
        callback_it->second(result);
        pending_callbacks_.erase(callback_it);
        
        if (result.success) {
            completed_requests_++;
        } else {
            failed_requests_++;
        }
    }
 
    return true;
#else
    return false;
#endif
}

References completed_requests_, deserialize_result(), failed_requests_, j, pending_callbacks_, and TAG_RESULT.

Referenced by run_master().

◆ run_master()

void tinyllama::DistributedCoordinator::run_master ( )

Run as master node, coordinating requests This function blocks and coordinates the distributed system.

Definition at line 317 of file distributed_coordinator.cpp.

                                        {
#ifdef HAS_MPI
    if (!mpi_initialized_ || rank_ != 0) {
        Logger::log("Only rank 0 can run as master", Logger::ERROR);
        return;
    }
 
    Logger::log("Starting master node", Logger::INFO);
 
    while (true) {
        // Receive results from workers
        if (receive_result_from_worker()) {
            // Result processed successfully
        } else {
            // No result available, sleep briefly
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
 
        // Send periodic heartbeats
        static auto last_heartbeat = std::chrono::steady_clock::now();
        auto now = std::chrono::steady_clock::now();
        if (now - last_heartbeat > std::chrono::seconds(5)) {
            send_heartbeat();
            last_heartbeat = now;
        }
    }
#endif
}

References Logger::ERROR, Logger::INFO, j, mpi_initialized_, rank_, receive_result_from_worker(), and send_heartbeat().

◆ run_worker()

void tinyllama::DistributedCoordinator::run_worker ( )

Process requests as a worker node This function blocks and processes incoming requests.

Definition at line 283 of file distributed_coordinator.cpp.

                                        {
#ifdef HAS_MPI
    if (!mpi_initialized_) {
        Logger::log("MPI not initialized", Logger::ERROR);
        return;
    }
 
    Logger::log("Starting worker node (rank " + std::to_string(rank_) + ")", Logger::INFO);
 
    while (true) {
        // Check for shutdown signal
        int flag;
        MPI_Status status;
        MPI_Iprobe(0, TAG_SHUTDOWN, MPI_COMM_WORLD, &flag, &status);
        if (flag) {
            int shutdown_signal;
            MPI_Recv(&shutdown_signal, 1, MPI_INT, 0, TAG_SHUTDOWN, MPI_COMM_WORLD, &status);
            Logger::log("Received shutdown signal", Logger::INFO);
            break;
        }
 
        // Check for incoming requests
        if (receive_and_process_request()) {
            // Request processed successfully
        } else {
            // No request available, sleep briefly
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
 
    Logger::log("Worker node shutting down", Logger::INFO);
#endif
}

References Logger::ERROR, Logger::INFO, j, mpi_initialized_, rank_, receive_and_process_request(), and TAG_SHUTDOWN.

◆ select_optimal_node()

int tinyllama::DistributedCoordinator::select_optimal_node ( const InferenceRequest & request )

private

Select optimal node for processing request.

Definition at line 241 of file distributed_coordinator.cpp.

                                                                               {
    if (request.target_node != -1 && request.target_node < size_) {
        return request.target_node;
    }
 
    // Simple load balancing - select node with least active requests
    int best_node = -1;
    int min_load = INT_MAX;
 
    for (const auto& node : cluster_nodes_) {
        if (node.rank == 0) continue; // Skip master node
        if (node.role == NodeRole::WORKER || node.role == NodeRole::HYBRID) {
            int current_load = node.active_requests.load();
            if (current_load < min_load) {
                min_load = current_load;
                best_node = node.rank;
            }
        }
    }
 
    return best_node;
}

References cluster_nodes_, HYBRID, j, size_, and WORKER.

Referenced by submit_request().

◆ send_heartbeat()

void tinyllama::DistributedCoordinator::send_heartbeat ( )

private

Send heartbeat to all nodes.

Definition at line 570 of file distributed_coordinator.cpp.

                                            {
#ifdef HAS_MPI
    for (int i = 1; i < size_; ++i) {
        int heartbeat = 1;
        MPI_Send(&heartbeat, 1, MPI_INT, i, TAG_HEARTBEAT, MPI_COMM_WORLD);
    }
#endif
}

References j, size_, and TAG_HEARTBEAT.

Referenced by run_master().

◆ send_request_to_worker()

bool tinyllama::DistributedCoordinator::send_request_to_worker	(	int	worker_rank,
		const InferenceRequest &	request
	)

private

Send request to worker node.

Definition at line 264 of file distributed_coordinator.cpp.

                                                                                                    {
#ifdef HAS_MPI
    std::vector<char> serialized = serialize_request(request);
    
    // Send size first
    int size = static_cast<int>(serialized.size());
    int result = MPI_Send(&size, 1, MPI_INT, worker_rank, TAG_REQUEST, MPI_COMM_WORLD);
    if (result != MPI_SUCCESS) {
        return false;
    }
 
    // Send data
    result = MPI_Send(serialized.data(), size, MPI_CHAR, worker_rank, TAG_REQUEST, MPI_COMM_WORLD);
    return result == MPI_SUCCESS;
#else
    return false;
#endif
}

References j, serialize_request(), and TAG_REQUEST.

Referenced by submit_request().

◆ send_result_to_master()

bool tinyllama::DistributedCoordinator::send_result_to_master ( const InferenceResult & result )

private

Send result back to master.

Definition at line 432 of file distributed_coordinator.cpp.

                                                                                {
#ifdef HAS_MPI
    std::vector<char> serialized = serialize_result(result);
    
    // Send size first
    int size = static_cast<int>(serialized.size());
    int mpi_result = MPI_Send(&size, 1, MPI_INT, 0, TAG_RESULT, MPI_COMM_WORLD);
    if (mpi_result != MPI_SUCCESS) {
        return false;
    }
 
    // Send data
    mpi_result = MPI_Send(serialized.data(), size, MPI_CHAR, 0, TAG_RESULT, MPI_COMM_WORLD);
    return mpi_result == MPI_SUCCESS;
#else
    return false;
#endif
}

References j, serialize_result(), and TAG_RESULT.

Referenced by receive_and_process_request().

◆ serialize_request()

std::vector< char > tinyllama::DistributedCoordinator::serialize_request ( const InferenceRequest & request )

private

Serialize request for MPI transmission.

Definition at line 451 of file distributed_coordinator.cpp.

                                                                                         {
    std::ostringstream oss;
    oss << request.request_id << "|"
        << request.prompt << "|"
        << request.max_tokens << "|"
        << request.temperature << "|"
        << request.top_k << "|"
        << request.top_p << "|"
        << request.system_prompt << "|"
        << request.target_node;
    
    std::string str = oss.str();
    return std::vector<char>(str.begin(), str.end());
}

References j, and str.

Referenced by send_request_to_worker().

◆ serialize_result()

std::vector< char > tinyllama::DistributedCoordinator::serialize_result ( const InferenceResult & result )

private

Serialize result for MPI transmission.

Definition at line 495 of file distributed_coordinator.cpp.

                                                                                      {
    std::ostringstream oss;
    oss << result.request_id << "|"
        << result.generated_text << "|"
        << result.tokens_generated << "|"
        << result.inference_time_ms << "|"
        << result.processing_node << "|"
        << (result.success ? 1 : 0) << "|"
        << result.error_message;
    
    std::string str = oss.str();
    return std::vector<char>(str.begin(), str.end());
}

References j, and str.

Referenced by send_result_to_master().

◆ shutdown()

void tinyllama::DistributedCoordinator::shutdown ( )

Shutdown MPI and cleanup resources.

Definition at line 61 of file distributed_coordinator.cpp.

                                      {
#ifdef HAS_MPI
    if (mpi_initialized_) {
        // Send shutdown signal to all nodes
        if (rank_ == 0) {
            for (int i = 1; i < size_; ++i) {
                int shutdown_signal = 1;
                MPI_Send(&shutdown_signal, 1, MPI_INT, i, TAG_SHUTDOWN, MPI_COMM_WORLD);
            }
        }
        
        MPI_Finalize();
        mpi_initialized_ = false;
        Logger::log("MPI finalized", Logger::INFO);
    }
#endif
}

References Logger::INFO, j, mpi_initialized_, rank_, size_, and TAG_SHUTDOWN.

Referenced by ~DistributedCoordinator().

◆ submit_request()

bool tinyllama::DistributedCoordinator::submit_request	(	const InferenceRequest &	request,
		std::function< void(const InferenceResult &)>	callback
	)

Submit an inference request (master node only)

Parameters

request	The inference request
callback	Callback function for result

Returns: true if request was submitted successfully

Definition at line 196 of file distributed_coordinator.cpp.

                                                                                              {
#ifdef HAS_MPI
    if (!mpi_initialized_ || rank_ != 0) {
        Logger::log("Only master node (rank 0) can submit requests", Logger::ERROR);
        return false;
    }
 
    // Store callback for later
    pending_callbacks_[request.request_id] = callback;
 
    // Select optimal worker node
    int target_node = select_optimal_node(request);
    if (target_node == -1) {
        Logger::log("No available worker nodes", Logger::ERROR);
        InferenceResult error_result;
        error_result.request_id = request.request_id;
        error_result.success = false;
        error_result.error_message = "No available worker nodes";
        callback(error_result);
        return false;
    }
 
    // Send request to worker
    bool sent = send_request_to_worker(target_node, request);
    if (!sent) {
        Logger::log("Failed to send request to worker node " + std::to_string(target_node), Logger::ERROR);
        InferenceResult error_result;
        error_result.request_id = request.request_id;
        error_result.success = false;
        error_result.error_message = "Failed to send request to worker";
        callback(error_result);
        return false;
    }
 
    total_requests_++;
    return true;
#else
    // Fallback to local processing if MPI not available
    InferenceResult result = process_inference_locally(request);
    callback(result);
    return result.success;
#endif
}

References Logger::ERROR, j, mpi_initialized_, pending_callbacks_, process_inference_locally(), rank_, tinyllama::DistributedCoordinator::InferenceResult::request_id, select_optimal_node(), send_request_to_worker(), and total_requests_.

Member Data Documentation

◆ cluster_nodes_

std::vector<NodeInfo> tinyllama::DistributedCoordinator::cluster_nodes_

private

Definition at line 150 of file distributed_coordinator.h.

Referenced by discover_cluster(), get_cluster_info(), get_load_stats(), handle_heartbeat(), receive_and_process_request(), and select_optimal_node().

◆ completed_requests_

std::atomic<int> tinyllama::DistributedCoordinator::completed_requests_ {0}

private

Definition at line 155 of file distributed_coordinator.h.

155{0};

Referenced by get_load_stats(), and receive_result_from_worker().

◆ failed_requests_

std::atomic<int> tinyllama::DistributedCoordinator::failed_requests_ {0}

private

Definition at line 156 of file distributed_coordinator.h.

156{0};

Referenced by get_load_stats(), and receive_result_from_worker().

◆ gpu_layers_

int tinyllama::DistributedCoordinator::gpu_layers_

private

Definition at line 144 of file distributed_coordinator.h.

Referenced by discover_cluster().

◆ model_path_

std::string tinyllama::DistributedCoordinator::model_path_

private

Definition at line 142 of file distributed_coordinator.h.

◆ mpi_initialized_

bool tinyllama::DistributedCoordinator::mpi_initialized_ = false

private

Definition at line 148 of file distributed_coordinator.h.

Referenced by discover_cluster(), initialize(), is_mpi_available(), run_master(), run_worker(), shutdown(), and submit_request().

◆ pending_callbacks_

std::map<std::string, std::function<void(const InferenceResult&)> > tinyllama::DistributedCoordinator::pending_callbacks_

private

Definition at line 151 of file distributed_coordinator.h.

Referenced by receive_result_from_worker(), and submit_request().

◆ rank_

int tinyllama::DistributedCoordinator::rank_ = -1

private

Definition at line 146 of file distributed_coordinator.h.

Referenced by discover_cluster(), get_rank(), initialize(), process_inference_locally(), receive_and_process_request(), run_master(), run_worker(), shutdown(), and submit_request().

◆ role_

NodeRole tinyllama::DistributedCoordinator::role_

private

Definition at line 141 of file distributed_coordinator.h.

Referenced by discover_cluster().

◆ size_

int tinyllama::DistributedCoordinator::size_ = -1

private

Definition at line 147 of file distributed_coordinator.h.

Referenced by discover_cluster(), get_size(), initialize(), select_optimal_node(), send_heartbeat(), and shutdown().

◆ TAG_HEARTBEAT

constexpr int tinyllama::DistributedCoordinator::TAG_HEARTBEAT = 3

staticconstexprprivate

Definition at line 161 of file distributed_coordinator.h.

Referenced by send_heartbeat().

◆ TAG_NODE_INFO

constexpr int tinyllama::DistributedCoordinator::TAG_NODE_INFO = 5

staticconstexprprivate

Definition at line 163 of file distributed_coordinator.h.

Referenced by discover_cluster().

◆ TAG_REQUEST

constexpr int tinyllama::DistributedCoordinator::TAG_REQUEST = 1

staticconstexprprivate

Definition at line 159 of file distributed_coordinator.h.

Referenced by receive_and_process_request(), and send_request_to_worker().

◆ TAG_RESULT

constexpr int tinyllama::DistributedCoordinator::TAG_RESULT = 2

staticconstexprprivate

Definition at line 160 of file distributed_coordinator.h.

Referenced by receive_result_from_worker(), and send_result_to_master().

◆ TAG_SHUTDOWN

constexpr int tinyllama::DistributedCoordinator::TAG_SHUTDOWN = 4

staticconstexprprivate

Definition at line 162 of file distributed_coordinator.h.

Referenced by run_worker(), and shutdown().

◆ tokenizer_path_

std::string tinyllama::DistributedCoordinator::tokenizer_path_

private

Definition at line 143 of file distributed_coordinator.h.

◆ total_requests_

std::atomic<int> tinyllama::DistributedCoordinator::total_requests_ {0}

private

Definition at line 154 of file distributed_coordinator.h.

154{0};

Referenced by get_load_stats(), and submit_request().

The documentation for this class was generated from the following files:

Classes

Public Types

Public Member Functions

Private Member Functions

Private Attributes

Static Private Attributes

Detailed Description

Member Enumeration Documentation

◆ NodeRole

Constructor & Destructor Documentation

◆ DistributedCoordinator()

◆ ~DistributedCoordinator()

Member Function Documentation

◆ deserialize_request()

◆ deserialize_result()

◆ discover_cluster()

◆ get_cluster_info()

◆ get_load_stats()

◆ get_rank()

◆ get_size()

◆ handle_heartbeat()

◆ initialize()

◆ is_mpi_available()

◆ process_inference_locally()

◆ receive_and_process_request()

◆ receive_result_from_worker()

◆ run_master()

◆ run_worker()

◆ select_optimal_node()

◆ send_heartbeat()

◆ send_request_to_worker()

◆ send_result_to_master()

◆ serialize_request()

◆ serialize_result()

◆ shutdown()

◆ submit_request()

Member Data Documentation

◆ cluster_nodes_

◆ completed_requests_

◆ failed_requests_

◆ gpu_layers_

◆ model_path_

◆ mpi_initialized_

◆ pending_callbacks_

◆ rank_

◆ role_

◆ size_

◆ TAG_HEARTBEAT

◆ TAG_NODE_INFO

◆ TAG_REQUEST

◆ TAG_RESULT

◆ TAG_SHUTDOWN

◆ tokenizer_path_

◆ total_requests_