向量数据库在实时推荐系统中的性能瓶颈与优化策略

realtime-recommendation-optimization/
├── config/
│   └── settings.py           # 应用配置（数据库连接、索引参数等）
├── core/
│   ├── vector_db_service.py  # 向量数据库连接与操作封装
│   └── recommendation_service.py # 推荐业务逻辑核心
├── models/
│   └── data_models.py        # 数据模型定义（Pydantic）
├── api/
│   └── app.py                # FastAPI 应用入口
├── simulation/
│   ├── data_generator.py     # 生成模拟电影数据与向量
│   └── load_client.py        # 模拟高并发请求的客户端
├── tests/
│   └── test_performance.py   # 性能测试与瓶颈分析脚本
├── requirements.txt          # Python 依赖
└── docker-compose.yml        # 启动 Milvus 服务

"""
项目配置文件
集中管理所有配置项，便于优化调整和实验对比。
"""
import os
from pydantic_settings import BaseSettings

class Settings(BaseSettings):
    # Milvus 连接配置
    MILVUS_HOST: str = os.getenv("MILVUS_HOST", "localhost")
    MILVUS_PORT: int = int(os.getenv("MILVUS_PORT", "19530"))
    
    # 集合（表）配置
    COLLECTION_NAME: str = "movie_recommendations"
    VECTOR_DIMENSION: int = 384  # 使用的 sentence transformer 模型维度
    VECTOR_FIELD: str = "embedding"
    PRIMARY_KEY_FIELD: str = "movie_id"
    TEXT_FIELD: str = "description"
    
    # 索引配置 (优化前：使用默认的FLAT精确搜索，速度慢但召回率100%)
    # 优化后：使用混合索引 HNSW + IVF_PQ
    INDEX_TYPE: str = os.getenv("INDEX_TYPE", "HNSW")  # 可选：FLAT, IVF_FLAT, IVF_SQ8, IVF_PQ, HNSW
    METRIC_TYPE: str = "IP"  # 内积 (余弦相似度归一化后等价于内积)
    INDEX_PARAMS: dict = {
        "HNSW": {"M": 16, "efConstruction": 200},  # HNSW 构建参数
        "IVF_PQ": {"nlist": 1024, "m": 32, "nbits": 8}, # IVF_PQ 参数
    }
    SEARCH_PARAMS: dict = {
        "HNSW": {"ef": 50},
        "IVF_PQ": {"nprobe": 20},
        "FLAT": {}
    }
    
    # 缓存配置 (优化项)
    ENABLE_CACHE: bool = True
    CACHE_TTL: int = 300  # 缓存过期时间（秒）
    CACHE_MAX_SIZE: int = 1000
    
    # 推荐参数
    TOP_K: int = 10  # 每次推荐返回的物品数量
    
    # 异步更新配置 (优化项)
    ENABLE_ASYNC_UPDATE: bool = True
    UPDATE_BATCH_SIZE: int = 100
    UPDATE_QUEUE_SIZE: int = 1000

    class Config:
        env_file = ".env"

settings = Settings()

"""
向量数据库服务层
封装所有与 Milvus 的交互，实现连接池、索引管理、搜索和异步更新。
"""
import threading
from typing import List, Optional, Dict, Any
from loguru import logger
from pymilvus import (
    connections,
    utility,
    Collection,
    CollectionSchema,
    FieldSchema,
    DataType,
    Hits
)
from queue import Queue
import numpy as np

from config.settings import settings

