eryon_rt/

orcha.rs

/*
    Appellation: orchestrator <module>
    Contrib: @FL03
*/
use crate::NodeMessage;
use crate::frag::FragTonnetz;
use crate::tasks::{Task, TaskStatus, TaskType};
use eryon::ops::Execute;
use eryon_actors::prelude::{Actor, Driver, Operator, OperatorKind, TriadDriver, VNode};
use rshyper::EdgeId;
use rstmt::nrt::Triad;
use scsys::Timestamp;
use std::collections::{HashMap, HashSet, VecDeque};

/// The orchestrator is responsible for managing communications between partitions,
/// ensuring synchronized state transitions, and facilitating coordinated operations
/// across the fragment.
#[derive(Clone, Debug)]
pub struct Orchestrator {
    /// Maps task IDs to their execution status
    pub(crate) task_statuses: HashMap<usize, TaskStatus>,
    /// Current active tasks running in the system
    pub(crate) active_tasks: HashSet<usize>,
    /// Keeps track of nodes that have recently communicated
    pub(crate) active_nodes: HashSet<EdgeId>,
    /// Message queue for inter-node communication
    pub(crate) message_queue: VecDeque<NodeMessage>,
    /// Next available task ID
    pub(crate) next_task_id: usize,
}

impl Default for Orchestrator {
    fn default() -> Self {
        Self::new()
    }
}

impl Orchestrator {
    /// Create a new orchestrator
    pub fn new() -> Self {
        Orchestrator {
            task_statuses: HashMap::new(),
            active_tasks: HashSet::new(),
            active_nodes: HashSet::new(),
            message_queue: VecDeque::new(),
            next_task_id: 0,
        }
    }
    /// Get the count of active nodes
    pub fn active_node_count(&self) -> usize {
        self.active_nodes.len()
    }
    /// Get the count of active tasks
    pub fn active_task_count(&self) -> usize {
        self.active_tasks.len()
    }
    /// Check if a node is active
    pub fn is_node_active(&self, node_id: EdgeId) -> bool {
        self.active_nodes.contains(&node_id)
    }
}

