AI应用开发中RDMA网络的迁移策略与风险控制

rdma_migration_demo/
├── config/
│   └── default.yaml        # 全局配置文件
├── core/
│   ├── __init__.py
│   ├── detector.py         # 环境探测器
│   ├── backend.py          # 网络后端抽象与实现
│   ├── monitor.py          # 性能监控器
│   ├── circuit_breaker.py  # 熔断器
│   └── controller.py       # 迁移控制器
├── scripts/
│   ├── start_server.py     # 启动服务器节点
│   └── start_worker.py     # 启动工作节点
├── tests/
│   └── test_compatibility.py # 兼容性测试
├── utils/
│   ├── __init__.py
│   └── logger.py           # 日志配置
├── requirements.txt        # Python依赖
├── run_migration.py        # 主入口：执行完整迁移演示
└── README.md               # 项目说明（此处仅为示意，输出时将省略）

# RDMA迁移演示配置
migration:
  strategy: "progressive" # progressive(渐进) / big_bang(一次性)
  batch_size: 1 # 渐进式迁移中每批迁移的节点数
  health_check_interval: 5 # 健康检查间隔(秒)

network:
  tcp:
    host: "0.0.0.0"
    base_port: 23456
  rdma:
    enabled: true
    provider: "verbs" # verbs / sdpa
    gid_index: 0 # 用于RoCE的GID索引
  fallback_enabled: true # 是否启用降级回退

monitoring:
  metrics_window: 10 # 指标收集时间窗口(秒)
  latency_threshold_ms: 50.0 # 延迟阈值，超过可能触发告警
  error_rate_threshold: 0.05 # 错误率阈值(5%)

circuit_breaker:
  failure_threshold: 10 # 连续失败次数阈值
  recovery_timeout: 30 # 熔断后恢复尝试间隔(秒)
  half_open_max_tests: 3 # 半开状态最大试探次数

logging:
  level: "INFO"
  file: "logs/migration.log"

import subprocess
import re
import logging
from typing import Dict, Optional, Tuple

class RDMAEnvDetector:
    """RDMA环境探测器"""
    
    def __init__(self):
        self.logger = logging.getLogger(self.__class__.__name__)
    
    def detect_all(self) -> Dict:
        """执行所有环境检测"""
        report = {
            "ibv_devices": self._detect_ibv_devices(),
            "libibverbs": self._check_libibverbs(),
            "driver_version": self._get_driver_version(),
            "network_interfaces": self._get_rdma_interfaces(),
            "overall_status": "UNKNOWN"
        }
        
        # 判断整体状态
        if report["ibv_devices"]["present"] and report["libibverbs"]["available"]:
            report["overall_status"] = "READY"
        elif not report["ibv_devices"]["present"]:
            report["overall_status"] = "NO_HARDWARE"
        else:
            report["overall_status"] = "LIB_MISSING"
            
        return report
    
    def _detect_ibv_devices(self) -> Dict:
        """检测IB Verbs设备"""
        try:
            # 尝试执行 `ibv_devices` 命令
            result = subprocess.run(['ibv_devices'], 
                                    capture_output=True, 
                                    text=True, 
                                    timeout=5)
            devices = []
            if result.returncode == 0:
                lines = result.stdout.strip().split('\n')
                # 跳过标题行
                for line in lines[1:]:
                    if line.strip():
                        parts = line.split()
                        if len(parts) >= 2:
                            devices.append({"device": parts[0], "node_guid": parts[1]})
            
            return {"present": len(devices) > 0, "devices": devices, "command_output": result.stdout}
        except (FileNotFoundError, subprocess.TimeoutExpired) as e:
            self.logger.warning(f"IB Verbs device detection failed: {e}")
            return {"present": False, "devices": [], "command_output": str(e)}
    
    def _check_libibverbs(self) -> Dict:
        """检查libibverbs库是否可用"""
        try:
            # 尝试导入pyverbs（Python绑定），如果失败则尝试ctypes加载
            try:
                from pyverbs import device
                return {"available": True, "method": "pyverbs", "version": "N/A"}
            except ImportError:
                import ctypes
                # 尝试加载libibverbs.so
                try:
                    lib = ctypes.CDLL('libibverbs.so')
                    return {"available": True, "method": "ctypes", "version": "N/A"}
                except OSError:
                    return {"available": False, "method": "ctypes", "error": "Library not found"}
        except Exception as e:
            return {"available": False, "method": "unknown", "error": str(e)}
    
    def _get_driver_version(self) -> Optional[str]:
        """获取驱动版本（简化示例，实际可能从文件读取）"""
        try:
            # 示例：读取Mellanox驱动版本
            with open('/sys/class/infiniband/mlx5_0/version', 'r') as f:
                return f.read().strip()
        except:
            return None
    
    def _get_rdma_interfaces(self) -> list:
        """获取配置了RDMA的网络接口（如RoCE）"""
        interfaces = []
        try:
            # 使用 ip 命令查找具有‘rdma'链接类型的设备
            result = subprocess.run(['ip', '-o', 'link', 'show', 'type', 'rdma'], 
                                    capture_output=True, 
                                    text=True, 
                                    timeout=5)
            for line in result.stdout.strip().split('\n'):
                if line:
                    # 解析输出，例如：`1: mlx5_0: <...> ...`
                    match = re.search(r'^\d+: (\w+):', line)
                    if match:
                        interfaces.append(match.group(1))
        except Exception as e:
            self.logger.debug(f"Failed to get RDMA interfaces: {e}")
        return interfaces

