A/B 测试与模型迭代

from enum import Enum
from dataclasses import dataclass
from typing import Dict, Optional

class ExperimentStatus(Enum):
    DRAFT = "draft"
    RUNNING = "running"
    STOPPED = "stopped"
    CONCLUDED = "concluded"


@dataclass
class ExperimentConfig:
    """A/B 测试实验配置"""
    experiment_id: str
    name: str
    hypothesis: str
    metric_primary: str
    metrics_secondary: list[str]
    traffic_split: Dict[str, float]  # {"control": 0.5, "variant": 0.5}
    min_sample_size: int
    duration_days: int
    status: ExperimentStatus = ExperimentStatus.DRAFT

    def validate(self) -> bool:
        total = sum(self.traffic_split.values())
        return abs(total - 1.0) < 1e-6 and self.min_sample_size > 0

import hashlib
from typing import List


def consistent_hash(user_id: str, buckets: int = 100) -> int:
    """为同一用户始终返回相同的分桶结果"""
    return int(hashlib.md5(user_id.encode()).hexdigest(), 16) % buckets


def assign_variant(user_id: str, config: ExperimentConfig) -> Optional[str]:
    """根据流量分配配置为用户分配实验变体"""
    if config.status != ExperimentStatus.RUNNING:
        return None
    bucket = consistent_hash(user_id)
    cumulative = 0.0
    for variant, proportion in config.traffic_split.items():
        cumulative += proportion * 100
        if bucket < cumulative:
            return variant
    return None

import math
from scipy.stats import norm


def calculate_sample_size(
    baseline_rate: float,
    min_detectable_effect: float,
    alpha: float = 0.05,
    power: float = 0.80
) -> int:
    """计算每组所需的最小样本量（比例检验）"""
    z_alpha = norm.ppf(1 - alpha / 2)
    z_beta = norm.ppf(power)
    p1 = baseline_rate
    p2 = baseline_rate * (1 + min_detectable_effect)
    p_bar = (p1 + p2) / 2
    numerator = (z_alpha * math.sqrt(2 * p_bar * (1 - p_bar)) +
                 z_beta * math.sqrt(p1 * (1 - p1) + p2 * (1 - p2)))  2
    denominator = (p2 - p1)  2
    return int(math.ceil(numerator / denominator))

# 基线转化率 10%，期望检测 5% 相对提升
n = calculate_sample_size(0.10, 0.05)
print(f"Each group needs: {n} samples")

import pandas as pd
import numpy as np


def generate_experiment_data(
    n_control: int,
    n_variant: int,
    conversion_rate_control: float = 0.10,
    conversion_rate_variant: float = 0.108
) -> pd.DataFrame:
    """生成模拟实验数据"""
    np.random.seed(42)
    control = pd.DataFrame({
        "user_id": [f"u_{i}" for i in range(n_control)],
        "group": "control",
        "converted": np.random.binomial(1, conversion_rate_control, n_control)
    })
    variant = pd.DataFrame({
        "user_id": [f"u_{i + n_control}" for i in range(n_variant)],
        "group": "variant",
        "converted": np.random.binomial(1, conversion_rate_variant, n_variant)
    })
    return pd.concat([control, variant], ignore_index=True)

data = generate_experiment_data(50000, 50000)
print(data.groupby("group")["converted"].mean())

from scipy import stats
import numpy as np


def analyze_ab_test(
    conversions_control: int,
    n_control: int,
    conversions_variant: int,
    n_variant: int,
    alpha: float = 0.05
) -> dict:
    """频率派 A/B 测试分析"""
    p1 = conversions_control / n_control
    p2 = conversions_variant / n_variant
    lift = (p2 - p1) / p1 * 100

    z_stat, p_value = stats.proportions_ztest(
        [conversions_variant, conversions_control],
        [n_variant, n_control],
        alternative="two-sided"
    )

    # 计算效应量 (Cohen's h)
    h = 2 * np.arcsin(np.sqrt(p2)) - 2 * np.arcsin(np.sqrt(p1))

    # 计算 95% 置信区间
    se = np.sqrt(p1 * (1 - p1) / n_control + p2 * (1 - p2) / n_variant)
    ci_low = (p2 - p1) - 1.96 * se
    ci_high = (p2 - p1) + 1.96 * se

    return {
        "p_value": round(p_value, 6),
        "significant": p_value < alpha,
        "lift_pct": round(lift, 2),
        "effect_size_h": round(h, 4),
        "ci_95": [round(ci_low, 6), round(ci_high, 6)]
    }