class VectorDBService:
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance._initialized = False
            return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self.collection = None
        self._update_queue = Queue(maxsize=settings.UPDATE_QUEUE_SIZE)
        self._init_connection()
        self._get_or_create_collection()
        self._create_index_if_needed()
        if settings.ENABLE_ASYNC_UPDATE:
            self._start_async_update_worker()
        self._initialized = True
        logger.info("VectorDBService initialized.")
    
    def _init_connection(self):
        """初始化到 Milvus 的连接（连接池由 pymilvus 内部管理）。"""
        try:
            connections.connect(
                alias="default",
                host=settings.MILVUS_HOST,
                port=settings.MILVUS_PORT
            )
            logger.info(f"Connected to Milvus at {settings.MILVUS_HOST}:{settings.MILVUS_PORT}")
        except Exception as e:
            logger.error(f"Failed to connect to Milvus: {e}")
            raise
    
    def _get_or_create_collection(self):
        """获取或创建电影集合。"""
        if utility.has_collection(settings.COLLECTION_NAME):
            self.collection = Collection(settings.COLLECTION_NAME)
            logger.info(f"Collection '{settings.COLLECTION_NAME}' loaded.")
        else:
            # 定义 Schema
            fields = [
                FieldSchema(name=settings.PRIMARY_KEY_FIELD, dtype=DataType.INT64, is_primary=True, auto_id=False),
                FieldSchema(name=settings.TEXT_FIELD, dtype=DataType.VARCHAR, max_length=1000),
                FieldSchema(name=settings.VECTOR_FIELD, dtype=DataType.FLOAT_VECTOR, dim=settings.VECTOR_DIMENSION)
            ]
            schema = CollectionSchema(fields=fields, description="Movie recommendations collection")
            self.collection = Collection(name=settings.COLLECTION_NAME, schema=schema)
            logger.info(f"Collection '{settings.COLLECTION_NAME}' created.")
    
    def _create_index_if_needed(self):
        """在指定字段上创建索引（如果尚未创建）。"""
        if len(self.collection.indexes) == 0:
            logger.info(f"Creating '{settings.INDEX_TYPE}' index on field '{settings.VECTOR_FIELD}'...")
            index_params = {
                "index_type": settings.INDEX_TYPE,
                "metric_type": settings.METRIC_TYPE,
                "params": settings.INDEX_PARAMS.get(settings.INDEX_TYPE, {})
            }
            self.collection.create_index(settings.VECTOR_FIELD, index_params)
            logger.info("Index created.")
        # 加载集合到内存以优化搜索性能
        self.collection.load()
    
    def search_similar_items(self, query_vector: List[float], top_k: int = None) -> List[Dict]:
        """
        核心 ANN 搜索函数。
        优化点：使用合适的搜索参数，确保内存中的集合已加载。
        """
        top_k = top_k or settings.TOP_K
        search_params = settings.SEARCH_PARAMS.get(settings.INDEX_TYPE, {})
        
        # 执行搜索
        results = self.collection.search(
            data=[query_vector],
            anns_field=settings.VECTOR_FIELD,
            param=search_params,
            limit=top_k,
            output_fields=[settings.PRIMARY_KEY_FIELD, settings.TEXT_FIELD]  # 指定需要返回的字段
        )
        
        # 格式化结果
        ret = []
        for hits in results:
            for hit in hits:
                ret.append({
                    "movie_id": hit.entity.get(settings.PRIMARY_KEY_FIELD),
                    "description": hit.entity.get(settings.TEXT_FIELD),
                    "score": hit.score
                })
        return ret
    
    def insert_items(self, data: List[Dict[str, Any]]):
        """
        插入或更新物品向量。
        优化：如果启用异步，则放入队列；否则同步插入。
        """
        if settings.ENABLE_ASYNC_UPDATE:
            try:
                self._update_queue.put_nowait(data)
                logger.debug(f"Update task queued. Queue size: {self._update_queue.qsize()}")
            except queue.Full:
                logger.warning("Update queue is full. Dropping data.")
                # 在真实场景中，这里可能需要更复杂的背压策略或写入临时存储
        else:
            self._sync_insert(data)
    
    def _sync_insert(self, data: List[Dict[str, Any]]):
        """同步插入数据。"""
        if not data:
            return
        # 提取并组织数据以符合 Milvus 插入格式
        movie_ids = [item[settings.PRIMARY_KEY_FIELD] for item in data]
        descriptions = [item[settings.TEXT_FIELD] for item in data]
        embeddings = [item[settings.VECTOR_FIELD] for item in data]
        
        insert_data = [movie_ids, descriptions, embeddings]
        try:
            mr = self.collection.insert(insert_data)
            logger.info(f"Synced insert {len(data)} items. Insert count: {mr.insert_count}")
        except Exception as e:
            logger.error(f"Failed to insert data: {e}")
    
    def _start_async_update_worker(self):
        """启动后台线程，批量处理队列中的更新请求。"""
        def worker():
            buffer = []
            while True:
                try:
                    data = self._update_queue.get(timeout=1.0)
                    buffer.extend(data)
                    # 当缓冲区达到批量大小或队列暂时为空时，执行插入
                    if len(buffer) >= settings.UPDATE_BATCH_SIZE or self._update_queue.empty():
                        if buffer:
                            self._sync_insert(buffer)
                            buffer.clear()
                    self._update_queue.task_done()
                except queue.Empty:
                    # 定期检查并清空缓冲区
                    if buffer:
                        self._sync_insert(buffer)
                        buffer.clear()
                except Exception as e:
                    logger.error(f"Async update worker error: {e}")
        
        thread = threading.Thread(target=worker, daemon=True, name="AsyncUpdateWorker")
        thread.start()
        logger.info("Async update worker started.")
    
    def get_collection_stats(self) -> Dict:
        """获取集合统计信息，用于监控。"""
        if not self.collection:
            return {}
        stats = self.collection.num_entities
        return {"entity_count": stats}