# 简化导出
detector = RDMAEnvDetector()

import socket
import pickle
import threading
import time
from abc import ABC, abstractmethod
from enum import Enum
import logging
import numpy as np
from typing import Any, Optional

# 尝试导入RDMA相关库，失败时设为None
try:
    from pyverbs import device, qp, cq, pd, mr
    import pyverbs.enums as e
    RDMA_AVAILABLE = True
except ImportError:
    RDMA_AVAILABLE = False
    # 创建虚拟类型以避免类型检查错误（仅当RDMA不可用时）
    class MockRDMA:
        class device: pass
        class qp: pass
        class cq: pass
        class pd: pass
        class mr: pass
    e = type('MockEnum', (), {})()

class BackendType(Enum):
    TCP = "tcp"
    RDMA = "rdma"

class NetworkBackend(ABC):
    """网络后端抽象基类"""
    
    def __init__(self, backend_type: BackendType):
        self.type = backend_type
        self.logger = logging.getLogger(f"{self.__class__.__name__}.{backend_type.value}")
        self.is_connected = False
        self.stats = {"tx_bytes": 0, "rx_bytes": 0, "errors": 0}
    
    @abstractmethod
    def connect(self, remote_addr: str, remote_port: int) -> bool:
        pass
    
    @abstractmethod
    def send(self, data: Any) -> bool:
        pass
    
    @abstractmethod
    def recv(self, timeout: float = None) -> Optional[Any]:
        pass
    
    @abstractmethod
    def close(self):
        pass
    
    def get_stats(self) -> dict:
        return self.stats.copy()

class TCPBackend(NetworkBackend):
    """TCP后端实现"""
    
    def __init__(self):
        super().__init__(BackendType.TCP)
        self.sock: Optional[socket.socket] = None
        self.lock = threading.Lock()
    
    def connect(self, remote_addr: str, remote_port: int) -> bool:
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.settimeout(10.0)
            self.sock.connect((remote_addr, remote_port))
            self.is_connected = True
            self.logger.info(f"TCP connected to {remote_addr}:{remote_port}")
            return True
        except Exception as e:
            self.logger.error(f"TCP connection failed: {e}")
            self.is_connected = False
            return False
    
    def send(self, data: Any) -> bool:
        if not self.is_connected or not self.sock:
            self.stats["errors"] += 1
            return False
        
        try:
            with self.lock:
                # 序列化数据并添加长度前缀
                serialized = pickle.dumps(data)
                length = len(serialized)
                header = length.to_bytes(4, byteorder='big')
                self.sock.sendall(header + serialized)
                self.stats["tx_bytes"] += length + 4
                return True
        except Exception as e:
            self.logger.error(f"TCP send failed: {e}")
            self.stats["errors"] += 1
            return False
    
    def recv(self, timeout: float = None) -> Optional[Any]:
        if not self.is_connected or not self.sock:
            return None
        
        old_timeout = self.sock.gettimeout()
        try:
            if timeout is not None:
                self.sock.settimeout(timeout)
            
            # 接收长度头
            header = self.sock.recv(4)
            if not header or len(header) != 4:
                return None
            length = int.from_bytes(header, byteorder='big')
            
            # 接收数据主体
            chunks = []
            bytes_received = 0
            while bytes_received < length:
                chunk = self.sock.recv(min(length - bytes_received, 4096))
                if not chunk:
                    break
                chunks.append(chunk)
                bytes_received += len(chunk)
            
            if bytes_received == length:
                data = pickle.loads(b''.join(chunks))
                self.stats["rx_bytes"] += bytes_received + 4
                return data
            else:
                self.logger.warning("Incomplete data received")
                return None
                
        except socket.timeout:
            self.logger.debug("TCP receive timeout")
            return None
        except Exception as e:
            self.logger.error(f"TCP receive failed: {e}")
            self.stats["errors"] += 1
            return None
        finally:
            self.sock.settimeout(old_timeout)
    
    def close(self):
        if self.sock:
            self.sock.close()
            self.sock = None
        self.is_connected = False
        self.logger.info("TCP backend closed.")

