10-实战项目案例¶

"""
电商用户行为分析系统 - 核心实现
"""
# 需要安装: pip install clickhouse-driver
import json
from datetime import datetime, timedelta
from typing import List, Dict, Optional
import pandas as pd
from sqlalchemy import create_engine, text
import redis
import clickhouse_driver

class ECommerceAnalytics:
    """电商分析系统主类"""

    def __init__(self):
        # MySQL连接
        self.mysql_engine = create_engine(
            'mysql+pymysql://user:pass@localhost/ecommerce'
        )

        # ClickHouse连接
        self.ch_client = clickhouse_driver.Client(
            host='localhost',
            database='ecommerce'
        )

        # Redis连接
        self.redis_client = redis.Redis(host='localhost', port=6379, db=0)

    def log_user_behavior(self, event: Dict):
        """
        记录用户行为

        双写策略：
        1. 写入ClickHouse用于分析
        2. 更新Redis实时统计
        """
        # 1. 写入ClickHouse
        self.ch_client.execute('''
            INSERT INTO user_behavior_log
            (event_id, user_id, session_id, event_type, product_id,
             category_id, brand_id, price, quantity, device_type,
             os, browser, ip_address, country, city, referrer, event_time)
            VALUES
        ''', [(
            event['event_id'],
            event['user_id'],
            event['session_id'],
            event['event_type'],
            event.get('product_id', 0),
            event.get('category_id', 0),
            event.get('brand_id', 0),
            event.get('price', 0),
            event.get('quantity', 0),
            event.get('device_type', 'pc'),
            event.get('os', ''),
            event.get('browser', ''),
            event.get('ip_address', '0.0.0.0'),
            event.get('country', ''),
            event.get('city', ''),
            event.get('referrer', ''),
            datetime.now()
        )])

        # 2. 更新Redis实时统计
        self._update_realtime_stats(event)

    def _update_realtime_stats(self, event: Dict):
        """更新实时统计（Redis）"""
        today = datetime.now().strftime('%Y-%m-%d')
        hour = datetime.now().strftime('%H')

        # 使用Redis Pipeline批量更新
        pipe = self.redis_client.pipeline()

        # 今日总PV
        pipe.incr(f'stats:{today}:pv')

        # 今日UV（使用HyperLogLog）
        pipe.pfadd(f'stats:{today}:uv', event['user_id'])

        # 事件类型统计
        event_type = event['event_type']
        pipe.hincrby(f'stats:{today}:events', event_type, 1)

        # 小时级统计
        pipe.hincrby(f'stats:{today}:{hour}:events', event_type, 1)

        # 设置过期时间（保留7天）
        pipe.expire(f'stats:{today}:pv', 7 * 24 * 3600)

        pipe.execute()

    def get_realtime_dashboard(self) -> Dict:
        """获取实时看板数据"""
        today = datetime.now().strftime('%Y-%m-%d')

        # 从Redis读取实时数据
        pv = int(self.redis_client.get(f'stats:{today}:pv') or 0)
        uv = self.redis_client.pfcount(f'stats:{today}:uv')
        events = self.redis_client.hgetall(f'stats:{today}:events')

        # 计算转化率
        views = int(events.get(b'view', 0))
        purchases = int(events.get(b'purchase', 0))
        conversion_rate = (purchases / views * 100) if views > 0 else 0

        return {
            'date': today,
            'pv': pv,
            'uv': uv,
            'events': {k.decode(): int(v) for k, v in events.items()},
            'conversion_rate': round(conversion_rate, 2)
        }

    def get_user_behavior_analysis(self,
                                   start_date: str,
                                   end_date: str) -> pd.DataFrame:
        """
        用户行为分析

        使用ClickHouse进行OLAP分析
        """
        query = '''
            SELECT
                event_date,
                event_type,
                count() as event_count,
                uniqExact(user_id) as unique_users,
                avg(price) as avg_price,
                sum(price * quantity) as total_revenue
            FROM user_behavior_log
            WHERE event_date BETWEEN %(start)s AND %(end)s
            GROUP BY event_date, event_type
            ORDER BY event_date, event_type
        '''

        return pd.read_sql(
            query,
            self.ch_client,
            params={'start': start_date, 'end': end_date}
        )

    def get_user_funnel_analysis(self, date: str) -> Dict:
        """
        用户漏斗分析

        计算从浏览到购买的转化漏斗
        """
        query = '''
            SELECT
                event_type,
                uniqExact(user_id) as user_count,
                count() as event_count
            FROM user_behavior_log
            WHERE event_date = %(date)s
            GROUP BY event_type
        '''

        result = self.ch_client.execute(query, {'date': date})

        # 构建漏斗数据
        funnel = {}
        for row in result:
            funnel[row[0]] = {
                'users': row[1],
                'events': row[2]
            }

        # 计算转化率
        views = funnel.get('view', {}).get('users', 0)
        clicks = funnel.get('click', {}).get('users', 0)
        carts = funnel.get('cart', {}).get('users', 0)
        purchases = funnel.get('purchase', {}).get('users', 0)

        return {
            'date': date,
            'funnel': {
                'view': views,
                'click': clicks,
                'cart': carts,
                'purchase': purchases
            },
            'conversion_rates': {
                'view_to_click': round(clicks / views * 100, 2) if views > 0 else 0,
                'click_to_cart': round(carts / clicks * 100, 2) if clicks > 0 else 0,
                'cart_to_purchase': round(purchases / carts * 100, 2) if carts > 0 else 0,
                'overall': round(purchases / views * 100, 2) if views > 0 else 0
            }
        }

    def get_user_rfm_analysis(self, analysis_date: str = None) -> pd.DataFrame:
        """
        RFM用户价值分析

        R: Recency - 最近一次购买距今天数
        F: Frequency - 购买频率
        M: Monetary - 购买金额
        """
        if analysis_date is None:
            analysis_date = datetime.now().strftime('%Y-%m-%d')

        query = '''
            SELECT
                user_id,
                dateDiff('day', max(event_time), %(date)s) as recency,
                count() as frequency,
                sum(price * quantity) as monetary
            FROM user_behavior_log
            WHERE event_type = 'purchase'
              AND event_time >= date_sub(%(date)s, 365)
            GROUP BY user_id
        '''

        df = pd.read_sql(
            query,
            self.ch_client,
            params={'date': analysis_date}
        )

        # RFM评分（使用分位数）
        df['R_score'] = pd.qcut(df['recency'], 5, labels=[5,4,3,2,1])
        df['F_score'] = pd.qcut(df['frequency'].rank(method='first'), 5, labels=[1,2,3,4,5])
        df['M_score'] = pd.qcut(df['monetary'], 5, labels=[1,2,3,4,5])

        # 综合RFM得分
        df['RFM_score'] = (df['R_score'].astype(str) +
                          df['F_score'].astype(str) +
                          df['M_score'].astype(str))

        # 用户分层
        def segment_user(row):
            if row['RFM_score'] in ['555', '554', '544', '545', '454', '455', '445']:
                return '重要价值客户'
            elif row['RFM_score'] in ['543', '444', '435', '355', '354', '345', '344', '335']:
                return '重要保持客户'
            elif row['RFM_score'] in ['512', '511', '422', '421', '412', '411', '311']:
                return '新客户'
            elif row['RFM_score'] in ['155', '154', '144', '214', '215', '115', '114']:
                return '重要挽留客户'
            else:
                return '一般客户'

        df['segment'] = df.apply(segment_user, axis=1)

        return df