"""
推荐服务核心逻辑
整合向量搜索、缓存和业务规则。
"""
import time
from typing import List, Dict, Any, Optional
from functools import lru_cache
from loguru import logger

from config.settings import settings
from core.vector_db_service import VectorDBService
from models.data_models import MovieItem

class RecommendationService:
    def __init__(self):
        self.vector_db = VectorDBService()
        self._cache = {}  # 简化内存缓存，生产环境应使用 Redis
        self._cache_timestamps = {}
    
    def _get_cache_key(self, movie_id: int, top_k: int) -> str:
        return f"{movie_id}:{top_k}"
    
    def _get_from_cache(self, cache_key: str) -> Optional[List[Dict]]:
        if not settings.ENABLE_CACHE:
            return None
        if cache_key in self._cache:
            if time.time() - self._cache_timestamps[cache_key] < settings.CACHE_TTL:
                logger.debug(f"Cache hit for key: {cache_key}")
                return self._cache[cache_key]
            else:
                # 缓存过期
                del self._cache[cache_key]
                del self._cache_timestamps[cache_key]
        return None
    
    def _set_to_cache(self, cache_key: str, results: List[Dict]):
        if not settings.ENABLE_CACHE:
            return
        # 简单的 LRU 驱逐策略
        if len(self._cache) >= settings.CACHE_MAX_SIZE:
            oldest_key = min(self._cache_timestamps, key=self._cache_timestamps.get)
            del self._cache[oldest_key]
            del self._cache_timestamps[oldest_key]
            logger.debug(f"Evicted cache key: {oldest_key}")
        
        self._cache[cache_key] = results
        self._cache_timestamps[cache_key] = time.time()
        logger.debug(f"Cached results for key: {cache_key}")
    
    @lru_cache(maxsize=1024)
    def _get_movie_vector_from_db(self, movie_id: int) -> Optional[List[float]]:
        """通过电影ID获取其向量。使用 LRU 缓存减少对 Milvus 的 primary key 查询。"""
        # 注意：这是一个简化实现。真实场景中，可能需要在 Milvus 中通过 query() 或创建索引来高效获取。
        # 这里假设我们有一个辅助的、支持主键查询的存储（如关系型数据库）来存放向量。
        # 为简化演示，我们假设此函数能从某个地方高效获取。
        # 本示例中，我们直接返回 None，并在上层调用时使用一个全局的映射（模拟）。
        # 在实际代码中，这里应有具体实现。
        return None
    
    def get_similar_movies(self, movie_id: int, top_k: int = None) -> List[Dict[str, Any]]:
        """
        推荐相似电影的主函数。
        优化点：多级缓存、高效向量获取、降级策略。
        """
        top_k = top_k or settings.TOP_K
        cache_key = self._get_cache_key(movie_id, top_k)
        
        # 1. 检查缓存
        cached_results = self._get_from_cache(cache_key)
        if cached_results is not None:
            return cached_results
        
        # 2. 获取查询电影的向量 (性能关键点，假设已优化)
        query_vector = self._get_movie_vector_from_db(movie_id)
        if query_vector is None:
            logger.warning(f"Vector for movie {movie_id} not found. Returning empty results.")
            # 可选降级：返回热门电影
            return self._get_fallback_recommendations(top_k)
        
        # 3. 执行向量数据库 ANN 搜索
        try:
            start_time = time.perf_counter()
            similar_items = self.vector_db.search_similar_items(query_vector, top_k)
            search_latency = time.perf_counter() - start_time
            logger.info(f"ANN search for movie {movie_id} took {search_latency:.3f}s")
            
            # 过滤掉查询电影本身（分数最高）
            filtered_items = [item for item in similar_items if item['movie_id'] != movie_id]
            
            # 4. 写入缓存
            self._set_to_cache(cache_key, filtered_items)
            
            return filtered_items
        except Exception as e:
            logger.error(f"ANN search failed for movie {movie_id}: {e}")
            # 降级策略
            return self._get_fallback_recommendations(top_k)
    
    def _get_fallback_recommendations(self, top_k: int) -> List[Dict]:
        """降级策略：返回一个预定义的热门电影列表。"""
        # 这里返回一个空列表作为示例。真实场景应从缓存或快速路径获取。
        return []
    
    def process_user_interaction(self, user_id: int, movie_id: int):
        """
        处理用户交互（如点击）。
        可能触发实时向量更新或用户偏好向量计算（本项目简化）。
        """
        logger.info(f"User {user_id} interacted with movie {movie_id}")
        # 在更复杂的系统中，这里可能会更新用户画像，或触发新一轮的实时推荐计算。
        # 本项目主要聚焦物品相似性推荐，因此此函数作为扩展点。
        pass