result = analyze_ab_test(5000, 50000, 5400, 50000)
print(result)

import numpy as np
from scipy.stats import beta


def bayesian_ab_analysis(
    conversions_a: int, n_a: int,
    conversions_b: int, n_b: int,
    n_samples: int = 100000
) -> dict:
    """贝叶斯 A/B 测试分析"""
    # 使用 Beta 先验 (Jeffreys prior)
    posterior_a = beta.rvs(
        conversions_a + 0.5, n_a - conversions_a + 0.5,
        size=n_samples
    )
    posterior_b = beta.rvs(
        conversions_b + 0.5, n_b - conversions_b + 0.5,
        size=n_samples
    )

    prob_b_better = (posterior_b > posterior_a).mean()
    expected_lift = ((posterior_b - posterior_a) / posterior_a).mean() * 100
    ci_lift = np.percentile(
        (posterior_b - posterior_a) / posterior_a * 100,
        [2.5, 97.5]
    )

    return {
        "prob_b_better": round(prob_b_better, 4),
        "expected_lift_pct": round(expected_lift, 2),
        "credible_interval_95": [round(ci_lift[0], 2), round(ci_lift[1], 2)]
    }

result = bayesian_ab_analysis(5000, 50000, 5400, 50000)
print(result)

import numpy as np
from typing import List


class ThompsonSampling:
    """Thompson Sampling 多臂老虎机"""

    def __init__(self, n_arms: int, alpha_prior: float = 1.0,
                 beta_prior: float = 1.0):
        self.n_arms = n_arms
        self.alpha = np.full(n_arms, alpha_prior)  #  successes + prior
        self.beta = np.full(n_arms, beta_prior)     #  failures + prior

    def select_arm(self) -> int:
        """从 Beta 后验采样，选择值最大的 arm"""
        samples = np.random.beta(self.alpha, self.beta)
        return int(np.argmax(samples))

    def update(self, arm: int, reward: float):
        """更新指定 arm 的后验分布"""
        if reward > 0:
            self.alpha[arm] += 1
        else:
            self.beta[arm] += 1

    def get_estimates(self) -> np.ndarray:
        """返回各 arm 的期望估计"""
        return self.alpha / (self.alpha + self.beta)


# 模拟: 3 个版本,真实转化率分别为 10%, 12%, 8%
ts = ThompsonSampling(3)
true_rates = [0.10, 0.12, 0.08]
for _ in range(10000):
    arm = ts.select_arm()
    reward = float(np.random.random() < true_rates[arm])
    ts.update(arm, reward)
print(f"Estimates: {ts.get_estimates()}")

import numpy as np
from typing import Tuple


class EpsilonGreedy:
    """Epsilon-Greedy 多臂老虎机"""

    def __init__(self, n_arms: int, epsilon: float = 0.1,
                 decay: float = 0.999):
        self.n_arms = n_arms
        self.epsilon = epsilon
        self.decay = decay
        self.counts = np.zeros(n_arms)
        self.values = np.zeros(n_arms)

    def select_arm(self) -> int:
        if np.random.random() < self.epsilon:
            return int(np.random.randint(self.n_arms))
        # 打破平局: 随机选择最优的
        max_val = np.max(self.values)
        best_arms = np.where(self.values == max_val)[0]
        return int(np.random.choice(best_arms))

    def update(self, arm: int, reward: float):
        self.counts[arm] += 1
        n = self.counts[arm]
        self.values[arm] += (reward - self.values[arm]) / n
        self.epsilon *= self.decay

    def regret(self, true_rates: np.ndarray) -> float:
        """计算累积 regret"""
        best_rate = np.max(true_rates)
        total_reward = np.sum(self.counts * self.values)
        optimal_reward = np.sum(self.counts) * best_rate
        return optimal_reward - total_reward