class RDMABackend(NetworkBackend):
    """RDMA后端简化实现 (基于Verbs)"""
    
    def __init__(self, gid_index: int = 0):
        super().__init__(BackendType.RDMA)
        if not RDMA_AVAILABLE:
            raise RuntimeError("RDMA libraries not available. Cannot initialize RDMABackend.")
        
        self.gid_index = gid_index
        self.context = None
        self.pd = None
        self.cq = None
        self.qp = None
        self.mr = None
        self.buffer = None
        self.buffer_size = 1024 * 1024  # 1MB 默认缓冲区
        self.remote_qp_num = None
        self.remote_lid = None
        self.remote_gid = None
        self.lock = threading.Lock()
    
    def connect(self, remote_addr: str, remote_port: int) -> bool:
        """简化连接流程：实际RDMA连接需要交换QP信息，这里仅为演示逻辑"""
        try:
            # 1. 获取设备列表并打开第一个设备
            dev_list = device.get_device_list()
            if not dev_list:
                self.logger.error("No RDMA devices found.")
                return False
            self.device = dev_list[0]
            self.context = self.device.open()
            
            # 2. 创建保护域、完成队列和队列对
            self.pd = pd.PD(self.context)
            self.cq = cq.CQ(self.context, cqe=10)  # 简化大小
            qp_init_attr = qp.QPInitAttr(cap=qp.QPCap(), qp_type=e.IBV_QPT_RC)
            self.qp = qp.QP(self.pd, qp_init_attr, e.IBV_QPT_RC)
            
            # 3. 分配并注册内存区域
            self.buffer = bytearray(self.buffer_size)
            self.mr = mr.MR(self.pd, self.buffer, 
                           access=e.IBV_ACCESS_LOCAL_WRITE | e.IBV_ACCESS_REMOTE_WRITE)
            
            # 4. 在实际应用中，这里需要通过TCP Socket交换QP信息（QP号，LID，GID等）
            # 此处模拟一个成功的连接建立
            self.logger.info(f"RDMA backend initialized (simulated connection to {remote_addr}).")
            self.is_connected = True
            return True
        except Exception as e:
            self.logger.error(f"RDMA backend initialization failed: {e}")
            self.is_connected = False
            return False
    
    def send(self, data: Any) -> bool:
        """简化的RDMA SEND操作"""
        if not self.is_connected:
            self.stats["errors"] += 1
            return False
        
        try:
            with self.lock:
                # 序列化数据
                serialized = pickle.dumps(data)
                length = len(serialized)
                
                if length > self.buffer_size:
                    self.logger.error(f"Data size {length} exceeds buffer {self.buffer_size}")
                    return False
                
                # 将数据复制到注册的内存区域
                self.buffer[:length] = serialized
                
                # **模拟**一次RDMA SEND操作。
                # 实际代码会构造SEND工作请求，提交到QP，并轮询CQ等待完成。
                # 此处用sleep模拟低延迟传输。
                time.sleep(0.001)  # 模拟1ms RDMA延迟
                
                self.stats["tx_bytes"] += length
                self.logger.debug(f"RDMA (simulated) sent {length} bytes.")
                return True
        except Exception as e:
            self.logger.error(f"RDMA (simulated) send failed: {e}")
            self.stats["errors"] += 1
            return False
    
    def recv(self, timeout: float = None) -> Optional[Any]:
        """简化的RDMA RECV操作"""
        # 在实际RDMA中，接收需要预置RECV WR。这里模拟接收固定数据。
        # 为演示目的，我们返回一个模拟的梯度张量。
        if not self.is_connected:
            return None
        time.sleep(0.001)  # 模拟接收延迟
        # 返回一个模拟的随机梯度
        sim_gradient = np.random.randn(100, 100).astype(np.float32)
        self.stats["rx_bytes"] += sim_gradient.nbytes
        return sim_gradient
    
    def close(self):
        # 按顺序清理RDMA资源
        resources = [self.qp, self.cq, self.mr, self.pd, self.context]
        for res in resources:
            if res:
                try:
                    res.close()
                except:
                    pass
        self.is_connected = False
        self.logger.info("RDMA backend closed.")

class BackendFactory:
    """后端工厂，根据配置和可用性创建合适的后端"""
    
    @staticmethod
    def create_backend(backend_type: BackendType, **kwargs) -> Optional[NetworkBackend]:
        try:
            if backend_type == BackendType.TCP:
                return TCPBackend()
            elif backend_type == BackendType.RDMA:
                if not RDMA_AVAILABLE:
                    raise RuntimeError("RDMA libraries not available.")
                gid_index = kwargs.get('gid_index', 0)
                return RDMABackend(gid_index=gid_index)
            else:
                raise ValueError(f"Unsupported backend type: {backend_type}")
        except Exception as e:
            logging.getLogger("BackendFactory").error(f"Failed to create backend {backend_type}: {e}")
            return None

import time
import threading
import psutil
from collections import deque
from dataclasses import dataclass
from typing import Deque, Dict, Optional
import logging