"""
Pydantic 数据模型
用于 API 请求/响应验证和数据结构定义。
"""
from pydantic import BaseModel, Field
from typing import List, Optional

class MovieItem(BaseModel):
    movie_id: int = Field(..., description="电影唯一ID")
    description: str = Field(..., description="电影描述文本")
    embedding: Optional[List[float]] = Field(None, description="电影描述对应的向量", min_items=384, max_items=384)

class RecommendationRequest(BaseModel):
    movie_id: int = Field(..., ge=1, description="查询的电影ID")
    top_k: Optional[int] = Field(10, ge=1, le=100, description="返回的推荐数量")

class RecommendationResponse(BaseModel):
    request_movie_id: int
    recommendations: List[MovieItem]
    cached: bool = Field(default=False, description="结果是否来自缓存")
    search_latency_ms: Optional[float] = Field(None, description="向量搜索耗时（毫秒）")

"""
FastAPI 应用主文件
提供推荐 RESTful API 端点。
"""
from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
import time
from loguru import logger

from config.settings import settings
from core.recommendation_service import RecommendationService
from models.data_models import RecommendationRequest, RecommendationResponse, MovieItem

# 全局服务实例
rec_service = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动
    global rec_service
    rec_service = RecommendationService()
    logger.info("RecommendationService started.")
    yield
    # 关闭
    logger.info("Shutting down...")
    # 可在此处添加清理逻辑

app = FastAPI(title="Real-time Movie Recommendation API", lifespan=lifespan)

@app.get("/health")
async def health_check():
    return {"status": "healthy", "milvus_index": settings.INDEX_TYPE}

@app.post("/recommend", response_model=RecommendationResponse)
async def get_recommendations(request: RecommendationRequest, background_tasks: BackgroundTasks):
    """
    实时推荐接口。

    1. 接收一个电影ID。
    2. 返回其最相似的 top_k 个电影。
    """
    start_time = time.perf_counter()
    
    if rec_service is None:
        raise HTTPException(status_code=503, detail="Service not initialized")
    
    # 调用推荐服务
    similar_items = rec_service.get_similar_movies(request.movie_id, request.top_k)
    
    # 记录处理延迟
    process_latency = (time.perf_counter() - start_time) * 1000  # 转毫秒
    
    # 构建响应 (简化，实际需判断是否来自缓存)
    resp = RecommendationResponse(
        request_movie_id=request.movie_id,
        recommendations=[
            MovieItem(movie_id=item['movie_id'], description=item['description'])
            for item in similar_items[:request.top_k]
        ],
        cached=False,  # 真实场景需要从服务层传递此信息
        search_latency_ms=process_latency
    )
    
    # 可选：将用户交互放入后台任务处理，不阻塞本次响应
    background_tasks.add_task(rec_service.process_user_interaction, user_id=0, movie_id=request.movie_id) # 简化user_id
    
    return resp

@app.get("/stats")
async def get_stats():
    """获取系统统计信息（用于监控）。"""
    if rec_service is None:
        return {}
    stats = rec_service.vector_db.get_collection_stats()
    return {
        "collection_stats": stats,
        "cache_enabled": settings.ENABLE_CACHE,
        "index_type": settings.INDEX_TYPE
    }

