osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
895 lines
29 KiB
Rust
895 lines
29 KiB
Rust
//! Merge planning and strategy recommendation system
|
|
|
|
use crate::algorithms::MergeUtils;
|
|
use crate::config::MergeStrategy;
|
|
use crate::error::{MergeError, Result};
|
|
use crate::strategies::StrategyUtils;
|
|
use crate::types::{Model, ModelMetadata};
|
|
use std::collections::HashMap;
|
|
use tracing::{debug, info};
|
|
|
|
/// Model merge planner and strategy recommender
|
|
pub struct MergePlanner {
|
|
compatibility_cache: HashMap<String, bool>,
|
|
}
|
|
|
|
impl MergePlanner {
|
|
/// Create a new merge planner
|
|
pub fn new() -> Self {
|
|
Self {
|
|
compatibility_cache: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Recommend the best merge strategy for given models
|
|
pub fn recommend_strategy(&self, models: &[ModelMetadata]) -> Result<MergeStrategy> {
|
|
if models.len() < 2 {
|
|
return Err(MergeError::algorithm(
|
|
"MergePlanner",
|
|
"At least 2 models required for strategy recommendation",
|
|
));
|
|
}
|
|
|
|
info!(
|
|
"Analyzing {} models for strategy recommendation",
|
|
models.len()
|
|
);
|
|
|
|
// Analyze model characteristics
|
|
let analysis = self.analyze_models(models)?;
|
|
|
|
// Generate strategy recommendation based on analysis
|
|
let recommended_strategy = self.select_optimal_strategy(&analysis)?;
|
|
|
|
info!("Recommended strategy: {:?}", recommended_strategy);
|
|
Ok(recommended_strategy)
|
|
}
|
|
|
|
/// Analyze model compatibility and characteristics
|
|
pub fn analyze_compatibility(&self, models: &[Model]) -> Result<CompatibilityReport> {
|
|
let mut report = CompatibilityReport::new();
|
|
|
|
if models.len() < 2 {
|
|
report.overall_compatible = true;
|
|
return Ok(report);
|
|
}
|
|
|
|
// Check pairwise compatibility
|
|
for i in 0..models.len() {
|
|
for j in (i + 1)..models.len() {
|
|
let compatibility = self.check_model_pair_compatibility(&models[i], &models[j])?;
|
|
|
|
if !compatibility.compatible {
|
|
report.overall_compatible = false;
|
|
report
|
|
.compatibility_issues
|
|
.extend(compatibility.issues.clone());
|
|
}
|
|
|
|
report.pairwise_compatibility.push(compatibility);
|
|
}
|
|
}
|
|
|
|
// Compute similarity matrix
|
|
report.similarity_matrix = StrategyUtils::compute_pairwise_similarities(models)?;
|
|
|
|
// Analyze diversity
|
|
let all_indices: Vec<usize> = (0..models.len()).collect();
|
|
report.diversity_score = StrategyUtils::compute_group_diversity(models, &all_indices)?;
|
|
|
|
// Estimate resource requirements
|
|
report.memory_requirements = self.estimate_merge_memory_requirements(models)?;
|
|
|
|
info!(
|
|
"Compatibility analysis: {}/{} compatible, diversity: {:.3}",
|
|
report
|
|
.pairwise_compatibility
|
|
.iter()
|
|
.filter(|c| c.compatible)
|
|
.count(),
|
|
report.pairwise_compatibility.len(),
|
|
report.diversity_score
|
|
);
|
|
|
|
Ok(report)
|
|
}
|
|
|
|
/// Plan merge execution with resource constraints
|
|
pub fn plan_merge_execution(
|
|
&self,
|
|
models: &[Model],
|
|
strategy: &MergeStrategy,
|
|
constraints: &ResourceConstraints,
|
|
) -> Result<ExecutionPlan> {
|
|
info!("Planning merge execution for {} models", models.len());
|
|
|
|
let mut plan = ExecutionPlan::new();
|
|
|
|
// Validate resource constraints
|
|
self.validate_resource_constraints(models, strategy, constraints)?;
|
|
|
|
// Determine execution phases
|
|
plan.phases = self.plan_execution_phases(models, strategy, constraints)?;
|
|
|
|
// Estimate execution time
|
|
plan.estimated_duration_ms = self.estimate_execution_time(models, strategy)?;
|
|
|
|
// Identify potential bottlenecks
|
|
plan.bottlenecks = self.identify_bottlenecks(models, strategy, constraints)?;
|
|
|
|
// Generate checkpoints
|
|
plan.checkpoints = self.plan_checkpoints(models, strategy)?;
|
|
|
|
// Risk assessment
|
|
plan.risks = self.assess_execution_risks(models, strategy)?;
|
|
|
|
info!(
|
|
"Execution plan created: {} phases, estimated duration: {}ms",
|
|
plan.phases.len(),
|
|
plan.estimated_duration_ms
|
|
);
|
|
|
|
Ok(plan)
|
|
}
|
|
|
|
/// Analyze models for strategy selection
|
|
fn analyze_models(&self, models: &[ModelMetadata]) -> Result<ModelAnalysis> {
|
|
let mut analysis = ModelAnalysis::new();
|
|
|
|
// Compute model statistics
|
|
for model in models {
|
|
let stats = ModelStats {
|
|
size_mb: model.size_bytes.unwrap_or(0) / (1024 * 1024),
|
|
complexity_score: self.estimate_model_complexity(model),
|
|
similarity_to_others: 0.0, // Will be computed below
|
|
training_quality: model.metrics.get("accuracy").copied().unwrap_or(0.5),
|
|
};
|
|
analysis.model_stats.push(stats);
|
|
}
|
|
|
|
// Compute average model size
|
|
analysis.avg_model_size_mb = analysis
|
|
.model_stats
|
|
.iter()
|
|
.map(|s| s.size_mb)
|
|
.sum::<usize>() as f32
|
|
/ models.len() as f32;
|
|
|
|
// Compute diversity (placeholder - would need actual model data)
|
|
analysis.diversity_score = 0.6; // Placeholder
|
|
|
|
// Determine merge complexity
|
|
analysis.merge_complexity = self.assess_merge_complexity(models);
|
|
|
|
// Check for special requirements
|
|
analysis.special_requirements = self.identify_special_requirements(models);
|
|
|
|
Ok(analysis)
|
|
}
|
|
|
|
/// Select optimal merge strategy based on analysis
|
|
fn select_optimal_strategy(&self, analysis: &ModelAnalysis) -> Result<MergeStrategy> {
|
|
debug!(
|
|
"Selecting strategy for analysis: {:?}",
|
|
analysis.merge_complexity
|
|
);
|
|
|
|
match analysis.merge_complexity {
|
|
MergeComplexity::Simple => {
|
|
// For simple merges, use efficient uniform averaging
|
|
if analysis.model_stats.len() > 5 {
|
|
Ok(MergeStrategy::ModelSoup(crate::config::ModelSoupConfig {
|
|
soup_method: crate::config::SoupMethod::Uniform,
|
|
max_models: 10,
|
|
..Default::default()
|
|
}))
|
|
} else {
|
|
Ok(MergeStrategy::Slerp(crate::config::SlerpConfig {
|
|
t: 0.5,
|
|
normalization: crate::config::NormalizationMethod::L2,
|
|
..Default::default()
|
|
}))
|
|
}
|
|
}
|
|
MergeComplexity::Moderate => {
|
|
// For moderate complexity, consider parameter conflicts
|
|
if analysis.diversity_score > 0.7 {
|
|
Ok(MergeStrategy::Ties(crate::config::TiesConfig {
|
|
sign_threshold: 0.6,
|
|
magnitude_threshold: 0.1,
|
|
density: 0.8,
|
|
..Default::default()
|
|
}))
|
|
} else {
|
|
Ok(MergeStrategy::TaskArithmetic(
|
|
crate::config::TaskArithmeticConfig {
|
|
scaling_factors: vec![1.0; analysis.model_stats.len() - 1],
|
|
preserve_signs: true,
|
|
..Default::default()
|
|
},
|
|
))
|
|
}
|
|
}
|
|
MergeComplexity::Complex => {
|
|
// For complex merges, use sophisticated methods
|
|
if analysis.avg_model_size_mb > 1000.0 {
|
|
Ok(MergeStrategy::Dare(crate::config::DareConfig {
|
|
drop_probability: 0.1,
|
|
adaptive_dropping: true,
|
|
importance_threshold: 0.05,
|
|
..Default::default()
|
|
}))
|
|
} else {
|
|
Ok(MergeStrategy::Fisher(crate::config::FisherConfig {
|
|
fisher_method: crate::config::FisherMethod::Empirical,
|
|
diagonal_fisher: true,
|
|
num_samples: 1000,
|
|
..Default::default()
|
|
}))
|
|
}
|
|
}
|
|
MergeComplexity::VeryComplex => {
|
|
// For very complex scenarios, use progressive merging
|
|
Ok(MergeStrategy::Progressive(
|
|
crate::config::ProgressiveConfig {
|
|
num_steps: 10,
|
|
validation_frequency: 2,
|
|
early_stopping_patience: Some(3),
|
|
enable_rollback: true,
|
|
..Default::default()
|
|
},
|
|
))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Check compatibility between two models
|
|
fn check_model_pair_compatibility(
|
|
&self,
|
|
model1: &Model,
|
|
model2: &Model,
|
|
) -> Result<PairwiseCompatibility> {
|
|
let mut compatibility = PairwiseCompatibility {
|
|
model1_name: model1.name.clone(),
|
|
model2_name: model2.name.clone(),
|
|
compatible: true,
|
|
similarity_score: 0.0,
|
|
issues: Vec::new(),
|
|
};
|
|
|
|
// Check architecture compatibility
|
|
if model1.architecture != model2.architecture {
|
|
compatibility.compatible = false;
|
|
compatibility.issues.push(format!(
|
|
"Architecture mismatch: {} vs {}",
|
|
model1.architecture.arch_type, model2.architecture.arch_type
|
|
));
|
|
}
|
|
|
|
// Check parameter compatibility
|
|
if model1.parameters.len() != model2.parameters.len() {
|
|
compatibility.compatible = false;
|
|
compatibility.issues.push(format!(
|
|
"Parameter count mismatch: {} vs {}",
|
|
model1.parameters.len(),
|
|
model2.parameters.len()
|
|
));
|
|
} else {
|
|
// Check individual parameters
|
|
for (name, param1) in &model1.parameters {
|
|
if let Some(param2) = model2.parameters.get(name) {
|
|
if param1.shape != param2.shape {
|
|
compatibility.compatible = false;
|
|
compatibility.issues.push(format!(
|
|
"Parameter shape mismatch in {}: {:?} vs {:?}",
|
|
name, param1.shape, param2.shape
|
|
));
|
|
}
|
|
} else {
|
|
compatibility.compatible = false;
|
|
compatibility
|
|
.issues
|
|
.push(format!("Parameter {name} missing in model2"));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute similarity if compatible
|
|
if compatibility.compatible {
|
|
compatibility.similarity_score = MergeUtils::compute_similarity(model1, model2)?;
|
|
}
|
|
|
|
Ok(compatibility)
|
|
}
|
|
|
|
/// Estimate model complexity from metadata
|
|
fn estimate_model_complexity(&self, metadata: &ModelMetadata) -> f32 {
|
|
let size_factor =
|
|
(metadata.size_bytes.unwrap_or(0) as f32 / (100 * 1024 * 1024) as f32).min(1.0);
|
|
let metric_factor = 1.0 - metadata.metrics.get("accuracy").unwrap_or(&0.5);
|
|
|
|
f32::midpoint(size_factor, metric_factor)
|
|
}
|
|
|
|
/// Assess overall merge complexity
|
|
fn assess_merge_complexity(&self, models: &[ModelMetadata]) -> MergeComplexity {
|
|
let avg_size_mb = models
|
|
.iter()
|
|
.map(|m| m.size_bytes.unwrap_or(0) / (1024 * 1024))
|
|
.sum::<usize>() as f32
|
|
/ models.len() as f32;
|
|
|
|
let model_count = models.len();
|
|
|
|
match (avg_size_mb, model_count) {
|
|
(size, count) if size < 100.0 && count <= 3 => MergeComplexity::Simple,
|
|
(size, count) if size < 500.0 && count <= 5 => MergeComplexity::Moderate,
|
|
(size, count) if size < 2000.0 || count <= 10 => MergeComplexity::Complex,
|
|
_ => MergeComplexity::VeryComplex,
|
|
}
|
|
}
|
|
|
|
/// Identify special requirements for models
|
|
fn identify_special_requirements(&self, models: &[ModelMetadata]) -> Vec<String> {
|
|
let mut requirements = Vec::new();
|
|
|
|
let total_size_gb = models
|
|
.iter()
|
|
.map(|m| m.size_bytes.unwrap_or(0))
|
|
.sum::<usize>() as f32
|
|
/ (1024.0 * 1024.0 * 1024.0);
|
|
|
|
if total_size_gb > 10.0 {
|
|
requirements.push("high_memory".to_string());
|
|
}
|
|
|
|
if models.len() > 8 {
|
|
requirements.push("parallel_processing".to_string());
|
|
}
|
|
|
|
// Check for diverse architectures (would need actual architecture comparison)
|
|
requirements.push("architecture_validation".to_string());
|
|
|
|
requirements
|
|
}
|
|
|
|
/// Validate resource constraints
|
|
fn validate_resource_constraints(
|
|
&self,
|
|
models: &[Model],
|
|
strategy: &MergeStrategy,
|
|
constraints: &ResourceConstraints,
|
|
) -> Result<()> {
|
|
// Memory validation
|
|
if let Some(max_memory_mb) = constraints.max_memory_mb {
|
|
let estimated_memory = StrategyUtils::estimate_memory_requirements(models, "merge")?;
|
|
let estimated_mb = estimated_memory / (1024 * 1024);
|
|
|
|
if estimated_mb > max_memory_mb {
|
|
return Err(MergeError::memory(estimated_mb));
|
|
}
|
|
}
|
|
|
|
// Time validation
|
|
if let Some(max_time_sec) = constraints.max_time_sec {
|
|
let estimated_time_ms = self.estimate_execution_time(models, strategy)?;
|
|
let estimated_sec = estimated_time_ms / 1000;
|
|
|
|
if estimated_sec > max_time_sec as u64 {
|
|
return Err(MergeError::timeout(estimated_time_ms));
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Plan execution phases
|
|
fn plan_execution_phases(
|
|
&self,
|
|
models: &[Model],
|
|
strategy: &MergeStrategy,
|
|
_constraints: &ResourceConstraints,
|
|
) -> Result<Vec<ExecutionPhase>> {
|
|
let mut phases = Vec::new();
|
|
|
|
// Phase 1: Preparation
|
|
phases.push(ExecutionPhase {
|
|
name: "Preparation".to_string(),
|
|
description: "Load and validate models".to_string(),
|
|
estimated_duration_ms: 5000 * models.len() as u64,
|
|
memory_requirement_mb: models
|
|
.iter()
|
|
.map(super::types::Model::memory_size)
|
|
.sum::<usize>()
|
|
/ (1024 * 1024),
|
|
dependencies: Vec::new(),
|
|
});
|
|
|
|
// Phase 2: Strategy-specific execution
|
|
let merge_duration = match strategy {
|
|
MergeStrategy::Ties(_) => 10000 + 2000 * models.len() as u64,
|
|
MergeStrategy::Dare(_) => 15000 + 3000 * models.len() as u64,
|
|
MergeStrategy::Slerp(_) => 8000 + 1500 * models.len() as u64,
|
|
MergeStrategy::TaskArithmetic(_) => 12000 + 2500 * models.len() as u64,
|
|
MergeStrategy::Fisher(_) => 20000 + 5000 * models.len() as u64,
|
|
MergeStrategy::ModelSoup(_) => 7000 + 1000 * models.len() as u64,
|
|
MergeStrategy::Frankenmerge(_) => 25000 + 4000 * models.len() as u64,
|
|
MergeStrategy::Progressive(_) => 30000 + 6000 * models.len() as u64,
|
|
};
|
|
|
|
phases.push(ExecutionPhase {
|
|
name: "Merging".to_string(),
|
|
description: format!("Execute {} merge strategy", strategy.name()),
|
|
estimated_duration_ms: merge_duration,
|
|
memory_requirement_mb: StrategyUtils::estimate_memory_requirements(
|
|
models,
|
|
strategy.name(),
|
|
)? / (1024 * 1024),
|
|
dependencies: vec!["Preparation".to_string()],
|
|
});
|
|
|
|
// Phase 3: Validation
|
|
phases.push(ExecutionPhase {
|
|
name: "Validation".to_string(),
|
|
description: "Validate merged model".to_string(),
|
|
estimated_duration_ms: 3000,
|
|
memory_requirement_mb: 100, // Minimal for validation
|
|
dependencies: vec!["Merging".to_string()],
|
|
});
|
|
|
|
Ok(phases)
|
|
}
|
|
|
|
/// Estimate execution time
|
|
fn estimate_execution_time(&self, models: &[Model], strategy: &MergeStrategy) -> Result<u64> {
|
|
let base_time = 10000; // 10 seconds base
|
|
let model_factor = models.len() as u64 * 2000; // 2 seconds per model
|
|
let param_factor = models
|
|
.iter()
|
|
.map(super::types::Model::parameter_count)
|
|
.sum::<usize>() as u64
|
|
/ 1000000; // 1ms per 1M params
|
|
|
|
let strategy_multiplier = match strategy {
|
|
MergeStrategy::Ties(_) => 2.0,
|
|
MergeStrategy::Dare(_) => 2.5,
|
|
MergeStrategy::Slerp(_) => 1.5,
|
|
MergeStrategy::TaskArithmetic(_) => 2.0,
|
|
MergeStrategy::Fisher(_) => 4.0,
|
|
MergeStrategy::ModelSoup(_) => 1.2,
|
|
MergeStrategy::Frankenmerge(_) => 3.0,
|
|
MergeStrategy::Progressive(_) => 5.0,
|
|
};
|
|
|
|
Ok(((base_time + model_factor + param_factor) as f32 * strategy_multiplier) as u64)
|
|
}
|
|
|
|
/// Estimate memory requirements for merging
|
|
fn estimate_merge_memory_requirements(&self, models: &[Model]) -> Result<usize> {
|
|
let total_model_memory: usize = models.iter().map(super::types::Model::memory_size).sum();
|
|
let overhead = (total_model_memory as f32 * 0.5) as usize; // 50% overhead
|
|
Ok(total_model_memory + overhead)
|
|
}
|
|
|
|
/// Identify potential bottlenecks
|
|
fn identify_bottlenecks(
|
|
&self,
|
|
models: &[Model],
|
|
strategy: &MergeStrategy,
|
|
constraints: &ResourceConstraints,
|
|
) -> Result<Vec<String>> {
|
|
let mut bottlenecks = Vec::new();
|
|
|
|
// Memory bottleneck
|
|
let estimated_memory = self.estimate_merge_memory_requirements(models)?;
|
|
if let Some(max_memory) = constraints.max_memory_mb
|
|
&& estimated_memory / (1024 * 1024) > max_memory * 9 / 10
|
|
{
|
|
bottlenecks.push("Memory usage near limit".to_string());
|
|
}
|
|
|
|
// CPU bottleneck for complex strategies
|
|
match strategy {
|
|
MergeStrategy::Fisher(_) | MergeStrategy::Progressive(_) => {
|
|
bottlenecks.push("CPU-intensive strategy".to_string());
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
// I/O bottleneck for many models
|
|
if models.len() > 10 {
|
|
bottlenecks.push("I/O bottleneck with many models".to_string());
|
|
}
|
|
|
|
Ok(bottlenecks)
|
|
}
|
|
|
|
/// Plan execution checkpoints
|
|
fn plan_checkpoints(&self, _models: &[Model], strategy: &MergeStrategy) -> Result<Vec<String>> {
|
|
let mut checkpoints = vec![
|
|
"Models loaded".to_string(),
|
|
"Compatibility validated".to_string(),
|
|
];
|
|
|
|
match strategy {
|
|
MergeStrategy::Progressive(_) => {
|
|
checkpoints.extend(vec![
|
|
"Phase 1 complete".to_string(),
|
|
"Phase 2 complete".to_string(),
|
|
"Phase 3 complete".to_string(),
|
|
]);
|
|
}
|
|
_ => {
|
|
checkpoints.push("Merge parameters computed".to_string());
|
|
}
|
|
}
|
|
|
|
checkpoints.extend(vec![
|
|
"Merge complete".to_string(),
|
|
"Validation passed".to_string(),
|
|
]);
|
|
|
|
Ok(checkpoints)
|
|
}
|
|
|
|
/// Assess execution risks
|
|
fn assess_execution_risks(
|
|
&self,
|
|
models: &[Model],
|
|
strategy: &MergeStrategy,
|
|
) -> Result<Vec<String>> {
|
|
let mut risks = Vec::new();
|
|
|
|
// Large model risk
|
|
let max_model_size = models
|
|
.iter()
|
|
.map(super::types::Model::memory_size)
|
|
.max()
|
|
.unwrap_or(0);
|
|
if max_model_size > 2 * 1024 * 1024 * 1024 {
|
|
// 2GB
|
|
risks.push("Large model may cause memory issues".to_string());
|
|
}
|
|
|
|
// Strategy-specific risks
|
|
match strategy {
|
|
MergeStrategy::Dare(_) => {
|
|
risks.push("Random dropping may affect model quality".to_string());
|
|
}
|
|
MergeStrategy::Fisher(_) => {
|
|
risks.push("Fisher information computation is resource-intensive".to_string());
|
|
}
|
|
MergeStrategy::Progressive(_) => {
|
|
risks.push("Progressive merge may fail at intermediate steps".to_string());
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
// Model diversity risk
|
|
if models.len() > 15 {
|
|
risks.push("Many models may reduce merge quality".to_string());
|
|
}
|
|
|
|
Ok(risks)
|
|
}
|
|
}
|
|
|
|
// Strategy name extension
|
|
impl MergeStrategy {
|
|
fn name(&self) -> &'static str {
|
|
match self {
|
|
Self::Ties(_) => "TIES",
|
|
Self::Dare(_) => "DARE",
|
|
Self::Slerp(_) => "SLERP",
|
|
Self::TaskArithmetic(_) => "TaskArithmetic",
|
|
Self::Fisher(_) => "Fisher",
|
|
Self::ModelSoup(_) => "ModelSoup",
|
|
Self::Frankenmerge(_) => "Frankenmerge",
|
|
Self::Progressive(_) => "Progressive",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model analysis results
|
|
#[derive(Debug)]
|
|
struct ModelAnalysis {
|
|
model_stats: Vec<ModelStats>,
|
|
avg_model_size_mb: f32,
|
|
diversity_score: f32,
|
|
merge_complexity: MergeComplexity,
|
|
special_requirements: Vec<String>,
|
|
}
|
|
|
|
impl ModelAnalysis {
|
|
fn new() -> Self {
|
|
Self {
|
|
model_stats: Vec::new(),
|
|
avg_model_size_mb: 0.0,
|
|
diversity_score: 0.0,
|
|
merge_complexity: MergeComplexity::Simple,
|
|
special_requirements: Vec::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Statistics for individual models
|
|
#[derive(Debug)]
|
|
struct ModelStats {
|
|
size_mb: usize,
|
|
complexity_score: f32,
|
|
similarity_to_others: f32,
|
|
training_quality: f32,
|
|
}
|
|
|
|
/// Merge complexity levels
|
|
#[derive(Debug, PartialEq)]
|
|
enum MergeComplexity {
|
|
Simple,
|
|
Moderate,
|
|
Complex,
|
|
VeryComplex,
|
|
}
|
|
|
|
/// Compatibility report for models
|
|
#[derive(Debug)]
|
|
pub struct CompatibilityReport {
|
|
pub overall_compatible: bool,
|
|
pub pairwise_compatibility: Vec<PairwiseCompatibility>,
|
|
pub compatibility_issues: Vec<String>,
|
|
pub similarity_matrix: Vec<Vec<f32>>,
|
|
pub diversity_score: f32,
|
|
pub memory_requirements: usize,
|
|
}
|
|
|
|
impl CompatibilityReport {
|
|
fn new() -> Self {
|
|
Self {
|
|
overall_compatible: true,
|
|
pairwise_compatibility: Vec::new(),
|
|
compatibility_issues: Vec::new(),
|
|
similarity_matrix: Vec::new(),
|
|
diversity_score: 0.0,
|
|
memory_requirements: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Pairwise model compatibility
|
|
#[derive(Debug)]
|
|
pub struct PairwiseCompatibility {
|
|
pub model1_name: String,
|
|
pub model2_name: String,
|
|
pub compatible: bool,
|
|
pub similarity_score: f32,
|
|
pub issues: Vec<String>,
|
|
}
|
|
|
|
/// Resource constraints for merge execution
|
|
#[derive(Debug)]
|
|
pub struct ResourceConstraints {
|
|
pub max_memory_mb: Option<usize>,
|
|
pub max_time_sec: Option<usize>,
|
|
pub max_gpu_memory_mb: Option<usize>,
|
|
pub parallel_workers: Option<usize>,
|
|
}
|
|
|
|
impl Default for ResourceConstraints {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_memory_mb: Some(8192), // 8GB default
|
|
max_time_sec: Some(3600), // 1 hour default
|
|
max_gpu_memory_mb: None,
|
|
parallel_workers: Some(4),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Merge execution plan
|
|
#[derive(Debug)]
|
|
pub struct ExecutionPlan {
|
|
pub phases: Vec<ExecutionPhase>,
|
|
pub estimated_duration_ms: u64,
|
|
pub bottlenecks: Vec<String>,
|
|
pub checkpoints: Vec<String>,
|
|
pub risks: Vec<String>,
|
|
}
|
|
|
|
impl ExecutionPlan {
|
|
fn new() -> Self {
|
|
Self {
|
|
phases: Vec::new(),
|
|
estimated_duration_ms: 0,
|
|
bottlenecks: Vec::new(),
|
|
checkpoints: Vec::new(),
|
|
risks: Vec::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Execution phase description
|
|
#[derive(Debug)]
|
|
pub struct ExecutionPhase {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub estimated_duration_ms: u64,
|
|
pub memory_requirement_mb: usize,
|
|
pub dependencies: Vec<String>,
|
|
}
|
|
|
|
impl Default for MergePlanner {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::types::{DataType, ModelArchitecture, ParameterTensor};
|
|
|
|
fn create_test_metadata(name: &str, size_mb: usize, accuracy: f32) -> ModelMetadata {
|
|
let mut metadata = ModelMetadata::default();
|
|
metadata.size_bytes = Some(size_mb * 1024 * 1024);
|
|
metadata.metrics.insert("accuracy".to_string(), accuracy);
|
|
metadata
|
|
}
|
|
|
|
fn create_test_model(name: &str) -> Model {
|
|
let arch = ModelArchitecture {
|
|
arch_type: "test".to_string(),
|
|
num_layers: 12,
|
|
hidden_dim: 768,
|
|
params: std::collections::HashMap::new(),
|
|
};
|
|
|
|
let mut model = Model::new(name.to_string(), arch);
|
|
let param = ParameterTensor::new(
|
|
"weight".to_string(),
|
|
vec![768, 768],
|
|
DataType::Float32,
|
|
vec![0.0; 768 * 768],
|
|
);
|
|
model.add_parameter(param);
|
|
model
|
|
}
|
|
|
|
#[test]
|
|
fn test_strategy_recommendation_simple() -> Result<()> {
|
|
let models = vec![
|
|
create_test_metadata("model1", 50, 0.8), // Small, good model
|
|
create_test_metadata("model2", 60, 0.75), // Small, good model
|
|
];
|
|
|
|
let planner = MergePlanner::new();
|
|
let strategy = planner.recommend_strategy(&models)?;
|
|
|
|
// Should recommend simple strategy for small models
|
|
match strategy {
|
|
MergeStrategy::Slerp(_) | MergeStrategy::ModelSoup(_) => (),
|
|
_ => panic!("Expected simple strategy for small models"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_strategy_recommendation_complex() -> Result<()> {
|
|
let models = vec![
|
|
create_test_metadata("model1", 2000, 0.9), // Large model
|
|
create_test_metadata("model2", 2500, 0.85), // Large model
|
|
create_test_metadata("model3", 1800, 0.88), // Large model
|
|
];
|
|
|
|
let planner = MergePlanner::new();
|
|
let strategy = planner.recommend_strategy(&models)?;
|
|
|
|
// Should recommend complex strategy for large models
|
|
match strategy {
|
|
MergeStrategy::Dare(_) | MergeStrategy::Fisher(_) | MergeStrategy::Progressive(_) => (),
|
|
_ => panic!("Expected complex strategy for large models"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_compatibility_analysis() -> Result<()> {
|
|
let model1 = create_test_model("model1");
|
|
let model2 = create_test_model("model2");
|
|
let models = vec![model1, model2];
|
|
|
|
let planner = MergePlanner::new();
|
|
let report = planner.analyze_compatibility(&models)?;
|
|
|
|
assert!(report.overall_compatible);
|
|
assert_eq!(report.pairwise_compatibility.len(), 1);
|
|
assert!(report.similarity_matrix.len() == 2);
|
|
assert!(report.diversity_score >= 0.0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_execution_planning() -> Result<()> {
|
|
let models = vec![create_test_model("model1"), create_test_model("model2")];
|
|
let strategy = MergeStrategy::Ties(crate::config::TiesConfig::default());
|
|
let constraints = ResourceConstraints::default();
|
|
|
|
let planner = MergePlanner::new();
|
|
let plan = planner.plan_merge_execution(&models, &strategy, &constraints)?;
|
|
|
|
assert!(!plan.phases.is_empty());
|
|
assert!(plan.estimated_duration_ms > 0);
|
|
assert!(!plan.checkpoints.is_empty());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_resource_constraint_validation() -> Result<()> {
|
|
let models = vec![create_test_model("model1")];
|
|
let strategy = MergeStrategy::Slerp(crate::config::SlerpConfig::default());
|
|
|
|
// Very restrictive constraints should fail
|
|
let restrictive_constraints = ResourceConstraints {
|
|
max_memory_mb: Some(1), // 1MB - too small
|
|
max_time_sec: Some(1), // 1 second - too short
|
|
..Default::default()
|
|
};
|
|
|
|
let planner = MergePlanner::new();
|
|
let result =
|
|
planner.validate_resource_constraints(&models, &strategy, &restrictive_constraints);
|
|
|
|
assert!(result.is_err()); // Should fail validation
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_merge_complexity_assessment() {
|
|
let planner = MergePlanner::new();
|
|
|
|
// Simple case
|
|
let simple_models = vec![create_test_metadata("model1", 50, 0.8)];
|
|
let simple_complexity = planner.assess_merge_complexity(&simple_models);
|
|
assert_eq!(simple_complexity, MergeComplexity::Simple);
|
|
|
|
// Complex case
|
|
let complex_models = vec![
|
|
create_test_metadata("model1", 2000, 0.9),
|
|
create_test_metadata("model2", 2500, 0.85),
|
|
create_test_metadata("model3", 1800, 0.88),
|
|
];
|
|
let complex_complexity = planner.assess_merge_complexity(&complex_models);
|
|
assert!(matches!(
|
|
complex_complexity,
|
|
MergeComplexity::Complex | MergeComplexity::VeryComplex
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_bottleneck_identification() -> Result<()> {
|
|
let large_models: Vec<Model> = (0..15)
|
|
.map(|i| create_test_model(&format!("model{}", i)))
|
|
.collect();
|
|
let strategy = MergeStrategy::Fisher(crate::config::FisherConfig::default());
|
|
let constraints = ResourceConstraints {
|
|
max_memory_mb: Some(100), // Very low memory limit
|
|
..Default::default()
|
|
};
|
|
|
|
let planner = MergePlanner::new();
|
|
let bottlenecks = planner.identify_bottlenecks(&large_models, &strategy, &constraints)?;
|
|
|
|
assert!(!bottlenecks.is_empty());
|
|
assert!(
|
|
bottlenecks
|
|
.iter()
|
|
.any(|b| b.contains("Memory") || b.contains("CPU") || b.contains("I/O"))
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
}
|