sglang_v0.5.2/sglang/sgl-router/py_test/integration/conftest.py

110 lines
2.8 KiB
Python

import os
import subprocess
import time
from pathlib import Path
from typing import Dict, Iterable, List, Tuple
import pytest
import requests
from ..fixtures.ports import find_free_port
from ..fixtures.router_manager import RouterManager
def pytest_configure(config):
config.addinivalue_line("markers", "integration: mark as router integration test")
@pytest.fixture
def router_manager() -> Iterable[RouterManager]:
mgr = RouterManager()
try:
yield mgr
finally:
mgr.stop_all()
def _spawn_mock_worker(args: List[str]) -> Tuple[subprocess.Popen, str, str]:
repo_root = Path(__file__).resolve().parents[2]
script = repo_root / "py_test" / "fixtures" / "mock_worker.py"
port = find_free_port()
worker_id = f"worker-{port}"
base_cmd = [
"python3",
str(script),
"--port",
str(port),
"--worker-id",
worker_id,
]
cmd = base_cmd + args
proc = subprocess.Popen(cmd)
url = f"http://127.0.0.1:{port}"
_wait_health(url)
return proc, url, worker_id
def _wait_health(url: str, timeout: float = 10.0):
start = time.time()
with requests.Session() as s:
while time.time() - start < timeout:
try:
r = s.get(f"{url}/health", timeout=1)
if r.status_code == 200:
return
except requests.RequestException:
pass
time.sleep(0.1)
raise TimeoutError(f"Mock worker at {url} did not become healthy")
@pytest.fixture
def mock_worker():
"""Start a single healthy mock worker; yields (process, url, worker_id)."""
proc, url, worker_id = _spawn_mock_worker([])
try:
yield proc, url, worker_id
finally:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout=3)
except subprocess.TimeoutExpired:
proc.kill()
@pytest.fixture
def mock_workers():
"""Factory to start N workers with custom args.
Usage:
procs, urls, ids = mock_workers(n=3, args=["--latency-ms", "5"]) # same args for all
...
"""
procs: List[subprocess.Popen] = []
def _start(n: int, args: List[str] | None = None):
args = args or []
new_procs: List[subprocess.Popen] = []
urls: List[str] = []
ids: List[str] = []
for _ in range(n):
p, url, wid = _spawn_mock_worker(args)
procs.append(p)
new_procs.append(p)
urls.append(url)
ids.append(wid)
return new_procs, urls, ids
try:
yield _start
finally:
for p in procs:
if p.poll() is None:
p.terminate()
try:
p.wait(timeout=3)
except subprocess.TimeoutExpired:
p.kill()