"""
模拟数据生成器
创建电影描述并利用 Sentence Transformer 生成向量。
"""
import numpy as np
from faker import Faker
from sentence_transformers import SentenceTransformer
import random
from loguru import logger
from tqdm import tqdm

from config.settings import settings
from core.vector_db_service import VectorDBService

fake = Faker()
# 使用轻量级模型
model = SentenceTransformer('all-MiniLM-L6-v2')  # 维度 384

def generate_movie_description() -> str:
    """生成虚构的电影描述。"""
    genres = ["Sci-Fi", "Drama", "Comedy", "Action", "Romance", "Thriller", "Documentary"]
    plot = fake.paragraph(nb_sentences=3)
    genre = random.choice(genres)
    return f"A {genre} film. {plot}"

def generate_and_insert_movies(num_movies: int = 10000, batch_size: int = 500):
    """
    生成电影数据并插入 Milvus。
    这是初始化数据集的步骤，不是实时更新。
    """
    logger.info(f"Generating {num_movies} movie records...")
    vector_db = VectorDBService()
    
    all_data = []
    for i in tqdm(range(num_movies), desc="Generating movies"):
        movie_id = i + 1
        description = generate_movie_description()
        # 生成向量
        embedding = model.encode(description).tolist()
        
        movie_data = {
            settings.PRIMARY_KEY_FIELD: movie_id,
            settings.TEXT_FIELD: description,
            settings.VECTOR_FIELD: embedding
        }
        all_data.append(movie_data)
        
        # 批量插入
        if len(all_data) >= batch_size:
            vector_db.insert_items(all_data)
            all_data.clear()
    
    # 插入剩余数据
    if all_data:
        vector_db.insert_items(all_data)
    
    logger.info(f"Data generation complete. Inserted {num_movies} movies.")
    # 注意：由于可能启用了异步更新，数据可能仍在队列中。在实际运行中，需要等待队列清空。
    # 这里为简化，我们假设同步插入或等待几秒。
    import time
    time.sleep(5)
    stats = vector_db.get_collection_stats()
    logger.info(f"Final collection stats: {stats}")

if __name__ == "__main__":
    generate_and_insert_movies(num_movies=5000)  # 生成5000部电影用于演示

"""
模拟高并发请求的客户端
用于压力测试和瓶颈分析。
"""
import asyncio
import aiohttp
import random
import time
from dataclasses import dataclass
from typing import List
import statistics
from loguru import logger

API_BASE_URL = "http://localhost:8000"  # 假设 FastAPI 运行在此地址

@dataclass
class TestResult:
    latency_ms: float
    status_code: int
    cached: bool

async def send_single_request(session: aiohttp.ClientSession, movie_id: int) -> TestResult:
    """发送单个推荐请求。"""
    url = f"{API_BASE_URL}/recommend"
    payload = {"movie_id": movie_id, "top_k": 10}
    start = time.perf_counter()
    try:
        async with session.post(url, json=payload) as resp:
            latency = (time.perf_counter() - start) * 1000
            resp_json = await resp.json()
            cached = resp_json.get('cached', False)
            return TestResult(latency_ms=latency, status_code=resp.status, cached=cached)
    except Exception as e:
        latency = (time.perf_counter() - start) * 1000
        logger.error(f"Request failed for movie {movie_id}: {e}")
        return TestResult(latency_ms=latency, status_code=0, cached=False)

