evolution_engine_auto_scale.py

"""
自动扩展的代码进化引擎
支持动态添加/移除 Worker，真正的灵活扩展
"""
import requests
import json
from datetime import datetime
import random
import concurrent.futures
import time
import threading


class EvolutionEngineAutoScale:
    def __init__(self, worker_urls=None, auto_discover=True):
        """
        初始化自动扩展引擎
        
        参数:
            worker_urls: 初始 Worker 列表（可选）
            auto_discover: 是否自动发现新 Worker
        """
        self.worker_urls = worker_urls or []
        self.available_workers = []
        self.auto_discover = auto_discover
        self.worker_stats = {}  # 记录每个 Worker 的统计信息
        
        # 启动 Worker 监控线程
        if auto_discover:
            self.monitor_thread = threading.Thread(
                target=self._monitor_workers,
                daemon=True
            )
            self.monitor_thread.start()
        
        print(f"🔍 初始化自动扩展引擎...")
        self.check_workers()
    
    def _monitor_workers(self):
        """后台监控 Worker 状态"""
        while True:
            time.sleep(10)  # 每 10 秒检查一次
            old_count = len(self.available_workers)
            self.check_workers(silent=True)
            new_count = len(self.available_workers)
            
            if new_count > old_count:
                print(f"\n✨ 检测到新 Worker！当前可用: {new_count} 个")
            elif new_count < old_count:
                print(f"\n⚠️ Worker 减少！当前可用: {new_count} 个")
    
    def add_worker(self, worker_url):
        """动态添加 Worker"""
        if worker_url not in self.worker_urls:
            self.worker_urls.append(worker_url)
            print(f"➕ 添加 Worker: {worker_url}")
            self.check_workers()
    
    def remove_worker(self, worker_url):
        """动态移除 Worker"""
        if worker_url in self.worker_urls:
            self.worker_urls.remove(worker_url)
        if worker_url in self.available_workers:
            self.available_workers.remove(worker_url)
        if worker_url in self.worker_stats:
            del self.worker_stats[worker_url]
        print(f"➖ 移除 Worker: {worker_url}")
    
    def check_workers(self, silent=False):
        """检查哪些 Worker 可用"""
        if not silent:
            print(f"🔍 检测 Worker...")
        
        new_available = []
        for url in self.worker_urls:
            try:
                response = requests.get(f"{url}/health", timeout=2)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('model_loaded'):
                        new_available.append(url)
                        
                        # 初始化统计信息
                        if url not in self.worker_stats:
                            self.worker_stats[url] = {
                                'requests': 0,
                                'failures': 0,
                                'total_time': 0
                            }
                        
                        if not silent:
                            print(f"  ✅ {url}")
                    else:
                        if not silent:
                            print(f"  ⚠️ {url} (模型未加载)")
            except:
                if not silent:
                    print(f"  ❌ {url} 不可用")
        
        self.available_workers = new_available
        self.available = len(self.available_workers) > 0
        
        if not silent:
            print(f"\n可用 Worker: {len(self.available_workers)}/{len(self.worker_urls)}")
    
    def get_best_worker(self):
        """选择最佳 Worker（基于负载均衡）"""
        if not self.available_workers:
            return None
        
        # 简单的负载均衡：选择请求数最少的 Worker
        best_worker = min(
            self.available_workers,
            key=lambda w: self.worker_stats.get(w, {}).get('requests', 0)
        )
        return best_worker
    
    def mutate_code(self, code, generation, best_time, worker_url=None):
        """生成代码变体"""
        if not self.available_workers:
            return code
        
        # 如果没有指定 Worker，选择最佳 Worker
        if worker_url is None:
            worker_url = self.get_best_worker()
        
        if worker_url is None:
            return code
        
        # 更新统计
        self.worker_stats[worker_url]['requests'] += 1
        
        prompt = f"""你是代码优化专家。请优化以下排序函数。

当前代码（第 {generation} 代，耗时 {best_time:.6f}秒）：
```python
{code}
```

优化要求：
1. 保持函数签名不变：def sort_array(arr)
2. 必须返回排序后的数组
3. 尝试改进算法效率
4. 可以尝试：改变算法逻辑、优化循环、使用更好的数据结构

只返回优化后的完整函数代码，不要解释：
```python"""

        start_time = time.time()
        try:
            response = requests.post(
                f"{worker_url}/generate",
                json={
                    "prompt": prompt,
                    "max_tokens": 300
                },
                timeout=60
            )
            
            elapsed = time.time() - start_time
            self.worker_stats[worker_url]['total_time'] += elapsed
            
            if response.status_code == 200:
                data = response.json()
                content = data['response']
                
                # 提取代码
                if "def sort_array" in content:
                    code_start = content.find("def sort_array")
                    code_end = content.find("```", code_start)
                    if code_end != -1:
                        new_code = content[code_start:code_end].strip()
                    else:
                        new_code = content[code_start:].strip()
                    
                    # 清理代码
                    lines = new_code.split('\n')
                    clean_lines = []
                    for line in lines:
                        if line.strip():
                            clean_lines.append(line)
                        if 'return' in line:
                            clean_lines.append(line)
                            break
                    
                    return '\n'.join(clean_lines) if clean_lines else code
                else:
                    return code
            else:
                self.worker_stats[worker_url]['failures'] += 1
                return code
                
        except Exception as e:
            self.worker_stats[worker_url]['failures'] += 1
            print(f"⚠️ Worker {worker_url} 失败: {e}")
            # 从可用列表中移除失败的 Worker
            if worker_url in self.available_workers:
                self.available_workers.remove(worker_url)
            return code
    
    def mutate_code_parallel(self, code, generation, best_time, num_variants=3):
        """并行生成多个代码变体（自动适应 Worker 数量）"""
        # 动态调整并发数
        actual_variants = min(num_variants, len(self.available_workers))
        if actual_variants == 0:
            return []
        
        variants = []
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=actual_variants) as executor:
            # 为每个变体分配一个 Worker
            futures = []
            for i in range(actual_variants):
                worker_url = self.available_workers[i % len(self.available_workers)]
                future = executor.submit(
                    self.mutate_code, code, generation, best_time, worker_url
                )
                futures.append(future)
            
            # 收集结果
            for future in concurrent.futures.as_completed(futures):
                try:
                    variant = future.result()
                    variants.append(variant)
                except Exception as e:
                    print(f"⚠️ 变体生成失败: {e}")
        
        return variants
    
    def print_stats(self):
        """打印 Worker 统计信息"""
        print("\n📊 Worker 统计:")
        print("=" * 60)
        for url, stats in self.worker_stats.items():
            status = "✅" if url in self.available_workers else "❌"
            avg_time = stats['total_time'] / stats['requests'] if stats['requests'] > 0 else 0
            print(f"{status} {url}")
            print(f"   请求: {stats['requests']}, 失败: {stats['failures']}, 平均耗时: {avg_time:.1f}s")
        print("=" * 60)
    
    def evolve(self, initial_code, generations=50, variants_per_gen=3):
        """核心进化循环（支持动态扩展）"""
        from code_executor import CodeExecutor
        
        executor = CodeExecutor()
        current_best = initial_code
        best_time = executor.measure_time(current_best)
        
        history = [{
            "generation": 0,
            "code": current_best,
            "time": best_time,
            "timestamp": datetime.now().isoformat(),
            "workers": len(self.available_workers)
        }]
        
        print(f"\n🧬 开始进化（自动扩展模式）...")
        print(f"初始 Worker 数量: {len(self.available_workers)}")
        print(f"Generation 0: {best_time:.6f}s (baseline)")
        
        for gen in range(1, generations + 1):
            # 检查 Worker 数量变化
            current_workers = len(self.available_workers)
            
            print(f"\n🔄 Generation {gen} (Worker: {current_workers})")
            
            if current_workers == 0:
                print("⚠️ 没有可用的 Worker，等待...")
                time.sleep(5)
                self.check_workers(silent=True)
                continue
            
            # 并行生成多个变体（自动适应 Worker 数量）
            variants = self.mutate_code_parallel(
                current_best, gen, best_time, variants_per_gen
            )
            
            if not variants:
                print("⚠️ 未能生成变体，继续...")
                continue
            
            # 测试所有变体
            improved = False
            for idx, variant in enumerate(variants):
                variant_time = executor.measure_time(variant)
                
                if variant_time is None:
                    print(f"  变体 {idx+1}: ❌ 执行失败")
                    continue
                
                print(f"  变体 {idx+1}: {variant_time:.6f}s", end=" ")
                
                if variant_time < best_time:
                    improvement = ((best_time - variant_time) / best_time) * 100
                    print(f"✅ 提升 {improvement:.1f}%")
                    
                    current_best = variant
                    best_time = variant_time
                    improved = True
                    
                    history.append({
                        "generation": gen,
                        "code": current_best,
                        "time": best_time,
                        "improvement": improvement,
                        "timestamp": datetime.now().isoformat(),
                        "workers": current_workers
                    })
                    
                    break
                else:
                    print("⚪ 无改进")
            
            if not improved:
                print(f"Generation {gen}: 无改进，继续探索...")
            else:
                print(f"✨ Generation {gen}: {best_time:.6f}s (best so far)")
        
        print(f"\n🎉 进化完成！")
        print(f"最终性能: {best_time:.6f}s")
        print(f"总提升: {((history[0]['time'] - best_time) / history[0]['time'] * 100):.1f}%")
        
        # 打印统计信息
        self.print_stats()
        
        with open('evolution_history.json', 'w') as f:
            json.dump(history, f, indent=2, ensure_ascii=False)
        
        return {
            "final_code": current_best,
            "final_time": best_time,
            "history": history
        }


if __name__ == "__main__":
    import sys
    
    print("🧬 代码进化实验室 - 自动扩展模式")
    print("=" * 50)
    
    # 配置初始 Worker 地址
    worker_urls = [
        "http://localhost:8001",  # 本地 Worker
        # 可以在运行时动态添加更多 Worker
    ]
    
    engine = EvolutionEngineAutoScale(worker_urls, auto_discover=True)
    
    if not engine.available:
        print("\n❌ 没有可用的 Worker")
        print("\n💡 启动 Worker:")
        print("   python mlx_worker.py")
        print("\n💡 或在运行过程中启动 Worker，系统会自动检测")
        sys.exit(1)
    
    print("\n💡 提示：")
    print("  - 可以在进化过程中启动新的 Worker")
    print("  - 系统会自动检测并使用新 Worker")
    print("  - Worker 失败会自动移除")
    print("")
    
    # 测试进化
    initial_code = """def sort_array(arr):
    n = len(arr)
    for i in range(n):
        for j in range(n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    return arr"""
    
    result = engine.evolve(initial_code, generations=10, variants_per_gen=3)
    
    print("\n" + "="*50)
    print("最终代码：")
    print("="*50)
    print(result['final_code'])