flexllm 0.3.3__py3-none-any.whl

flexllm/async_api/core.py

@@ -0,0 +1,373 @@
+import asyncio
+from asyncio import Queue
+import itertools
+
+import time
+from typing import Any, Dict, Iterable, List, Optional, Callable, AsyncIterator, AsyncGenerator, Tuple, Union
+from aiohttp import ClientSession, TCPConnector, ClientTimeout
+from contextlib import asynccontextmanager
+
+from ..utils.core import async_retry
+from .interface import RequestResult
+from .progress import ProgressTracker, ProgressBarConfig
+
+from dataclasses import dataclass
+
+
+@dataclass
+class StreamingResult:
+    completed_requests: List[RequestResult]
+    progress: Optional[ProgressTracker]
+    is_final: bool
+
+
+class RateLimiter:
+    """
+    速率限制器
+
+    Args:
+        max_qps: 每秒最大请求数
+        use_bucket: 是否使用漏桶算法（aiolimiter），默认 True
+                    False 时使用简单的锁+sleep 实现
+    """
+
+    def __init__(self, max_qps: Optional[float] = None, use_bucket: bool = True):
+        self.max_qps = max_qps
+        self._use_bucket = use_bucket
+
+        if max_qps:
+            if use_bucket:
+                from aiolimiter import AsyncLimiter
+                self._limiter = AsyncLimiter(max_qps, 1)
+            else:
+                self._lock = asyncio.Lock()
+                self._last_request_time = 0
+                self._min_interval = 1 / max_qps
+
+    async def acquire(self):
+        if not self.max_qps:
+            return
+        if self._use_bucket:
+            await self._limiter.acquire()
+        else:
+            async with self._lock:
+                elapsed = time.time() - self._last_request_time
+                if elapsed < self._min_interval:
+                    await asyncio.sleep(self._min_interval - elapsed)
+                self._last_request_time = time.time()
+
+
+class ConcurrentRequester:
+    """
+    并发请求管理器
+
+    Example
+    -------
+
+    requester = ConcurrentRequester(
+        concurrency_limit=5,
+        max_qps=10,
+        timeout=0.7,
+    )
+
+    request_params = [
+        {
+            'json': {
+                'messages': [{"role": "user", "content": "讲个笑话" }],
+                'model': "qwen2.5:latest",
+            },
+            'headers': {'Content-Type': 'application/json'}
+        } for i in range(10)
+    ]
+
+    # 执行并发请求
+    results, tracker = await requester.process_requests(
+        request_params=request_params,
+        url='http://localhost:11434/v1/chat/completions',
+        method='POST',
+        show_progress=True
+    )
+    """
+
+    def __init__(
+            self,
+            concurrency_limit: int,
+            max_qps: Optional[float] = None,
+            timeout: Optional[float] = None,
+            retry_times: int = 3,
+            retry_delay: float = 0.3
+    ):
+        self._concurrency_limit = concurrency_limit
+        if timeout:
+            self._timeout = ClientTimeout(total=timeout, connect=min(10., timeout))
+        else:
+            self._timeout = None
+        self._rate_limiter = RateLimiter(max_qps)
+        self._semaphore = asyncio.Semaphore(concurrency_limit)
+        self.retry_times = retry_times
+        self.retry_delay = retry_delay
+
+    @asynccontextmanager
+    async def _get_session(self):
+        connector = TCPConnector(limit=self._concurrency_limit+10, limit_per_host=0, force_close=False)
+        async with ClientSession(timeout=self._timeout, connector=connector, trust_env=True) as session:
+            yield session
+
+    @staticmethod
+    async def _make_requests( session: ClientSession,method: str, url: str,  **kwargs):
+        async with session.request(method, url, **kwargs) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return response, data
+
+    async def make_requests(self, session: ClientSession,method: str, url: str,  **kwargs):
+        return await async_retry(self.retry_times, self.retry_delay)(self._make_requests)(session,method, url, **kwargs)
+
+    async def _send_single_request(
+            self,
+            session: ClientSession,
+            request_id: int,
+            url: str,
+            method: str = 'POST',
+            meta: dict = None,
+            **kwargs
+    ) -> RequestResult:
+        """发送单个请求"""
+        async with self._semaphore:
+            try:
+                # todo: 速率限制也许需要优化
+                await self._rate_limiter.acquire()
+
+                start_time = time.time()
+                response, data = await self.make_requests(session, method, url, **kwargs)
+                latency = time.time() - start_time
+
+                if response.status != 200:
+                    error_info = {
+                        'status_code': response.status,
+                        'response_data': data,
+                        'error': f"HTTP {response.status}"
+                    }
+                    return RequestResult(
+                        request_id=request_id,
+                        data=error_info,
+                        status='error',
+                        meta=meta,
+                        latency=latency
+                    )
+
+                return RequestResult(
+                    request_id=request_id,
+                    data=data,
+                    status="success",
+                    meta=meta,
+                    latency=latency
+                )
+
+            except asyncio.TimeoutError as e:
+                return RequestResult(
+                    request_id=request_id,
+                    data={'error': 'Timeout error', 'detail': str(e)},
+                    status='error',
+                    meta=meta,
+                    latency=time.time() - start_time
+                )
+            except Exception as e:
+                return RequestResult(
+                    request_id=request_id,
+                    data={'error': e.__class__.__name__, 'detail': str(e)},
+                    status='error',
+                    meta=meta,
+                    latency=time.time() - start_time
+                )
+
+    async def process_with_concurrency_window(
+            self,
+            items: Iterable,
+            process_func: Callable,
+            concurrency_limit: int,
+            progress: Optional[ProgressTracker] = None,
+            batch_size: int = 1,
+    ) -> AsyncGenerator[StreamingResult, Any]:
+        """
+        使用滑动窗口方式处理并发任务，支持流式返回结果
+
+        Args:
+            items: 待处理的项目迭代器
+            process_func: 处理单个项目的异步函数，接收item和项目item_id作为参数
+            concurrency_limit: 并发限制数量,也是窗口大小
+            progress: 可选的进度跟踪器
+            batch_size: 每次yield返回的最小完成请求数量
+
+        Yields:
+             生成 StreamingResult 对象序列
+        """
+
+        async def handle_completed_tasks(done_tasks, batch, is_final=False):
+            """内部函数处理已完成的任务"""
+            for task in done_tasks:
+                result = await task
+                if progress:
+                    progress.update(result)
+                batch.append(result)
+
+            if len(batch) >= batch_size or (is_final and batch):
+                if is_final and progress:
+                    progress.summary()
+                yield StreamingResult(
+                    completed_requests=sorted(batch, key=lambda x: x.request_id),
+                    progress=progress,
+                    is_final=is_final
+                )
+                batch.clear()
+
+        item_id = 0
+        active_tasks = set()
+        completed_batch = []
+
+        # 处理输入项目
+        for item in items:
+            if len(active_tasks) >= concurrency_limit:
+                done, active_tasks = await asyncio.wait(
+                    active_tasks,
+                    return_when=asyncio.FIRST_COMPLETED
+                )
+                async for result in handle_completed_tasks(done, completed_batch):
+                    yield result
+
+            active_tasks.add(asyncio.create_task(process_func(item, item_id)))
+            item_id += 1
+
+        # 处理剩余任务
+        if active_tasks:
+            done, _ = await asyncio.wait(active_tasks)
+            async for result in handle_completed_tasks(done, completed_batch, is_final=True):
+                yield result
+
+    async def _stream_requests(
+            self,
+            queue: Queue,
+            request_params: Iterable[Dict[str, Any]],
+            url: str,
+            method: str = 'POST',
+            total_requests: Optional[int] = None,
+            show_progress: bool = True,
+            batch_size: Optional[int] = None
+    ) :
+        """
+        流式处理批量请求，实时返回已完成的结果
+
+        Args:
+            request_params: 请求参数列表
+            url: 请求URL
+            method: 请求方法
+            total_requests: 总请求数量
+            show_progress: 是否显示进度
+            batch_size: 每次yield返回的最小完成请求数量
+        """
+        progress = None
+        if batch_size is None:
+            batch_size = self._concurrency_limit
+        if total_requests is None and show_progress:
+            request_params, params_for_counting = itertools.tee(request_params)
+            total_requests = sum(1 for _ in params_for_counting)
+
+        if show_progress and total_requests is not None:
+            progress = ProgressTracker(
+                total_requests,
+                concurrency=self._concurrency_limit,
+                config=ProgressBarConfig()
+            )
+
+        async with self._get_session() as session:
+            async for result in self.process_with_concurrency_window(
+                    items=request_params,
+                    process_func=lambda params, request_id: self._send_single_request(
+                        session=session,
+                        request_id=request_id,
+                        url=url,
+                        method=method,
+                        meta=params.pop('meta', None),
+                        **params
+                    ),
+                    concurrency_limit=self._concurrency_limit,
+                    progress=progress,
+                    batch_size=batch_size,
+            ):
+                await queue.put(result)
+
+        await queue.put(None)
+
+    async def aiter_stream_requests(self,
+                                  request_params: Iterable[Dict[str, Any]],
+                                  url: str,
+                                  method: str = 'POST',
+                                  total_requests: Optional[int] = None,
+                                  show_progress: bool = True,
+                                  batch_size: Optional[int] = None
+                                  ):
+        queue = Queue()
+        task = asyncio.create_task(self._stream_requests(queue,
+                                                         request_params=request_params,
+                                                         url=url,
+                                                         method=method,
+                                                         total_requests=total_requests,
+                                                         show_progress=show_progress,
+                                                         batch_size=batch_size))
+        try:
+            while True:
+                result = await queue.get()
+                if result is None:
+                    break
+                yield result
+        finally:
+            if not task.done():
+                task.cancel()
+
+
+    async def process_requests(
+            self,
+            request_params: Iterable[Dict[str, Any]],
+            url: str,
+            method: str = 'POST',
+            total_requests: Optional[int] = None,
+            show_progress: bool = True
+    ) -> Tuple[List[RequestResult], Optional[ProgressTracker]]:
+        """
+        处理批量请求
+
+        Returns:
+            Tuple[list[RequestResult], Optional[ProgressTracker]]:
+            请求结果列表和进度跟踪器（如果启用了进度显示）
+        """
+        progress = None
+        if total_requests is None and show_progress:
+            request_params, params_for_counting = itertools.tee(request_params)
+            total_requests = sum(1 for _ in params_for_counting)
+
+        if show_progress and total_requests is not None:
+            progress = ProgressTracker(
+                total_requests,
+                concurrency=self._concurrency_limit,
+                config=ProgressBarConfig()
+            )
+
+        results = []
+        async with self._get_session() as session:
+            async for result in self.process_with_concurrency_window(
+                items=request_params,
+                process_func=lambda params, request_id: self._send_single_request(
+                    session=session,
+                    request_id=request_id,
+                    url=url,
+                    method=method,
+                    meta=params.pop('meta', None),
+                    **params
+                ),
+                concurrency_limit=self._concurrency_limit,
+                progress=progress,
+            ):
+                results.extend(result.completed_requests)
+        # sort
+        results = sorted(results, key=lambda x: x.request_id)
+        return results, progress