@dataclass
class NetworkMetric:
    timestamp: float
    backend_type: str
    latency_ms: float  # 单次操作延迟（示例）
    throughput_mbps: float  # 吞吐量
    error_count: int
    cpu_percent: float
    mem_percent: float

class PerformanceMonitor:
    """性能监控器，收集系统与网络指标"""
    
    def __init__(self, window_size: int = 100):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.window_size = window_size
        self.metrics: Deque[NetworkMetric] = deque(maxlen=window_size)
        self.lock = threading.Lock()
        self._running = False
        self._collector_thread: Optional[threading.Thread] = None
        
        # 用于计算增量
        self._last_net_io = psutil.net_io_counters()
        self._last_net_time = time.time()
    
    def start(self):
        """启动后台指标收集"""
        if self._running:
            return
        self._running = True
        self._collector_thread = threading.Thread(target=self._collect_loop, daemon=True)
        self._collector_thread.start()
        self.logger.info("Performance monitor started.")
    
    def stop(self):
        """停止指标收集"""
        self._running = False
        if self._collector_thread:
            self._collector_thread.join(timeout=2)
        self.logger.info("Performance monitor stopped.")
    
    def _collect_loop(self):
        """收集循环，定期采集系统级指标"""
        while self._running:
            try:
                # 收集系统指标
                cpu = psutil.cpu_percent(interval=0.1)
                mem = psutil.virtual_memory().percent
                
                # 计算网络吞吐量（全局，非RDMA特定）
                net_io = psutil.net_io_counters()
                current_time = time.time()
                time_diff = current_time - self._last_net_time
                if time_diff > 0:
                    bytes_sent = net_io.bytes_sent - self._last_net_io.bytes_sent
                    bytes_recv = net_io.bytes_recv - self._last_net_io.bytes_recv
                    throughput_mbps = (bytes_sent + bytes_recv) * 8 / time_diff / 1_000_000
                else:
                    throughput_mbps = 0.0
                
                self._last_net_io = net_io
                self._last_net_time = current_time
                
                # 创建一个通用的系统指标样本（不关联特定后端）
                metric = NetworkMetric(
                    timestamp=current_time,
                    backend_type="system",
                    latency_ms=0.0,
                    throughput_mbps=throughput_mbps,
                    error_count=0,
                    cpu_percent=cpu,
                    mem_percent=mem
                )
                
                with self.lock:
                    self.metrics.append(metric)
                    
            except Exception as e:
                self.logger.error(f"Error in metric collection loop: {e}")
            time.sleep(5)  # 每5秒收集一次系统指标
    
    def record_operation(self, backend_type: str, latency_ms: float, error_occurred: bool):
        """记录一次网络操作的详细指标"""
        try:
            cpu = psutil.cpu_percent(interval=None)
            mem = psutil.virtual_memory().percent
            
            metric = NetworkMetric(
                timestamp=time.time(),
                backend_type=backend_type,
                latency_ms=latency_ms,
                throughput_mbps=0.0,  # 单次操作不计算吞吐
                error_count=1 if error_occurred else 0,
                cpu_percent=cpu,
                mem_percent=mem
            )
            with self.lock:
                self.metrics.append(metric)
        except Exception as e:
            self.logger.error(f"Failed to record operation metric: {e}")
    
    def get_recent_metrics(self, backend_type: Optional[str] = None, limit: int = 20) -> list:
        """获取最近收集的指标"""
        with self.lock:
            metrics_list = list(self.metrics)
        
        if backend_type:
            filtered = [m for m in metrics_list if m.backend_type == backend_type]
        else:
            filtered = metrics_list
        
        return filtered[-limit:]
    
    def calculate_statistics(self, backend_type: str, window_seconds: float = 30) -> Dict:
        """计算指定时间窗口内的统计信息"""
        now = time.time()
        window_start = now - window_seconds
        
        with self.lock:
            relevant_metrics = [
                m for m in self.metrics 
                if m.backend_type == backend_type and m.timestamp >= window_start
            ]
        
        if not relevant_metrics:
            return {}
        
        latencies = [m.latency_ms for m in relevant_metrics if m.latency_ms > 0]
        errors = sum(m.error_count for m in relevant_metrics)
        total_ops = len(relevant_metrics)
        
        stats = {
            "sample_count": total_ops,
            "avg_latency_ms": sum(latencies) / len(latencies) if latencies else 0,
            "p95_latency_ms": sorted(latencies)[int(len(latencies) * 0.95)] if latencies else 0,
            "error_rate": errors / total_ops if total_ops > 0 else 0,
            "time_window_s": window_seconds
        }
        return stats

import time
from enum import Enum
import logging
from typing import Callable, Optional

class CircuitState(Enum):
    CLOSED = "CLOSED"   # 正常状态，请求通过
    OPEN = "OPEN"       # 熔断状态，请求快速失败
    HALF_OPEN = "HALF_OPEN" # 半开状态，试探性允许部分请求通过