from dataclasses import dataclass
from typing import Dict, List
import time


@dataclass
class CanaryStage:
    """灰度发布的一个阶段"""
    traffic_pct: float
    min_duration_hours: int
    success_criteria: Dict[str, float]


class CanaryDeployment:
    """管理灰度发布流程"""

    def __init__(self, model_name: str):
        self.model_name = model_name
        self.stages: List[CanaryStage] = [
            CanaryStage(1, 24, {"error_rate": 0.01, "latency_p99": 500}),
            CanaryStage(5, 24, {"error_rate": 0.01, "latency_p99": 500}),
            CanaryStage(25, 48, {"error_rate": 0.005, "latency_p99": 450}),
            CanaryStage(50, 48, {"error_rate": 0.005, "latency_p99": 450}),
            CanaryStage(100, 24, {"error_rate": 0.005, "latency_p99": 400}),
        ]
        self.current_stage = 0

    def check_stage_health(self, metrics: Dict[str, float]) -> bool:
        """检查当前阶段的健康指标"""
        criteria = self.stages[self.current_stage].success_criteria
        return all(
            metrics.get(k, float("inf")) <= v
            for k, v in criteria.items()
        )

    def promote(self) -> bool:
        """推进到下一个阶段"""
        if self.current_stage >= len(self.stages) - 1:
            return False
        self.current_stage += 1
        return True

# Istio VirtualService 灰度流量配置
apiVersion: networking.istio.io/v1beta3
kind: VirtualService
metadata:
  name: model-serving-route
spec:
  hosts:
    - model-api.example.com
  http:
    - route:
        - destination:
            host: model-api
            subset: stable
          weight: 95
        - destination:
            host: model-api
            subset: canary
          weight: 5
      timeout: 2s
      retries:
        attempts: 3
        perTryTimeout: 500ms
---
apiVersion: networking.istio.io/v1beta3
kind: DestinationRule
metadata:
  name: model-api-versions
spec:
  host: model-api.example.com
  subsets:
    - name: stable
      labels:
        version: v1.2.0
    - name: canary
      labels:
        version: v1.3.0

import time
from enum import Enum
from typing import Optional
import threading


class ModelState(Enum):
    ACTIVE = "active"
    ROLLED_BACK = "rolled_back"
    ROLLING_BACK = "rolling_back"


class ModelRollbackManager:
    """模型回滚管理器"""

    def __init__(self, model_name: str):
        self.model_name = model_name
        self.state = ModelState.ACTIVE
        self.current_version: Optional[str] = None
        self.previous_version: Optional[str] = None
        self.rollback_thresholds = {
            "error_rate": 0.05,
            "latency_p99_ms": 1000,
            "conversion_drop_pct": -10,
        }
        self._lock = threading.Lock()

    def check_and_trigger_rollback(self, metrics: dict) -> bool:
        """检查指标，必要时触发回滚"""
        with self._lock:
            if self.state == ModelState.ROLLING_BACK:
                return False
            triggered = False
            if metrics.get("error_rate", 0) > self.rollback_thresholds["error_rate"]:
                triggered = True
            if metrics.get("latency_p99_ms", 0) > self.rollback_thresholds["latency_p99_ms"]:
                triggered = True
            if metrics.get("conversion_drop_pct", 0) < self.rollback_thresholds["conversion_drop_pct"]:
                triggered = True
            if triggered:
                return self._execute_rollback()
            return False

    def _execute_rollback(self) -> bool:
        if not self.previous_version:
            return False
        self.state = ModelState.ROLLING_BACK
        # 切换流量到旧版本
        time.sleep(1)  # 模拟流量切换
        self.current_version, self.previous_version = (
            self.previous_version, self.current_version
        )
        self.state = ModelState.ROLLED_BACK
        return True