flexllm/async_api/interface.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass
+class RequestResult:
+    """请求结果的数据类"""
+    request_id: int
+    data: Any
+    status: str
+    latency: float
+    meta: dict = None

flexllm/async_api/progress.py

@@ -0,0 +1,277 @@
+from enum import Enum
+from rich.console import Console
+from rich.markdown import Markdown
+from typing import Dict, List, Optional, TYPE_CHECKING
+import time
+from dataclasses import dataclass
+import statistics
+from collections import defaultdict
+
+if TYPE_CHECKING:
+    from .interface import RequestResult
+
+
+class ProgressBarStyle(Enum):
+    SOLID = ("█", "─", "⚡")  # 实心样式
+    BLANK = ("▉", " ", "⚡")
+    GRADIENT = ("▰", "▱", "⚡")  # 渐变样式
+    BLOCKS = ("▣", "▢", "⚡")  # 方块样式
+    ARROW = ("━", "─", "⚡")  # 箭头样式
+    DOTS = ("⣿", "⣀", "⚡")  # 点状样式
+    PIPES = ("┃", "┆", "⚡")  # 管道样式
+    STARS = ("★", "☆", "⚡")  # 星星样式
+
+
+@dataclass
+class ProgressBarConfig:
+    bar_length: int = 30
+    show_percentage: bool = True
+    show_speed: bool = True
+    show_counts: bool = True
+    show_time_stats: bool = True
+    style: ProgressBarStyle = ProgressBarStyle.BLANK
+    use_colors: bool = True
+
+
+class ProgressTracker:
+    def __init__(
+        self,
+        total_requests: int,
+        concurrency=1,
+        config: Optional[ProgressBarConfig] = None,
+    ):
+        self.console = Console()
+
+        # 统计信息
+        self.success_count = 0
+        self.error_count = 0
+        self.latencies: List[float] = []
+        self.errors: Dict[str, int] = defaultdict(int)  # 统计不同类型的错误
+
+        self.total_requests = total_requests
+        self.concurrency = concurrency
+        self.config = config or ProgressBarConfig()
+        self.completed_requests = 0
+        self.success_count = 0
+        self.error_count = 0
+        self.start_time = time.time()
+        self.latencies = []
+        self.errors = {}
+        self.last_speed_update = time.time()
+        self.recent_latencies = []  # 用于计算实时速度
+
+        # ANSI颜色代码
+        self.colors = {
+            "green": "\033[92m",
+            "yellow": "\033[93m",
+            "red": "\033[91m",
+            "blue": "\033[94m",
+            "purple": "\033[95m",
+            "cyan": "\033[96m",
+            "reset": "\033[0m",
+        }
+
+    def _format_time(self, seconds: float) -> str:
+        """格式化时间显示"""
+        if seconds > 3600:
+            return f"{seconds / 3600:.1f}h"
+        if seconds > 60:
+            return f"{seconds / 60:.1f}m"
+        return f"{seconds:.1f}s"
+
+    @staticmethod
+    def _format_speed(speed: float) -> str:
+        """格式化速度显示"""
+        # if speed >= 1:
+        return f"{speed:.1f} req/s"
+        # return f'{speed*1000:.0f} req/ms'
+
+    def _get_colored_text(self, text: str, color: str) -> str:
+        """添加颜色到文本"""
+        if self.config.use_colors:
+            return f"{self.colors[color]}{text}{self.colors['reset']}"
+        return text
+
+    def _calculate_speed(self) -> float:
+        """计算实际吞吐量（已完成请求数 / 已用时间）"""
+        elapsed = time.time() - self.start_time
+        if elapsed <= 0:
+            return 0
+        return self.completed_requests / elapsed
+
+    def update(self, result: "RequestResult") -> None:
+        """
+        更新进度和统计信息
+
+        Args:
+            result: 请求结果
+        """
+        self.completed_requests += 1
+        self.latencies.append(result.latency)
+        self.recent_latencies.append(result.latency)
+
+        # 只保留最近的30个请求用于计算速度
+        if len(self.recent_latencies) > 30:
+            self.recent_latencies.pop(0)
+
+        if result.status == "success":
+            self.success_count += 1
+        else:
+            self.error_count += 1
+            # 安全地获取错误类型，处理 result.data 为 None 的情况
+            if result.data and isinstance(result.data, dict):
+                error_type = result.data.get("error", "unknown")
+            else:
+                error_type = "unknown"
+            self.errors[error_type] = self.errors.get(error_type, 0) + 1
+
+        current_time = time.time()
+        total_time = current_time - self.start_time
+        progress = self.completed_requests / self.total_requests
+
+        # 计算统计信息
+        speed = self._calculate_speed()
+        avg_latency = statistics.mean(self.latencies) if self.latencies else 0
+        remaining_requests = self.total_requests - self.completed_requests
+        estimated_remaining_time = (
+            avg_latency * remaining_requests / self.concurrency
+            if avg_latency > 0
+            else 0
+        )
+
+        # 创建进度条
+        style = self.config.style.value
+        filled_length = int(self.config.bar_length * progress)
+        bar = style[0] * filled_length + style[1] * (
+            self.config.bar_length - filled_length
+        )
+
+        # 构建输出组件
+        components = []
+
+        # 进度条和百分比
+        progress_text = f"[{self._get_colored_text(bar, 'blue')}]"
+        if self.config.show_percentage:
+            progress_text += (
+                f" {self._get_colored_text(f'{progress * 100:.1f}%', 'green')}"
+            )
+        components.append(progress_text)
+
+        # 请求计数
+        if self.config.show_counts:
+            counts = f"({self.completed_requests}/{self.total_requests})"
+            components.append(self._get_colored_text(counts, "yellow"))
+
+        # 速度信息
+        if self.config.show_speed:
+            speed_text = f"{style[2]} {self._format_speed(speed)}"
+            components.append(self._get_colored_text(speed_text, "cyan"))
+
+        # 时间统计
+        if self.config.show_time_stats:
+            time_stats = (
+                f"avg: {self._format_time(avg_latency)} "
+                f"total: {self._format_time(total_time)} "
+                f"eta: {self._format_time(estimated_remaining_time)}"
+            )
+            components.append(self._get_colored_text(time_stats, "purple"))
+
+        # 打印进度 - 修复Windows编码问题
+        try:
+            print("\r" + " ".join(components), end="", flush=True)
+        except UnicodeEncodeError:
+            # Windows GBK编码兼容处理
+            safe_components = []
+            for comp in components:
+                # 替换有问题的Unicode字符
+                safe_comp = comp.replace("⚡", "*").replace("█", "#").replace("─", "-")
+                safe_comp = (
+                    safe_comp.replace("▉", "|").replace("▰", "=").replace("▱", "-")
+                )
+                safe_comp = (
+                    safe_comp.replace("▣", "[").replace("▢", "]").replace("━", "=")
+                )
+                safe_comp = (
+                    safe_comp.replace("┃", "|")
+                    .replace("┆", ":")
+                    .replace("★", "*")
+                    .replace("☆", "+")
+                )
+                safe_comp = safe_comp.replace("⣿", "#").replace("⣀", ".")
+                safe_components.append(safe_comp)
+            print("\r" + " ".join(safe_components), end="", flush=True)
+
+    def summary(self, show_p999=False, print_to_console=True) -> str:
+        """打印请求汇总信息"""
+        total_time = time.time() - self.start_time
+        avg_latency = sum(self.latencies) / len(self.latencies) if self.latencies else 0
+        throughput = self.success_count / total_time if total_time > 0 else 0
+
+        # 计算延迟分位数
+        sorted_latencies = sorted(self.latencies)
+        p50 = p95 = p99 = 0
+        if sorted_latencies:
+            p50 = sorted_latencies[int(len(sorted_latencies) * 0.5)]
+            p95 = sorted_latencies[int(len(sorted_latencies) * 0.95)]
+            p99 = sorted_latencies[int(len(sorted_latencies) * 0.99)]
+            p995 = sorted_latencies[int(len(sorted_latencies) * 0.995)]
+            p999 = sorted_latencies[int(len(sorted_latencies) * 0.999)]
+
+        p99_str = f"> - P99 延迟: {p99:.2f} 秒"
+        p999_str = f"""> - P99 延迟: {p99:.2f} 秒
+> - P995 延迟: {p995:.2f} 秒
+> - P999 延迟: {p999:.2f} 秒"""
+        p99_or_p999_str = p999_str if show_p999 else p99_str
+
+        summary = f"""
+                                   请求统计                                    
+
+| 总体情况                                                                   
+|  - 总请求数: {self.total_requests}                                                             
+|  - 成功请求数: {self.success_count}                                                           
+|  - 失败请求数: {self.error_count}                                                           
+|  - 成功率: {(self.success_count / self.total_requests * 100):.2f}%                                                         
+
+| 性能指标                                                                   
+|  - 平均延迟: {avg_latency:.2f} 秒                                                       
+|  - P50 延迟: {p50:.2f} 秒                                                       
+|  - P95 延迟: {p95:.2f} 秒                                                       
+|  - P99 延迟: {p99:.2f} 秒                                                       
+|  - 吞吐量: {throughput:.2f} 请求/秒                                                
+|  - 总运行时间: {total_time:.2f} 秒                                               
+
+"""
+        # 如果有错误，添加错误统计
+        if self.errors:
+            summary += "| 错误分布                                                                   \n"
+            for error_type, count in self.errors.items():
+                percentage = count / self.error_count * 100
+                summary += f"|  - {error_type}: {count} ({percentage:.1f}%)                            \n"
+
+        summary += "-" * 76
+        if print_to_console:
+            print()  # 打印空行
+            try:
+                # 尝试使用Rich输出，如果失败则使用普通print
+                self.console.print(summary)
+            except UnicodeEncodeError:
+                # 在Windows GBK环境下，如果出现编码错误，使用普通print
+                print(summary)
+        return summary
+
+
+if __name__ == "__main__":
+    from .interface import RequestResult
+
+    config = ProgressBarConfig()
+    tracker = ProgressTracker(100, 1, config)
+    for i in range(100):
+        time.sleep(0.1)
+        tracker.update(
+            result=RequestResult(
+                request_id=i,
+                data=None,
+                status="success",
+                latency=0.1,
+            )
+        )