class CircuitBreaker:
    """
    熔断器模式实现。
    当失败次数超过阈值，熔断电路，经过恢复期后进入半开状态试探。
    """
    
    def __init__(self, 
                 name: str,
                 failure_threshold: int = 5,
                 recovery_timeout: float = 30.0,
                 half_open_max_success: int = 3):
        self.name = name
        self.logger = logging.getLogger(f"CircuitBreaker.{name}")
        
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.half_open_max_success = half_open_max_success
        
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_failure_time: Optional[float] = None
        self.last_state_change_time = time.time()
        
    def call(self, func: Callable, *args, **kwargs):
        """
        保护一个函数调用。根据熔断器状态决定是执行、快速失败还是试探。
        """
        # 检查是否需要从OPEN状态恢复
        if self.state == CircuitState.OPEN:
            if self._open_expired():
                self._enter_half_open()
            else:
                raise CircuitBreakerOpenError(f"Circuit '{self.name}' is OPEN. Fast fail.")
        
        # 执行被保护的调用
        try:
            result = func(*args, **kwargs)
            self._on_success()
            return result
        except Exception as e:
            self._on_failure()
            raise e
    
    def _on_success(self):
        """调用成功时的处理"""
        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            self.logger.debug(f"Half-open success {self.success_count}/{self.half_open_max_success}")
            if self.success_count >= self.half_open_max_success:
                self._enter_closed()
        else: # CLOSED state
            self.failure_count = 0  # 连续成功则重置失败计数
    
    def _on_failure(self):
        """调用失败时的处理"""
        self.failure_count += 1
        self.last_failure_time = time.time()
        self.logger.warning(f"Failure recorded. Count: {self.failure_count}/{self.failure_threshold}")
        
        if self.state == CircuitState.HALF_OPEN:
            # 半开状态下任何失败都立刻重新熔断
            self._enter_open()
        elif self.state == CircuitState.CLOSED:
            if self.failure_count >= self.failure_threshold:
                self._enter_open()
    
    def _enter_closed(self):
        """进入闭合状态"""
        old_state = self.state
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.success_count = 0
        self.last_state_change_time = time.time()
        self.logger.info(f"State changed from {old_state} to {self.state}")
    
    def _enter_open(self):
        """进入熔断状态"""
        old_state = self.state
        self.state = CircuitState.OPEN
        self.last_state_change_time = time.time()
        self.logger.warning(f"State changed from {old_state} to {self.state}. Circuit tripped!")
    
    def _enter_half_open(self):
        """进入半开状态"""
        old_state = self.state
        self.state = CircuitState.HALF_OPEN
        self.success_count = 0
        self.last_state_change_time = time.time()
        self.logger.info(f"State changed from {old_state} to {self.state}. Testing recovery.")
    
    def _open_expired(self) -> bool:
        """检查OPEN状态是否已过恢复超时时间"""
        if self.state != CircuitState.OPEN:
            return False
        open_duration = time.time() - self.last_state_change_time
        return open_duration >= self.recovery_timeout
    
    def get_status(self) -> dict:
        """获取当前状态信息"""
        return {
            "name": self.name,
            "state": self.state.value,
            "failure_count": self.failure_count,
            "success_count": self.success_count,
            "last_state_change": self.last_state_change_time,
            "is_open": self.state == CircuitState.OPEN
        }

class CircuitBreakerOpenError(Exception):
    """熔断器打开时抛出的异常"""
    pass

import time
import yaml
from typing import Dict, List, Optional
import logging
from enum import Enum

from .detector import detector
from .backend import BackendFactory, BackendType, NetworkBackend
from .monitor import PerformanceMonitor
from .circuit_breaker import CircuitBreaker

class MigrationStage(Enum):
    INITIAL = "initial"          # 初始状态，TCP运行
    DETECTION = "detection"      # 环境探测
    VALIDATION = "validation"    # 验证RDMA连通性
    DUAL_STACK = "dual_stack"    # 双栈运行（TCP+RDMA）
    RDMA_ONLY = "rdma_only"      # 仅RDMA运行
    ROLLBACK = "rollback"        # 回滚至TCP

