use std::time::{Instant, Duration};
use std::collections::{HashMap, VecDeque};
use serde::{Serialize, Deserialize};
use crate::device::{RiDevice, RiDeviceType, RiDeviceCapabilities};
#[derive(Debug, Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiDeviceDiscoveryEngine {
fingerprints: HashMap<String, DeviceFingerprint>,
discovery_history: VecDeque<DiscoveryRecord>,
confidence_threshold: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct DeviceFingerprint {
device_type: RiDeviceType,
capabilities: RiDeviceCapabilities,
identification_patterns: Vec<IdentificationPattern>,
confidence_score: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct IdentificationPattern {
field: String,
pattern: String,
weight: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct DiscoveryRecord {
timestamp: u64, device_id: String,
device_info: HashMap<String, String>,
identified_type: Option<RiDeviceType>,
confidence: f64,
}
impl Default for RiDeviceDiscoveryEngine {
fn default() -> Self {
Self::new()
}
}
impl RiDeviceDiscoveryEngine {
#[cfg(not(feature = "pyo3"))]
pub fn new() -> Self {
let mut engine = Self {
fingerprints: HashMap::new(),
discovery_history: VecDeque::with_capacity(1000),
confidence_threshold: 0.7,
};
engine.initialize_default_fingerprints();
engine
}
#[cfg(not(feature = "pyo3"))]
pub fn discover_devices(&mut self, scan_results: Vec<DeviceScanResult>) -> Vec<RiDevice> {
let mut discovered_devices = Vec::with_capacity(4);
for scan_result in scan_results {
if let Some(device) = self.identify_device(scan_result) {
discovered_devices.push(device);
}
}
discovered_devices
}
fn identify_device(&mut self, scan_result: DeviceScanResult) -> Option<RiDevice> {
let device_info = scan_result.device_info;
let mut best_match: Option<(String, f64)> = None;
for (fingerprint_id, fingerprint) in &self.fingerprints {
let confidence = self.calculate_match_confidence(&device_info, fingerprint);
if confidence > self.confidence_threshold {
match best_match {
None => best_match = Some((fingerprint_id.clone(), confidence)),
Some((_, best_confidence)) if confidence > best_confidence => {
best_match = Some((fingerprint_id.clone(), confidence));
}
_ => {}
}
}
}
if let Some((fingerprint_id, confidence)) = best_match {
let fingerprint = match self.fingerprints.get(&fingerprint_id) {
Some(f) => f,
None => return None,
};
let device = RiDevice::new(
scan_result.device_name.clone(),
fingerprint.device_type,
).with_capabilities(fingerprint.capabilities.clone());
self.record_discovery(
scan_result.device_id,
device_info,
Some(fingerprint.device_type),
confidence,
);
Some(device)
} else {
self.record_discovery(
scan_result.device_id,
device_info,
None,
0.0,
);
None
}
}
fn calculate_match_confidence(&self, device_info: &HashMap<String, String>, fingerprint: &DeviceFingerprint) -> f64 {
let mut total_weight = 0.0;
let mut matched_weight = 0.0;
for pattern in &fingerprint.identification_patterns {
total_weight += pattern.weight;
if let Some(value) = device_info.get(&pattern.field) {
if self.matches_pattern(value, &pattern.pattern) {
matched_weight += pattern.weight;
}
}
}
if total_weight > 0.0 {
matched_weight / total_weight
} else {
0.0
}
}
fn matches_pattern(&self, value: &str, pattern: &str) -> bool {
value.to_lowercase().contains(&pattern.to_lowercase())
}
fn record_discovery(&mut self, device_id: String, device_info: HashMap<String, String>, identified_type: Option<RiDeviceType>, confidence: f64) {
let record = DiscoveryRecord {
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or(Duration::from_secs(0))
.as_secs(),
device_id,
device_info,
identified_type,
confidence,
};
self.discovery_history.push_back(record);
if self.discovery_history.len() > 1000 {
self.discovery_history.pop_front();
}
}
fn initialize_default_fingerprints(&mut self) {
let gpu_fingerprint = DeviceFingerprint {
device_type: RiDeviceType::GPU,
capabilities: RiDeviceCapabilities {
memory_gb: Some(16.0),
compute_units: Some(512),
storage_gb: Some(500.0),
bandwidth_gbps: Some(900.0),
custom_capabilities: vec![("cuda_support".to_string(), "true".to_string()), ("tensor_cores".to_string(), "true".to_string())].into_iter().collect(),
},
identification_patterns: vec![
IdentificationPattern {
field: "device_name".to_string(),
pattern: "nvidia".to_string(),
weight: 0.4,
},
IdentificationPattern {
field: "driver".to_string(),
pattern: "cuda".to_string(),
weight: 0.6,
},
],
confidence_score: 0.9,
};
let tpu_fingerprint = DeviceFingerprint {
device_type: RiDeviceType::Custom, capabilities: RiDeviceCapabilities {
memory_gb: Some(32.0),
compute_units: Some(128),
storage_gb: Some(1000.0),
bandwidth_gbps: Some(1200.0),
custom_capabilities: vec![("tpu_support".to_string(), "true".to_string()), ("ml_accelerator".to_string(), "true".to_string())].into_iter().collect(),
},
identification_patterns: vec![
IdentificationPattern {
field: "device_name".to_string(),
pattern: "tpu".to_string(),
weight: 0.8,
},
IdentificationPattern {
field: "vendor".to_string(),
pattern: "google".to_string(),
weight: 0.2,
},
],
confidence_score: 0.95,
};
self.fingerprints.insert("gpu".to_string(), gpu_fingerprint);
self.fingerprints.insert("tpu".to_string(), tpu_fingerprint); }
#[cfg(not(feature = "pyo3"))]
pub fn get_discovery_stats(&self) -> DiscoveryStats {
let total_attempts = self.discovery_history.len();
let successful_identifications = self.discovery_history
.iter()
.filter(|record| record.identified_type.is_some())
.count();
let avg_confidence = if total_attempts > 0 {
self.discovery_history
.iter()
.map(|record| record.confidence)
.sum::<f64>() / total_attempts as f64
} else {
0.0
};
DiscoveryStats {
total_attempts,
successful_identifications,
success_rate: if total_attempts > 0 {
successful_identifications as f64 / total_attempts as f64
} else {
0.0
},
avg_confidence,
}
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiDeviceDiscoveryEngine {
#[new]
pub fn new() -> Self {
let mut engine = Self {
fingerprints: HashMap::new(),
discovery_history: VecDeque::with_capacity(1000),
confidence_threshold: 0.7,
};
engine.initialize_default_fingerprints();
engine
}
pub fn discover_devices(&mut self, scan_results: Vec<DeviceScanResult>) -> Vec<RiDevice> {
let mut discovered_devices = Vec::with_capacity(4);
for scan_result in scan_results {
if let Some(device) = self.identify_device(scan_result) {
discovered_devices.push(device);
}
}
discovered_devices
}
pub fn get_confidence_threshold(&self) -> f64 {
self.confidence_threshold
}
pub fn set_confidence_threshold(&mut self, threshold: f64) {
self.confidence_threshold = threshold;
}
pub fn get_discovery_stats(&self) -> DiscoveryStats {
let total_attempts = self.discovery_history.len();
let successful_identifications = self.discovery_history
.iter()
.filter(|r| r.identified_type.is_some())
.count();
let avg_confidence = if total_attempts > 0 {
self.discovery_history
.iter()
.map(|record| record.confidence)
.sum::<f64>() / total_attempts as f64
} else {
0.0
};
DiscoveryStats {
total_attempts,
successful_identifications,
success_rate: if total_attempts > 0 {
successful_identifications as f64 / total_attempts as f64
} else {
0.0
},
avg_confidence,
}
}
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct DeviceScanResult {
pub device_id: String,
pub device_name: String,
pub device_info: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct DiscoveryStats {
pub total_attempts: usize,
pub successful_identifications: usize,
pub success_rate: f64,
pub avg_confidence: f64,
}
#[derive(Debug, Clone)]
pub struct RiResourceScheduler {
performance_history: HashMap<String, Vec<PerformanceRecord>>,
device_loads: HashMap<String, f64>,
policies: Vec<SchedulingPolicy>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct PerformanceRecord {
timestamp: u64, latency_ms: f64,
throughput: f64,
error_rate: f64,
utilization: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SchedulingPolicy {
name: String,
priority: u8,
conditions: Vec<PolicyCondition>,
action: PolicyAction,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct PolicyCondition {
metric: String,
operator: String,
threshold: f64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
enum PolicyAction {
PreferDevice(String),
AvoidDevice(String),
LoadBalance,
PriorityBased,
}
impl Default for RiResourceScheduler {
fn default() -> Self {
Self::new()
}
}
impl RiResourceScheduler {
pub fn new() -> Self {
Self {
performance_history: HashMap::new(),
device_loads: HashMap::new(),
policies: Vec::new(),
}
}
pub fn schedule_resource(
&mut self,
request: &ResourceRequest,
available_devices: &[RiDevice],
) -> Option<String> {
let suitable_devices: Vec<&RiDevice> = available_devices
.iter()
.filter(|device| self.meets_requirements(device, request))
.collect();
if suitable_devices.is_empty() {
return None;
}
let mut device_scores: Vec<(String, f64)> = suitable_devices
.iter()
.map(|device| {
let score = self.calculate_device_score(device, request);
(device.id().to_string(), score)
})
.collect();
device_scores.sort_by(|a, b| {
b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
});
device_scores.first().map(|(device_id, _)| device_id.clone())
}
fn meets_requirements(&self, device: &RiDevice, request: &ResourceRequest) -> bool {
let capabilities = device.capabilities();
if let Some(required_memory) = request.required_memory_gb {
if let Some(available_memory) = capabilities.memory_gb {
if available_memory < required_memory {
return false;
}
} else {
return false; }
}
if let Some(required_compute) = request.required_compute_units {
if let Some(available_compute) = capabilities.compute_units {
if available_compute < required_compute {
return false;
}
} else {
return false; }
}
if let Some(required_bandwidth) = request.required_bandwidth_gbps {
if let Some(available_bandwidth) = capabilities.bandwidth_gbps {
if available_bandwidth < required_bandwidth {
return false;
}
} else {
return false; }
}
for (required_key, required_value) in &request.required_custom_capabilities {
if let Some(available_value) = capabilities.custom_capabilities.get(required_key) {
if available_value != required_value {
return false;
}
} else {
return false; }
}
true
}
fn calculate_device_score(&self, device: &RiDevice, request: &ResourceRequest) -> f64 {
let device_id = device.id();
let base_score = 100.0;
let mut score = base_score;
if let Some(&load) = self.device_loads.get(device_id) {
score -= load.powi(2) * 70.0;
}
if let Some(performance_records) = self.performance_history.get(device_id) {
if !performance_records.is_empty() {
let recent_records = &performance_records[performance_records.len().saturating_sub(10)..];
let record_count = recent_records.len();
if record_count > 0 {
let record_count_f64 = record_count as f64;
let avg_latency = recent_records
.iter()
.map(|record| record.latency_ms)
.sum::<f64>() / record_count_f64;
let avg_throughput = recent_records
.iter()
.map(|record| record.throughput)
.sum::<f64>() / record_count_f64;
let avg_error_rate = recent_records
.iter()
.map(|record| record.error_rate)
.sum::<f64>() / record_count_f64;
let avg_utilization = recent_records
.iter()
.map(|record| record.utilization)
.sum::<f64>() / record_count_f64;
score += (100.0 - avg_latency.min(100.0)) * 0.2;
score += (avg_throughput.min(100.0)) * 0.2;
score -= (avg_error_rate.min(1.0)) * 20.0;
let utilization_bonus = if (0.6..=0.8).contains(&avg_utilization) {
10.0
} else if (0.4..=0.9).contains(&avg_utilization) {
5.0
} else {
0.0
};
score += utilization_bonus;
}
}
}
let health_score = device.health_score() as f64;
score += (health_score / 100.0) * 20.0;
if device.is_responsive(300) { score += 10.0;
}
for policy in &self.policies {
score = self.apply_policy_score_adjustment(device, request, policy, score);
}
score.clamp(0.0, 200.0)
}
fn apply_policy_score_adjustment(
&self,
device: &RiDevice,
request: &ResourceRequest,
policy: &SchedulingPolicy,
current_score: f64,
) -> f64 {
let conditions_met = policy.conditions.iter().all(|condition| {
self.evaluate_condition(device, request, condition)
});
if conditions_met {
match &policy.action {
PolicyAction::PreferDevice(preferred_device) => {
if device.id() == preferred_device {
current_score + 30.0 } else {
current_score
}
}
PolicyAction::AvoidDevice(avoided_device) => {
if device.id() == avoided_device {
current_score - 30.0 } else {
current_score
}
}
PolicyAction::LoadBalance => {
let load_penalty = self.device_loads.get(device.id()).unwrap_or(&0.0) * 20.0;
current_score - load_penalty
}
PolicyAction::PriorityBased => {
let priority_bonus = (request.priority as f64 / 10.0) * 15.0;
current_score + priority_bonus
}
}
} else {
current_score
}
}
fn evaluate_condition(
&self,
device: &RiDevice,
request: &ResourceRequest,
condition: &PolicyCondition,
) -> bool {
let value = match condition.metric.as_str() {
"device_type" => {
match device.device_type() {
RiDeviceType::GPU => 1.0,
RiDeviceType::Memory => 2.0,
RiDeviceType::CPU => 3.0,
RiDeviceType::Storage => 4.0,
RiDeviceType::Network => 5.0,
RiDeviceType::Sensor => 6.0,
RiDeviceType::Actuator => 7.0,
RiDeviceType::Custom => 8.0,
}
}
"request_priority" => request.priority as f64,
"time_of_day" => {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
(now % 86400) as f64 / 3600.0
}
_ => 0.0,
};
match condition.operator.as_str() {
">" => value > condition.threshold,
">=" => value >= condition.threshold,
"<" => value < condition.threshold,
"<=" => value <= condition.threshold,
"==" => (value - condition.threshold).abs() < 0.001,
_ => false,
}
}
pub fn record_performance(
&mut self,
device_id: &str,
latency_ms: f64,
throughput: f64,
error_rate: f64,
utilization: f64,
) {
let record = PerformanceRecord {
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
latency_ms,
throughput,
error_rate,
utilization,
};
self.performance_history
.entry(device_id.to_string())
.or_default()
.push(record);
if let Some(history) = self.performance_history.get_mut(device_id) {
if history.len() > 100 {
history.remove(0);
}
}
}
pub fn update_device_load(&mut self, device_id: &str, load: f64) {
self.device_loads.insert(device_id.to_string(), load.clamp(0.0, 1.0));
}
pub fn add_policy(&mut self, policy: SchedulingPolicy) {
self.policies.push(policy);
self.policies.sort_by(|a, b| b.priority.cmp(&a.priority));
}
}
#[derive(Debug, Clone)]
pub struct ResourceRequest {
pub request_id: String,
pub required_memory_gb: Option<f64>,
pub required_compute_units: Option<usize>,
pub required_bandwidth_gbps: Option<f64>,
pub required_custom_capabilities: HashMap<String, String>,
pub priority: u8,
pub deadline: Option<Instant>,
}