from dataclasses import dataclass
from typing import List
import json


@dataclass
class RollbackEvent:
    """回滚事件记录"""
    model_name: str
    from_version: str
    to_version: str
    trigger_reason: str
    metrics_snapshot: dict
    timestamp: str
    operator: str


class RollbackHistory:
    """维护回滚历史并生成分析报告"""

    def __init__(self):
        self.events: List[RollbackEvent] = []

    def record(self, event: RollbackEvent):
        self.events.append(event)

    def get_rollback_rate(self, model_name: str) -> float:
        total = sum(1 for e in self.events if e.model_name == model_name)
        return total / max(1, 10)  # 假设 10 次发布

    def common_reasons(self) -> List[str]:
        reasons = [e.trigger_reason for e in self.events]
        from collections import Counter
        return [r for r, _ in Counter(reasons).most_common(5)]

    def generate_report(self) -> dict:
        return {
            "total_rollbacks": len(self.events),
            "common_reasons": self.common_reasons(),
            "recent_events": [
                json.dumps({
                    "model": e.model_name,
                    "from": e.from_version,
                    "to": e.to_version,
                    "reason": e.trigger_reason
                })
                for e in self.events[-5:]
            ]
        }

import mlflow
from mlflow.tracking import MlflowClient


class MLflowModelRegistry:
    """基于 MLflow 的模型注册与部署管理"""

    def __init__(self, tracking_uri: str = "http://localhost:5000"):
        self.client = MlflowClient(tracking_uri=tracking_uri)

    def register_model(
        self,
        model_name: str,
        run_id: str,
        description: str = ""
    ) -> str:
        """注册模型并返回版本号"""
        result = self.client.create_model_version(
            name=model_name,
            source=f"runs:/{run_id}/model",
            run_id=run_id,
            description=description
        )
        return result.version

    def promote_to_staging(self, model_name: str, version: str):
        """将模型版本推进到 Staging"""
        self.client.transition_model_version_stage(
            name=model_name, version=version, stage="Staging"
        )

    def promote_to_production(self, model_name: str, version: str):
        """将模型版本推进到 Production"""
        self.client.transition_model_version_stage(
            name=model_name, version=version, stage="Production"
        )

    def get_production_version(self, model_name: str) -> str:
        versions = self.client.get_latest_versions(
            model_name, stages=["Production"]
        )
        return versions[0].version if versions else None

import mlflow
import requests
from typing import Optional


class ABTestModelRouter:
    """集成 MLflow 的 A/B 测试模型路由器"""

    def __init__(self, model_name: str, experiment_id: str):
        self.model_name = model_name
        self.experiment_id = experiment_id
        self.client = mlflow.tracking.MlflowClient()
        self._model_cache = {}

    def get_model_uri(self, stage: str = "Production") -> Optional[str]:
        """获取指定阶段模型的 Serving URI"""
        versions = self.client.get_latest_versions(
            self.model_name, stages=[stage]
        )
        if not versions:
            return None
        return versions[0].source

    def predict(self, user_id: str, features: dict) -> dict:
        """根据实验配置路由到对应模型并返回预测结果"""
        variant = self._get_assigned_variant(user_id)
        if variant == "control":
            uri = self.get_model_uri("Production")
        else:
            uri = self.get_model_uri("Staging")

        if not uri:
            raise RuntimeError(f"No model found for variant: {variant}")

        response = requests.post(
            f"{uri}/invocations",
            json={"inputs": [features]},
            timeout=5
        )
        prediction = response.json()
        self._log_prediction(user_id, variant, prediction)
        return prediction

    def _log_prediction(self, user_id, variant, prediction):
        pass  # 记录到实验分析系统

    def _get_assigned_variant(self, user_id: str) -> str:
        import hashlib
        bucket = int(hashlib.md5(user_id.encode()).hexdigest(), 16) % 100
        return "variant" if bucket < 10 else "control"  # 10% 流量