class MigrationController:
    """迁移过程的总控制器"""
    
    def __init__(self, config_path: str = "config/default.yaml"):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.load_config(config_path)
        
        self.stage = MigrationStage.INITIAL
        self.backend_primary: Optional[NetworkBackend] = None  # 当前主用后端
        self.backend_secondary: Optional[NetworkBackend] = None # 备用后端（用于回滚）
        self.monitor = PerformanceMonitor()
        self.circuit_breaker = CircuitBreaker(
            name="RDMA_Migration",
            failure_threshold=self.config['circuit_breaker']['failure_threshold'],
            recovery_timeout=self.config['circuit_breaker']['recovery_timeout']
        )
        
        # 迁移节点列表（模拟）
        self.nodes_to_migrate = [f"node-{i}" for i in range(3)]
        self.migrated_nodes: List[str] = []
        
        self._should_stop = False
    
    def load_config(self, config_path: str):
        """加载配置文件"""
        try:
            with open(config_path, 'r') as f:
                self.config = yaml.safe_load(f)
            self.logger.info(f"Config loaded from {config_path}")
        except Exception as e:
            self.logger.error(f"Failed to load config: {e}")
            # 提供默认配置
            self.config = {
                'migration': {'strategy': 'progressive', 'batch_size': 1},
                'network': {'rdma': {'enabled': True}, 'fallback_enabled': True},
                'circuit_breaker': {'failure_threshold': 10, 'recovery_timeout': 30}
            }
    
    def start_migration(self):
        """启动完整的迁移流程"""
        self.logger.info("=== Starting RDMA Migration Process ===")
        self.monitor.start()
        
        try:
            self.run_stage_detection()
            if self._should_stop:
                return
                
            self.run_stage_validation()
            if self._should_stop:
                return
                
            # 根据策略选择渐进式或一次性迁移
            if self.config['migration']['strategy'] == 'progressive':
                self.run_stage_progressive_migration()
            else:
                self.run_stage_big_bang_migration()
                
        except Exception as e:
            self.logger.critical(f"Migration process failed critically: {e}")
            self.trigger_rollback("Critical failure")
        finally:
            self.monitor.stop()
            self.cleanup()
    
    def run_stage_detection(self):
        """执行阶段：环境探测"""
        self.set_stage(MigrationStage.DETECTION)
        self.logger.info("--- Stage: Environment Detection ---")
        
        report = detector.detect_all()
        self.logger.info(f"Detection Report: {report['overall_status']}")
        
        if report['overall_status'] != 'READY':
            self.logger.error(f"RDMA environment not ready: {report['overall_status']}. Aborting migration.")
            self._should_stop = True
            return
        
        self.logger.info("RDMA environment check passed.")
        # 初始化主用后端为TCP（初始状态）
        self.backend_primary = BackendFactory.create_backend(BackendType.TCP)
        if not self.backend_primary:
            self.logger.error("Failed to initialize primary TCP backend. Abort.")
            self._should_stop = True
    
    def run_stage_validation(self):
        """执行阶段：RDMA连通性验证"""
        self.set_stage(MigrationStage.VALIDATION)
        self.logger.info("--- Stage: RDMA Connectivity Validation ---")
        
        # 在实际应用中，这里会与目标节点建立测试连接。
        # 此处简化：尝试创建RDMA后端并"连接"。
        test_rdma_backend = BackendFactory.create_backend(
            BackendType.RDMA, 
            gid_index=self.config['network']['rdma'].get('gid_index', 0)
        )
        
        if not test_rdma_backend:
            self.logger.error("Failed to create test RDMA backend. Aborting migration.")
            self._should_stop = True
            return
        
        # 模拟连接测试（假设连接成功）
        test_success = test_rdma_backend.connect("127.0.0.1", 9999)
        test_rdma_backend.close()
        
        if test_success:
            self.logger.info("RDMA connectivity validation passed.")
        else:
            self.logger.error("RDMA connectivity validation failed.")
            self._should_stop = True
    
    def run_stage_progressive_migration(self):
        """执行阶段：渐进式迁移"""
        self.logger.info("--- Stage: Progressive Migration ---")
        batch_size = self.config['migration']['batch_size']
        
        for i, node in enumerate(self.nodes_to_migrate):
            if self._should_stop:
                break
                
            self.logger.info(f"Migrating batch starting with node: {node} (Batch {i+1}/{(len(self.nodes_to_migrate)+batch_size-1)//batch_size})")
            
            # 1. 为目标批次节点切换到双栈模式
            self.set_stage(MigrationStage.DUAL_STACK)
            self.logger.info(f"Entering dual-stack mode for node {node}...")
            
            # 创建RDMA后端作为新的主用，TCP作为备用
            rdma_backend = BackendFactory.create_backend(
                BackendType.RDMA,
                gid_index=self.config['network']['rdma'].get('gid_index', 0)
            )
            if not rdma_backend or not rdma_backend.connect("simulated_remote", 20000):
                self.logger.error(f"Failed to establish RDMA backend for {node}. Skipping batch.")
                continue
            
            # 切换：原主用TCP变为备用，RDMA变为主用
            self.backend_secondary = self.backend_primary  # TCP降级为备用
            self.backend_primary = rdma_backend
            
            # 2. 监控运行一段时间
            self._run_validation_workload(node)
            
            # 3. 检查监控指标，决定是否继续
            stats = self.monitor.calculate_statistics(backend_type="rdma", window_seconds=15)
            if stats.get('error_rate', 1.0) > self.config['monitoring']['error_rate_threshold']:
                self.logger.warning(f"High error rate detected for {node}. Triggering rollback for this batch.")
                # 回滚这个批次：主用切回TCP
                self.backend_primary.close()
                self.backend_primary = self.backend_secondary
                self.backend_secondary = None
            else:
                # 迁移成功，关闭备用TCP连接
                self.logger.info(f"Migration successful for node {node}.")
                if self.backend_secondary:
                    self.backend_secondary.close()
                    self.backend_secondary = None
                self.migrated_nodes.append(node)
                self.set_stage(MigrationStage.RDMA_ONLY)
                
            time.sleep(2)  # 批次间隔
        
        self.logger.info(f"Progressive migration completed. Migrated nodes: {self.migrated_nodes}")
    
    def run_stage_big_bang_migration(self):
        """执行阶段：一次性迁移（高风险）"""
        self.logger.info("--- Stage: Big Bang Migration ---")
        self.set_stage(MigrationStage.DUAL_STACK) # 短暂进入双栈
        
        # 一次性为所有节点切换
        rdma_backend = BackendFactory.create_backend(
            BackendType.RDMA,
            gid_index=self.config['network']['rdma'].get('gid_index', 0)
        )
        if not rdma_backend:
            self.logger.error("Big bang migration failed: cannot create RDMA backend.")
            self.trigger_rollback("RDMA backend creation failed")
            return
        
        # 保留旧TCP后端作为备用
        old_tcp_backend = self.backend_primary
        self.backend_primary = rdma_backend
        self.backend_secondary = old_tcp_backend
        
        # 执行验证工作负载，用熔断器保护
        try:
            self.circuit_breaker.call(self._run_validation_workload, "all_nodes")
        except Exception as e:
            self.logger.error(f"Big bang migration failed during validation: {e}")
            self.trigger_rollback(f"Validation failed: {e}")
            return
        
        # 检查熔断器状态
        if self.circuit_breaker.get_status()['is_open']:
            self.logger.error("Circuit breaker opened during big bang migration.")
            self.trigger_rollback("Circuit breaker tripped")
        else:
            self.logger.info("Big bang migration successful.")
            self.set_stage(MigrationStage.RDMA_ONLY)
            if self.backend_secondary:
                self.backend_secondary.close()
                self.backend_secondary = None
    
    def _run_validation_workload(self, node_id: str):
        """运行验证工作负载（模拟梯度通信）"""
        self.logger.info(f"Running validation workload for {node_id}...")
        workloads = 5  # 发送5个模拟梯度
        for i in range(workloads):
            if self._should_stop:
                break
            start_time = time.time()
            
            # 模拟梯度数据
            import numpy as np
            simulated_gradient = np.random.randn(50, 50).astype(np.float32)
            error_occurred = False
            
            try:
                # 使用熔断器保护发送操作
                success = self.circuit_breaker.call(self.backend_primary.send, simulated_gradient)
                if not success:
                    error_occurred = True
                    self.logger.warning(f"Workload send failed for {node_id}, iter {i}")
            except Exception as e:
                error_occurred = True
                self.logger.error(f"Workload send raised exception: {e}")
            
            latency_ms = (time.time() - start_time) * 1000
            # 记录指标
            backend_type_str = self.backend_primary.type.value if self.backend_primary else "unknown"
            self.monitor.record_operation(backend_type_str, latency_ms, error_occurred)
            
            time.sleep(0.5)  # 模拟工作间隔
        self.logger.info(f"Validation workload for {node_id} finished.")
    
    def trigger_rollback(self, reason: str):
        """触发回滚流程"""
        self.logger.warning(f"!!! Triggering Rollback !!! Reason: {reason}")
        self.set_stage(MigrationStage.ROLLBACK)
        
        if self.backend_secondary:
            # 如果有备用TCP后端，切换回它
            self.logger.info("Switching back to TCP backend.")
            if self.backend_primary:
                self.backend_primary.close()
            self.backend_primary = self.backend_secondary
            self.backend_secondary = None
        elif self.backend_primary and self.backend_primary.type == BackendType.RDMA:
            # 没有备用，但主用是RDMA，需要新建一个TCP后端
            self.logger.info("Creating new TCP backend for rollback.")
            self.backend_primary.close()
            self.backend_primary = BackendFactory.create_backend(BackendType.TCP)
        
        self.logger.info("Rollback completed. Service should now be on TCP.")
        self._should_stop = True
    
    def set_stage(self, new_stage: MigrationStage):
        """更新迁移阶段"""
        old_stage = self.stage
        self.stage = new_stage
        self.logger.info(f"Migration stage changed: {old_stage.value} -> {new_stage.value}")
    
    def stop(self):
        """外部请求停止迁移"""
        self.logger.info("Migration stop requested.")
        self._should_stop = True
    
    def cleanup(self):
        """清理资源"""
        self.logger.info("Cleaning up resources...")
        for backend in [self.backend_primary, self.backend_secondary]:
            if backend:
                try:
                    backend.close()
                except:
                    pass
        self.logger.info("Cleanup done.")
    
    def get_status(self) -> Dict:
        """获取控制器当前状态"""
        return {
            "stage": self.stage.value,
            "primary_backend": self.backend_primary.type.value if self.backend_primary else None,
            "secondary_backend": self.backend_secondary.type.value if self.backend_secondary else None,
            "migrated_nodes": self.migrated_nodes,
            "circuit_breaker": self.circuit_breaker.get_status()
        }