async def run_concurrent_test(total_requests: int, concurrent_users: int, max_movie_id: int):
    """
    运行并发测试。
    total_requests: 总请求数
    concurrent_users: 并发用户数（协程数）
    max_movie_id: 电影ID范围上限
    """
    logger.info(f"Starting load test: {total_requests} reqs, {concurrent_users} concurrent users.")
    
    connector = aiohttp.TCPConnector(limit=0)  # 不限制连接数
    timeout = aiohttp.ClientTimeout(total=30)
    
    async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
        # 创建任务队列
        tasks = []
        request_count = 0
        
        # 使用信号量控制并发度
        semaphore = asyncio.Semaphore(concurrent_users)
        
        async def limited_request(movie_id):
            async with semaphore:
                return await send_single_request(session, movie_id)
        
        # 准备任务列表
        for _ in range(total_requests):
            movie_id = random.randint(1, max_movie_id)
            task = asyncio.create_task(limited_request(movie_id))
            tasks.append(task)
        
        # 等待所有任务完成并收集结果
        results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # 分析结果
    successful_results = []
    failed_count = 0
    cache_hits = 0
    
    for r in results:
        if isinstance(r, Exception):
            failed_count += 1
            continue
        if r.status_code == 200:
            successful_results.append(r)
            if r.cached:
                cache_hits += 1
        else:
            failed_count += 1
    
    # 计算指标
    latencies = [r.latency_ms for r in successful_results]
    if latencies:
        avg_latency = statistics.mean(latencies)
        p50_latency = statistics.median(latencies)
        p90_latency = np.percentile(latencies, 90) if len(latencies) >= 10 else None
        p99_latency = np.percentile(latencies, 99) if len(latencies) >= 100 else None
        throughput = len(successful_results) / (max(latencies) / 1000) if latencies else 0
    else:
        avg_latency = p50_latency = p90_latency = p99_latency = throughput = 0
    
    logger.info("\n" + "="*50)
    logger.info("LOAD TEST RESULTS")
    logger.info("="*50)
    logger.info(f"Total Requests: {total_requests}")
    logger.info(f"Successful: {len(successful_results)}")
    logger.info(f"Failed: {failed_count}")
    logger.info(f"Cache Hit Rate: {cache_hits/len(successful_results)*100:.1f}% ({cache_hits}/{len(successful_results)})")
    logger.info(f"Average Latency: {avg_latency:.2f} ms")
    logger.info(f"P50 Latency: {p50_latency:.2f} ms")
    logger.info(f"P90 Latency: {p90_latency:.2f} ms" if p90_latency else "P90: N/A")
    logger.info(f"P99 Latency: {p99_latency:.2f} ms" if p99_latency else "P99: N/A")
    logger.info(f"Estimated Throughput: {throughput:.2f} req/s")
    logger.info("="*50)
    
    return {
        "successful": len(successful_results),
        "failed": failed_count,
        "cache_hit_rate": cache_hits/len(successful_results) if successful_results else 0,
        "avg_latency_ms": avg_latency,
        "p99_latency_ms": p99_latency,
        "throughput_rps": throughput
    }

if __name__ == "__main__":
    import sys
    import numpy as np
    total = int(sys.argv[1]) if len(sys.argv) > 1 else 1000
    concurrent = int(sys.argv[2]) if len(sys.argv) > 2 else 50
    max_id = int(sys.argv[3]) if len(sys.argv) > 3 else 5000
    
    asyncio.run(run_concurrent_test(total, concurrent, max_id))

"""
性能测试与瓶颈分析脚本
对比不同配置下的性能表现。
"""
import subprocess
import time
import json
import asyncio
from simulation.load_client import run_concurrent_test
from loguru import logger
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # 用于无头环境
import matplotlib.pyplot as plt

def start_fastapi_server():
    """启动 FastAPI 服务器（在后台）。"""
    # 注意：这是一个示意函数。在实际测试中，你可能需要更复杂的进程管理。
    # 例如使用 subprocess.Popen，并确保在测试结束后终止进程。
    logger.info("Starting FastAPI server...")
    # 假设服务器已在另一个终端启动，或使用如下命令：
    # proc = subprocess.Popen(["uvicorn", "api.app:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "4"])
    # return proc
    return None

