[model-gateway] rename py_test to e2e_test (#16454)
This commit is contained in:
Simo Lin
@@ -0,0 +1 @@
|
||||
"""Shared fixtures for router integration tests."""
|
||||
@@ -0,0 +1,236 @@
|
||||
"""
|
||||
Generate self-signed certificates for mTLS integration testing.
|
||||
Creates a Certificate Authority (CA), server certificates, and client certificates.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import ipaddress
|
||||
from pathlib import Path
|
||||
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.x509.oid import NameOID
|
||||
|
||||
|
||||
def generate_private_key():
|
||||
"""Generate an RSA private key."""
|
||||
return rsa.generate_private_key(
|
||||
public_exponent=65537,
|
||||
key_size=2048,
|
||||
)
|
||||
|
||||
|
||||
def generate_ca_certificate():
|
||||
"""Generate a self-signed CA certificate."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = issuer = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "Test CA"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(issuer)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
|
||||
.add_extension(
|
||||
x509.BasicConstraints(ca=True, path_length=None),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_cert_sign=True,
|
||||
crl_sign=True,
|
||||
key_encipherment=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.sign(private_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def generate_server_certificate(ca_key, ca_cert):
|
||||
"""Generate a server certificate signed by the CA."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "localhost"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(ca_cert.subject)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
|
||||
.add_extension(
|
||||
x509.SubjectAlternativeName(
|
||||
[
|
||||
x509.DNSName("localhost"),
|
||||
x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_encipherment=True,
|
||||
key_cert_sign=False,
|
||||
crl_sign=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.ExtendedKeyUsage(
|
||||
[
|
||||
x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def generate_client_certificate(ca_key, ca_cert):
|
||||
"""Generate a client certificate signed by the CA."""
|
||||
private_key = generate_private_key()
|
||||
|
||||
subject = x509.Name(
|
||||
[
|
||||
x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
|
||||
x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.LOCALITY_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"),
|
||||
x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Test"),
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, "test-client"),
|
||||
]
|
||||
)
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(ca_cert.subject)
|
||||
.public_key(private_key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.utcnow())
|
||||
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
|
||||
.add_extension(
|
||||
x509.KeyUsage(
|
||||
digital_signature=True,
|
||||
key_encipherment=True,
|
||||
key_cert_sign=False,
|
||||
crl_sign=False,
|
||||
content_commitment=False,
|
||||
data_encipherment=False,
|
||||
key_agreement=False,
|
||||
encipher_only=False,
|
||||
decipher_only=False,
|
||||
),
|
||||
critical=True,
|
||||
)
|
||||
.add_extension(
|
||||
x509.ExtendedKeyUsage(
|
||||
[
|
||||
x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
|
||||
]
|
||||
),
|
||||
critical=False,
|
||||
)
|
||||
.sign(ca_key, hashes.SHA256())
|
||||
)
|
||||
|
||||
return private_key, cert
|
||||
|
||||
|
||||
def save_key(key, path: Path):
|
||||
"""Save private key to PEM file."""
|
||||
with open(path, "wb") as f:
|
||||
f.write(
|
||||
key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.TraditionalOpenSSL,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def save_cert(cert, path: Path):
|
||||
"""Save certificate to PEM file."""
|
||||
with open(path, "wb") as f:
|
||||
f.write(cert.public_bytes(serialization.Encoding.PEM))
|
||||
|
||||
|
||||
def generate_all_certificates(output_dir: Path):
|
||||
"""Generate all certificates and keys for mTLS testing."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print("==> Generating CA certificate...")
|
||||
ca_key, ca_cert = generate_ca_certificate()
|
||||
save_key(ca_key, output_dir / "ca-key.pem")
|
||||
save_cert(ca_cert, output_dir / "ca-cert.pem")
|
||||
|
||||
print("==> Generating server certificate...")
|
||||
server_key, server_cert = generate_server_certificate(ca_key, ca_cert)
|
||||
save_key(server_key, output_dir / "server-key.pem")
|
||||
save_cert(server_cert, output_dir / "server-cert.pem")
|
||||
|
||||
print("==> Generating client certificate...")
|
||||
client_key, client_cert = generate_client_certificate(ca_key, ca_cert)
|
||||
save_key(client_key, output_dir / "client-key.pem")
|
||||
save_cert(client_cert, output_dir / "client-cert.pem")
|
||||
|
||||
print(f"==> Certificates generated successfully in {output_dir}")
|
||||
print()
|
||||
print("Files created:")
|
||||
print(" - ca-cert.pem : CA certificate (for verifying server/client certs)")
|
||||
print(" - ca-key.pem : CA private key")
|
||||
print(" - server-cert.pem : Server certificate")
|
||||
print(" - server-key.pem : Server private key")
|
||||
print(" - client-cert.pem : Client certificate")
|
||||
print(" - client-key.pem : Client private key")
|
||||
print()
|
||||
print("Test server can use: server-cert.pem + server-key.pem")
|
||||
print("Test router can use: client-cert.pem + client-key.pem + ca-cert.pem")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
script_dir = Path(__file__).parent
|
||||
certs_dir = script_dir / "test_certs"
|
||||
generate_all_certificates(certs_dir)
|
||||
@@ -0,0 +1,288 @@
|
||||
"""
|
||||
Lightweight mock worker HTTP server for router integration tests.
|
||||
|
||||
Implements minimal endpoints used by the router:
|
||||
- GET /health, /health_generate
|
||||
- POST /generate, /v1/completions, /v1/chat/completions
|
||||
- POST /flush_cache
|
||||
- GET /get_server_info, /get_model_info, /v1/models
|
||||
|
||||
Behavior knobs are controlled via CLI flags to simulate failures, latency, and load.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Optional
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse
|
||||
|
||||
# Global state (per-process)
|
||||
_inflight = 0
|
||||
_failures_seen = 0
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--host", default="127.0.0.1")
|
||||
p.add_argument("--port", type=int, required=True)
|
||||
p.add_argument("--worker-id", default=None)
|
||||
p.add_argument("--latency-ms", type=int, default=0)
|
||||
p.add_argument("--timeout", action="store_true")
|
||||
p.add_argument("--status-code", type=int, default=200)
|
||||
p.add_argument("--fail-first-n", type=int, default=0)
|
||||
p.add_argument("--random-fail-rate", type=float, default=0.0)
|
||||
p.add_argument("--require-api-key", action="store_true")
|
||||
p.add_argument("--api-key", default=None)
|
||||
p.add_argument("--max-payload-bytes", type=int, default=10 * 1024 * 1024)
|
||||
p.add_argument("--stream", action="store_true")
|
||||
p.add_argument("--dp-size", type=int, default=1)
|
||||
p.add_argument("--crash-on-request", action="store_true")
|
||||
p.add_argument("--health-fail-after-ms", type=int, default=0)
|
||||
# TLS/mTLS configuration
|
||||
p.add_argument(
|
||||
"--ssl-certfile", type=str, default=None, help="Path to SSL certificate file"
|
||||
)
|
||||
p.add_argument("--ssl-keyfile", type=str, default=None, help="Path to SSL key file")
|
||||
p.add_argument(
|
||||
"--ssl-ca-certs",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to CA certificates for client verification",
|
||||
)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def _extract_worker_id(args: argparse.Namespace) -> str:
|
||||
if args.worker_id:
|
||||
return str(args.worker_id)
|
||||
# default to port (unique enough for tests)
|
||||
return f"worker-{args.port}"
|
||||
|
||||
|
||||
def create_app(args: argparse.Namespace) -> FastAPI:
|
||||
app = FastAPI()
|
||||
worker_id = _extract_worker_id(args)
|
||||
start_ts = time.time()
|
||||
crashed = {"done": False}
|
||||
|
||||
async def maybe_delay():
|
||||
if args.latency_ms > 0:
|
||||
await asyncio.sleep(args.latency_ms / 1000.0)
|
||||
|
||||
def should_fail() -> Optional[int]:
|
||||
global _failures_seen
|
||||
# Fail first N requests (500)
|
||||
if args.fail_first_n > 0 and _failures_seen < args.fail_first_n:
|
||||
_failures_seen += 1
|
||||
return 500
|
||||
# Random failure probability (500)
|
||||
if args.random_fail_rate > 0.0 and random.random() < args.random_fail_rate:
|
||||
return 500
|
||||
# Forced status code override (non-200) for all responses
|
||||
if args.status_code != 200:
|
||||
return int(args.status_code)
|
||||
return None
|
||||
|
||||
def check_api_key(request: Request):
|
||||
if not args.require_api_key:
|
||||
return
|
||||
auth = request.headers.get("Authorization")
|
||||
if not auth or not auth.startswith("Bearer "):
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
key = auth.split(" ", 1)[1]
|
||||
if args.api_key and key != args.api_key:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
|
||||
@asynccontextmanager
|
||||
async def track_inflight():
|
||||
global _inflight
|
||||
_inflight += 1
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_inflight -= 1
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
if (
|
||||
args.health_fail_after_ms
|
||||
and (time.time() - start_ts) * 1000.0 >= args.health_fail_after_ms
|
||||
):
|
||||
return PlainTextResponse("bad", status_code=500)
|
||||
return PlainTextResponse("ok", status_code=200)
|
||||
|
||||
@app.get("/health_generate")
|
||||
async def health_generate():
|
||||
return PlainTextResponse("ok", status_code=200)
|
||||
|
||||
@app.post("/flush_cache")
|
||||
async def flush_cache():
|
||||
return PlainTextResponse("ok", status_code=200)
|
||||
|
||||
@app.get("/get_model_info")
|
||||
async def get_model_info():
|
||||
return JSONResponse({"model": "mock", "vocab_size": 32000})
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def list_models():
|
||||
return JSONResponse({"data": [{"id": "mock", "object": "model"}]})
|
||||
|
||||
@app.get("/get_server_info")
|
||||
async def get_server_info(request: Request):
|
||||
# Enforce API key on server info when required (used by dp_aware probing)
|
||||
check_api_key(request)
|
||||
return JSONResponse(
|
||||
{
|
||||
"worker_id": worker_id,
|
||||
"load_in_flight": _inflight,
|
||||
"cache": {"size": 0, "hit_rate": 0.0},
|
||||
"dp_size": int(args.dp_size),
|
||||
}
|
||||
)
|
||||
|
||||
@app.get("/get_load")
|
||||
async def get_load(request: Request):
|
||||
check_api_key(request)
|
||||
# Return format matching real workers: array of load info per DP rank
|
||||
return JSONResponse(
|
||||
[
|
||||
{
|
||||
"dp_rank": 0,
|
||||
"num_reqs": _inflight,
|
||||
"num_waiting_reqs": 0,
|
||||
"num_tokens": _inflight,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
def make_json_response(obj: dict, status_code: int = 200) -> JSONResponse:
|
||||
resp = JSONResponse(obj, status_code=status_code)
|
||||
resp.headers["X-Worker-Id"] = worker_id
|
||||
return resp
|
||||
|
||||
async def handle_text_request(request: Request):
|
||||
# Authorization
|
||||
check_api_key(request)
|
||||
|
||||
# Payload limit
|
||||
body = await request.body()
|
||||
if len(body) > args.max_payload_bytes:
|
||||
return make_json_response({"error": "payload too large"}, status_code=413)
|
||||
|
||||
# Simulate crash on first request
|
||||
if args.crash_on_request and not crashed["done"]:
|
||||
crashed["done"] = True
|
||||
os._exit(1)
|
||||
|
||||
# Optional timeout (simulate hang)
|
||||
if args.timeout:
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
# Optional latency
|
||||
await maybe_delay()
|
||||
|
||||
# Optional failures
|
||||
fail_code = should_fail()
|
||||
if fail_code is not None and fail_code != 200:
|
||||
return make_json_response(
|
||||
{"error": f"mock failure {fail_code}"}, status_code=fail_code
|
||||
)
|
||||
|
||||
# Build response echoing minimal shape
|
||||
try:
|
||||
data = await request.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
data = {}
|
||||
|
||||
received_headers = {k.lower(): v for k, v in request.headers.items()}
|
||||
|
||||
now = time.time()
|
||||
ret = {
|
||||
"id": f"cmpl-{int(now*1000)}",
|
||||
"object": "text_completion",
|
||||
"created": int(now),
|
||||
"model": "mock",
|
||||
"choices": [
|
||||
{
|
||||
"text": "ok",
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"worker_id": worker_id,
|
||||
"echo": data,
|
||||
"received_headers": received_headers,
|
||||
}
|
||||
return make_json_response(ret, status_code=200)
|
||||
|
||||
async def handle_stream_request(request: Request):
|
||||
check_api_key(request)
|
||||
|
||||
async def gen():
|
||||
# minimal 2-chunk stream then [DONE]
|
||||
for i in range(2):
|
||||
await asyncio.sleep(0.01)
|
||||
chunk = {
|
||||
"choices": [{"delta": {"content": "x"}}],
|
||||
"worker_id": worker_id,
|
||||
}
|
||||
yield f"data: {json.dumps(chunk)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
headers = {"X-Worker-Id": worker_id}
|
||||
return StreamingResponse(gen(), media_type="text/event-stream", headers=headers)
|
||||
|
||||
@app.post("/generate")
|
||||
async def generate(request: Request):
|
||||
async with track_inflight():
|
||||
if args.stream:
|
||||
return await handle_stream_request(request)
|
||||
return await handle_text_request(request)
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def completions(request: Request):
|
||||
async with track_inflight():
|
||||
if args.stream:
|
||||
return await handle_stream_request(request)
|
||||
return await handle_text_request(request)
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat_completions(request: Request):
|
||||
async with track_inflight():
|
||||
if args.stream:
|
||||
return await handle_stream_request(request)
|
||||
return await handle_text_request(request)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
app = create_app(args)
|
||||
# Handle SIGTERM gracefully for fast test teardown
|
||||
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
|
||||
|
||||
# Configure SSL if certificates are provided
|
||||
ssl_config = {}
|
||||
if args.ssl_certfile and args.ssl_keyfile:
|
||||
ssl_config["ssl_certfile"] = args.ssl_certfile
|
||||
ssl_config["ssl_keyfile"] = args.ssl_keyfile
|
||||
# If CA certs provided, require client certificates (mTLS)
|
||||
if args.ssl_ca_certs:
|
||||
ssl_config["ssl_ca_certs"] = args.ssl_ca_certs
|
||||
ssl_config["ssl_cert_reqs"] = 2 # ssl.CERT_REQUIRED
|
||||
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="warning", **ssl_config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,8 @@
|
||||
import socket
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
"""Return an available TCP port on localhost."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
@@ -0,0 +1,248 @@
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from .ports import find_free_port
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcHandle:
|
||||
process: subprocess.Popen
|
||||
url: str
|
||||
|
||||
|
||||
class RouterManager:
|
||||
"""Helper to spawn a router process and interact with admin endpoints."""
|
||||
|
||||
def __init__(self):
|
||||
self._children: List[subprocess.Popen] = []
|
||||
|
||||
def start_router(
|
||||
self,
|
||||
worker_urls: Optional[List[str]] = None,
|
||||
policy: str = "round_robin",
|
||||
port: Optional[int] = None,
|
||||
extra: Optional[Dict] = None,
|
||||
# PD options
|
||||
pd_disaggregation: bool = False,
|
||||
prefill_urls: Optional[List[tuple]] = None,
|
||||
decode_urls: Optional[List[str]] = None,
|
||||
prefill_policy: Optional[str] = None,
|
||||
decode_policy: Optional[str] = None,
|
||||
) -> ProcHandle:
|
||||
worker_urls = worker_urls or []
|
||||
port = port or find_free_port()
|
||||
cmd = [
|
||||
"python3",
|
||||
"-m",
|
||||
"sglang_router.launch_router",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--policy",
|
||||
policy,
|
||||
]
|
||||
# Avoid Prometheus port collisions by assigning a free port per router
|
||||
prom_port = find_free_port()
|
||||
cmd.extend(
|
||||
["--prometheus-port", str(prom_port), "--prometheus-host", "127.0.0.1"]
|
||||
)
|
||||
if worker_urls:
|
||||
cmd.extend(["--worker-urls", *worker_urls])
|
||||
|
||||
# PD routing configuration
|
||||
if pd_disaggregation:
|
||||
cmd.append("--pd-disaggregation")
|
||||
if prefill_urls:
|
||||
for url, bport in prefill_urls:
|
||||
if bport is None:
|
||||
cmd.extend(["--prefill", url, "none"])
|
||||
else:
|
||||
cmd.extend(["--prefill", url, str(bport)])
|
||||
if decode_urls:
|
||||
for url in decode_urls:
|
||||
cmd.extend(["--decode", url])
|
||||
if prefill_policy:
|
||||
cmd.extend(["--prefill-policy", prefill_policy])
|
||||
if decode_policy:
|
||||
cmd.extend(["--decode-policy", decode_policy])
|
||||
|
||||
# Map supported extras to CLI flags (subset for integration)
|
||||
if extra:
|
||||
flag_map = {
|
||||
"max_payload_size": "--max-payload-size",
|
||||
"dp_aware": "--dp-aware",
|
||||
"api_key": "--api-key",
|
||||
# Health/monitoring
|
||||
"worker_startup_check_interval": "--worker-startup-check-interval",
|
||||
# Cache-aware tuning
|
||||
"cache_threshold": "--cache-threshold",
|
||||
"balance_abs_threshold": "--balance-abs-threshold",
|
||||
"balance_rel_threshold": "--balance-rel-threshold",
|
||||
# Retry
|
||||
"retry_max_retries": "--retry-max-retries",
|
||||
"retry_initial_backoff_ms": "--retry-initial-backoff-ms",
|
||||
"retry_max_backoff_ms": "--retry-max-backoff-ms",
|
||||
"retry_backoff_multiplier": "--retry-backoff-multiplier",
|
||||
"retry_jitter_factor": "--retry-jitter-factor",
|
||||
"disable_retries": "--disable-retries",
|
||||
# Circuit breaker
|
||||
"cb_failure_threshold": "--cb-failure-threshold",
|
||||
"cb_success_threshold": "--cb-success-threshold",
|
||||
"cb_timeout_duration_secs": "--cb-timeout-duration-secs",
|
||||
"cb_window_duration_secs": "--cb-window-duration-secs",
|
||||
"disable_circuit_breaker": "--disable-circuit-breaker",
|
||||
# Rate limiting
|
||||
"max_concurrent_requests": "--max-concurrent-requests",
|
||||
"queue_size": "--queue-size",
|
||||
"queue_timeout_secs": "--queue-timeout-secs",
|
||||
"rate_limit_tokens_per_second": "--rate-limit-tokens-per-second",
|
||||
# mTLS configuration
|
||||
"client_cert_path": "--client-cert-path",
|
||||
"client_key_path": "--client-key-path",
|
||||
"ca_cert_paths": "--ca-cert-paths",
|
||||
}
|
||||
for k, v in extra.items():
|
||||
if v is None:
|
||||
continue
|
||||
flag = flag_map.get(k)
|
||||
if not flag:
|
||||
continue
|
||||
if isinstance(v, bool):
|
||||
if v:
|
||||
cmd.append(flag)
|
||||
elif isinstance(v, list):
|
||||
# Handle list arguments (e.g., ca_cert_paths)
|
||||
if v: # Only add if list is not empty
|
||||
cmd.append(flag)
|
||||
cmd.extend([str(item) for item in v])
|
||||
else:
|
||||
cmd.extend([flag, str(v)])
|
||||
|
||||
proc = subprocess.Popen(cmd)
|
||||
self._children.append(proc)
|
||||
url = f"http://127.0.0.1:{port}"
|
||||
self._wait_health(url)
|
||||
return ProcHandle(process=proc, url=url)
|
||||
|
||||
def _wait_health(self, base_url: str, timeout: float = 30.0):
|
||||
start = time.time()
|
||||
with requests.Session() as s:
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = s.get(f"{base_url}/health", timeout=2)
|
||||
if r.status_code == 200:
|
||||
return
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.2)
|
||||
raise TimeoutError(f"Router at {base_url} did not become healthy")
|
||||
|
||||
def add_worker(self, base_url: str, worker_url: str, timeout: float = 30.0) -> None:
|
||||
r = requests.post(f"{base_url}/workers", json={"url": worker_url})
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"add_worker failed: {r.status_code} {r.text}" # ACCEPTED status
|
||||
|
||||
payload = r.json()
|
||||
worker_id = payload.get("worker_id")
|
||||
assert worker_id, f"add_worker did not return worker_id: {payload}"
|
||||
|
||||
# Poll until worker is actually added and healthy
|
||||
start = time.time()
|
||||
with requests.Session() as s:
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = s.get(f"{base_url}/workers/{worker_id}", timeout=2)
|
||||
if r.status_code == 200:
|
||||
data = r.json()
|
||||
# Check if registration job failed
|
||||
job_status = data.get("job_status")
|
||||
if job_status and job_status.get("state") == "failed":
|
||||
raise RuntimeError(
|
||||
f"Worker registration failed: {job_status.get('message', 'Unknown error')}"
|
||||
)
|
||||
# Check if worker is healthy and registered (not just in job queue)
|
||||
if data.get("is_healthy", False):
|
||||
return
|
||||
# Worker not ready yet, continue polling
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.1)
|
||||
raise TimeoutError(
|
||||
f"Worker {worker_url} was not added and healthy after {timeout}s"
|
||||
)
|
||||
|
||||
def remove_worker(
|
||||
self, base_url: str, worker_url: str, timeout: float = 30.0
|
||||
) -> None:
|
||||
# Resolve worker_id from the current registry snapshot
|
||||
r_list = requests.get(f"{base_url}/workers")
|
||||
assert (
|
||||
r_list.status_code == 200
|
||||
), f"list_workers failed: {r_list.status_code} {r_list.text}"
|
||||
workers = r_list.json().get("workers", [])
|
||||
worker_id = next(
|
||||
(w.get("id") for w in workers if w.get("url") == worker_url), None
|
||||
)
|
||||
assert (
|
||||
worker_id
|
||||
), f"could not find worker_id for url={worker_url}. workers={workers}"
|
||||
|
||||
r = requests.delete(f"{base_url}/workers/{worker_id}")
|
||||
assert (
|
||||
r.status_code == 202
|
||||
), f"remove_worker failed: {r.status_code} {r.text}" # ACCEPTED status
|
||||
|
||||
# Poll until worker is actually removed (GET returns 404) or timeout
|
||||
start = time.time()
|
||||
last_status = None
|
||||
with requests.Session() as s:
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = s.get(f"{base_url}/workers/{worker_id}", timeout=2)
|
||||
if r.status_code == 404:
|
||||
# Worker successfully removed
|
||||
return
|
||||
elif r.status_code == 200:
|
||||
# Check if removal job failed
|
||||
data = r.json()
|
||||
job_status = data.get("job_status")
|
||||
if job_status:
|
||||
last_status = job_status
|
||||
if job_status.get("state") == "failed":
|
||||
raise RuntimeError(
|
||||
f"Worker removal failed: {job_status.get('message', 'Unknown error')}"
|
||||
)
|
||||
# Worker still being processed, continue polling
|
||||
except requests.RequestException:
|
||||
pass
|
||||
time.sleep(0.1)
|
||||
|
||||
# Provide detailed timeout error with last known status
|
||||
error_msg = f"Worker {worker_url} was not removed after {timeout}s"
|
||||
if last_status:
|
||||
error_msg += f". Last job status: {last_status}"
|
||||
raise TimeoutError(error_msg)
|
||||
|
||||
def list_workers(self, base_url: str) -> list[str]:
|
||||
r = requests.get(f"{base_url}/workers")
|
||||
assert r.status_code == 200, f"list_workers failed: {r.status_code} {r.text}"
|
||||
data = r.json()
|
||||
# Extract URLs from WorkerInfo objects
|
||||
workers = data.get("workers", [])
|
||||
return [w["url"] for w in workers]
|
||||
|
||||
def stop_all(self):
|
||||
for p in self._children:
|
||||
if p.poll() is None:
|
||||
p.terminate()
|
||||
try:
|
||||
p.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
p.kill()
|
||||
self._children.clear()
|
||||
Reference in New Issue
Block a user