#!/usr/bin/env python3
"""
RDMA迁移演示主入口脚本。
"""
import sys
import time
import logging
from utils.logger import setup_logging
from core.controller import MigrationController

def main():
    # 设置日志
    setup_logging()
    logger = logging.getLogger("Main")
    
    logger.info("RDMA Migration Demo Starting...")
    
    # 初始化控制器
    controller = MigrationController("config/default.yaml")
    
    # 启动迁移流程（在主线程运行，方便捕获键盘中断）
    try:
        controller.start_migration()
    except KeyboardInterrupt:
        logger.info("\nMigration interrupted by user.")
        controller.stop()
    except Exception as e:
        logger.critical(f"Unhandled exception in main: {e}", exc_info=True)
        sys.exit(1)
    
    # 打印最终状态
    final_status = controller.get_status()
    logger.info("=== Migration Process Finished ===")
    logger.info(f"Final Stage: {final_status['stage']}")
    logger.info(f"Primary Backend: {final_status['primary_backend']}")
    logger.info(f"Migrated Nodes: {final_status['migrated_nodes']}")
    
    # 保持一段时间以便观察
    logger.info("Demo will exit in 5 seconds...")
    time.sleep(5)

if __name__ == "__main__":
    main()

import logging
import sys
import os
from logging.handlers import RotatingFileHandler