def run_performance_scenarios():
    """
    运行多个性能测试场景，对比不同优化策略。
    """
    scenarios = [
        {"name": "Baseline (FLAT index, No Cache)", "env_vars": {"INDEX_TYPE": "FLAT", "ENABLE_CACHE": "False"}, "workers": 1},
        {"name": "HNSW Index Only", "env_vars": {"INDEX_TYPE": "HNSW", "ENABLE_CACHE": "False"}, "workers": 1},
        {"name": "HNSW + Cache", "env_vars": {"INDEX_TYPE": "HNSW", "ENABLE_CACHE": "True"}, "workers": 1},
        {"name": "IVF_PQ Index + Cache", "env_vars": {"INDEX_TYPE": "IVF_PQ", "ENABLE_CACHE": "True"}, "workers": 1},
        {"name": "HNSW + Cache + 4 Workers", "env_vars": {"INDEX_TYPE": "HNSW", "ENABLE_CACHE": "True"}, "workers": 4},
    ]
    
    results = []
    base_env = {}
    
    for scenario in scenarios:
        logger.info(f"\n{'='*60}")
        logger.info(f"Running scenario: {scenario['name']}")
        logger.info('='*60)
        
        # 在实际中，这里需要重启服务以应用新的环境变量/配置。
        # 为简化，我们假设通过修改 config/settings.py 或环境变量，并重启服务。
        # 本示例中，我们直接运行客户端，并假设服务端配置已手动调整好。
        
        # 运行负载测试
        time.sleep(5)  # 等待服务稳定
        test_result = asyncio.run(
            run_concurrent_test(total_requests=2000, concurrent_users=100, max_movie_id=5000)
        )
        test_result['scenario'] = scenario['name']
        results.append(test_result)
    
    # 输出对比报告
    df = pd.DataFrame(results)
    print("\n\nPERFORMANCE COMPARISON REPORT")
    print("="*80)
    print(df.to_string(index=False))
    
    # 绘制对比图
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    ax = axes[0, 0]
    df.plot.bar(x='scenario', y='avg_latency_ms', ax=ax, legend=False, color='skyblue')
    ax.set_title('Average Latency (ms)')
    ax.set_ylabel('ms')
    ax.tick_params(axis='x', rotation=45)
    
    ax = axes[0, 1]
    df.plot.bar(x='scenario', y='p99_latency_ms', ax=ax, legend=False, color='lightcoral')
    ax.set_title('P99 Latency (ms)')
    ax.set_ylabel('ms')
    ax.tick_params(axis='x', rotation=45)
    
    ax = axes[1, 0]
    df.plot.bar(x='scenario', y='throughput_rps', ax=ax, legend=False, color='lightgreen')
    ax.set_title('Throughput (req/s)')
    ax.set_ylabel('requests per second')
    ax.tick_params(axis='x', rotation=45)
    
    ax = axes[1, 1]
    df.plot.bar(x='scenario', y='cache_hit_rate', ax=ax, legend=False, color='gold')
    ax.set_title('Cache Hit Rate')
    ax.set_ylabel('rate')
    ax.tick_params(axis='x', rotation=45)
    
    plt.tight_layout()
    plt.savefig('performance_comparison.png', dpi=150)
    logger.info("Performance comparison chart saved to 'performance_comparison.png'")

if __name__ == "__main__":
    run_performance_scenarios()

fastapi==0.104.1
uvicorn[standard]==0.24.0
pymilvus==2.3.6
sentence-transformers==2.2.2
numpy==1.24.3
pydantic==2.5.0
pydantic-settings==2.1.0
faker==19.6.2
loguru==0.7.2
aiohttp==3.9.1
tqdm==4.66.1
pandas==2.1.4
matplotlib==3.8.2
locust==2.20.1  # 可选，另一种负载测试工具

version: '3.5'

services:
  etcd:
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:

      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    volumes:

      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd

  minio:
    container_name: milvus-minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: minioadmin
      MINIO_SECRET_KEY: minioadmin
    volumes:

      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
    command: minio server /minio_data
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
      interval: 30s
      timeout: 20s
      retries: 3

  standalone:
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.3.6
    command: ["milvus", "run", "standalone"]
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:

      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
    ports:

      - "19530:19530"
      - "9091:9091"
    depends_on:

      - "etcd"
      - "minio"

networks:
  default:
    name: milvus

# 在项目根目录下
docker-compose up -d
# 等待所有服务就绪，可以通过 docker-compose logs -f 查看日志

python -m venv venv
# 在 Windows 上: venv\Scripts\activate
source venv/bin/activate  # 在 Linux/Mac 上

pip install -r requirements.txt

python -m simulation.data_generator

uvicorn api.app:app --host 0.0.0.0 --port 8000 --reload

python -m simulation.load_client 1000 50 5000
# 参数: <总请求数> <并发用户数> <最大电影ID>

python -m simulation.load_client 1000 50 5000

python -m tests.test_performance

# tests/test_basic.py
import sys
sys.path.append('.')
from core.recommendation_service import RecommendationService

def test_service_initialization():
    service = RecommendationService()
    assert service.vector_db is not None
    print("Service initialization test passed.")

if __name__ == "__main__":
    test_service_initialization()

# 健康检查
curl http://localhost:8000/health

# 获取推荐 (假设电影ID 42存在)
curl -X POST "http://localhost:8000/recommend" \
     -H "Content-Type: application/json" \
     -d '{"movie_id": 42, "top_k": 5}'

# 查看系统统计
curl http://localhost:8000/stats