# 使用示例
analytics = ECommerceAnalytics()

# 记录用户行为
event = {
    'event_id': 'uuid-123',
    'user_id': 10001,
    'session_id': 'sess-abc',
    'event_type': 'purchase',
    'product_id': 5001,
    'category_id': 10,
    'brand_id': 100,
    'price': 299.00,
    'quantity': 2,
    'device_type': 'mobile'
}
analytics.log_user_behavior(event)

# 获取实时看板
dashboard = analytics.get_realtime_dashboard()
print(f"今日PV: {dashboard['pv']}, UV: {dashboard['uv']}")

# RFM分析
rfm_df = analytics.get_user_rfm_analysis()
print(rfm_df.groupby('segment').size())

"""
智能推荐系统 - 核心实现
"""
import numpy as np
from typing import List, Dict, Tuple
from sqlalchemy import create_engine, text
from pgvector.sqlalchemy import Vector
import redis
import hashlib
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel  # Pydantic数据验证模型
import random

app = FastAPI(title="推荐系统API")

class RecommendationRequest(BaseModel):
    user_id: int
    context: Dict = {}
    num_recommendations: int = 10

class RecommendationSystem:
    """推荐系统主类"""

    def __init__(self):
        # PostgreSQL连接
        self.engine = create_engine('postgresql://user:pass@localhost/recommendation')

        # Redis连接
        self.redis = redis.Redis(host='localhost', port=6379, db=0)

        # 缓存配置
        self.cache_ttl = 3600  # 1小时

    def get_recommendations(self,
                           user_id: int,
                           context: Dict = None,
                           num_results: int = 10) -> List[Dict]:
        """
        获取推荐结果

        流程：
        1. 检查A/B测试分组
        2. 尝试从缓存获取
        3. 根据分组选择算法
        4. 生成推荐
        5. 缓存结果
        """
        context = context or {}

        # 1. 检查A/B测试
        ab_group = self._get_ab_test_group(user_id)
        algorithm = self._select_algorithm(ab_group)

        # 2. 检查缓存
        cache_key = f"rec:{user_id}:{algorithm}:{hash(str(context))}"
        cached = self.redis.get(cache_key)
        if cached:
            return json.loads(cached)

        # 3. 生成推荐
        if algorithm == 'collaborative_filtering':
            recommendations = self._collaborative_filtering(user_id, num_results)
        elif algorithm == 'content_based':
            recommendations = self._content_based(user_id, num_results)
        elif algorithm == 'vector_similarity':
            recommendations = self._vector_similarity(user_id, num_results)
        else:
            recommendations = self._hybrid_recommendation(user_id, num_results)

        # 4. 应用业务规则
        recommendations = self._apply_business_rules(recommendations, context)

        # 5. 记录推荐日志
        self._log_recommendation(user_id, context, recommendations, algorithm)

        # 6. 缓存结果
        self.redis.setex(cache_key, self.cache_ttl, json.dumps(recommendations))

        return recommendations

    def _get_ab_test_group(self, user_id: int) -> str:
        """获取用户A/B测试分组"""
        # 检查是否有正在运行的测试
        with self.engine.connect() as conn:
            result = conn.execute(text('''
                SELECT test_id, traffic_split
                FROM ab_tests
                WHERE status = 'running'
                LIMIT 1
            '''))
            test = result.fetchone()

            if not test:
                return 'control'

            test_id, traffic_split = test

            # 检查用户是否已分配
            result = conn.execute(text('''
                SELECT group_name FROM ab_test_assignments
                WHERE user_id = :user_id AND test_id = :test_id
            '''), {'user_id': user_id, 'test_id': test_id})

            assignment = result.fetchone()
            if assignment:
                return assignment.group_name

            # 新用户分配
            group = 'A' if random.random() < traffic_split else 'B'
            conn.execute(text('''
                INSERT INTO ab_test_assignments (user_id, test_id, group_name)
                VALUES (:user_id, :test_id, :group)
            '''), {'user_id': user_id, 'test_id': test_id, 'group': group})
            conn.commit()

            return group

    def _select_algorithm(self, ab_group: str) -> str:
        """根据A/B分组选择算法"""
        algorithms = {
            'A': 'vector_similarity',
            'B': 'collaborative_filtering',
            'control': 'hybrid'
        }
        return algorithms.get(ab_group, 'hybrid')

    def _vector_similarity(self, user_id: int, num_results: int) -> List[Dict]:
        """
        基于向量相似度的推荐

        使用pgvector进行相似度搜索
        """
        with self.engine.connect() as conn:
            # 获取用户向量
            result = conn.execute(text('''
                SELECT user_vector FROM user_profiles WHERE user_id = :user_id
            '''), {'user_id': user_id})

            row = result.fetchone()
            if not row or row.user_vector is None:
                # 新用户，返回热门商品
                return self._get_popular_products(num_results)

            user_vector = row.user_vector

            # 向量相似度搜索
            result = conn.execute(text('''
                SELECT
                    p.product_id,
                    p.product_vector <=> :user_vector as distance,
                    p.price,
                    p.popularity_score
                FROM product_features p
                WHERE p.product_vector IS NOT NULL
                ORDER BY p.product_vector <=> :user_vector
                LIMIT :limit
            '''), {'user_vector': str(user_vector), 'limit': num_results * 2})

            recommendations = []
            for row in result:
                recommendations.append({
                    'product_id': row.product_id,
                    'score': 1 - float(row.distance),  # 转换为相似度
                    'price': float(row.price),
                    'algorithm': 'vector_similarity'
                })

            return recommendations[:num_results]

    def _collaborative_filtering(self, user_id: int, num_results: int) -> List[Dict]:
        """
        协同过滤推荐

        基于相似用户的购买行为
        """
        with self.engine.connect() as conn:
            # 找到相似用户
            result = conn.execute(text('''
                WITH similar_users AS (
                    SELECT
                        b.user_id as other_user,
                        COUNT(*) as common_products
                    FROM user_behavior_sequences a
                    JOIN user_behavior_sequences b ON a.product_id = b.product_id
                    WHERE a.user_id = :user_id
                      AND b.user_id != :user_id
                      AND a.event_type = 'purchase'
                      AND b.event_type = 'purchase'
                    GROUP BY b.user_id
                    ORDER BY common_products DESC
                    LIMIT 20
                )
                SELECT
                    ub.product_id,
                    COUNT(*) as score,
                    AVG(p.price) as avg_price
                FROM user_behavior_sequences ub
                JOIN similar_users su ON ub.user_id = su.other_user
                JOIN product_features p ON ub.product_id = p.product_id
                WHERE ub.event_type = 'purchase'
                  AND ub.product_id NOT IN (
                      SELECT product_id FROM user_behavior_sequences
                      WHERE user_id = :user_id AND event_type = 'purchase'
                  )
                GROUP BY ub.product_id
                ORDER BY score DESC
                LIMIT :limit
            '''), {'user_id': user_id, 'limit': num_results})

            recommendations = []
            for row in result:
                recommendations.append({
                    'product_id': row.product_id,
                    'score': row.score,
                    'price': float(row.avg_price),
                    'algorithm': 'collaborative_filtering'
                })

            return recommendations

    def _content_based(self, user_id: int, num_results: int) -> List[Dict]:
        """基于内容的推荐"""
        with self.engine.connect() as conn:
            # 获取用户偏好
            result = conn.execute(text('''
                SELECT preferences FROM user_profiles WHERE user_id = :user_id
            '''), {'user_id': user_id})

            row = result.fetchone()
            if not row:
                return self._get_popular_products(num_results)

            preferences = row.preferences or {}
            preferred_categories = preferences.get('categories', [])

            # 基于偏好的推荐
            result = conn.execute(text('''
                SELECT
                    product_id,
                    price,
                    popularity_score,
                    attributes
                FROM product_features
                WHERE category_id = ANY(:categories)
                ORDER BY popularity_score DESC
                LIMIT :limit
            '''), {'categories': preferred_categories, 'limit': num_results})

            recommendations = []
            for row in result:
                recommendations.append({
                    'product_id': row.product_id,
                    'score': row.popularity_score,
                    'price': float(row.price),
                    'algorithm': 'content_based'
                })

            return recommendations

    def _hybrid_recommendation(self, user_id: int, num_results: int) -> List[Dict]:
        """混合推荐"""
        # 获取多种算法的推荐
        vector_recs = self._vector_similarity(user_id, num_results // 2)
        cf_recs = self._collaborative_filtering(user_id, num_results // 2)

        # 合并并去重
        seen_products = set()
        recommendations = []

        for rec in vector_recs + cf_recs:
            if rec['product_id'] not in seen_products:
                seen_products.add(rec['product_id'])
                rec['algorithm'] = 'hybrid'
                recommendations.append(rec)

        return recommendations[:num_results]

    def _apply_business_rules(self,
                             recommendations: List[Dict],
                             context: Dict) -> List[Dict]:
        """应用业务规则"""
        # 过滤已购买商品
        # 提升新品曝光
        # 价格区间过滤
        # 库存过滤
        return recommendations

    def _get_popular_products(self, num_results: int) -> List[Dict]:
        """获取热门商品（冷启动）"""
        with self.engine.connect() as conn:
            result = conn.execute(text('''
                SELECT product_id, price, popularity_score
                FROM product_features
                ORDER BY popularity_score DESC
                LIMIT :limit
            '''), {'limit': num_results})

            return [
                {
                    'product_id': row.product_id,
                    'score': row.popularity_score,
                    'price': float(row.price),
                    'algorithm': 'popular'
                }
                for row in result
            ]

    def _log_recommendation(self,
                           user_id: int,
                           context: Dict,
                           recommendations: List[Dict],
                           algorithm: str):
        """记录推荐日志"""
        import uuid

        with self.engine.connect() as conn:
            conn.execute(text('''
                INSERT INTO recommendation_logs
                (user_id, request_id, context, recommendations, algorithm)
                VALUES (:user_id, :request_id, :context, :recommendations, :algorithm)
            '''), {
                'user_id': user_id,
                'request_id': str(uuid.uuid4()),
                'context': json.dumps(context),
                'recommendations': json.dumps(recommendations),
                'algorithm': algorithm
            })
            conn.commit()

# FastAPI路由
rec_system = RecommendationSystem()

@app.post("/recommendations")
async def get_recommendations(request: RecommendationRequest):  # async def定义异步函数；用await调用
    """推荐API"""
    try:
        recommendations = rec_system.get_recommendations(
            user_id=request.user_id,
            context=request.context,
            num_results=request.num_recommendations
        )

        return {
            'user_id': request.user_id,
            'recommendations': recommendations,
            'total': len(recommendations)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/feedback")
async def record_feedback(user_id: int, product_id: int, action: str):
    """记录用户反馈（点击、购买等）"""
    # 更新用户画像
    # 更新推荐模型
    return {'status': 'success'}

"""
MLOps模型管理平台 - 核心实现
"""
import os
import json
import hashlib
import shutil
from datetime import datetime
from typing import List, Dict, Optional
from sqlalchemy import create_engine, text
import boto3
from botocore.exceptions import ClientError

class MLOpsPlatform:
    """MLOps平台主类"""

    def __init__(self, db_url: str, storage_endpoint: str, storage_bucket: str):
        self.engine = create_engine(db_url)
        self.s3_client = boto3.client(
            's3',
            endpoint_url=storage_endpoint,
            aws_access_key_id='minioadmin',
            aws_secret_access_key='minioadmin'
        )
        self.bucket = storage_bucket
        self._ensure_bucket_exists()

    def _ensure_bucket_exists(self):
        """确保存储桶存在"""
        try:  # try/except捕获异常
            self.s3_client.head_bucket(Bucket=self.bucket)
        except ClientError:
            self.s3_client.create_bucket(Bucket=self.bucket)

    # ==================== 数据集管理 ====================

    def register_dataset(self,
                        name: str,
                        version: str,
                        local_path: str,
                        description: str = "",
                        tags: Dict = None) -> int:
        """注册数据集"""
        # 计算数据集信息
        size_bytes = os.path.getsize(local_path)

        # 上传到对象存储
        storage_path = f"datasets/{name}/{version}/data"
        self.s3_client.upload_file(local_path, self.bucket, storage_path)

        # 记录到数据库
        with self.engine.connect() as conn:
            result = conn.execute(text('''
                INSERT INTO datasets
                (name, version, description, storage_path, size_bytes, tags, created_by)
                VALUES (:name, :version, :desc, :path, :size, :tags, :user)
                RETURNING dataset_id
            '''), {
                'name': name,
                'version': version,
                'desc': description,
                'path': storage_path,
                'size': size_bytes,
                'tags': json.dumps(tags or {}),  # json.dumps将Python对象转为JSON字符串
                'user': 'system'
            })
            conn.commit()

            dataset_id = result.fetchone()[0]

        print(f"✅ 数据集注册成功: {name} v{version} (ID: {dataset_id})")
        return dataset_id

    def get_dataset(self, name: str, version: str = None) -> Optional[Dict]:
        """获取数据集信息"""
        with self.engine.connect() as conn:
            if version:
                result = conn.execute(text('''
                    SELECT * FROM datasets
                    WHERE name = :name AND version = :version
                '''), {'name': name, 'version': version})
            else:
                result = conn.execute(text('''
                    SELECT * FROM datasets
                    WHERE name = :name
                    ORDER BY created_at DESC LIMIT 1
                '''), {'name': name})

            row = result.fetchone()
            if row:
                return {
                    'dataset_id': row.dataset_id,
                    'name': row.name,
                    'version': row.version,
                    'storage_path': row.storage_path,
                    'size_bytes': row.size_bytes,
                    'tags': json.loads(row.tags) if row.tags else {}  # json.loads将JSON字符串转为Python对象
                }
            return None

    def download_dataset(self, dataset_id: int, local_path: str):
        """下载数据集"""
        dataset = self.get_dataset_by_id(dataset_id)
        if not dataset:
            raise ValueError(f"数据集不存在: {dataset_id}")

        self.s3_client.download_file(
            self.bucket,
            dataset['storage_path'],
            local_path
        )
        print(f"✅ 数据集下载完成: {local_path}")

    # ==================== 实验管理 ====================

    def create_experiment(self, name: str, description: str = "", tags: Dict = None) -> int:
        """创建实验"""
        with self.engine.connect() as conn:
            result = conn.execute(text('''
                INSERT INTO experiments (name, description, tags, created_by)
                VALUES (:name, :desc, :tags, :user)
                RETURNING experiment_id
            '''), {
                'name': name,
                'desc': description,
                'tags': json.dumps(tags or {}),
                'user': 'system'
            })
            conn.commit()

            exp_id = result.fetchone()[0]

        print(f"✅ 实验创建成功: {name} (ID: {exp_id})")
        return exp_id

    def start_run(self, experiment_id: int, run_name: str = None) -> str:
        """开始运行"""
        import uuid
        run_id = uuid.uuid4().hex[:16]  # 切片操作：[start:end:step]提取子序列

        with self.engine.connect() as conn:
            conn.execute(text('''
                INSERT INTO runs (run_id, experiment_id, tags)
                VALUES (:run_id, :exp_id, :tags)
            '''), {
                'run_id': run_id,
                'exp_id': experiment_id,
                'tags': json.dumps({'mlflow.runName': run_name} if run_name else {})
            })
            conn.commit()

        print(f"▶️ 运行开始: {run_id}")
        return run_id

    def log_param(self, run_id: str, key: str, value: any):
        """记录参数"""
        with self.engine.connect() as conn:
            conn.execute(text('''
                INSERT INTO run_params (run_id, param_key, param_value)
                VALUES (:run_id, :key, :value)
                ON CONFLICT (run_id, param_key) DO UPDATE
                SET param_value = EXCLUDED.param_value
            '''), {
                'run_id': run_id,
                'key': key,
                'value': str(value)
            })
            conn.commit()

    def log_metric(self, run_id: str, key: str, value: float, step: int = None):
        """记录指标"""
        with self.engine.connect() as conn:
            conn.execute(text('''
                INSERT INTO run_metrics (run_id, metric_key, metric_value, step)
                VALUES (:run_id, :key, :value, :step)
            '''), {
                'run_id': run_id,
                'key': key,
                'value': value,
                'step': step
            })
            conn.commit()

    def log_artifact(self, run_id: str, local_path: str, artifact_path: str = None):
        """记录Artifact"""
        artifact_name = artifact_path or os.path.basename(local_path)
        storage_path = f"artifacts/{run_id}/{artifact_name}"

        # 上传到存储
        if os.path.isfile(local_path):
            self.s3_client.upload_file(local_path, self.bucket, storage_path)
            size_bytes = os.path.getsize(local_path)
        else:
            # 上传目录
            for root, dirs, files in os.walk(local_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    rel_path = os.path.relpath(file_path, local_path)
                    s3_path = f"{storage_path}/{rel_path}"
                    self.s3_client.upload_file(file_path, self.bucket, s3_path)
            size_bytes = sum(
                os.path.getsize(os.path.join(dirpath, filename))
                for dirpath, dirnames, filenames in os.walk(local_path)
                for filename in filenames
            )

        # 记录到数据库
        with self.engine.connect() as conn:
            conn.execute(text('''
                INSERT INTO run_artifacts (run_id, artifact_name, artifact_type, storage_path, size_bytes)
                VALUES (:run_id, :name, :type, :path, :size)
            '''), {
                'run_id': run_id,
                'name': artifact_name,
                'type': 'file',
                'path': storage_path,
                'size': size_bytes
            })
            conn.commit()

        print(f"📎 Artifact已记录: {artifact_name}")

    def end_run(self, run_id: str, status: str = "COMPLETED"):
        """结束运行"""
        with self.engine.connect() as conn:
            conn.execute(text('''
                UPDATE runs
                SET status = :status, end_time = NOW()
                WHERE run_id = :run_id
            '''), {'status': status, 'run_id': run_id})
            conn.commit()

        print(f"✅ 运行结束: {run_id} ({status})")

    # ==================== 模型管理 ====================

    def register_model(self,
                      name: str,
                      run_id: str,
                      model_path: str,
                      description: str = "",
                      tags: Dict = None) -> Dict:
        """注册模型"""
        # 确保模型存在
        with self.engine.connect() as conn:
            result = conn.execute(text('''
                SELECT 1 FROM registered_models WHERE name = :name
            '''), {'name': name})

            if not result.fetchone():
                conn.execute(text('''
                    INSERT INTO registered_models (name, description, tags)
                    VALUES (:name, :desc, :tags)
                '''), {
                    'name': name,
                    'desc': description,
                    'tags': json.dumps(tags or {})
                })

            # 获取版本号
            result = conn.execute(text('''
                SELECT COALESCE(MAX(version), 0) + 1
                FROM model_versions mv
                JOIN registered_models rm ON mv.model_id = rm.model_id
                WHERE rm.name = :name
            '''), {'name': name})

            version = result.fetchone()[0]

            # 上传模型
            model_storage_path = f"models/{name}/{version}/model"
            if os.path.isdir(model_path):
                for root, dirs, files in os.walk(model_path):
                    for file in files:
                        file_path = os.path.join(root, file)
                        rel_path = os.path.relpath(file_path, model_path)
                        s3_path = f"{model_storage_path}/{rel_path}"
                        self.s3_client.upload_file(file_path, self.bucket, s3_path)
            else:
                self.s3_client.upload_file(model_path, self.bucket, model_storage_path)

            # 记录版本
            result = conn.execute(text('''
                INSERT INTO model_versions
                (model_id, version, run_id, storage_path, model_format, training_metrics)
                SELECT model_id, :version, :run_id, :path, :format, :metrics
                FROM registered_models WHERE name = :name
                RETURNING version_id
            '''), {
                'name': name,
                'version': version,
                'run_id': run_id,
                'path': model_storage_path,
                'format': 'sklearn',
                'metrics': '{}'
            })
            conn.commit()

            version_id = result.fetchone()[0]

        print(f"✅ 模型注册成功: {name} v{version}")
        return {'model_name': name, 'version': version, 'version_id': version_id}

    def transition_model_stage(self, name: str, version: int, stage: str):
        """转换模型阶段"""
        with self.engine.connect() as conn:
            # 如果要设置为Production，降级其他Production版本
            if stage == 'PRODUCTION':
                conn.execute(text('''
                    UPDATE model_versions mv
                    SET status = 'ARCHIVED'
                    FROM registered_models rm
                    WHERE mv.model_id = rm.model_id
                      AND rm.name = :name
                      AND mv.status = 'PRODUCTION'
                '''), {'name': name})

            conn.execute(text('''
                UPDATE model_versions mv
                SET status = :stage
                FROM registered_models rm
                WHERE mv.model_id = rm.model_id
                  AND rm.name = :name
                  AND mv.version = :version
            '''), {'name': name, 'version': version, 'stage': stage})
            conn.commit()

        print(f"🔄 模型 {name} v{version} 阶段变更为: {stage}")

    def load_model(self, name: str, version: int = None, stage: str = None):
        """加载模型"""
        with self.engine.connect() as conn:
            if version:
                result = conn.execute(text('''
                    SELECT mv.storage_path, mv.model_format
                    FROM model_versions mv
                    JOIN registered_models rm ON mv.model_id = rm.model_id
                    WHERE rm.name = :name AND mv.version = :version
                '''), {'name': name, 'version': version})
            elif stage:
                result = conn.execute(text('''
                    SELECT mv.storage_path, mv.model_format
                    FROM model_versions mv
                    JOIN registered_models rm ON mv.model_id = rm.model_id
                    WHERE rm.name = :name AND mv.status = :stage
                    ORDER BY mv.version DESC LIMIT 1
                '''), {'name': name, 'stage': stage})
            else:
                result = conn.execute(text('''
                    SELECT mv.storage_path, mv.model_format
                    FROM model_versions mv
                    JOIN registered_models rm ON mv.model_id = rm.model_id
                    WHERE rm.name = :name
                    ORDER BY mv.version DESC LIMIT 1
                '''), {'name': name})

            row = result.fetchone()
            if not row:
                raise ValueError(f"模型不存在: {name}")

            # 下载模型
            local_path = f"/tmp/models/{name}"
            os.makedirs(local_path, exist_ok=True)

            # 列出并下载所有文件
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket,
                Prefix=row.storage_path
            )

            for obj in response.get('Contents', []):
                rel_path = os.path.relpath(obj['Key'], row.storage_path)
                file_path = os.path.join(local_path, rel_path)
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                self.s3_client.download_file(self.bucket, obj['Key'], file_path)

            # 加载模型
            import joblib
            model_file = os.path.join(local_path, 'model.pkl')
            return joblib.load(model_file)

# 使用示例
mlops = MLOpsPlatform(
    db_url='postgresql://user:pass@localhost/mlops',
    storage_endpoint='http://localhost:9000',
    storage_bucket='mlops-storage'
)

# 注册数据集
dataset_id = mlops.register_dataset(
    name='customer_churn',
    version='1.0.0',
    local_path='data/customer_churn.csv',
    description='客户流失预测数据集',
    tags={'domain': 'finance', 'type': 'classification'}
)

# 创建实验并记录运行
exp_id = mlops.create_experiment('churn_prediction', '客户流失预测模型')
run_id = mlops.start_run(exp_id, run_name='xgboost_baseline')

# 记录参数和指标
mlops.log_param(run_id, 'model', 'xgboost')
mlops.log_param(run_id, 'n_estimators', 100)
mlops.log_param(run_id, 'max_depth', 6)

for epoch in range(10):
    mlops.log_metric(run_id, 'train_auc', 0.8 + epoch * 0.02, step=epoch)
    mlops.log_metric(run_id, 'val_auc', 0.78 + epoch * 0.015, step=epoch)

mlops.log_metric(run_id, 'test_auc', 0.95)

# 保存模型并注册
mlops.log_artifact(run_id, 'model.pkl')
model_info = mlops.register_model(
    name='churn_predictor',
    run_id=run_id,
    model_path='model.pkl',
    description='XGBoost客户流失预测模型'
)

# 部署到生产
mlops.transition_model_stage('churn_predictor', model_info['version'], 'PRODUCTION')

mlops.end_run(run_id)