def setup_logging(log_file='logs/migration.log', level=logging.INFO):
    """配置日志记录"""
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    
    # 格式化器
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    
    # 控制台处理器
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)
    console_handler.setLevel(level)
    
    # 文件处理器（轮转）
    file_handler = RotatingFileHandler(
        log_file, maxBytes=10*1024*1024, backupCount=5
    )
    file_handler.setFormatter(formatter)
    file_handler.setLevel(level)
    
    # 获取根日志记录器并配置
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)  # 根记录器捕获所有级别
    root_logger.addHandler(console_handler)
    root_logger.addHandler(file_handler)
    
    # 抑制一些库的噪音日志
    logging.getLogger("urllib3").setLevel(logging.WARNING)
    logging.getLogger("pyverbs").setLevel(logging.WARNING)

# 核心依赖
numpy>=1.21.0
pyyaml>=6.0
psutil>=5.9.0

# RDMA支持 (可选，仅在有RDMA硬件的环境安装)
# pyverbs  # 通常通过系统包管理器安装，如 `apt-get install python3-pyverbs`
# 为演示，我们将其标记为可选

# 开发与测试
pytest>=7.0.0

# 克隆项目或创建上述文件结构后，进入项目根目录
cd rdma_migration_demo

# 安装Python依赖 (使用虚拟环境推荐)
python3 -m venv venv
source venv/bin/activate  # Linux/macOS
# venv\Scripts\activate   # Windows

pip install -r requirements.txt

# 确保在项目根目录下
python run_migration.py

# 运行简单的兼容性测试（检查环境探测逻辑）
python -m pytest tests/test_compatibility.py -v

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from core.detector import RDMAEnvDetector
from core.backend import BackendFactory, BackendType

def test_detector_smoke():
    """测试探测器基本功能不崩溃"""
    detector = RDMAEnvDetector()
    report = detector.detect_all()
    assert 'overall_status' in report
    print(f"Detection report: {report['overall_status']}")

def test_tcp_backend_creation():
    """测试TCP后端工厂"""
    backend = BackendFactory.create_backend(BackendType.TCP)
    assert backend is not None
    assert backend.type == BackendType.TCP
    print("TCP backend creation test passed.")

def test_rdma_backend_creation_graceful():
    """测试RDMA后端工厂在无库情况下的优雅降级"""
    # 这个测试取决于环境。我们主要检查工厂不会崩溃。
    try:
        backend = BackendFactory.create_backend(BackendType.RDMA)
        # 如果创建成功，则应有有效对象；如果失败，则返回None或抛出异常被捕获。
        if backend:
            assert backend.type == BackendType.RDMA
            print("RDMA backend creation test passed (lib available).")
        else:
            print("RDMA backend creation returned None (lib likely unavailable).")
    except RuntimeError as e:
        if "not available" in str(e):
            print(f"RDMA backend creation gracefully handled missing lib: {e}")
        else:
            raise

if __name__ == "__main__":
    test_detector_smoke()
    test_tcp_backend_creation()
    test_rdma_backend_creation_graceful()
    print("All compatibility tests completed.")