3 changed files with 656 additions and 277 deletions
-
6weed/mount/ml/batch_optimizer.go
-
367weed/mount/ml/optimization_engine_test.go
-
14weed/mount/ml/phase4_integration_test.go
@ -0,0 +1,367 @@ |
|||||
|
package ml |
||||
|
|
||||
|
import ( |
||||
|
"testing" |
||||
|
) |
||||
|
|
||||
|
// TestOptimizationEngine_Basic tests the basic functionality of the optimization engine
|
||||
|
func TestOptimizationEngine_Basic(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
if engine == nil { |
||||
|
t.Fatal("Should create optimization engine") |
||||
|
} |
||||
|
|
||||
|
if !engine.enabled { |
||||
|
t.Error("Engine should be enabled") |
||||
|
} |
||||
|
|
||||
|
// Check that default rules and strategies are loaded
|
||||
|
if len(engine.rules) == 0 { |
||||
|
t.Error("Should have default rules loaded") |
||||
|
} |
||||
|
|
||||
|
if len(engine.strategies) == 0 { |
||||
|
t.Error("Should have default strategies loaded") |
||||
|
} |
||||
|
|
||||
|
t.Logf("Engine initialized with %d rules, %d strategies", len(engine.rules), len(engine.strategies)) |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_RuleEvaluation tests rule evaluation
|
||||
|
func TestOptimizationEngine_RuleEvaluation(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
// Create test context for sequential access of a large model file
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/models/large_model.pth", |
||||
|
FileSize: 2 * 1024 * 1024 * 1024, // 2GB
|
||||
|
FileType: "model", |
||||
|
AccessPattern: SequentialAccess, |
||||
|
AccessFrequency: 10, |
||||
|
Framework: "pytorch", |
||||
|
WorkloadType: "training", |
||||
|
} |
||||
|
|
||||
|
// Apply optimizations
|
||||
|
result := engine.OptimizeAccess(context) |
||||
|
|
||||
|
if result == nil { |
||||
|
t.Fatal("Should return optimization result") |
||||
|
} |
||||
|
|
||||
|
if !result.Applied { |
||||
|
t.Error("Should apply optimizations for large model file with sequential access") |
||||
|
} |
||||
|
|
||||
|
if result.Confidence < 0.5 { |
||||
|
t.Errorf("Expected confidence >= 0.5, got %.2f", result.Confidence) |
||||
|
} |
||||
|
|
||||
|
if len(result.Optimizations) == 0 { |
||||
|
t.Error("Should have applied optimizations") |
||||
|
} |
||||
|
|
||||
|
t.Logf("Applied %d optimizations with confidence %.2f", |
||||
|
len(result.Optimizations), result.Confidence) |
||||
|
|
||||
|
for i, opt := range result.Optimizations { |
||||
|
t.Logf("Optimization %d: type=%s, target=%s", i+1, opt.Type, opt.Target) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_FrameworkDetection tests framework detection
|
||||
|
func TestOptimizationEngine_FrameworkDetection(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
testCases := []struct { |
||||
|
filePath string |
||||
|
expectedFramework string |
||||
|
}{ |
||||
|
{"/models/model.pth", "pytorch"}, |
||||
|
{"/models/model.pt", "pytorch"}, |
||||
|
{"/models/saved_model.pb", "tensorflow"}, |
||||
|
{"/models/model.h5", "tensorflow"}, |
||||
|
{"/models/checkpoint.ckpt", "tensorflow"}, |
||||
|
{"/data/dataset.tfrecord", "tensorflow"}, |
||||
|
{"/unknown/file.bin", ""}, |
||||
|
} |
||||
|
|
||||
|
for _, tc := range testCases { |
||||
|
framework := engine.detectFramework(tc.filePath, nil) |
||||
|
|
||||
|
if tc.expectedFramework == "" { |
||||
|
if framework != "" { |
||||
|
t.Errorf("File %s: expected no framework detection, got %s", tc.filePath, framework) |
||||
|
} |
||||
|
} else { |
||||
|
if framework != tc.expectedFramework { |
||||
|
t.Errorf("File %s: expected framework %s, got %s", |
||||
|
tc.filePath, tc.expectedFramework, framework) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_FileTypeDetection tests file type detection
|
||||
|
func TestOptimizationEngine_FileTypeDetection(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
testCases := []struct { |
||||
|
filePath string |
||||
|
expectedType string |
||||
|
}{ |
||||
|
{"/models/model.pth", "model"}, |
||||
|
{"/data/dataset.csv", "dataset"}, |
||||
|
{"/configs/config.yaml", "config"}, |
||||
|
{"/logs/training.log", "log"}, |
||||
|
{"/unknown/file.bin", "unknown"}, |
||||
|
} |
||||
|
|
||||
|
for _, tc := range testCases { |
||||
|
fileType := engine.detectFileType(tc.filePath) |
||||
|
|
||||
|
if fileType != tc.expectedType { |
||||
|
t.Errorf("File %s: expected type %s, got %s", |
||||
|
tc.filePath, tc.expectedType, fileType) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_ConditionEvaluation tests condition evaluation
|
||||
|
func TestOptimizationEngine_ConditionEvaluation(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/models/test.pth", |
||||
|
FileSize: 5 * 1024 * 1024, // 5MB
|
||||
|
FileType: "model", |
||||
|
AccessPattern: SequentialAccess, |
||||
|
Framework: "pytorch", |
||||
|
} |
||||
|
|
||||
|
// Test various condition types
|
||||
|
testConditions := []struct { |
||||
|
condition RuleCondition |
||||
|
expected bool |
||||
|
}{ |
||||
|
{ |
||||
|
condition: RuleCondition{ |
||||
|
Type: "file_pattern", |
||||
|
Property: "extension", |
||||
|
Operator: "equals", |
||||
|
Value: ".pth", |
||||
|
}, |
||||
|
expected: true, |
||||
|
}, |
||||
|
{ |
||||
|
condition: RuleCondition{ |
||||
|
Type: "file_context", |
||||
|
Property: "size", |
||||
|
Operator: "greater_than", |
||||
|
Value: 1024 * 1024, // 1MB
|
||||
|
}, |
||||
|
expected: true, |
||||
|
}, |
||||
|
{ |
||||
|
condition: RuleCondition{ |
||||
|
Type: "access_pattern", |
||||
|
Property: "pattern_type", |
||||
|
Operator: "equals", |
||||
|
Value: "sequential", |
||||
|
}, |
||||
|
expected: true, |
||||
|
}, |
||||
|
{ |
||||
|
condition: RuleCondition{ |
||||
|
Type: "workload_context", |
||||
|
Property: "framework", |
||||
|
Operator: "equals", |
||||
|
Value: "tensorflow", |
||||
|
}, |
||||
|
expected: false, |
||||
|
}, |
||||
|
} |
||||
|
|
||||
|
for i, tc := range testConditions { |
||||
|
result := engine.evaluateCondition(tc.condition, context) |
||||
|
if result != tc.expected { |
||||
|
t.Errorf("Condition %d: expected %v, got %v", i+1, tc.expected, result) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_PluginSystem tests the plugin system
|
||||
|
func TestOptimizationEngine_PluginSystem(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
// Register a test plugin
|
||||
|
plugin := NewPyTorchPlugin() |
||||
|
err := engine.RegisterPlugin(plugin) |
||||
|
if err != nil { |
||||
|
t.Fatalf("Failed to register plugin: %v", err) |
||||
|
} |
||||
|
|
||||
|
// Verify plugin is registered
|
||||
|
if _, exists := engine.plugins["pytorch"]; !exists { |
||||
|
t.Error("PyTorch plugin should be registered") |
||||
|
} |
||||
|
|
||||
|
// Test framework detection through plugin
|
||||
|
confidence := plugin.DetectFramework("/models/test.pth", nil) |
||||
|
if confidence < 0.5 { |
||||
|
t.Errorf("Expected high confidence for .pth file, got %.2f", confidence) |
||||
|
} |
||||
|
|
||||
|
// Test optimization hints
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/models/test.pth", |
||||
|
FileSize: 100 * 1024 * 1024, // 100MB
|
||||
|
FileType: "model", |
||||
|
Framework: "pytorch", |
||||
|
} |
||||
|
|
||||
|
hints := plugin.GetOptimizationHints(context) |
||||
|
if len(hints) == 0 { |
||||
|
t.Error("Plugin should provide optimization hints") |
||||
|
} |
||||
|
|
||||
|
t.Logf("Plugin provided %d optimization hints", len(hints)) |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_UsagePatterns tests usage pattern learning
|
||||
|
func TestOptimizationEngine_UsagePatterns(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/models/training_model.pth", |
||||
|
FileSize: 50 * 1024 * 1024, // 50MB
|
||||
|
FileType: "model", |
||||
|
AccessPattern: SequentialAccess, |
||||
|
Framework: "pytorch", |
||||
|
WorkloadType: "training", |
||||
|
} |
||||
|
|
||||
|
// Apply optimization multiple times to build usage patterns
|
||||
|
for i := 0; i < 5; i++ { |
||||
|
result := engine.OptimizeAccess(context) |
||||
|
if result == nil { |
||||
|
t.Fatalf("Optimization %d failed", i+1) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// Check that usage patterns are being tracked
|
||||
|
if len(engine.usagePatterns) == 0 { |
||||
|
t.Error("Should have learned usage patterns") |
||||
|
} |
||||
|
|
||||
|
// Verify pattern characteristics
|
||||
|
for patternKey, pattern := range engine.usagePatterns { |
||||
|
t.Logf("Learned pattern: %s (frequency=%d, success_rate=%.2f)", |
||||
|
patternKey, pattern.Frequency, pattern.SuccessRate) |
||||
|
|
||||
|
if pattern.Frequency < 1 { |
||||
|
t.Errorf("Pattern %s should have frequency >= 1", patternKey) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_Metrics tests metrics collection
|
||||
|
func TestOptimizationEngine_Metrics(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
metrics := engine.GetMetrics() |
||||
|
|
||||
|
if metrics == nil { |
||||
|
t.Fatal("Should return metrics") |
||||
|
} |
||||
|
|
||||
|
expectedKeys := []string{"enabled", "rules_count", "templates_count", "strategies_count"} |
||||
|
|
||||
|
for _, key := range expectedKeys { |
||||
|
if _, exists := metrics[key]; !exists { |
||||
|
t.Errorf("Metrics should contain key: %s", key) |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
if metrics["enabled"] != true { |
||||
|
t.Error("Metrics should show engine as enabled") |
||||
|
} |
||||
|
|
||||
|
t.Logf("Engine metrics: %+v", metrics) |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_ConfigurationDriven tests configuration-driven optimization
|
||||
|
func TestOptimizationEngine_ConfigurationDriven(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
defer engine.Shutdown() |
||||
|
|
||||
|
// Test that the engine can apply optimizations based on its loaded configuration
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/data/dataset.csv", |
||||
|
FileSize: 10 * 1024 * 1024, // 10MB
|
||||
|
FileType: "dataset", |
||||
|
AccessPattern: SequentialAccess, |
||||
|
Framework: "", |
||||
|
WorkloadType: "training", |
||||
|
BatchSize: 32, |
||||
|
} |
||||
|
|
||||
|
result := engine.OptimizeAccess(context) |
||||
|
|
||||
|
if result == nil { |
||||
|
t.Fatal("Should return optimization result") |
||||
|
} |
||||
|
|
||||
|
// The engine should make intelligent decisions based on context
|
||||
|
if result.Applied && len(result.Optimizations) > 0 { |
||||
|
t.Logf("Successfully applied %d optimizations", len(result.Optimizations)) |
||||
|
|
||||
|
for _, opt := range result.Optimizations { |
||||
|
if opt.Type == "" || opt.Target == "" { |
||||
|
t.Error("Optimization should have valid type and target") |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
if len(result.Recommendations) > 0 { |
||||
|
t.Logf("Generated %d recommendations", len(result.Recommendations)) |
||||
|
for _, rec := range result.Recommendations { |
||||
|
t.Logf("Recommendation: %s", rec) |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
// TestOptimizationEngine_Shutdown tests proper shutdown
|
||||
|
func TestOptimizationEngine_Shutdown(t *testing.T) { |
||||
|
engine := NewOptimizationEngine(true) |
||||
|
|
||||
|
if !engine.enabled { |
||||
|
t.Error("Engine should start enabled") |
||||
|
} |
||||
|
|
||||
|
engine.Shutdown() |
||||
|
|
||||
|
if engine.enabled { |
||||
|
t.Error("Engine should be disabled after shutdown") |
||||
|
} |
||||
|
|
||||
|
// Test that optimization doesn't work after shutdown
|
||||
|
context := &OptimizationContext{ |
||||
|
FilePath: "/test.pth", |
||||
|
FileSize: 1024, |
||||
|
} |
||||
|
|
||||
|
result := engine.OptimizeAccess(context) |
||||
|
if result.Applied { |
||||
|
t.Error("Should not apply optimizations after shutdown") |
||||
|
} |
||||
|
} |
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue