add instrumented benchmarking with timing breakdown · zzstoatzz.io/prefect-server@a8b6d48

+194 -159

scripts/benchmark

··· 1 1 #!/usr/bin/env -S uv run --script --quiet 2 2 # /// script 3 3 # requires-python = ">=3.12" 4 - # dependencies = ["httpx", "rich"] 4 + # dependencies = ["httpx", "rich", "psutil"] 5 5 # /// 6 6 """ 7 - benchmark runner for prefect-server 7 + Instrumented benchmark runner for prefect-server. 8 8 9 - compares performance between zig and python prefect servers by running 10 - the same workload against each and measuring execution time. 9 + Compares zig vs python servers with: 10 + - Timing breakdown by test section 11 + - Memory usage tracking 12 + - Request count per section 11 13 12 - usage: 13 - ./scripts/benchmark --server zig test_flow.py 14 - ./scripts/benchmark --server python test_flow.py 15 - ./scripts/benchmark --compare test_flow.py 16 - ./scripts/benchmark --compare scripts/test-with-client 17 - 18 - the workload script must respect PREFECT_API_URL environment variable. 14 + Usage: 15 + ./scripts/benchmark --server zig 16 + ./scripts/benchmark --server python 17 + ./scripts/benchmark --compare 18 + ./scripts/benchmark --compare --iterations 3 19 19 """ 20 20 21 21 import argparse 22 + import json 22 23 import os 23 24 import shutil 24 25 import signal 25 26 import subprocess 26 27 import sys 27 28 import time 28 - from dataclasses import dataclass 29 + from dataclasses import dataclass, field 29 30 from pathlib import Path 30 31 from typing import Literal 31 32 32 33 import httpx 34 + import psutil 33 35 from rich.console import Console 34 36 from rich.table import Table 35 37 ··· 38 40 SERVER_PORT = 4200 39 41 API_URL = f"http://localhost:{SERVER_PORT}/api" 40 42 HEALTH_ENDPOINT = f"{API_URL}/health" 43 + 44 + 45 + @dataclass 46 + class SectionResult: 47 + name: str 48 + duration_ms: float 49 + requests: int 50 + passed: bool 41 51 42 52 43 53 @dataclass 44 54 class BenchmarkResult: 45 55 server: str 46 - workload: str 47 - duration_ms: float 48 56 success: bool 57 + total_duration_ms: float 58 + total_requests: int 59 + memory_mb: float 60 + sections: list[SectionResult] = field(default_factory=list) 49 61 error: str | None = None 50 62 51 63 52 64 def wait_for_health(timeout: float = 30.0) -> bool: 53 - """wait for server health endpoint to respond""" 65 + """Wait for server health endpoint.""" 54 66 start = time.time() 55 67 while time.time() - start < timeout: 56 68 try: ··· 64 76 65 77 66 78 def kill_port(port: int) -> None: 67 - """kill any process listening on the given port""" 79 + """Kill any process on the given port.""" 68 80 try: 69 - result = subprocess.run( 70 - ["lsof", "-ti", f":{port}"], 71 - capture_output=True, 72 - text=True, 73 - ) 81 + result = subprocess.run(["lsof", "-ti", f":{port}"], capture_output=True, text=True) 74 82 if result.stdout.strip(): 75 - pids = result.stdout.strip().split("\n") 76 - for pid in pids: 83 + for pid in result.stdout.strip().split("\n"): 77 84 try: 78 85 os.kill(int(pid), signal.SIGTERM) 79 86 except (ProcessLookupError, ValueError): ··· 83 90 pass 84 91 85 92 93 + def get_process_memory_mb(pid: int) -> float: 94 + """Get memory usage of a process in MB.""" 95 + try: 96 + proc = psutil.Process(pid) 97 + return proc.memory_info().rss / (1024 * 1024) 98 + except (psutil.NoSuchProcess, psutil.AccessDenied): 99 + return 0.0 100 + 101 + 86 102 class ServerManager: 87 - """manages starting and stopping prefect servers""" 103 + """Manages server lifecycle.""" 88 104 89 105 def __init__(self, server_type: Literal["zig", "python"]): 90 106 self.server_type = server_type ··· 92 108 self.project_root = Path(__file__).parent.parent 93 109 94 110 def start(self) -> bool: 95 - """start the server, returns True if successful""" 111 + """Start the server.""" 96 112 kill_port(SERVER_PORT) 97 113 98 114 if self.server_type == "zig": 99 115 return self._start_zig() 100 - else: 101 - return self._start_python() 116 + return self._start_python() 102 117 103 118 def _start_zig(self) -> bool: 104 - """start the zig server""" 105 119 binary = self.project_root / "zig-out" / "bin" / "prefect-server" 106 120 if not binary.exists(): 107 121 console.print(f"[red]error: zig binary not found at {binary}[/red]") 108 - console.print("[yellow]run 'zig build' first[/yellow]") 109 122 return False 110 123 111 - # clean database for consistent benchmarks 124 + # clean db 112 125 db_path = self.project_root / "prefect.db" 113 126 if db_path.exists(): 114 127 db_path.unlink() ··· 130 143 return True 131 144 132 145 def _start_python(self) -> bool: 133 - """start the python prefect server""" 134 - # check if uvx is available 135 146 if not shutil.which("uvx"): 136 147 console.print("[red]error: uvx not found[/red]") 137 148 return False 138 149 139 - # use a temp database for python server too 140 150 temp_db = self.project_root / ".benchmark-python.db" 141 151 if temp_db.exists(): 142 152 temp_db.unlink() ··· 153 163 stderr=subprocess.DEVNULL, 154 164 ) 155 165 156 - # python server takes longer to start 157 166 if not wait_for_health(timeout=60.0): 158 167 self.stop() 159 168 return False 160 169 return True 161 170 171 + def get_memory_mb(self) -> float: 172 + """Get server memory usage.""" 173 + if self.process: 174 + return get_process_memory_mb(self.process.pid) 175 + return 0.0 176 + 162 177 def stop(self) -> None: 163 - """stop the server""" 164 178 if self.process: 165 179 self.process.terminate() 166 180 try: ··· 168 182 except subprocess.TimeoutExpired: 169 183 self.process.kill() 170 184 self.process = None 171 - 172 - # ensure port is free 173 185 kill_port(SERVER_PORT) 174 186 175 187 176 - def run_workload(workload_path: str, project_root: Path) -> tuple[float, bool, str | None]: 177 - """ 178 - run a workload script and return (duration_ms, success, error_message) 179 - """ 180 - path = Path(workload_path) 181 - if not path.is_absolute(): 182 - path = project_root / path 183 - if not path.exists(): 184 - return 0, False, f"workload not found: {workload_path}" 188 + def run_workload(project_root: Path) -> dict | None: 189 + """Run test-api-sequence and return JSON results.""" 190 + workload = project_root / "scripts" / "test-api-sequence" 185 191 186 192 env = os.environ.copy() 187 193 env["PREFECT_API_URL"] = API_URL 188 194 env["PREFECT_LOGGING_LEVEL"] = "WARNING" 189 - 190 - # clear uv's virtual env vars so nested uv run uses the project's venv 191 195 env.pop("VIRTUAL_ENV", None) 192 196 env.pop("UV_RUN_RECURSION_DEPTH", None) 193 197 194 - # determine how to run the script 195 - if path.suffix == ".py": 196 - # run in project context so venv dependencies are available 197 - cmd = ["uv", "run", "--project", str(project_root), "python", str(path)] 198 - else: 199 - # assume it's an executable script (like scripts/test-with-client) 200 - cmd = [str(path)] 201 - 202 - start = time.perf_counter() 203 198 try: 204 199 result = subprocess.run( 205 - cmd, 200 + [str(workload), "--json"], 206 201 cwd=project_root, 207 202 env=env, 208 203 capture_output=True, 209 204 text=True, 210 - timeout=300, # 5 minute timeout 205 + timeout=300, 211 206 ) 212 - duration_ms = (time.perf_counter() - start) * 1000 213 207 214 208 if result.returncode != 0: 215 - return duration_ms, False, result.stderr or result.stdout 216 - return duration_ms, True, None 209 + # try to parse JSON anyway (might have partial results) 210 + pass 217 211 218 - except subprocess.TimeoutExpired: 219 - duration_ms = (time.perf_counter() - start) * 1000 220 - return duration_ms, False, "timeout after 5 minutes" 221 - except Exception as e: 222 - duration_ms = (time.perf_counter() - start) * 1000 223 - return duration_ms, False, str(e) 212 + # find JSON in output (might have other output mixed in) 213 + for line in result.stdout.strip().split("\n"): 214 + if line.startswith("{"): 215 + return json.loads(line) 224 216 217 + return None 225 218 226 - def benchmark_single( 227 - server_type: Literal["zig", "python"], 228 - workload: str, 229 - iterations: int = 1, 230 - ) -> list[BenchmarkResult]: 231 - """run benchmark against a single server""" 219 + except (subprocess.TimeoutExpired, json.JSONDecodeError) as e: 220 + console.print(f"[red]workload error: {e}[/red]") 221 + return None 222 + 223 + 224 + def benchmark_single(server_type: Literal["zig", "python"], iterations: int = 1) -> list[BenchmarkResult]: 225 + """Run benchmark against a single server.""" 232 226 results = [] 233 227 234 228 console.print(f"\n[bold]benchmarking {server_type} server[/bold]") ··· 244 238 if not manager.start(): 245 239 results.append(BenchmarkResult( 246 240 server=server_type, 247 - workload=workload, 248 - duration_ms=0, 249 241 success=False, 242 + total_duration_ms=0, 243 + total_requests=0, 244 + memory_mb=0, 250 245 error="failed to start server", 251 246 )) 252 247 console.print("[red]failed[/red]") ··· 254 249 255 250 console.print("running workload...", end=" ") 256 251 257 - duration_ms, success, error = run_workload(workload, manager.project_root) 252 + # measure memory before workload 253 + memory_before = manager.get_memory_mb() 254 + 255 + data = run_workload(manager.project_root) 256 + 257 + # measure memory after workload 258 + memory_after = manager.get_memory_mb() 259 + memory_mb = max(memory_before, memory_after) 260 + 258 261 manager.stop() 259 262 263 + if data is None: 264 + results.append(BenchmarkResult( 265 + server=server_type, 266 + success=False, 267 + total_duration_ms=0, 268 + total_requests=0, 269 + memory_mb=memory_mb, 270 + error="workload failed", 271 + )) 272 + console.print("[red]failed[/red]") 273 + continue 274 + 275 + sections = [ 276 + SectionResult( 277 + name=s["name"], 278 + duration_ms=s["duration_ms"], 279 + requests=s["requests"], 280 + passed=s["passed"], 281 + ) 282 + for s in data.get("sections", []) 283 + ] 284 + 260 285 results.append(BenchmarkResult( 261 286 server=server_type, 262 - workload=workload, 263 - duration_ms=duration_ms, 264 - success=success, 265 - error=error, 287 + success=data.get("passed", False), 288 + total_duration_ms=data.get("total_duration_ms", 0), 289 + total_requests=data.get("total_requests", 0), 290 + memory_mb=memory_mb, 291 + sections=sections, 266 292 )) 267 293 268 - if success: 269 - console.print(f"[green]{duration_ms:.0f}ms[/green]") 270 - else: 271 - console.print(f"[red]failed: {error}[/red]") 294 + console.print(f"[green]{data.get('total_duration_ms', 0):.0f}ms[/green] ({memory_mb:.1f}MB)") 272 295 273 296 return results 274 297 275 298 276 299 def print_comparison(zig_results: list[BenchmarkResult], python_results: list[BenchmarkResult]): 277 - """print a comparison table""" 278 - table = Table(title="benchmark results") 279 - table.add_column("server", style="cyan") 280 - table.add_column("duration", justify="right") 281 - table.add_column("status", justify="center") 300 + """Print detailed comparison.""" 282 301 283 - def avg_duration(results: list[BenchmarkResult]) -> float: 302 + def avg(results: list[BenchmarkResult], key: str) -> float: 284 303 successful = [r for r in results if r.success] 285 304 if not successful: 286 305 return 0 287 - return sum(r.duration_ms for r in successful) / len(successful) 306 + return sum(getattr(r, key) for r in successful) / len(successful) 288 307 289 - def format_duration(ms: float) -> str: 290 - if ms < 1000: 291 - return f"{ms:.0f}ms" 292 - return f"{ms/1000:.2f}s" 308 + def avg_section(results: list[BenchmarkResult], section_name: str) -> float: 309 + successful = [r for r in results if r.success] 310 + if not successful: 311 + return 0 312 + total = 0 313 + count = 0 314 + for r in successful: 315 + for s in r.sections: 316 + if s.name == section_name: 317 + total += s.duration_ms 318 + count += 1 319 + return total / count if count else 0 293 320 294 - zig_avg = avg_duration(zig_results) 295 - python_avg = avg_duration(python_results) 321 + # summary table 322 + console.print() 323 + summary = Table(title="summary") 324 + summary.add_column("metric", style="cyan") 325 + summary.add_column("zig", justify="right") 326 + summary.add_column("python", justify="right") 327 + summary.add_column("zig advantage", justify="right") 296 328 297 - zig_success = sum(1 for r in zig_results if r.success) 298 - python_success = sum(1 for r in python_results if r.success) 329 + zig_time = avg(zig_results, "total_duration_ms") 330 + python_time = avg(python_results, "total_duration_ms") 331 + speedup = python_time / zig_time if zig_time > 0 else 0 299 332 300 - table.add_row( 301 - "zig", 302 - format_duration(zig_avg), 303 - f"{zig_success}/{len(zig_results)} ok", 333 + summary.add_row( 334 + "time", 335 + f"{zig_time:.0f}ms", 336 + f"{python_time:.0f}ms", 337 + f"[green]{speedup:.1f}x faster[/green]" if speedup > 1 else f"[yellow]{1/speedup:.1f}x slower[/yellow]", 304 338 ) 305 - table.add_row( 306 - "python", 307 - format_duration(python_avg), 308 - f"{python_success}/{len(python_results)} ok", 339 + 340 + zig_mem = avg(zig_results, "memory_mb") 341 + python_mem = avg(python_results, "memory_mb") 342 + mem_ratio = python_mem / zig_mem if zig_mem > 0 else 0 343 + 344 + summary.add_row( 345 + "memory", 346 + f"{zig_mem:.1f}MB", 347 + f"{python_mem:.1f}MB", 348 + f"[green]{mem_ratio:.1f}x smaller[/green]" if mem_ratio > 1 else f"[yellow]{1/mem_ratio:.1f}x larger[/yellow]", 309 349 ) 310 350 311 - console.print() 312 - console.print(table) 351 + zig_reqs = avg(zig_results, "total_requests") 352 + python_reqs = avg(python_results, "total_requests") 353 + summary.add_row("requests", f"{zig_reqs:.0f}", f"{python_reqs:.0f}", "") 313 354 314 - # show speedup if both succeeded 315 - if zig_avg > 0 and python_avg > 0: 316 - if zig_avg < python_avg: 317 - speedup = python_avg / zig_avg 318 - console.print(f"\n[bold green]zig is {speedup:.1f}x faster[/bold green]") 319 - elif python_avg < zig_avg: 320 - speedup = zig_avg / python_avg 321 - console.print(f"\n[bold yellow]python is {speedup:.1f}x faster[/bold yellow]") 355 + console.print(summary) 356 + 357 + # section breakdown 358 + if zig_results and zig_results[0].success and python_results and python_results[0].success: 359 + console.print() 360 + breakdown = Table(title="timing breakdown by section") 361 + breakdown.add_column("section", style="cyan") 362 + breakdown.add_column("zig", justify="right") 363 + breakdown.add_column("python", justify="right") 364 + breakdown.add_column("speedup", justify="right") 365 + 366 + section_names = [s.name for s in zig_results[0].sections] 367 + for name in section_names: 368 + zig_ms = avg_section(zig_results, name) 369 + python_ms = avg_section(python_results, name) 370 + section_speedup = python_ms / zig_ms if zig_ms > 0 else 0 371 + 372 + speedup_str = "" 373 + if section_speedup > 1.5: 374 + speedup_str = f"[green]{section_speedup:.1f}x[/green]" 375 + elif section_speedup > 1: 376 + speedup_str = f"[dim]{section_speedup:.1f}x[/dim]" 377 + elif section_speedup > 0: 378 + speedup_str = f"[yellow]{section_speedup:.1f}x[/yellow]" 379 + 380 + breakdown.add_row(name, f"{zig_ms:.1f}ms", f"{python_ms:.1f}ms", speedup_str) 381 + 382 + console.print(breakdown) 383 + 384 + # final verdict 385 + if zig_time > 0 and python_time > 0: 386 + if speedup > 1: 387 + console.print(f"\n[bold green]zig is {speedup:.1f}x faster overall[/bold green]") 322 388 else: 323 - console.print("\n[bold]both servers performed equally[/bold]") 389 + console.print(f"\n[bold yellow]python is {1/speedup:.1f}x faster overall[/bold yellow]") 324 390 325 391 326 392 def main(): 327 - parser = argparse.ArgumentParser( 328 - description="benchmark prefect servers", 329 - formatter_class=argparse.RawDescriptionHelpFormatter, 330 - epilog=""" 331 - examples: 332 - ./scripts/benchmark --server zig test_flow.py 333 - ./scripts/benchmark --server python test_flow.py 334 - ./scripts/benchmark --compare test_flow.py 335 - ./scripts/benchmark --compare test_flow.py --iterations 3 336 - """, 337 - ) 338 - parser.add_argument( 339 - "workload", 340 - help="path to workload script (must respect PREFECT_API_URL)", 341 - ) 342 - parser.add_argument( 343 - "--server", 344 - choices=["zig", "python"], 345 - help="server to benchmark (use --compare to benchmark both)", 346 - ) 347 - parser.add_argument( 348 - "--compare", 349 - action="store_true", 350 - help="benchmark both servers and compare results", 351 - ) 352 - parser.add_argument( 353 - "--iterations", 354 - type=int, 355 - default=1, 356 - help="number of iterations per server (default: 1)", 357 - ) 393 + parser = argparse.ArgumentParser(description="benchmark prefect servers") 394 + parser.add_argument("--server", choices=["zig", "python"], help="server to benchmark") 395 + parser.add_argument("--compare", action="store_true", help="compare both servers") 396 + parser.add_argument("--iterations", type=int, default=1, help="iterations per server") 358 397 359 398 args = parser.parse_args() 360 399 ··· 362 401 parser.error("must specify --server or --compare") 363 402 364 403 console.print(f"[bold]prefect-server benchmark[/bold]") 365 - console.print(f"workload: {args.workload}") 366 404 if args.iterations > 1: 367 405 console.print(f"iterations: {args.iterations}") 368 406 369 407 if args.compare: 370 - zig_results = benchmark_single("zig", args.workload, args.iterations) 371 - python_results = benchmark_single("python", args.workload, args.iterations) 408 + zig_results = benchmark_single("zig", args.iterations) 409 + python_results = benchmark_single("python", args.iterations) 372 410 print_comparison(zig_results, python_results) 373 411 else: 374 - results = benchmark_single(args.server, args.workload, args.iterations) 375 - if results: 376 - successful = [r for r in results if r.success] 377 - if successful: 378 - avg = sum(r.duration_ms for r in successful) / len(successful) 379 - console.print(f"\n[bold]average: {avg:.0f}ms[/bold]") 412 + results = benchmark_single(args.server, args.iterations) 413 + if results and results[0].success: 414 + console.print(f"\n[bold]average: {sum(r.total_duration_ms for r in results) / len(results):.0f}ms[/bold]") 380 415 381 416 382 417 if __name__ == "__main__":