impl Orchestrator {
    /// Get all active node IDs
    pub fn active_node_ids(&self) -> Vec<EdgeId> {
        self.active_nodes.iter().copied().collect()
    }
    /// Coordinate knowledge sharing between nodes
    pub fn coordinate_knowledge_sharing<D>(
        &mut self,
        fragment: &mut FragTonnetz<D>,
    ) -> crate::Result<()>
    where
        D: TriadDriver,
    {
        // 1. Identify knowledge-bearing nodes
        let nodes_with_knowledge: Vec<_> = fragment
            .partitions()
            .iter()
            .filter(|(_, node)| node.has_surface_network() && node.total_features() > 0)
            .map(|(&id, _)| id)
            .collect();

        if nodes_with_knowledge.len() <= 1 {
            return Ok(()); // Nothing to coordinate
        }

        // 2. Extract patterns once per node - this is our main optimization
        let mut extracted_patterns = Vec::with_capacity(nodes_with_knowledge.len());

        for &node_id in &nodes_with_knowledge {
            if let Some(node) = fragment.get_vnode(&node_id) {
                match node.extract_knowledge_patterns() {
                    Ok(patterns) if !patterns.is_empty() => {
                        extracted_patterns.push((node_id, patterns));
                    }
                    _ => continue,
                }
            }
        }

        // 3. Simple filtering: Share each source with at most 5 target nodes
        // This limits overall complexity without complex musical calculations
        let max_targets_per_source = 5;

        // 4. Perform transfers using the extracted patterns
        for (source_id, patterns) in extracted_patterns {
            // Skip empty pattern sets
            if patterns.is_empty() {
                continue;
            }

            // Choose target nodes (up to max_targets_per_source)
            let mut target_count = 0;

            for &target_id in &nodes_with_knowledge {
                if target_id != source_id && target_count < max_targets_per_source {
                    if let Some(target_node) = fragment.get_vnode_mut(&target_id) {
                        if let Err(_e) =
                            target_node.integrate_external_knowledge(&patterns, &source_id)
                        {
                            #[cfg(feature = "tracing")]
                            tracing::warn!("Failed to integrate knowledge: {:?}", _e);
                        }

                        target_count += 1;
                    }
                }
            }
        }

        // 5. Record coordination event
        for node in fragment.partitions_mut().values_mut() {
            node.store_mut().record_event(
                "knowledge_coordination",
                Some(vec![nodes_with_knowledge.len()]),
            );
        }

        Ok(())
    }
    /// Create decision task based on current system state
    pub fn create_adaptive_task<D>(&self, fragment: &FragTonnetz<D>) -> Task
    where
        D: Driver<Triad>,
    {
        // Get system status
        let capabilities = fragment.calculate_node_capabilities();
        let avg_memory = capabilities
            .values()
            .map(|cap| cap.memory_usage)
            .sum::<usize>() as f32
            / capabilities.len().max(1) as f32;

        // Decide what task is most needed
        if avg_memory > 5000.0 {
            // If memory usage is high
            Task::new(
                self.next_task_id,
                8,
                TaskType::OptimizeMemory { max_features: 1000 },
            )
        } else {
            // Check if we need to coordinate learning
            let learning_nodes = capabilities
                .values()
                .filter(|cap| cap.learning_capability)
                .count();

            if learning_nodes > 1 {
                Task::new(self.next_task_id, 6, TaskType::CoordinateLearning)
            } else {
                // Default to computational optimization
                Task::new(self.next_task_id, 5, TaskType::ComputeTransformations)
            }
        }
    }
    /// Create a task for coordinating learning
    pub fn create_coordinate_learning_task(&self, priority: usize) -> Task {
        Task::new(self.next_task_id, priority, TaskType::CoordinateLearning)
    }
    /// Create a task for node mode change
    pub fn create_node_mode_change_task(&self, node_id: EdgeId, mode: OperatorKind) -> Task {
        Task::new(
            self.next_task_id,
            7,
            TaskType::ChangeNodeMode(node_id, mode),
        )
    }
    /// Execute a task on the fragment
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, fragment, task), name = "execute", target = "orchestrator")
    )]
    pub fn execute_task<D>(
        &mut self,
        fragment: &mut FragTonnetz<D>,
        task: &Task,
    ) -> crate::Result<()>
    where
        D: TriadDriver,
    {
        match &task.kind {
            TaskType::Transform {
                node_id,
                transformation,
            } => {
                if let Some(node) = fragment.get_vnode_mut(node_id) {
                    node.transform(*transformation)?;
                } else {
                    return Err(crate::RuntimeError::NodeNotFound);
                }
            }
            TaskType::BatchTransform(operations) => {
                for (node_id, transforms) in operations {
                    if let Some(node) = fragment.get_vnode_mut(node_id) {
                        node.transform_batch(transforms)?;
                    } else {
                        return Err(crate::RuntimeError::NodeNotFound);
                    }
                }
            }
            TaskType::ComputeTransformations => {
                fragment.compute_transformations();
            }
            TaskType::SharePatterns => {
                self.share_patterns(fragment)?;
            }
            TaskType::CoordinateLearning => {
                fragment.coordinate_learning()?;
            }

            TaskType::OptimizeMemory { max_features } => {
                // First collect nodes that need optimization
                let node_stats: Vec<(EdgeId, usize)> = fragment
                    .partitions()
                    .iter()
                    .map(|(id, node)| (*id, node.total_features()))
                    .filter(|(_, count)| *count > max_features / 2) // Only optimize if significant features
                    .collect();

                let mut _optimized = 0;
                let mut _pruned = 0;

                // Perform optimization on selected nodes
                for (node_id, feature_count) in node_stats {
                    if let Some(node) = fragment.get_vnode_mut(&node_id) {
                        // Adjust max features based on node's actual count
                        let node_max = std::cmp::min(*max_features, feature_count * 2 / 3);

                        match node.optimize_memory(node_max) {
                            Ok(stats) => {
                                _optimized += 1;
                                _pruned += stats.pruned_count();

                                // Record optimization event
                                node.store_mut().record_event(
                                    "memory_optimized",
                                    Some(vec![stats.features_before(), stats.features_after()]),
                                );
                            }
                            Err(_e) => {
                                #[cfg(feature = "tracing")]
                                tracing::error!(
                                    "Failed to optimize memory for node {node_id}: {_e:?}"
                                );
                            }
                        }
                    }
                }
                #[cfg(feature = "tracing")]
                tracing::info!(
                    "Memory optimization: {} nodes optimized, {} features pruned",
                    _optimized,
                    _pruned
                );
                // Record overall memory optimization with timestamp
                let timestamp = Timestamp::<u64>::now().to_string();

                fragment.partitions_mut().values_mut().for_each(|node| {
                    node.store_mut()
                        .set_property("last_memory_optimization", &timestamp);
                });
            }
            TaskType::BalanceResources => {
                fragment.balance_resources()?;
            }
            TaskType::ChangeNodeMode(node_id, mode) => {
                fragment.set_node_operator(*node_id, *mode)?;
            }
            TaskType::ProcessAgentProposals => {
                self.process_agent_proposals(fragment)?;
            }
            _ => todo!("TaskType is not currently supported"),
        }

        Ok(())
    }
    /// Get task status
    pub fn get_task_status(&self, task_id: usize) -> Option<&TaskStatus> {
        self.task_statuses.get(&task_id)
    }
    /// Get all tasks that can be executed in parallel
    pub fn get_parallel_tasks(&self) -> Vec<usize> {
        self.task_statuses
            .iter()
            .filter_map(|(&id, status)| {
                if *status == TaskStatus::Queued {
                    Some(id)
                } else {
                    None
                }
            })
            .collect()
    }
    /// Initialize the orchestrator with the fragment
    pub fn initialize<D>(&mut self, fragment: &FragTonnetz<D>)
    where
        D: Driver<Triad>,
    {
        // Register all partitions as active nodes
        for &edge_id in fragment.partitions().keys() {
            self.active_nodes.insert(edge_id);
        }
    }
    /// Mark a node as active
    pub fn mark_node_active(&mut self, node_id: EdgeId) {
        self.active_nodes.insert(node_id);
    }
    /// Mark a node as inactive
    pub fn mark_node_inactive(&mut self, node_id: EdgeId) {
        self.active_nodes.remove(&node_id);
    }
    /// Monitor node health and resource utilization
    pub fn monitor_node_health<D>(&mut self, fragment: &FragTonnetz<D>) -> Vec<EdgeId>
    where
        D: Driver<Triad>,
    {
        let mut unhealthy_nodes = Vec::new();
        let capabilities = fragment.calculate_node_capabilities();

        // Find nodes with excessive resource usage
        for (&id, cap) in &capabilities {
            if cap.memory_usage > 10000 {
                // Arbitrary high memory threshold
                unhealthy_nodes.push(id);
            }
        }

        unhealthy_nodes
    }
    /// Optimize operator distribution across the fragment
    pub fn optimize_operator_distribution<D>(
        &mut self,
        fragment: &mut FragTonnetz<D>,
    ) -> crate::Result<()>
    where
        D: Driver<Triad>,
    {
        fragment.optimize_operator_distribution()
    }
    /// Process proposed transformations from agents
    pub fn process_agent_proposals<D>(
        &mut self,
        fragment: &mut FragTonnetz<D>,
    ) -> crate::Result<Vec<usize>>
    where
        D: TriadDriver,
    {
        // Collect proposals from agent nodes
        let proposals = fragment.collect_proposed_transformations();
        if proposals.is_empty() {
            return Ok(Vec::new());
        }

        // Create tasks from proposals
        let mut tasks = Vec::new();
        for (node_id, transformation) in proposals {
            let task_id = self.register_task();

            let task = Task::new(
                task_id,
                5,
                TaskType::Transform {
                    node_id,
                    transformation,
                },
            );

            tasks.push(task);
        }

        // Execute the tasks
        let mut completed = Vec::new();
        for task in tasks {
            let task_id = task.id;
            self.update_task_status(task_id, TaskStatus::Running);

            match self.execute_task(fragment, &task) {
                Ok(_) => {
                    self.update_task_status(task_id, TaskStatus::Completed);
                    completed.push(task_id);
                }
                Err(e) => {
                    self.update_task_status(task_id, TaskStatus::Failed(format!("{:?}", e)));
                }
            }
        }

        Ok(completed)
    }
    /// Process messages in the queue
    pub fn process_messages<D>(&mut self, fragment: &mut FragTonnetz<D>) -> crate::Result<()>
    where
        D: TriadDriver,
    {
        while let Some(message) = self.message_queue.pop_front() {
            match &message {
                NodeMessage::TransformRequest {
                    target, transform, ..
                } => {
                    // Apply transformation to target node
                    if let Some(node) = fragment.get_vnode_mut(target) {
                        node.transform(*transform)?;
                    } else {
                        return Err(crate::RuntimeError::NodeNotFound);
                    }
                }
                NodeMessage::StateSync {
                    source,
                    state_hash: _,
                } => {
                    // For now, just mark the node as active
                    self.active_nodes.insert(*source);
                }
                NodeMessage::PatternShare {
                    source,
                    pattern,
                    importance: _,
                } => {
                    // Share with other nodes (except the source)
                    for (&edge_id, node) in fragment.partitions_mut() {
                        if edge_id != *source {
                            let _ = node.learn_transformation_sequence(pattern);
                        }
                    }
                }
            }
        }

        Ok(())
    }
    /// Register a new task and return its ID
    pub fn register_task(&mut self) -> usize {
        let task_id = self.next_task_id;
        self.next_task_id += 1;

        self.task_statuses.insert(task_id, TaskStatus::Queued);
        task_id
    }
    /// Share patterns between nodes
    pub fn share_patterns<D>(&mut self, fragment: &mut FragTonnetz<D>) -> crate::Result<()>
    where
        D: Driver<Triad>,
    {
        // Get all nodes
        let node_ids: Vec<_> = fragment.partitions().keys().cloned().collect();

        if node_ids.len() <= 1 {
            return Ok(()); // Nothing to share
        }

        // Share patterns from each node with others
        for i in 0..node_ids.len() {
            let source_id = node_ids[i];

            // Skip if source doesn't allow sharing
            if let Some(source_node) = fragment.get_vnode(&source_id) {
                if !<Operator<f32> as Actor<D, f32>>::allows_pattern_sharing(source_node.operator())
                {
                    continue;
                }
            } else {
                continue;
            }

            for j in 0..node_ids.len() {
                if i == j {
                    continue;
                }

                let target_id = node_ids[j];

                // Skip if target doesn't allow sharing
                if let Some(target_node) = fragment.get_vnode(&target_id) {
                    if !<Operator<f32> as Actor<D, f32>>::allows_pattern_sharing(
                        target_node.operator(),
                    ) {
                        continue;
                    }
                } else {
                    continue;
                }
                // If both allow sharing, do the sharing
                if let (Some(source), Some(target)) = (
                    fragment.clone().get_vnode(&source_id),
                    fragment.get_vnode_mut(&target_id),
                ) {
                    // Create temporary copy
                    let mut source_copy = VNode::from_driver(source.driver().clone());
                    source_copy.store_mut().merge(source.store());

                    source_copy.share_patterns(target)?;
                }
            }
        }

        Ok(())
    }

    /// Schedule tasks for execution
    pub fn schedule_tasks<D>(
        &mut self,
        fragment: &mut FragTonnetz<D>,
        tasks: Vec<Task>,
    ) -> crate::Result<Vec<usize>>
    where
        D: TriadDriver,
    {
        let mut task_ids = Vec::with_capacity(tasks.len());

        for task in tasks {
            let task_id = self.register_task();
            self.update_task_status(task_id, TaskStatus::Running);

            match self.execute_task(fragment, &task) {
                Ok(_) => {
                    self.update_task_status(task_id, TaskStatus::Completed);
                }
                Err(e) => {
                    self.update_task_status(task_id, TaskStatus::Failed(format!("{:?}", e)));
                }
            }

            task_ids.push(task_id);
        }

        Ok(task_ids)
    }
    /// Send a message between nodes
    pub fn send_message(&mut self, message: NodeMessage) {
        self.message_queue.push_back(message);
    }
    /// Update the status of a task
    pub fn update_task_status(&mut self, task_id: usize, status: TaskStatus) {
        match status {
            TaskStatus::Running => {
                self.active_tasks.insert(task_id);
            }
            TaskStatus::Completed | TaskStatus::Failed(_) => {
                self.active_tasks.remove(&task_id);
            }
            _ => {}
        }

        self.task_statuses.insert(task_id, status);
    }
}

impl<D> Execute<(Task, &mut FragTonnetz<D>)> for Orchestrator
where
    D: TriadDriver,
{
    type Output = crate::Result<()>;

    fn execute(mut self, params: (Task, &mut FragTonnetz<D>)) -> Self::Output {
        let (task, fragment) = params;
        self.update_task_status(task.id, TaskStatus::Running);

        match self.execute_task(fragment, &task) {
            Ok(_) => {
                self.update_task_status(task.id, TaskStatus::Completed);
                Ok(())
            }
            Err(err) => {
                self.update_task_status(task.id, TaskStatus::Failed(format!("{:?}", err)));
                Err(err)
            }
        }
    }
}