Comet ML：实验管理与追踪

"""Comet ML实验管理与可复现研究工具。 本模块演示Comet ML的实验追踪、Optimizer超参搜索、模型Registry管理和团队协作功能，展示可复现研究的完整流程。 依赖: comet_ml >= 3.3 (模拟实现，无需实际安装)""" import osimport jsonimport randomfrom dataclasses import dataclass, fieldfrom typing import Dict, List, Optional, Any  @dataclassclass Experiment:    """Comet ML实验数据结构。"""    experiment_key: str    name: str    parameters: Dict[str, Any] = field(default_factory=dict)    metrics: Dict[str, float] = field(default_factory=dict)    code_version: str = ""    environment: Dict[str, str] = field(default_factory=dict)    status: str = "running"  class CometMLTracker:    """Comet ML实验追踪器，模拟核心API功能。"""     def __init__(self, workspace: str, project_name: str):        self.workspace = workspace        self.project_name = project_name        self.experiments: List[Experiment] = []        self._exp_counter = 0     def start_experiment(self, name: str = "") -> Experiment:        """创建并启动新实验，自动记录环境快照。"""        self._exp_counter += 1        exp_key = f"exp_{self._exp_counter:04d}"        experiment = Experiment(            experiment_key=exp_key,            name=name or f"experiment_{self._exp_counter}",            code_version=self._get_code_version(),            environment=self._capture_environment()        )        self.experiments.append(experiment)        print(f"[Comet] Started experiment: {exp_key} ({experiment.name})")        return experiment     def log_parameters(self, exp: Experiment, params: Dict):        """记录超参数，支持嵌套字典和列表。"""        for key, value in params.items():            exp.parameters[key] = value        print(f"[Comet] Logged {len(params)} parameters to {exp.experiment_key}")     def log_metric(self, exp: Experiment, name: str, value: float, step: int = 0):        """记录指标值，支持按step追踪训练曲线。"""        metric_key = f"{name}_step_{step}"        exp.metrics[metric_key] = value        if step == 0:            exp.metrics[name] = value     def log_metrics(self, exp: Experiment, metrics: Dict, step: int = 0):        """批量记录指标。"""        for name, value in metrics.items():            self.log_metric(exp, name, value, step)     def end_experiment(self, exp: Experiment):        """结束实验并标记状态。"""        exp.status = "finished"        print(f"[Comet] Experiment {exp.experiment_key} finished. "              f"Final metrics: {exp.metrics.get('accuracy', 'N/A')}")     def _get_code_version(self) -> str:        """获取当前代码版本（模拟Git commit hash）。"""        return f"commit_{random.randint(10000, 99999)}"     def _capture_environment(self) -> Dict[str, str]:        """捕获执行环境快照确保可复现性。"""        return {            "python_version": "3.11.5",            "framework": "PyTorch 2.1",            "cuda_version": "12.1",            "gpu": "NVIDIA A100",            "os": "Linux 5.15.0"        }     def compare_experiments(self, exp_keys: List[str]) -> str:        """对比多个实验的指标差异。"""        results = []        for key in exp_keys:            exp = next((e for e in self.experiments if e.experiment_key == key), None)            if exp:                results.append(f"{key}: acc={exp.metrics.get('accuracy', 'N/A')}")        return " | ".join(results)  class Optimizer:    """Comet ML Optimizer超参搜索模块。"""     def __init__(self, algorithm: str = "bayes"):        self.algorithm = algorithm        self.trials: List[Dict] = []     def search(self, config: Dict, train_fn, max_trials: int = 50):        """执行超参搜索，支持Bayesian/Grid/Random策略。"""        print(f"[Optimizer] Starting {self.algorithm} search "              f"with {max_trials} trials")        best_score = 0.0        best_params = None        for trial in range(max_trials):            params = self._sample_parameters(config, trial)            score = train_fn(params)            self.trials.append({"trial": trial, "params": params, "score": score})            if score > best_score:                best_score = score                best_params = params                print(f"[Optimizer] Trial {trial}: score={score:.4f} (NEW BEST)")        print(f"[Optimizer] Best score: {best_score:.4f}")        return best_params, best_score     def _sample_parameters(self, config: Dict, trial: int) -> Dict:        """根据搜索策略采样超参数。"""        params = {}        for name, spec in config.items():            if spec["type"] == "float":                low, high = spec["min"], spec["max"]                params[name] = random.uniform(low, high)            elif spec["type"] == "int":                low, high = spec["min"], spec["max"]                params[name] = random.randint(low, high)            elif spec["type"] == "categorical":                params[name] = random.choice(spec["values"])        return params  class ModelRegistry:    """Comet ML模型Registry，管理模型全生命周期。"""     def __init__(self):        self.models: Dict[str, Dict] = {}     def register_model(self, name: str, version: str,                       experiment_key: str, stage: str = "staging"):        """注册模型版本并关联训练实验。"""        if name not in self.models:            self.models[name] = {}        self.models[name][version] = {            "experiment_key": experiment_key,            "stage": stage,            "artifact_path": f"models/{name}/{version}"        }        print(f"[Registry] Registered {name} v{version} "              f"(stage: {stage}, exp: {experiment_key})")     def transition_stage(self, name: str, version: str, new_stage: str):        """转换模型阶段标记（staging -> production）。"""        if name in self.models and version in self.models[name]:            old_stage = self.models[name][version]["stage"]            self.models[name][version]["stage"] = new_stage            print(f"[Registry] {name} v{version}: {old_stage} -> {new_stage}")     def get_model_info(self, name: str, version: str) -> Dict:        """查询模型信息，返回关联的实验和阶段。"""        return self.models.get(name, {}).get(version, {})  def simulate_training(params: Dict) -> float:    """模拟训练函数，返回模型准确率。"""    lr = params.get("learning_rate", 0.001)    batch_size = params.get("batch_size", 32)    base_acc = 0.85    lr_bonus = min(lr * 50, 0.08)    bs_bonus = min(batch_size / 512, 0.05)    noise = random.uniform(-0.02, 0.02)    return min(base_acc + lr_bonus + bs_bonus + noise, 0.99)  if __name__ == "__main__":    # 1. 初始化Comet ML追踪器    tracker = CometMLTracker(workspace="ai-research-team", project_name="llm-finetuning")     # 2. 创建实验并记录参数    exp = tracker.start_experiment(name="bert-finetune-v1")    tracker.log_parameters(exp, {        "learning_rate": 0.0001,        "batch_size": 32,        "epochs": 10,        "model_name": "bert-base-chinese"    })     # 3. 模拟训练过程并记录指标    for epoch in range(10):        accuracy = 0.80 + epoch * 0.015 + random.uniform(-0.01, 0.01)        loss = 0.50 - epoch * 0.04 + random.uniform(-0.02, 0.02)        tracker.log_metrics(exp, {"accuracy": accuracy, "loss": loss}, step=epoch)     tracker.end_experiment(exp)     # 4. 使用Optimizer搜索最优超参数    optimizer = Optimizer(algorithm="bayes")    search_config = {        "learning_rate": {"type": "float", "min": 0.0001, "max": 0.01},        "batch_size": {"type": "categorical", "values": [16, 32, 64, 128]}    }    best_params, best_score = optimizer.search(        search_config, simulate_training, max_trials=20    )     # 5. 注册最优模型到Registry    registry = ModelRegistry()    registry.register_model("bert-classifier", "1.0.0", exp.experiment_key, "staging")    registry.transition_stage("bert-classifier", "1.0.0", "production")     # 6. 查询模型信息验证可追溯性    model_info = registry.get_model_info("bert-classifier", "1.0.0")    print(f"\nModel info: {json.dumps(model_info, indent=2)}")    print(f"Total experiments: {len(tracker.experiments)}")    print(f"Optimizer trials: {len(optimizer.trials)}")