+383 -271

scripts/test-api-sequence

··· 4 4 # dependencies = ["httpx", "rich"] 5 5 # /// 6 6 """ 7 - Test the exact API call sequence for a flow run against prefect-zig. 7 + Instrumented API test suite for prefect-server. 8 8 9 - This mimics what the Prefect Python client does, step by step: 10 - 1. POST /flows/ - create/get flow 11 - 2. POST /flow_runs/ - create flow run with PENDING state 12 - 3. GET /flow_runs/{id} - read flow run 13 - 4. POST /flow_runs/{id}/set_state - transition to RUNNING 14 - 5. (execute user code) 15 - 6. POST /flow_runs/{id}/set_state - transition to COMPLETED or FAILED 9 + Runs the full API sequence and reports timing breakdown per section. 10 + Use --json for machine-readable output (for benchmark script). 16 11 """ 17 12 13 + import json as json_lib 18 14 import os 19 15 import sys 16 + import time 20 17 import uuid 21 - from datetime import datetime, timezone 18 + from dataclasses import dataclass, field 19 + from typing import Callable 22 20 23 21 import httpx 24 22 from rich.console import Console 25 23 from rich.panel import Panel 24 + from rich.table import Table 26 25 27 26 console = Console() 28 27 BASE_URL = os.environ.get("PREFECT_API_URL", "http://localhost:4200/api") 28 + QUIET = "--json" in sys.argv or "--quiet" in sys.argv 29 29 30 30 31 - def iso_now() -> str: 32 - return datetime.now(timezone.utc).isoformat() 31 + @dataclass 32 + class TestResult: 33 + name: str 34 + passed: bool 35 + duration_ms: float 36 + requests: int = 0 37 + error: str | None = None 33 38 34 39 35 - def test_flow_sequence(flow_name: str = "test-flow", should_fail: bool = False): 36 - """Run through the exact sequence of API calls the Prefect client makes.""" 40 + class CountingClient(httpx.Client): 41 + """HTTP client that counts requests.""" 37 42 38 - console.print(Panel(f"testing flow sequence: {flow_name} (fail={should_fail})", style="blue")) 39 - console.print(f"server: {BASE_URL}\n") 43 + def __init__(self, *args, **kwargs): 44 + super().__init__(*args, **kwargs) 45 + self.request_count = 0 40 46 41 - with httpx.Client(base_url=BASE_URL, timeout=10) as client: 42 - # step 1: create/get flow 43 - console.print("[bold]1. POST /flows/[/bold]") 44 - resp = client.post("/flows/", json={"name": flow_name}) 45 - if resp.status_code not in (200, 201): 46 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 47 - return False 48 - flow = resp.json() 49 - console.print(f" flow_id: {flow.get('id')}") 47 + def request(self, *args, **kwargs): 48 + self.request_count += 1 49 + return super().request(*args, **kwargs) 50 50 51 - # validate flow response has required fields 52 - for field in ["id", "name", "created"]: 53 - if field not in flow: 54 - console.print(f"[red]FAIL[/red]: missing field '{field}' in flow response") 55 - return False 56 51 57 - # step 2: create flow run 58 - console.print("\n[bold]2. POST /flow_runs/[/bold]") 59 - flow_run_create = { 60 - "flow_id": flow["id"], 61 - "name": f"run-{uuid.uuid4().hex[:8]}", 62 - "parameters": {"x": 1, "y": 2}, 63 - "state": { 64 - "type": "PENDING", 65 - "name": "Pending", 66 - "timestamp": iso_now(), 67 - "message": None, 68 - }, 69 - } 70 - resp = client.post("/flow_runs/", json=flow_run_create) 71 - if resp.status_code not in (200, 201): 72 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 73 - return False 74 - flow_run = resp.json() 75 - flow_run_id = flow_run.get("id") 76 - console.print(f" flow_run_id: {flow_run_id}") 77 - console.print(f" state: {flow_run.get('state_type')}") 52 + def run_test(name: str, test_fn: Callable[[CountingClient], bool]) -> TestResult: 53 + """Run a test function with timing and request counting.""" 54 + if not QUIET: 55 + console.print(Panel(f"testing {name}", style="blue")) 78 56 79 - # validate flow run response 80 - for field in ["id", "flow_id", "name", "state_type", "state"]: 81 - if field not in flow_run: 82 - console.print(f"[red]FAIL[/red]: missing field '{field}' in flow_run response") 83 - return False 57 + client = CountingClient(base_url=BASE_URL, timeout=10) 58 + start = time.perf_counter() 84 59 85 - # step 3: read flow run (refresh before execution) 86 - console.print("\n[bold]3. GET /flow_runs/{id}[/bold]") 87 - resp = client.get(f"/flow_runs/{flow_run_id}") 88 - if resp.status_code != 200: 89 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 90 - return False 91 - flow_run = resp.json() 92 - console.print(f" state: {flow_run.get('state_type')}") 60 + try: 61 + passed = test_fn(client) 62 + duration_ms = (time.perf_counter() - start) * 1000 63 + return TestResult( 64 + name=name, 65 + passed=passed, 66 + duration_ms=duration_ms, 67 + requests=client.request_count, 68 + ) 69 + except Exception as e: 70 + duration_ms = (time.perf_counter() - start) * 1000 71 + return TestResult( 72 + name=name, 73 + passed=False, 74 + duration_ms=duration_ms, 75 + requests=client.request_count, 76 + error=str(e), 77 + ) 78 + finally: 79 + client.close() 93 80 94 - # step 4: set state to RUNNING 95 - console.print("\n[bold]4. POST /flow_runs/{id}/set_state (RUNNING)[/bold]") 96 - set_state_running = { 97 - "state": { 98 - "type": "RUNNING", 99 - "name": "Running", 100 - "timestamp": iso_now(), 101 - "message": None, 102 - }, 103 - "force": False, 104 - } 105 - resp = client.post(f"/flow_runs/{flow_run_id}/set_state", json=set_state_running) 106 - if resp.status_code != 200: 107 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 108 - return False 109 - result = resp.json() 110 - console.print(f" status: {result.get('status')}") 111 - console.print(f" state.type: {result.get('state', {}).get('type')}") 112 81 113 - # validate orchestration result 114 - if result.get("status") != "ACCEPT": 115 - console.print(f"[yellow]WARN[/yellow]: expected ACCEPT, got {result.get('status')}") 82 + # ---------- test functions ---------- 116 83 117 - # step 5: simulate user code execution 118 - console.print("\n[bold]5. (execute user code)[/bold]") 119 - if should_fail: 120 - console.print(" simulating failure...") 121 - error_message = "Flow run encountered an exception: ValueError('test error')" 122 - else: 123 - console.print(" simulating success...") 124 - error_message = None 125 84 126 - # step 6: set final state 127 - if should_fail: 128 - console.print("\n[bold]6. POST /flow_runs/{id}/set_state (FAILED)[/bold]") 129 - final_state = { 130 - "state": { 131 - "type": "FAILED", 132 - "name": "Failed", 133 - "timestamp": iso_now(), 134 - "message": error_message, 135 - }, 136 - "force": False, 137 - } 138 - else: 139 - console.print("\n[bold]6. POST /flow_runs/{id}/set_state (COMPLETED)[/bold]") 140 - final_state = { 141 - "state": { 142 - "type": "COMPLETED", 143 - "name": "Completed", 144 - "timestamp": iso_now(), 145 - "message": None, 146 - }, 147 - "force": False, 148 - } 149 - 150 - resp = client.post(f"/flow_runs/{flow_run_id}/set_state", json=final_state) 151 - if resp.status_code != 200: 152 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 153 - return False 154 - result = resp.json() 155 - console.print(f" status: {result.get('status')}") 156 - console.print(f" state.type: {result.get('state', {}).get('type')}") 157 - 158 - # verify final state 159 - console.print("\n[bold]7. GET /flow_runs/{id} (verify final state)[/bold]") 160 - resp = client.get(f"/flow_runs/{flow_run_id}") 161 - if resp.status_code != 200: 162 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 163 - return False 164 - flow_run = resp.json() 165 - final_type = flow_run.get("state_type") 166 - expected_type = "FAILED" if should_fail else "COMPLETED" 167 - 168 - if final_type == expected_type: 169 - console.print(f" [green]state: {final_type} (correct)[/green]") 170 - else: 171 - console.print(f" [red]state: {final_type} (expected {expected_type})[/red]") 172 - return False 173 - 174 - return True 175 - 176 - 177 - def test_admin_endpoints(): 85 + def test_admin(client: CountingClient) -> bool: 178 86 """Test admin/health endpoints.""" 179 - console.print(Panel("testing admin endpoints", style="blue")) 180 - 181 - with httpx.Client(base_url=BASE_URL, timeout=10) as client: 182 - # health 87 + # health 88 + if not QUIET: 183 89 console.print("[bold]GET /health[/bold]") 184 - resp = client.get("/health") 185 - if resp.status_code != 200: 90 + resp = client.get("/health") 91 + if resp.status_code != 200: 92 + if not QUIET: 186 93 console.print(f"[red]FAIL[/red]: {resp.status_code}") 187 - return False 188 - console.print(f" {resp.json()}") 94 + return False 95 + if not QUIET: 96 + console.print(f" {resp.text}") 189 97 190 - # version 98 + # version 99 + if not QUIET: 191 100 console.print("[bold]GET /admin/version[/bold]") 192 - resp = client.get("/admin/version") 193 - if resp.status_code != 200: 101 + resp = client.get("/admin/version") 102 + if resp.status_code != 200: 103 + if not QUIET: 194 104 console.print(f"[red]FAIL[/red]: {resp.status_code}") 195 - return False 105 + return False 106 + if not QUIET: 196 107 console.print(f" {resp.json()}") 197 108 198 - # csrf-token 109 + # csrf-token 110 + if not QUIET: 199 111 console.print("[bold]GET /csrf-token[/bold]") 200 - resp = client.get("/csrf-token") 201 - if resp.status_code != 200: 112 + resp = client.get("/csrf-token", params={"client": "test-client"}) 113 + if resp.status_code == 200: 114 + if not QUIET: 115 + console.print(f" token received") 116 + elif resp.status_code == 422: 117 + if not QUIET: 118 + console.print(f" csrf protection disabled (ok)") 119 + else: 120 + if not QUIET: 202 121 console.print(f"[red]FAIL[/red]: {resp.status_code}") 203 - return False 204 - console.print(f" token received") 122 + return False 205 123 206 124 return True 207 125 208 126 209 - def test_filter_endpoints(): 210 - """Test filter endpoints.""" 211 - console.print(Panel("testing filter endpoints", style="blue")) 127 + def test_flow_run(client: CountingClient, should_fail: bool = False) -> bool: 128 + """Test flow run lifecycle.""" 129 + suffix = "fail" if should_fail else "success" 130 + if not QUIET: 131 + console.print(f"server: {BASE_URL}\n") 132 + 133 + # create flow 134 + if not QUIET: 135 + console.print("[bold]1. POST /flows/[/bold]") 136 + resp = client.post("/flows/", json={"name": f"bench-flow-{suffix}"}) 137 + if resp.status_code not in (200, 201): 138 + if not QUIET: 139 + console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 140 + return False 141 + flow = resp.json() 142 + if not QUIET: 143 + console.print(f" flow_id: {flow.get('id')}") 144 + 145 + # create flow run 146 + if not QUIET: 147 + console.print("\n[bold]2. POST /flow_runs/[/bold]") 148 + resp = client.post("/flow_runs/", json={ 149 + "flow_id": flow["id"], 150 + "name": f"run-{uuid.uuid4().hex[:8]}", 151 + "state": {"type": "PENDING", "name": "Pending"}, 152 + }) 153 + if resp.status_code not in (200, 201): 154 + if not QUIET: 155 + console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 156 + return False 157 + flow_run = resp.json() 158 + flow_run_id = flow_run.get("id") 159 + if not QUIET: 160 + console.print(f" flow_run_id: {flow_run_id}") 212 161 213 - with httpx.Client(base_url=BASE_URL, timeout=10) as client: 214 - # flows/filter 215 - console.print("[bold]POST /flows/filter[/bold]") 216 - resp = client.post("/flows/filter", json={}) 217 - if resp.status_code != 200: 162 + # read flow run 163 + if not QUIET: 164 + console.print("\n[bold]3. GET /flow_runs/{id}[/bold]") 165 + resp = client.get(f"/flow_runs/{flow_run_id}") 166 + if resp.status_code != 200: 167 + if not QUIET: 218 168 console.print(f"[red]FAIL[/red]: {resp.status_code}") 219 - return False 220 - flows = resp.json() 221 - console.print(f" {len(flows)} flows") 169 + return False 170 + if not QUIET: 171 + console.print(f" state: {resp.json().get('state_type')}") 222 172 223 - # flow_runs/filter 224 - console.print("[bold]POST /flow_runs/filter[/bold]") 225 - resp = client.post("/flow_runs/filter", json={}) 226 - if resp.status_code != 200: 173 + # set RUNNING 174 + if not QUIET: 175 + console.print("\n[bold]4. POST /flow_runs/{id}/set_state (RUNNING)[/bold]") 176 + resp = client.post(f"/flow_runs/{flow_run_id}/set_state", json={ 177 + "state": {"type": "RUNNING", "name": "Running"}, 178 + "force": False, 179 + }) 180 + if resp.status_code not in (200, 201): 181 + if not QUIET: 227 182 console.print(f"[red]FAIL[/red]: {resp.status_code}") 228 - return False 229 - runs = resp.json() 230 - console.print(f" {len(runs)} flow runs") 183 + return False 184 + if not QUIET: 185 + console.print(f" status: {resp.json().get('status')}") 231 186 232 - # task_runs/filter 233 - console.print("[bold]POST /task_runs/filter[/bold]") 234 - resp = client.post("/task_runs/filter", json={}) 235 - if resp.status_code != 200: 187 + # set final state 188 + final_type = "FAILED" if should_fail else "COMPLETED" 189 + final_name = "Failed" if should_fail else "Completed" 190 + if not QUIET: 191 + console.print(f"\n[bold]5. POST /flow_runs/{{id}}/set_state ({final_type})[/bold]") 192 + resp = client.post(f"/flow_runs/{flow_run_id}/set_state", json={ 193 + "state": {"type": final_type, "name": final_name}, 194 + "force": False, 195 + }) 196 + if resp.status_code not in (200, 201): 197 + if not QUIET: 236 198 console.print(f"[red]FAIL[/red]: {resp.status_code}") 237 - return False 238 - tasks = resp.json() 239 - console.print(f" {len(tasks)} task runs") 199 + return False 200 + if not QUIET: 201 + console.print(f" status: {resp.json().get('status')}") 202 + 203 + # verify 204 + if not QUIET: 205 + console.print("\n[bold]6. GET /flow_runs/{id} (verify)[/bold]") 206 + resp = client.get(f"/flow_runs/{flow_run_id}") 207 + if resp.status_code != 200: 208 + return False 209 + actual_type = resp.json().get("state_type") 210 + if actual_type != final_type: 211 + if not QUIET: 212 + console.print(f"[red]FAIL[/red]: expected {final_type}, got {actual_type}") 213 + return False 214 + if not QUIET: 215 + console.print(f" [green]state: {actual_type} (correct)[/green]") 240 216 241 217 return True 242 218 243 219 244 - def test_task_run_sequence(): 245 - """Test task run API sequence.""" 246 - console.print(Panel("testing task run sequence", style="blue")) 247 - 248 - with httpx.Client(base_url=BASE_URL, timeout=10) as client: 249 - # create task run 220 + def test_task_run(client: CountingClient) -> bool: 221 + """Test task run lifecycle.""" 222 + # create 223 + if not QUIET: 250 224 console.print("[bold]POST /task_runs/[/bold]") 251 - task_run_create = { 252 - "task_key": "test-task", 253 - "dynamic_key": f"key-{uuid.uuid4().hex[:8]}", 254 - "name": f"task-{uuid.uuid4().hex[:8]}", 255 - "state": {"type": "PENDING", "name": "Pending"}, 256 - } 257 - resp = client.post("/task_runs/", json=task_run_create) 258 - if resp.status_code not in (200, 201): 259 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 260 - return False 261 - task_run = resp.json() 262 - task_run_id = task_run.get("id") 225 + resp = client.post("/task_runs/", json={ 226 + "task_key": "bench-task", 227 + "dynamic_key": f"key-{uuid.uuid4().hex[:8]}", 228 + "name": f"task-{uuid.uuid4().hex[:8]}", 229 + "state": {"type": "PENDING", "name": "Pending"}, 230 + }) 231 + if resp.status_code not in (200, 201): 232 + if not QUIET: 233 + console.print(f"[red]FAIL[/red]: {resp.status_code}") 234 + return False 235 + task_run_id = resp.json().get("id") 236 + if not QUIET: 263 237 console.print(f" task_run_id: {task_run_id}") 264 238 265 - # read task run 266 - console.print(f"[bold]GET /task_runs/{task_run_id}[/bold]") 267 - resp = client.get(f"/task_runs/{task_run_id}") 268 - if resp.status_code != 200: 269 - console.print(f"[red]FAIL[/red]: {resp.status_code}") 270 - return False 271 - console.print(f" state: {resp.json().get('state_type')}") 239 + # read 240 + resp = client.get(f"/task_runs/{task_run_id}") 241 + if resp.status_code != 200: 242 + return False 272 243 273 - # set state to RUNNING 274 - console.print(f"[bold]POST /task_runs/{task_run_id}/set_state (RUNNING)[/bold]") 275 - resp = client.post(f"/task_runs/{task_run_id}/set_state", json={ 276 - "state": {"type": "RUNNING", "name": "Running"} 277 - }) 278 - if resp.status_code != 200: 279 - console.print(f"[red]FAIL[/red]: {resp.status_code}") 280 - return False 281 - console.print(f" status: {resp.json().get('status')}") 244 + # RUNNING 245 + resp = client.post(f"/task_runs/{task_run_id}/set_state", json={ 246 + "state": {"type": "RUNNING", "name": "Running"}, 247 + "force": False, 248 + }) 249 + if resp.status_code not in (200, 201): 250 + return False 251 + if not QUIET: 252 + console.print(f" -> RUNNING: {resp.json().get('status')}") 282 253 283 - # set state to COMPLETED 284 - console.print(f"[bold]POST /task_runs/{task_run_id}/set_state (COMPLETED)[/bold]") 285 - resp = client.post(f"/task_runs/{task_run_id}/set_state", json={ 286 - "state": {"type": "COMPLETED", "name": "Completed"} 287 - }) 254 + # COMPLETED 255 + resp = client.post(f"/task_runs/{task_run_id}/set_state", json={ 256 + "state": {"type": "COMPLETED", "name": "Completed"}, 257 + "force": False, 258 + }) 259 + if resp.status_code not in (200, 201): 260 + return False 261 + if not QUIET: 262 + console.print(f" -> COMPLETED: {resp.json().get('status')}") 263 + 264 + return True 265 + 266 + 267 + def test_filters(client: CountingClient) -> bool: 268 + """Test filter endpoints.""" 269 + for endpoint, label in [ 270 + ("/flows/filter", "flows"), 271 + ("/flow_runs/filter", "flow_runs"), 272 + ("/task_runs/filter", "task_runs"), 273 + ]: 274 + resp = client.post(endpoint, json={}) 288 275 if resp.status_code != 200: 289 - console.print(f"[red]FAIL[/red]: {resp.status_code}") 276 + if not QUIET: 277 + console.print(f"[red]FAIL[/red]: {endpoint} {resp.status_code}") 290 278 return False 291 - console.print(f" status: {resp.json().get('status')}") 279 + if not QUIET: 280 + console.print(f" {label}: {len(resp.json())} items") 292 281 293 282 return True 294 283 295 284 296 - def test_logs_endpoint(): 285 + def test_logs(client: CountingClient) -> bool: 297 286 """Test logs endpoint.""" 298 - console.print(Panel("testing logs endpoint", style="blue")) 287 + from datetime import datetime, timezone 288 + 289 + logs = [ 290 + {"level": 20, "message": "test log 1", "name": "test", "timestamp": datetime.now(timezone.utc).isoformat()}, 291 + {"level": 30, "message": "test log 2", "name": "test", "timestamp": datetime.now(timezone.utc).isoformat()}, 292 + ] 293 + resp = client.post("/logs/", json=logs) 294 + if resp.status_code not in (200, 201, 204): 295 + if not QUIET: 296 + console.print(f"[red]FAIL[/red]: {resp.status_code}") 297 + return False 298 + if not QUIET: 299 + console.print(f" {len(logs)} logs sent") 300 + return True 301 + 302 + 303 + def test_blocks(client: CountingClient) -> bool: 304 + """Test blocks API (types, schemas, documents).""" 305 + slug = f"bench-block-{uuid.uuid4().hex[:8]}" 306 + 307 + # create block type 308 + if not QUIET: 309 + console.print("[bold]block_types[/bold]") 310 + resp = client.post("/block_types/", json={ 311 + "name": f"Bench Block {slug}", 312 + "slug": slug, 313 + "description": "benchmark block type", 314 + }) 315 + if resp.status_code not in (200, 201): 316 + if not QUIET: 317 + console.print(f"[red]FAIL[/red]: create block_type {resp.status_code}") 318 + return False 319 + block_type = resp.json() 320 + block_type_id = block_type.get("id") 321 + if not QUIET: 322 + console.print(f" created: {block_type_id}") 323 + 324 + # get by slug 325 + resp = client.get(f"/block_types/slug/{slug}") 326 + if resp.status_code != 200: 327 + return False 328 + 329 + # create schema 330 + if not QUIET: 331 + console.print("[bold]block_schemas[/bold]") 332 + resp = client.post("/block_schemas/", json={ 333 + "block_type_id": block_type_id, 334 + "fields": {"properties": {"value": {"type": "string"}}}, 335 + "capabilities": ["test"], 336 + "version": "1.0.0", 337 + }) 338 + if resp.status_code not in (200, 201): 339 + if not QUIET: 340 + console.print(f"[red]FAIL[/red]: create block_schema {resp.status_code}") 341 + return False 342 + block_schema = resp.json() 343 + block_schema_id = block_schema.get("id") 344 + checksum = block_schema.get("checksum") 345 + if not QUIET: 346 + console.print(f" created: {block_schema_id}") 299 347 300 - with httpx.Client(base_url=BASE_URL, timeout=10) as client: 301 - console.print("[bold]POST /logs/[/bold]") 302 - logs = [ 303 - {"level": 20, "message": "test log 1", "name": "test", "timestamp": iso_now()}, 304 - {"level": 30, "message": "test log 2", "name": "test", "timestamp": iso_now()}, 305 - ] 306 - resp = client.post("/logs/", json=logs) 307 - if resp.status_code not in (200, 201, 204): 308 - console.print(f"[red]FAIL[/red]: {resp.status_code} {resp.text}") 348 + # get by checksum 349 + resp = client.get(f"/block_schemas/checksum/{checksum}") 350 + if resp.status_code != 200: 351 + return False 352 + 353 + # create document 354 + if not QUIET: 355 + console.print("[bold]block_documents[/bold]") 356 + doc_name = f"bench-doc-{uuid.uuid4().hex[:8]}" 357 + resp = client.post("/block_documents/", json={ 358 + "name": doc_name, 359 + "block_type_id": block_type_id, 360 + "block_schema_id": block_schema_id, 361 + "data": {"value": "secret-value"}, 362 + }) 363 + if resp.status_code not in (200, 201): 364 + if not QUIET: 365 + console.print(f"[red]FAIL[/red]: create block_document {resp.status_code}") 366 + return False 367 + block_doc = resp.json() 368 + block_doc_id = block_doc.get("id") 369 + if not QUIET: 370 + console.print(f" created: {block_doc_id}") 371 + 372 + # get by id 373 + resp = client.get(f"/block_documents/{block_doc_id}") 374 + if resp.status_code != 200: 375 + return False 376 + 377 + # get by slug/name 378 + resp = client.get(f"/block_types/slug/{slug}/block_documents/name/{doc_name}") 379 + if resp.status_code != 200: 380 + return False 381 + 382 + # update 383 + resp = client.patch(f"/block_documents/{block_doc_id}", json={"data": {"value": "updated"}}) 384 + if resp.status_code != 204: 385 + return False 386 + if not QUIET: 387 + console.print(f" updated") 388 + 389 + # filters 390 + for endpoint in ["/block_types/filter", "/block_schemas/filter", "/block_documents/filter"]: 391 + resp = client.post(endpoint, json={}) 392 + if resp.status_code != 200: 309 393 return False 310 - console.print(f" {len(logs)} logs sent") 394 + 395 + # delete 396 + resp = client.delete(f"/block_documents/{block_doc_id}") 397 + if resp.status_code != 204: 398 + return False 399 + if not QUIET: 400 + console.print(f" deleted") 311 401 312 402 return True 313 403 314 404 315 405 def main(): 316 - console.print("\n[bold cyan]prefect-server API test[/bold cyan]\n") 406 + json_output = "--json" in sys.argv 407 + 408 + if not QUIET: 409 + console.print("\n[bold cyan]prefect-server API test[/bold cyan]\n") 410 + 411 + results: list[TestResult] = [] 317 412 318 - results = [] 413 + # run all tests 414 + results.append(run_test("admin", test_admin)) 415 + results.append(run_test("flow_run (success)", lambda c: test_flow_run(c, should_fail=False))) 416 + results.append(run_test("flow_run (failure)", lambda c: test_flow_run(c, should_fail=True))) 417 + results.append(run_test("task_run", test_task_run)) 418 + results.append(run_test("filters", test_filters)) 419 + results.append(run_test("logs", test_logs)) 420 + results.append(run_test("blocks", test_blocks)) 319 421 320 - # admin endpoints 321 - results.append(("admin", test_admin_endpoints())) 322 - console.print() 422 + total_duration = sum(r.duration_ms for r in results) 423 + total_requests = sum(r.requests for r in results) 424 + all_passed = all(r.passed for r in results) 323 425 324 - # flow sequence (happy path) 325 - results.append(("flow (success)", test_flow_sequence("happy-flow", should_fail=False))) 326 - console.print() 426 + if json_output: 427 + # machine-readable output for benchmark script 428 + output = { 429 + "passed": all_passed, 430 + "total_duration_ms": total_duration, 431 + "total_requests": total_requests, 432 + "sections": [ 433 + { 434 + "name": r.name, 435 + "passed": r.passed, 436 + "duration_ms": r.duration_ms, 437 + "requests": r.requests, 438 + "error": r.error, 439 + } 440 + for r in results 441 + ], 442 + } 443 + print(json_lib.dumps(output)) 444 + else: 445 + # human-readable output 446 + console.print("\n" + "=" * 60) 327 447 328 - # flow sequence (failure path) 329 - results.append(("flow (failure)", test_flow_sequence("sad-flow", should_fail=True))) 330 - console.print() 448 + table = Table(title="timing breakdown") 449 + table.add_column("section", style="cyan") 450 + table.add_column("time", justify="right") 451 + table.add_column("reqs", justify="right") 452 + table.add_column("status", justify="center") 331 453 332 - # task run sequence 333 - results.append(("task_run", test_task_run_sequence())) 334 - console.print() 454 + for r in results: 455 + status = "[green]✓[/green]" if r.passed else "[red]✗[/red]" 456 + table.add_row(r.name, f"{r.duration_ms:.1f}ms", str(r.requests), status) 335 457 336 - # filter endpoints 337 - results.append(("filter", test_filter_endpoints())) 338 - console.print() 458 + table.add_row("", "", "", "", style="dim") 459 + table.add_row("[bold]total[/bold]", f"[bold]{total_duration:.1f}ms[/bold]", f"[bold]{total_requests}[/bold]", "") 339 460 340 - # logs endpoint 341 - results.append(("logs", test_logs_endpoint())) 342 - console.print() 461 + console.print(table) 343 462 344 - # summary 345 - console.print("=" * 50) 346 - all_passed = all(r[1] for r in results) 347 - for name, passed in results: 348 - status = "[green]✓[/green]" if passed else "[red]✗[/red]" 349 - console.print(f" {status} {name}") 463 + if all_passed: 464 + console.print("\n[bold green]all tests passed[/bold green]") 465 + else: 466 + console.print("\n[bold red]some tests failed[/bold red]") 350 467 351 - if all_passed: 352 - console.print("\n[bold green]all tests passed[/bold green]") 353 - sys.exit(0) 354 - else: 355 - console.print("\n[bold red]some tests failed[/bold red]") 356 - sys.exit(1) 468 + sys.exit(0 if all_passed else 1) 357 469 358 470 359 471 if __name__ == "__main__":

+2 -2

src/services/event_persister.zig

··· 3 3 const log = @import("../logging.zig"); 4 4 const messaging = @import("../utilities/messaging.zig"); 5 5 const db = @import("../db/sqlite.zig"); 6 - const common = @import("../api/common.zig"); 6 + const time_util = @import("../utilities/time.zig"); 7 7 8 8 const BATCH_SIZE: usize = 100; 9 9 const FLUSH_INTERVAL_MS: u64 = 1000; // 1 second ··· 111 111 var success_count: usize = 0; 112 112 var truncated_count: usize = 0; 113 113 var ts_buf: [32]u8 = undefined; 114 - const now_ts = common.getTimestamp(&ts_buf); 114 + const now_ts = time_util.timestamp(&ts_buf); 115 115 116 116 for (batch) |event| { 117 117 if (event.truncated) truncated_count += 1;