diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index 3dd8b045..f5dec4ce 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -302,3 +302,5 @@ async def ashutdown(self) -> None: for lm in self._lm_cache.values(): await lm.close() self._lm_cache.clear() + if hasattr(self._orchestrator, "shutdown"): + self._orchestrator.shutdown() diff --git a/its_hub/core/orchestrator.py b/its_hub/core/orchestrator.py index 9cc01e29..7bacded0 100644 --- a/its_hub/core/orchestrator.py +++ b/its_hub/core/orchestrator.py @@ -2,6 +2,7 @@ import contextlib import logging import threading +from concurrent.futures import ThreadPoolExecutor from its_hub._rust import _PyLMOrchestrator from its_hub.api import ( @@ -20,24 +21,26 @@ class _ThreadSafeAsyncSemaphore: limit is respected across event loops and threads, and wraps acquire/release for use in async contexts without blocking the event loop. - FIXME: When a task wants to acquire the semaphore, it submits self._sem.acquire (a blocking call) - to the default ThreadPoolExecutor. The default thread pool can be exhausted when max_concurrency - is low (e.g., 2-4) with large batches. Our default max_concurrency=32 should help avoid the issue - but needs further investigation, if necessary. + A dedicated ThreadPoolExecutor is used so that blocked acquires do not + exhaust the default executor and starve unrelated async operations. """ def __init__(self, value: int): self._sem = threading.Semaphore(value) + self._executor = ThreadPoolExecutor( + max_workers=value, thread_name_prefix="lm-sem" + ) async def acquire(self): loop = asyncio.get_running_loop() - # Run blocking acquire in the default executor so the event loop - # stays responsive while waiting for a slot. - await loop.run_in_executor(None, self._sem.acquire) + await loop.run_in_executor(self._executor, self._sem.acquire) def release(self): self._sem.release() + def shutdown(self): + self._executor.shutdown(wait=True) + async def __aenter__(self): await self.acquire() return self @@ -69,6 +72,10 @@ def __init__(self, max_concurrency: int = 32): else None ) + def shutdown(self): + if self._semaphore is not None: + self._semaphore.shutdown() + async def agenerate( self, lm: AbstractLanguageModel, diff --git a/its_hub/integration/iaas/envoy_config.yaml b/its_hub/integration/iaas/envoy_config.yaml index 67ceacef..452aeacf 100644 --- a/its_hub/integration/iaas/envoy_config.yaml +++ b/its_hub/integration/iaas/envoy_config.yaml @@ -115,7 +115,8 @@ static_resources: prefix: "/" headers: - name: X-ITS-Route - exact_match: "its-service" + string_match: + exact: "its-service" route: cluster: iaas_upstream timeout: 300s # **CUSTOMIZE**: Timeout for ITS processing diff --git a/pyproject.toml b/pyproject.toml index b6ce4246..24220cf2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,6 +74,7 @@ dev = [ "jupytext>=1.15.0", "jupyter>=1.0.0", "scipy>=1.11.0", + "pyyaml>=6.0", ] # Experimental - not officially supported in MVP diff --git a/tests/e2e/test_iaas_envoy_e2e.py b/tests/e2e/test_iaas_envoy_e2e.py new file mode 100644 index 00000000..64a3ad1c --- /dev/null +++ b/tests/e2e/test_iaas_envoy_e2e.py @@ -0,0 +1,300 @@ +""" +End-to-end test for IaaS + Envoy integration. + +Tests the full stack: Client -> Envoy -> ext_proc router -> IaaS -> LLM + +Starts the IaaS service, ext_proc router, and (optionally) Envoy, then +runs requests through the stack and verifies correct routing and responses. + +Usage: + # With real LLM endpoint: + python tests/e2e/test_iaas_envoy_e2e.py \\ + --endpoint http://localhost:8100/v1 \\ + --model_name Qwen/Qwen2.5-Math-7B-Instruct + + # With built-in mock LLM (no external dependencies except Envoy): + python tests/e2e/test_iaas_envoy_e2e.py --mock-llm + + # Skip Envoy (test IaaS service only): + python tests/e2e/test_iaas_envoy_e2e.py --mock-llm --skip-envoy +""" + +import argparse +import shutil +import sys + +from tests.e2e.utils.iaas_helpers import ( + configure_iaas, + find_free_port, + http_get, + http_post, + start_envoy, + start_iaas_stack, + start_mock_llm, + stop_processes, +) + +# --------------------------------------------------------------------------- +# Result tracker +# --------------------------------------------------------------------------- + + +class _Result: + def __init__(self): + self.passed = 0 + self.failed = 0 + self.skipped = 0 + self.details = [] + + def ok(self, name): + self.passed += 1 + self.details.append(f" PASS {name}") + print(f" PASS {name}") + + def fail(self, name, reason): + self.failed += 1 + self.details.append(f" FAIL {name}: {reason}") + print(f" FAIL {name}: {reason}") + + def skip(self, name, reason): + self.skipped += 1 + self.details.append(f" SKIP {name}: {reason}") + print(f" SKIP {name}: {reason}") + + def summary(self): + total = self.passed + self.failed + self.skipped + print(f"\n{'='*60}") + print(f"Results: {self.passed}/{total} passed, {self.failed} failed, {self.skipped} skipped") + if self.failed: + print("\nFailed tests:") + for d in self.details: + if "FAIL" in d: + print(d) + print(f"{'='*60}") + return self.failed == 0 + + +# --------------------------------------------------------------------------- +# Tests: IaaS direct +# --------------------------------------------------------------------------- + + +def check_iaas_direct(iaas_url, llm_endpoint, model_name, api_key, result): + """Test IaaS service directly (without Envoy).""" + print("\n--- IaaS Direct Tests ---") + + # Configure + try: + configure_iaas(iaas_url, llm_endpoint, model_name, api_key) + result.ok("iaas_configure") + except RuntimeError as e: + result.fail("iaas_configure", str(e)) + return + + # Models endpoint + status, body = http_get(f"{iaas_url}/v1/models") + if status == 200 and body.get("data") and body["data"][0]["id"] == model_name: + result.ok("iaas_models") + else: + result.fail("iaas_models", f"unexpected: {body}") + + # Chat completion via body budget + status, body = http_post(f"{iaas_url}/v1/chat/completions", { + "model": model_name, + "messages": [{"role": "user", "content": "What is 2+2?"}], + "budget": 3, + }) + if status == 200 and body.get("choices"): + content = body["choices"][0]["message"]["content"] + if content: + result.ok("iaas_chat_completion_body_budget") + else: + result.fail("iaas_chat_completion_body_budget", "empty content") + else: + result.fail("iaas_chat_completion_body_budget", f"status {status}: {body}") + + # Chat completion via header budget + status, body = http_post( + f"{iaas_url}/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": "What is 3+3?"}], + }, + headers={"X-ITS-Budget": "2"}, + ) + if status == 200 and body.get("choices"): + result.ok("iaas_chat_completion_header_budget") + else: + result.fail("iaas_chat_completion_header_budget", f"status {status}: {body}") + + # Chat completion with header overrides + status, body = http_post( + f"{iaas_url}/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": "What is 5+5?"}], + }, + headers={ + "X-ITS-Budget": "2", + "X-ITS-Endpoint": llm_endpoint, + "X-ITS-API-Key": api_key, + }, + ) + if status == 200 and body.get("choices"): + result.ok("iaas_header_overrides") + else: + result.fail("iaas_header_overrides", f"status {status}: {body}") + + +# --------------------------------------------------------------------------- +# Tests: Envoy-routed +# --------------------------------------------------------------------------- + + +def check_envoy_routed(envoy_url, iaas_url, llm_endpoint, model_name, api_key, result): + """Test requests routed through Envoy.""" + print("\n--- Envoy-Routed Tests ---") + + # Configure IaaS first + try: + configure_iaas(iaas_url, llm_endpoint, model_name, api_key) + except RuntimeError: + result.fail("envoy_precondition", "could not configure IaaS") + return + + # ITS request through Envoy (should route to IaaS) + status, body = http_post( + f"{envoy_url}/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": "What is 7+7?"}], + }, + headers={ + "X-ITS-Budget": "2", + "X-ITS-Endpoint": llm_endpoint, + "X-ITS-API-Key": api_key, + }, + ) + if status == 200 and body.get("choices"): + result.ok("envoy_its_request") + else: + result.fail("envoy_its_request", f"status {status}: {body}") + + # Non-ITS request through Envoy (should pass through to LLM) + status, body = http_post( + f"{envoy_url}/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": "Direct pass-through"}], + }, + ) + if status == 200 and body.get("choices"): + result.ok("envoy_passthrough") + else: + result.fail("envoy_passthrough", f"status {status}: {body}") + + # Verify ITS headers are stripped on pass-through + status, body = http_post( + f"{envoy_url}/v1/chat/completions", + { + "model": model_name, + "messages": [{"role": "user", "content": "Stray header test"}], + }, + headers={"X-ITS-Endpoint": "http://should-be-stripped/v1"}, + ) + if status == 200: + if "its_headers_received" not in body: + result.skip( + "envoy_stray_header_stripped", + "upstream does not report received headers (use --mock-llm)", + ) + elif body["its_headers_received"]: + result.fail("envoy_stray_header_stripped", f"ITS headers reached LLM: {body['its_headers_received']}") + else: + result.ok("envoy_stray_header_stripped") + else: + result.fail("envoy_stray_header_stripped", f"status {status}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser( + description="E2E tests for IaaS + Envoy integration", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument("--endpoint", help="LLM endpoint URL (e.g., http://localhost:8100/v1)") + p.add_argument("--model_name", default="mock-model", help="Model name at the endpoint") + p.add_argument("--api_key", default="NO_API_KEY", help="API key for the LLM endpoint") + p.add_argument("--mock-llm", action="store_true", help="Start a built-in mock LLM server") + p.add_argument("--skip-envoy", action="store_true", help="Skip Envoy tests (test IaaS only)") + return p.parse_args() + + +def main(): + args = parse_args() + result = _Result() + processes = [] + servers = [] + envoy_tmpdir = None + + try: + # --- Resolve LLM endpoint --- + if args.mock_llm: + llm_port = find_free_port() + servers.append(start_mock_llm(llm_port)) + llm_endpoint = f"http://127.0.0.1:{llm_port}/v1" + model_name = "mock-model" + print(f"Mock LLM started on port {llm_port}") + elif args.endpoint: + llm_endpoint = args.endpoint + model_name = args.model_name + from urllib.parse import urlparse + llm_port = urlparse(llm_endpoint).port or 80 + else: + print("Error: provide --endpoint or --mock-llm") + sys.exit(1) + + api_key = args.api_key + + # --- Start IaaS stack --- + stack_procs, iaas_url, ext_proc_port = start_iaas_stack(llm_port) + processes.extend(stack_procs) + + # --- Run IaaS direct tests --- + check_iaas_direct(iaas_url, llm_endpoint, model_name, api_key, result) + + # --- Envoy tests --- + if args.skip_envoy: + result.skip("envoy_tests", "skipped via --skip-envoy") + elif not shutil.which("envoy"): + result.skip("envoy_tests", "envoy binary not found in PATH") + else: + from urllib.parse import urlparse + iaas_port = urlparse(iaas_url).port + envoy_result = start_envoy(ext_proc_port, iaas_port, llm_port) + if envoy_result is None: + result.fail("envoy_tests", "Envoy or ext_proc failed to start") + else: + envoy_proc, envoy_url, envoy_tmpdir, _ = envoy_result + processes.append(("envoy", envoy_proc)) + check_envoy_routed(envoy_url, iaas_url, llm_endpoint, model_name, api_key, result) + + finally: + print("\nShutting down services...") + stop_processes(processes) + for server in servers: + server.shutdown() + if envoy_tmpdir: + import shutil as _shutil + _shutil.rmtree(envoy_tmpdir, ignore_errors=True) + + success = result.summary() + sys.exit(0 if success else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/test_iaas_envoy_perf.py b/tests/e2e/test_iaas_envoy_perf.py new file mode 100644 index 00000000..333f4ec8 --- /dev/null +++ b/tests/e2e/test_iaas_envoy_perf.py @@ -0,0 +1,374 @@ +""" +Performance test for IaaS + Envoy integration. + +Measures latency overhead and throughput for the ITS gateway stack. +Compares direct LLM access vs IaaS vs Envoy+IaaS routing. + +Prerequisites: + - LLM endpoint running (or use --mock-llm for a built-in mock) + - IaaS service running (started automatically with --start-services) + - Envoy running (optional, started automatically with --start-services) + +Usage: + # Quick test with mock LLM (measures gateway overhead only): + python tests/e2e/test_iaas_envoy_perf.py --mock-llm --start-services + + # Against running services: + python tests/e2e/test_iaas_envoy_perf.py \\ + --llm-url http://localhost:8100/v1 \\ + --iaas-url http://localhost:8109 \\ + --envoy-url http://localhost:8108 \\ + --model_name Qwen/Qwen2.5-Math-7B-Instruct \\ + --concurrency 10 --num-requests 50 + + # Vary budget to see scaling: + python tests/e2e/test_iaas_envoy_perf.py --mock-llm --start-services \\ + --budgets 1,4,8,16 +""" + +import argparse +import asyncio +import statistics +import sys +import time + +try: + import aiohttp + HAS_AIOHTTP = True +except ImportError: + HAS_AIOHTTP = False + +from tests.e2e.utils.iaas_helpers import ( + configure_iaas, + find_free_port, + http_post, + start_envoy, + start_iaas_stack, + start_mock_llm, + stop_processes, +) + +# --------------------------------------------------------------------------- +# Async benchmark +# --------------------------------------------------------------------------- + + +def _compute_stats(latencies, errors, error_details, wall_time, num_requests): + if error_details: + print(" Error details (first 5):") + for detail in error_details[:5]: + print(f" {detail}") + + if not latencies: + return None + + latencies.sort() + return { + "count": len(latencies), + "errors": errors, + "wall_time_s": round(wall_time, 2), + "rps": round(len(latencies) / wall_time, 1), + "p50_ms": round(latencies[len(latencies) // 2], 1), + "p95_ms": round(latencies[int(len(latencies) * 0.95)], 1), + "p99_ms": round(latencies[int(len(latencies) * 0.99)], 1), + "mean_ms": round(statistics.mean(latencies), 1), + "min_ms": round(latencies[0], 1), + "max_ms": round(latencies[-1], 1), + } + + +async def benchmark_endpoint(url, model_name, num_requests, concurrency, budget=None, headers=None, timeout_s=120): + """Send num_requests concurrent requests to an HTTP endpoint and collect latencies.""" + sem = asyncio.Semaphore(concurrency) + latencies = [] + errors = 0 + completed = 0 + + payload = { + "model": model_name, + "messages": [{"role": "user", "content": "What is 2+2?"}], + } + if budget is not None: + payload["budget"] = budget + + req_headers = {"Content-Type": "application/json"} + if headers: + req_headers.update(headers) + + error_details = [] + + async def _single_request(session, i): + nonlocal errors, completed + async with sem: + start = time.perf_counter() + try: + async with session.post(url, json=payload, headers=req_headers, timeout=aiohttp.ClientTimeout(total=timeout_s)) as resp: + body = await resp.read() + elapsed = (time.perf_counter() - start) * 1000 + if resp.status == 200: + latencies.append(elapsed) + else: + errors += 1 + error_details.append(f"req {i}: HTTP {resp.status}: {body[:200]}") + except Exception as e: + errors += 1 + error_details.append(f"req {i}: {type(e).__name__}: {e}") + completed += 1 + if completed % 10 == 0 or completed == num_requests: + print(f" {completed}/{num_requests} done", flush=True) + + try: + async with aiohttp.ClientSession() as session: + wall_start = time.perf_counter() + tasks = [asyncio.create_task(_single_request(session, i)) for i in range(num_requests)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=timeout_s) + wall_time = time.perf_counter() - wall_start + except TimeoutError: + wall_time = timeout_s + cancelled = num_requests - completed + errors += cancelled + print(f" TIMEOUT after {timeout_s}s ({completed}/{num_requests} completed, {cancelled} cancelled)") + + return _compute_stats(latencies, errors, error_details, wall_time, num_requests) + + +async def benchmark_algorithm(llm_url, model_name, api_key, num_requests, concurrency, budget, timeout_s=120): + """Benchmark using SelfConsistency algorithm directly (no IaaS service). + + This is the true baseline: same algorithm and orchestrator as IaaS, but + without the FastAPI/Envoy service layers. Overhead of IaaS vs this baseline + isolates the HTTP service cost. + """ + from its_hub.core.algorithms.self_consistency import SelfConsistency + from its_hub.core.lms.openai_lm import OpenAICompatibleLanguageModel + + sem = asyncio.Semaphore(concurrency) + latencies = [] + errors = 0 + completed = 0 + error_details = [] + + lm = OpenAICompatibleLanguageModel( + endpoint=llm_url, api_key=api_key, model_name=model_name, + ) + alg = SelfConsistency() + + async def _single_request(i): + nonlocal errors, completed + async with sem: + start = time.perf_counter() + try: + await alg.ainfer( + lm, "What is 2+2?", budget=budget, return_response_only=True, + ) + elapsed = (time.perf_counter() - start) * 1000 + latencies.append(elapsed) + except Exception as e: + errors += 1 + error_details.append(f"req {i}: {type(e).__name__}: {e}") + completed += 1 + if completed % 10 == 0 or completed == num_requests: + print(f" {completed}/{num_requests} done", flush=True) + + try: + wall_start = time.perf_counter() + tasks = [asyncio.create_task(_single_request(i)) for i in range(num_requests)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=timeout_s) + wall_time = time.perf_counter() - wall_start + except TimeoutError: + wall_time = timeout_s + cancelled = num_requests - completed + errors += cancelled + print(f" TIMEOUT after {timeout_s}s ({completed}/{num_requests} completed, {cancelled} cancelled)") + finally: + await lm.close() + + return _compute_stats(latencies, errors, error_details, wall_time, num_requests) + + +def print_stats(label, stats): + if stats is None: + print(f" {label}: no results (all requests failed)") + return + print(f" {label}:") + print(f" Requests: {stats['count']} ok, {stats['errors']} errors") + print(f" Wall time: {stats['wall_time_s']}s ({stats['rps']} req/s)") + print(f" Latency: p50={stats['p50_ms']}ms p95={stats['p95_ms']}ms p99={stats['p99_ms']}ms") + print(f" Range: min={stats['min_ms']}ms mean={stats['mean_ms']}ms max={stats['max_ms']}ms") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser( + description="Performance test for IaaS + Envoy integration", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + g = p.add_argument_group("endpoints") + g.add_argument("--llm-url", help="Direct LLM endpoint (e.g., http://localhost:8100/v1)") + g.add_argument("--iaas-url", help="IaaS service URL (e.g., http://localhost:8109)") + g.add_argument("--envoy-url", help="Envoy gateway URL (e.g., http://localhost:8108)") + g.add_argument("--model_name", default="mock-model") + g.add_argument("--api_key", default="NO_API_KEY") + + g = p.add_argument_group("auto-start") + g.add_argument("--mock-llm", action="store_true", help="Start built-in mock LLM") + g.add_argument("--mock-latency-ms", type=int, default=10, help="Mock LLM response latency (default: 10ms)") + g.add_argument("--start-services", action="store_true", help="Auto-start IaaS + ext_proc + Envoy") + + g = p.add_argument_group("benchmark config") + g.add_argument("--num-requests", type=int, default=50, help="Total requests per benchmark (default: 50)") + g.add_argument("--concurrency", type=int, default=10, help="Max concurrent requests (default: 10)") + g.add_argument("--budgets", default="1,4", help="Comma-separated budgets to test (default: 1,4)") + g.add_argument("--warmup", type=int, default=3, help="Warmup requests before benchmark (default: 3)") + + return p.parse_args() + + +def main(): + args = parse_args() + budgets = [int(b) for b in args.budgets.split(",")] + processes = [] + servers = [] + envoy_tmpdir = None + + if not HAS_AIOHTTP: + print("Error: aiohttp is required. Install with: pip install aiohttp") + sys.exit(1) + + try: + # --- Resolve endpoints --- + llm_url = args.llm_url + iaas_url = args.iaas_url + envoy_url = args.envoy_url + llm_port = None + + if args.mock_llm: + llm_port = find_free_port() + servers.append(start_mock_llm(llm_port, latency_ms=args.mock_latency_ms)) + llm_url = f"http://127.0.0.1:{llm_port}/v1" + print(f"Mock LLM started on port {llm_port} (latency={args.mock_latency_ms}ms)") + + if args.start_services: + if not llm_url: + print("Error: --start-services requires --llm-url or --mock-llm") + sys.exit(1) + + from urllib.parse import urlparse + if llm_port is None: + llm_port = urlparse(llm_url).port or 80 + + stack_procs, iaas_url, ext_proc_port = start_iaas_stack(llm_port) + processes.extend(stack_procs) + + iaas_port = urlparse(iaas_url).port + envoy_result = start_envoy(ext_proc_port, iaas_port, llm_port) + if envoy_result: + envoy_proc, envoy_url, envoy_tmpdir, _ = envoy_result + processes.append(("envoy", envoy_proc)) + else: + print("Envoy not available, skipping Envoy benchmarks") + + # --- Configure IaaS --- + if iaas_url: + if not llm_url: + print("Error: --iaas-url requires --llm-url or --mock-llm") + sys.exit(1) + configure_iaas(iaas_url, llm_url, args.model_name, args.api_key) + print("IaaS configured") + + # --- Warmup --- + print(f"\nWarming up ({args.warmup} requests per endpoint)...") + for _ in range(args.warmup): + if llm_url: + http_post(f"{llm_url}/chat/completions", { + "model": args.model_name, + "messages": [{"role": "user", "content": "warmup"}], + }) + if iaas_url: + http_post(f"{iaas_url}/v1/chat/completions", { + "model": args.model_name, + "messages": [{"role": "user", "content": "warmup"}], + "budget": 1, + }) + + # --- Benchmarks --- + print(f"\nBenchmark: {args.num_requests} requests, concurrency={args.concurrency}") + print("=" * 60) + + for budget in budgets: + print(f"\n--- budget={budget} ---") + + baseline_p50 = None + + if llm_url: + print(f"\n[Algorithm direct, budget={budget}]") + stats = asyncio.run(benchmark_algorithm( + llm_url, args.model_name, args.api_key, + args.num_requests, args.concurrency, budget=budget, + )) + print_stats(f"algorithm(budget={budget})", stats) + baseline_p50 = stats["p50_ms"] if stats else None + + if iaas_url: + print(f"\n[IaaS, budget={budget}]") + stats = asyncio.run(benchmark_endpoint( + f"{iaas_url}/v1/chat/completions", + args.model_name, + args.num_requests, + args.concurrency, + budget=budget, + )) + print_stats(f"iaas(budget={budget})", stats) + if stats and baseline_p50 is not None: + overhead = stats["p50_ms"] - baseline_p50 + print(f" Overhead vs direct: {overhead:+.1f}ms (p50)") + + if envoy_url: + print(f"\n[Envoy -> IaaS, budget={budget}]") + stats = asyncio.run(benchmark_endpoint( + f"{envoy_url}/v1/chat/completions", + args.model_name, + args.num_requests, + args.concurrency, + headers={ + "X-ITS-Budget": str(budget), + "X-ITS-Endpoint": llm_url, + "X-ITS-API-Key": args.api_key, + }, + )) + print_stats(f"envoy(budget={budget})", stats) + if stats and baseline_p50 is not None: + overhead = stats["p50_ms"] - baseline_p50 + print(f" Overhead vs direct: {overhead:+.1f}ms (p50)") + + if envoy_url: + print("\n--- pass-through ---") + print("\n[Envoy pass-through (no ITS)]") + stats = asyncio.run(benchmark_endpoint( + f"{envoy_url}/v1/chat/completions", + args.model_name, + args.num_requests, + args.concurrency, + )) + print_stats("envoy-passthrough", stats) + + print("\n" + "=" * 60) + print("Done.") + + finally: + print("\nShutting down services...") + stop_processes(processes) + for server in servers: + server.shutdown() + if envoy_tmpdir: + import shutil + shutil.rmtree(envoy_tmpdir, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/utils/iaas_helpers.py b/tests/e2e/utils/iaas_helpers.py new file mode 100644 index 00000000..8988f7fc --- /dev/null +++ b/tests/e2e/utils/iaas_helpers.py @@ -0,0 +1,327 @@ +"""Shared helpers for IaaS + Envoy e2e and performance tests.""" + +import json +import os +import shutil +import socket +import subprocess +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.request +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from tests.conftest import find_free_port + + +def wait_for_port(port, host="127.0.0.1", timeout=15): + """Wait until a port is accepting connections.""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + with socket.create_connection((host, port), timeout=1): + return True + except OSError: + time.sleep(0.2) + return False + + +def wait_for_http(url, timeout=15): + """Wait until an HTTP endpoint returns 200.""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + with urllib.request.urlopen(url, timeout=2) as resp: + if resp.status == 200: + return True + except Exception: + time.sleep(0.3) + return False + + +# --------------------------------------------------------------------------- +# HTTP helpers (urllib-based, no extra dependencies) +# --------------------------------------------------------------------------- + + +def _decode_json(raw): + try: + return json.loads(raw) + except (json.JSONDecodeError, UnicodeDecodeError): + return {"raw": raw.decode(errors="replace")[:500]} + + +def http_post(url, data, headers=None, timeout=30): + req_headers = {"Content-Type": "application/json"} + if headers: + req_headers.update(headers) + + body = json.dumps(data).encode() + req = urllib.request.Request(url, data=body, headers=req_headers, method="POST") + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.status, _decode_json(resp.read()) + except urllib.error.HTTPError as e: + return e.code, _decode_json(e.read()) + except urllib.error.URLError as e: + return 0, {"error": str(e)} + + +def http_get(url, timeout=10): + try: + with urllib.request.urlopen(url, timeout=timeout) as resp: + return resp.status, _decode_json(resp.read()) + except urllib.error.HTTPError as e: + return e.code, _decode_json(e.read()) + except urllib.error.URLError as e: + return 0, {"error": str(e)} + + +# --------------------------------------------------------------------------- +# Mock LLM server +# --------------------------------------------------------------------------- + + +class MockLLMHandler(BaseHTTPRequestHandler): + """OpenAI-compatible mock LLM with configurable latency.""" + + latency_ms = 0 + + def do_POST(self): + if self.path == "/v1/chat/completions": + content_length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(content_length)) if content_length else {} + + if self.latency_ms: + time.sleep(self.latency_ms / 1000.0) + + messages = body.get("messages", []) + user_msg = messages[-1]["content"] if messages else "unknown" + + its_headers = { + k: "" if k.lower() == "x-its-api-key" else v + for k, v in self.headers.items() + if k.lower().startswith("x-its-") + } + + response = { + "id": "mock-llm-001", + "object": "chat.completion", + "created": int(time.time()), + "model": body.get("model", "mock-model"), + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": f"Mock LLM response to: {user_msg}", + }, + "finish_reason": "stop", + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "its_headers_received": its_headers, + } + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response).encode()) + else: + self.send_response(404) + self.end_headers() + + def do_GET(self): + if self.path == "/health": + self.send_response(200) + self.end_headers() + self.wfile.write(b"OK") + elif self.path == "/v1/models": + response = {"data": [{"id": "mock-model", "object": "model"}]} + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response).encode()) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, format, *args): + pass + + +def start_mock_llm(port, latency_ms=0): + MockLLMHandler.latency_ms = latency_ms + server = ThreadingHTTPServer(("127.0.0.1", port), MockLLMHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return server + + +# --------------------------------------------------------------------------- +# Envoy config generation +# --------------------------------------------------------------------------- + + +_PROD_ENVOY_CONFIG = os.path.join( + os.path.dirname(__file__), os.pardir, os.pardir, os.pardir, + "its_hub", "integration", "iaas", "envoy_config.yaml", +) + + +def generate_envoy_config(envoy_port, ext_proc_port, iaas_port, llm_port): + """Load the production envoy_config.yaml and adapt it for local testing.""" + import yaml + + admin_port = find_free_port() + + with open(_PROD_ENVOY_CONFIG) as f: + config = yaml.safe_load(f) + + # --- Admin: simple, no access log --- + config["admin"] = { + "address": { + "socket_address": {"address": "127.0.0.1", "port_value": admin_port}, + }, + } + + # --- Listener port --- + listener = config["static_resources"]["listeners"][0] + sock_addr = listener["address"]["socket_address"] + sock_addr["port_value"] = envoy_port + sock_addr.pop("protocol", None) + + # --- HTTP connection manager tweaks --- + hcm = listener["filter_chains"][0]["filters"][0]["typed_config"] + hcm.pop("access_log", None) + + for hf in hcm.get("http_filters", []): + if "ext_proc" in hf.get("name", ""): + ep = hf["typed_config"] + ep["failure_mode_allow"] = False + ep["message_timeout"] = "10s" + ep["grpc_service"]["timeout"] = "10s" + + for vh in hcm.get("route_config", {}).get("virtual_hosts", []): + for route in vh.get("routes", []): + if "route" in route: + route["route"]["timeout"] = "60s" + + # --- Clusters: set test ports, remove health checks, use STATIC --- + port_map = { + "ext_proc_cluster": ext_proc_port, + "iaas_upstream": iaas_port, + "llm_upstream": llm_port, + } + for cluster in config["static_resources"]["clusters"]: + name = cluster["name"] + if name in port_map: + cluster["type"] = "STATIC" + cluster.pop("dns_lookup_family", None) + cluster.pop("health_checks", None) + ep = cluster["load_assignment"]["endpoints"][0]["lb_endpoints"][0]["endpoint"] + ep["address"]["socket_address"]["address"] = "127.0.0.1" + ep["address"]["socket_address"]["port_value"] = port_map[name] + + return yaml.dump(config, default_flow_style=False, sort_keys=False), admin_port + + +# --------------------------------------------------------------------------- +# Service lifecycle +# --------------------------------------------------------------------------- + + +def start_iaas_stack(llm_port): + """Start ext_proc + IaaS service. Returns (processes, iaas_url, ext_proc_port). + + Caller is responsible for calling stop_processes(processes) in a finally block. + """ + processes = [] + + ext_proc_port = find_free_port() + ext_proc_proc = subprocess.Popen( + [sys.executable, "-c", + f"import sys; sys.argv = ['its-iaas-ext-proc', '--port', '{ext_proc_port}']; " + f"from its_hub.integration.iaas.grpc_server import main; main()"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ) + processes.append(("ext_proc", ext_proc_proc)) + + iaas_port = find_free_port() + iaas_proc = subprocess.Popen( + [sys.executable, "-c", + f"import sys; sys.argv = ['its-iaas', '--port', '{iaas_port}']; " + f"from its_hub.integration.iaas.app_server import main; main()"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ) + processes.append(("iaas", iaas_proc)) + + iaas_url = f"http://127.0.0.1:{iaas_port}" + + print(f"Waiting for IaaS service on port {iaas_port}...") + if not wait_for_http(f"{iaas_url}/docs", timeout=20): + stop_processes(processes) + raise RuntimeError("IaaS service failed to start (health check timeout)") + print(f"IaaS service ready on port {iaas_port}") + + return processes, iaas_url, ext_proc_port + + +def start_envoy(ext_proc_port, iaas_port, llm_port): + """Start Envoy with generated config. Returns (process, envoy_url, tmpdir, admin_port) or None.""" + if not shutil.which("envoy"): + return None + + print(f"Waiting for ext_proc on port {ext_proc_port}...") + if not wait_for_port(ext_proc_port, timeout=15): + print("Warning: ext_proc not ready, skipping Envoy") + return None + print(f"ext_proc ready on port {ext_proc_port}") + + envoy_port = find_free_port() + tmpdir = tempfile.mkdtemp(prefix="its_envoy_") + config_path = os.path.join(tmpdir, "envoy.yaml") + config_text, admin_port = generate_envoy_config(envoy_port, ext_proc_port, iaas_port, llm_port) + with open(config_path, "w") as f: + f.write(config_text) + + envoy_proc = subprocess.Popen( + ["envoy", "-c", config_path, "--log-level", "warn"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ) + + print(f"Waiting for Envoy on port {envoy_port}...") + if not wait_for_port(envoy_port, timeout=15): + envoy_proc.terminate() + envoy_proc.wait() + shutil.rmtree(tmpdir, ignore_errors=True) + print("Warning: Envoy failed to start") + return None + + print(f"Envoy ready on port {envoy_port} (admin: {admin_port})") + time.sleep(1) + return envoy_proc, f"http://127.0.0.1:{envoy_port}", tmpdir, admin_port + + +def configure_iaas(iaas_url, llm_endpoint, model_name, api_key): + """Configure IaaS service. Raises on failure.""" + status, body = http_post(f"{iaas_url}/configure", { + "endpoint": llm_endpoint, + "api_key": api_key, + "model": model_name, + "alg": "self-consistency", + "regex_patterns": [r"\\boxed{([^}]+)}"], + }) + if status != 200: + raise RuntimeError(f"Failed to configure IaaS (status {status}): {body}") + + +def stop_processes(processes): + """Terminate and wait for all managed processes.""" + for name, proc in reversed(processes): + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + print(f" {name} stopped") diff --git a/tests/test_iaas_ext_processor.py b/tests/test_iaas_ext_processor.py new file mode 100644 index 00000000..8bee3c44 --- /dev/null +++ b/tests/test_iaas_ext_processor.py @@ -0,0 +1,310 @@ +"""Tests for the IaaS ext_proc router (integration/iaas/ext_processor.py). + +The IaaS ext_proc is a lightweight gRPC service that makes routing decisions: +- X-ITS-Budget present → route to IaaS service (set X-ITS-Route header) +- No X-ITS-Budget → pass through to upstream LLM (strip stray ITS headers) +""" + +# ruff: noqa: I001 +from unittest.mock import AsyncMock, MagicMock + +import pytest + +try: + import grpc + import its_hub.integration.proto # noqa: F401 + from envoy.config.core.v3 import base_pb2 + from envoy.service.ext_proc.v3 import external_processor_pb2 as ext_proc_pb2 + + HAS_GRPC = True +except ImportError: + HAS_GRPC = False + +pytestmark = pytest.mark.skipif(not HAS_GRPC, reason="ext_proc deps not installed") + + +# --------------------------------------------------------------------------- +# Proto helpers +# --------------------------------------------------------------------------- + + +def _header(key, value): + return base_pb2.HeaderValue(key=key, raw_value=value.encode("utf-8")) + + +def _request_headers(*extra_headers): + """Build a ProcessingRequest with request_headers.""" + headers = [ + _header(":path", "/v1/chat/completions"), + _header(":method", "POST"), + _header("content-type", "application/json"), + ] + for k, v in extra_headers: + headers.append(_header(k, v)) + return ext_proc_pb2.ProcessingRequest( + request_headers=ext_proc_pb2.HttpHeaders( + headers=base_pb2.HeaderMap(headers=headers) + ) + ) + + +def _response_headers(): + return ext_proc_pb2.ProcessingRequest( + response_headers=ext_proc_pb2.HttpHeaders( + headers=base_pb2.HeaderMap(headers=[]) + ) + ) + + +def _response_body(): + return ext_proc_pb2.ProcessingRequest( + response_body=ext_proc_pb2.HttpBody(body=b'{"choices": []}', end_of_stream=True) + ) + + +async def _async_iter(items): + for item in items: + yield item + + +async def _run_process(processor, requests, context=None): + if context is None: + context = MagicMock(spec=grpc.ServicerContext) + context.peer.return_value = "test-peer" + responses = [] + async for resp in processor.Process(_async_iter(requests), context): + responses.append(resp) + return responses + + +def _make_processor(): + from its_hub.integration.iaas.ext_processor import ExternalProcessorService + + return ExternalProcessorService() + + +# --------------------------------------------------------------------------- +# Tests: routing decisions +# --------------------------------------------------------------------------- + + +class TestIaaSExtProcRouting: + @pytest.mark.asyncio + async def test_budget_header_routes_to_iaas(self): + """X-ITS-Budget present → route to IaaS with X-ITS-Route header.""" + from its_hub.integration.iaas.ext_processor import _ROUTE_TO_IAAS + + processor = _make_processor() + request = _request_headers(("x-its-budget", "5")) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + assert responses[0] == _ROUTE_TO_IAAS + + resp_headers = responses[0].request_headers.response + assert resp_headers.clear_route_cache is True + + set_headers = resp_headers.header_mutation.set_headers + assert len(set_headers) == 1 + assert set_headers[0].header.key == "X-ITS-Route" + assert set_headers[0].header.raw_value == b"its-service" + + @pytest.mark.asyncio + async def test_budget_with_other_its_headers_routes_to_iaas(self): + """Multiple ITS headers with budget → still routes to IaaS.""" + from its_hub.integration.iaas.ext_processor import _ROUTE_TO_IAAS + + processor = _make_processor() + request = _request_headers( + ("x-its-budget", "3"), + ("x-its-endpoint", "http://llm/v1"), + ("x-its-api-key", "sk-test"), + ) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + assert responses[0] == _ROUTE_TO_IAAS + + @pytest.mark.asyncio + async def test_no_its_headers_passes_through(self): + """No ITS headers → pass through unchanged.""" + from its_hub.integration.iaas.ext_processor import _PASS_THROUGH + + processor = _make_processor() + request = _request_headers() + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + assert responses[0] == _PASS_THROUGH + + @pytest.mark.asyncio + async def test_its_headers_without_budget_strips_and_passes_through(self): + """ITS headers present but no X-ITS-Budget → strip stray headers, pass through.""" + processor = _make_processor() + request = _request_headers( + ("x-its-endpoint", "http://llm/v1"), + ("x-its-api-key", "sk-test"), + ) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + resp = responses[0] + mutation = resp.request_headers.response.header_mutation + assert "x-its-endpoint" in mutation.remove_headers + assert "x-its-api-key" in mutation.remove_headers + assert resp.request_headers.response.clear_route_cache is False + + @pytest.mark.asyncio + async def test_single_stray_its_header_stripped(self): + """A single stray X-ITS-Endpoint (no budget) → strip it.""" + processor = _make_processor() + request = _request_headers(("x-its-endpoint", "http://llm/v1")) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + mutation = responses[0].request_headers.response.header_mutation + assert "x-its-endpoint" in mutation.remove_headers + + +# --------------------------------------------------------------------------- +# Tests: response phases +# --------------------------------------------------------------------------- + + +class TestIaaSExtProcResponsePhases: + @pytest.mark.asyncio + async def test_response_headers_continue(self): + """response_headers phase → CONTINUE.""" + processor = _make_processor() + responses = await _run_process(processor, [_response_headers()]) + + assert len(responses) == 1 + status = responses[0].response_headers.response.status + assert status == ext_proc_pb2.CommonResponse.CONTINUE + + @pytest.mark.asyncio + async def test_response_body_continue(self): + """response_body phase → CONTINUE.""" + processor = _make_processor() + responses = await _run_process(processor, [_response_body()]) + + assert len(responses) == 1 + status = responses[0].response_body.response.status + assert status == ext_proc_pb2.CommonResponse.CONTINUE + + +# --------------------------------------------------------------------------- +# Tests: full request lifecycle +# --------------------------------------------------------------------------- + + +class TestIaaSExtProcLifecycle: + @pytest.mark.asyncio + async def test_full_its_request_lifecycle(self): + """Full lifecycle: request_headers → response_headers → response_body.""" + from its_hub.integration.iaas.ext_processor import _ROUTE_TO_IAAS + + processor = _make_processor() + requests = [ + _request_headers(("x-its-budget", "5")), + _response_headers(), + _response_body(), + ] + responses = await _run_process(processor, requests) + + assert len(responses) == 3 + assert responses[0] == _ROUTE_TO_IAAS + assert responses[1].response_headers.response.status == ext_proc_pb2.CommonResponse.CONTINUE + assert responses[2].response_body.response.status == ext_proc_pb2.CommonResponse.CONTINUE + + @pytest.mark.asyncio + async def test_full_passthrough_lifecycle(self): + """Full lifecycle without ITS headers: all phases CONTINUE.""" + from its_hub.integration.iaas.ext_processor import _PASS_THROUGH + + processor = _make_processor() + requests = [ + _request_headers(), + _response_headers(), + _response_body(), + ] + responses = await _run_process(processor, requests) + + assert len(responses) == 3 + assert responses[0] == _PASS_THROUGH + assert responses[1].HasField("response_headers") + assert responses[2].HasField("response_body") + + +# --------------------------------------------------------------------------- +# Tests: error handling +# --------------------------------------------------------------------------- + + +class TestIaaSExtProcErrors: + @pytest.mark.asyncio + async def test_stream_error_aborts_with_internal(self): + """Exception in request stream → abort with INTERNAL status.""" + processor = _make_processor() + context = AsyncMock(spec=grpc.aio.ServicerContext) + context.peer.return_value = "test-peer" + + async def _error_iter(): + raise RuntimeError("connection lost") + yield # noqa: unreachable — makes this an async generator + + responses = [] + async for resp in processor.Process(_error_iter(), context): + responses.append(resp) + + assert len(responses) == 0 + context.abort.assert_awaited_once_with( + grpc.StatusCode.INTERNAL, "connection lost" + ) + + @pytest.mark.asyncio + async def test_stream_error_logged(self, caplog): + """Stream error is logged at ERROR level.""" + processor = _make_processor() + context = AsyncMock(spec=grpc.aio.ServicerContext) + context.peer.return_value = "test-peer" + + async def _error_iter(): + raise ValueError("bad request") + yield # noqa: unreachable + + with caplog.at_level("ERROR"): + async for _ in processor.Process(_error_iter(), context): + pass + + assert any("Stream error" in record.message for record in caplog.records) + + +# --------------------------------------------------------------------------- +# Tests: case insensitivity +# --------------------------------------------------------------------------- + + +class TestIaaSExtProcCaseHandling: + @pytest.mark.asyncio + async def test_mixed_case_budget_header_routes(self): + """X-ITS-Budget with mixed case still routes (keys are lowered).""" + from its_hub.integration.iaas.ext_processor import _ROUTE_TO_IAAS + + processor = _make_processor() + request = _request_headers(("X-ITS-Budget", "5")) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + assert responses[0] == _ROUTE_TO_IAAS + + @pytest.mark.asyncio + async def test_uppercase_its_headers_stripped_on_passthrough(self): + """Mixed-case ITS headers without budget → lowered keys are stripped.""" + processor = _make_processor() + request = _request_headers(("X-ITS-Endpoint", "http://llm/v1")) + responses = await _run_process(processor, [request]) + + assert len(responses) == 1 + mutation = responses[0].request_headers.response.header_mutation + assert "x-its-endpoint" in mutation.remove_headers diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index e5c09655..0f08826f 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -391,6 +391,33 @@ async def test_concurrency_not_artificially_limited(self, orchestrator_cls): await orch.agenerate(lm, _make_batch(n)) assert lm.peak == n + @pytest.mark.asyncio + async def test_many_concurrent_agenerate_calls_do_not_deadlock(self, orchestrator_cls): + """Many concurrent agenerate() callers must not exhaust the default thread pool. + + Simulates the gateway scenario: N requests arrive simultaneously, + each calling agenerate() with a small batch. The orchestrator's + semaphore acquires must not starve the default executor that the + LM itself may need (e.g. for DNS, TLS, or other blocking I/O). + """ + orch = orchestrator_cls(max_concurrency=2) + + class ExecutorDependentLM: + """LM that uses the default executor, like a real HTTP client would.""" + async def agenerate_single(self, messages, loop=None, **kwargs): + _loop = loop or asyncio.get_running_loop() + await _loop.run_in_executor(None, lambda: None) + await asyncio.sleep(0.01) + return {"role": "assistant", "content": "ok"} + + lm = ExecutorDependentLM() + tasks = [ + asyncio.create_task(orch.agenerate(lm, _make_batch(1))) + for _ in range(50) + ] + results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=30) + assert len(results) == 50 + @pytest.mark.asyncio async def test_shared_semaphore_across_sequential_calls(self, orchestrator_cls): """Two overlapping agenerate calls share the same semaphore."""