Coverage for agentos/benchmarks/runner.py: 47%

1"""v0.80 — 性能基准测试运行器：延迟/吞吐/并发。""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import json 

7import time 

8from collections.abc import Awaitable, Callable 

9from dataclasses import dataclass, field 

10from typing import Any 

11 

12 

13@dataclass 

14class BenchmarkScenario: 

15 """单个基准测试场景。""" 

16 

17 name: str 

18 description: str = "" 

19 setup: Callable[[], Any] | None = None 

20 teardown: Callable[[Any], None] | None = None 

21 

22 

23@dataclass 

24class BenchmarkConfig: 

25 """基准测试配置。""" 

26 

27 warmup_iterations: int = 3 

28 measure_iterations: int = 10 

29 concurrency_levels: list[int] = field(default_factory=lambda: [1, 4, 8]) 

30 timeout_per_run: float = 30.0 

31 

32 

33@dataclass 

34class _LatencyStats: 

35 min_ms: float = 0 

36 max_ms: float = 0 

37 avg_ms: float = 0 

38 p50_ms: float = 0 

39 p95_ms: float = 0 

40 p99_ms: float = 0 

41 

42 @staticmethod 

43 def compute(latencies_ms: list[float]) -> _LatencyStats: 

44 if not latencies_ms: 

45 return _LatencyStats() 

46 s = sorted(latencies_ms) 

47 n = len(s) 

48 return _LatencyStats( 

49 min_ms=s[0], 

50 max_ms=s[-1], 

51 avg_ms=sum(s) / n, 

52 p50_ms=s[int(n * 0.5)], 

53 p95_ms=s[int(n * 0.95)] if int(n * 0.95) < n else s[-1], 

54 p99_ms=s[int(n * 0.99)] if int(n * 0.99) < n else s[-1], 

55 ) 

56 

57 

58@dataclass 

59class BenchmarkReport: 

60 """基准测试报告。""" 

61 

62 scenario: str = "" 

63 description: str = "" 

64 config: BenchmarkConfig = field(default_factory=BenchmarkConfig) 

65 results: list[dict[str, Any]] = field(default_factory=list) 

66 summary: str = "" 

67 

68 def to_json(self) -> str: 

69 return json.dumps( 

70 { 

71 "scenario": self.scenario, 

72 "description": self.description, 

73 "results": self.results, 

74 "summary": self.summary, 

75 }, 

76 indent=2, 

77 ensure_ascii=False, 

78 ) 

79 

80 def to_markdown(self) -> str: 

81 lines = [ 

82 f"# Benchmark: {self.scenario}", 

83 "", 

84 f"_{self.description}_", 

85 "", 

86 "| 并发 | 总调用 | 总耗时(s) | 吞吐(QPS) | 平均延迟(ms) | P50(ms) | P95(ms) | P99(ms) | 成功率 |", 

87 "|------|--------|-----------|-----------|-------------|---------|---------|---------|--------|", 

88 ] 

89 for r in self.results: 

90 lines.append( 

91 f"| {r['concurrency']} | {r['total_calls']} | {r['total_time_s']:.2f} | " 

92 f"{r['throughput_qps']:.1f} | {r['latency_stats']['avg_ms']:.1f} | " 

93 f"{r['latency_stats']['p50_ms']:.1f} | {r['latency_stats']['p95_ms']:.1f} | " 

94 f"{r['latency_stats']['p99_ms']:.1f} | {r['success_rate']*100:.0f}% |" 

95 ) 

96 if self.summary: 

97 lines.extend(["", f"> {self.summary}"]) 

98 return "\n".join(lines) 

99 

100 

101class BenchmarkRunner: 

102 """基准测试运行器。""" 

103 

104 def __init__(self, config: BenchmarkConfig | None = None): 

105 self.config = config or BenchmarkConfig() 

106 

107 async def run( 

108 self, 

109 scenario: BenchmarkScenario, 

110 callable_fn: Callable[[], Any], 

111 async_callable_fn: Callable[[], Awaitable[Any]] | None = None, 

112 ) -> BenchmarkReport: 

113 """运行基准测试。 

114 

115 Args: 

116 scenario: 测试场景。 

117 callable_fn: 同步测试函数。 

118 async_callable_fn: 异步测试函数（用于并发测试）。 

119 """ 

120 setup_state = scenario.setup() if scenario.setup else None 

121 

122 results: list[dict[str, Any]] = [] 

123 

124 for concurrency in self.config.concurrency_levels: 

125 total_calls = concurrency * self.config.measure_iterations 

126 total_start = time.perf_counter() 

127 success = 0 

128 latencies_ms: list[float] = [] 

129 

130 async def _one_call(): 

131 nonlocal success 

132 t0 = time.perf_counter() 

133 try: 

134 if async_callable_fn: 

135 await async_callable_fn() 

136 else: 

137 callable_fn() 

138 success += 1 

139 except Exception: 

140 pass 

141 latencies_ms.append((time.perf_counter() - t0) * 1000) 

142 

143 # warmup 

144 for _ in range(self.config.warmup_iterations): 

145 try: 

146 callable_fn() 

147 except Exception: 

148 pass 

149 

150 # measure 

151 tasks = [_one_call() for _ in range(total_calls)] 

152 await asyncio.gather(*tasks) 

153 

154 total_time = time.perf_counter() - total_start 

155 stats = _LatencyStats.compute(latencies_ms) 

156 

157 results.append( 

158 { 

159 "concurrency": concurrency, 

160 "total_calls": total_calls, 

161 "total_time_s": round(total_time, 3), 

162 "throughput_qps": round(total_calls / total_time, 1) if total_time > 0 else 0, 

163 "latency_stats": { 

164 "min_ms": round(stats.min_ms, 2), 

165 "max_ms": round(stats.max_ms, 2), 

166 "avg_ms": round(stats.avg_ms, 2), 

167 "p50_ms": round(stats.p50_ms, 2), 

168 "p95_ms": round(stats.p95_ms, 2), 

169 "p99_ms": round(stats.p99_ms, 2), 

170 }, 

171 "success_rate": round(success / total_calls, 4) if total_calls else 0, 

172 } 

173 ) 

174 

175 if scenario.teardown and setup_state is not None: 

176 scenario.teardown(setup_state) 

177 

178 avg_throughput = sum(r["throughput_qps"] for r in results) / len(results) if results else 0 

179 return BenchmarkReport( 

180 scenario=scenario.name, 

181 description=scenario.description, 

182 config=self.config, 

183 results=results, 

184 summary=f"平均吞吐: {avg_throughput:.1f} QPS | 并发级别: {self.config.concurrency_levels}", 

185 ) 

186 

187 

188async def run_benchmark( 

189 scenario_name: str, 

190 callable_fn: Callable[[], Any], 

191 config: BenchmarkConfig | None = None, 

192) -> BenchmarkReport: 

193 """便捷函数：运行一次基准测试并返回 Markdown 报告。""" 

194 runner = BenchmarkRunner(config) 

195 scenario = BenchmarkScenario(name=scenario_name) 

196 return await runner.run(scenario, callable_fn)