jesse-uxandGrok (agent loop)
a54f75e5e3
Consult, rectification turns, and report generation stop before a credit reservation when the product caps are hit. A full heavy-compute gate waits up to 8 seconds, and one consultation retries that 429 inside the same request. Co-Authored-By: Grok (agent loop)
343 lines
12 KiB
Python
343 lines
12 KiB
Python
"""Fail-fast bounded concurrency for heavy API compute endpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import threading
|
|
import time
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
|
|
from scripts import api_heavy_compute_gate as gate
|
|
from scripts.api_heavy_compute_gate import (
|
|
DEFAULT_CONCURRENCY,
|
|
HEAVY_COMPUTE_PATHS,
|
|
HeavyComputeBusy,
|
|
acquire_heavy_compute_slot,
|
|
heavy_compute_waiter_count,
|
|
is_heavy_compute_path,
|
|
release_heavy_compute_slot,
|
|
reset_heavy_compute_gate,
|
|
)
|
|
from scripts.jyotish_api_server import (
|
|
DEFAULT_ALLOWED_HOSTS,
|
|
DEFAULT_ALLOWED_ORIGINS,
|
|
JyotishAPIHandler,
|
|
)
|
|
|
|
|
|
class _FakeHeaders(dict):
|
|
def get(self, key, default=None):
|
|
return super().get(key, default)
|
|
|
|
|
|
class _FakeServer:
|
|
allowed_origins = DEFAULT_ALLOWED_ORIGINS
|
|
allowed_hosts = DEFAULT_ALLOWED_HOSTS
|
|
|
|
|
|
class _PostCaptureHandler(JyotishAPIHandler):
|
|
def __init__(self, path: str, payload: dict) -> None:
|
|
raw = json.dumps(payload).encode("utf-8")
|
|
self.headers = _FakeHeaders(
|
|
{
|
|
"Content-Length": str(len(raw)),
|
|
"Content-Type": "application/json",
|
|
}
|
|
)
|
|
self.server = _FakeServer()
|
|
self.path = path
|
|
self.rfile = BytesIO(raw)
|
|
self.wfile = BytesIO()
|
|
self.status_code = None
|
|
self.response_headers = []
|
|
self.client_address = ("test-heavy-compute", 0)
|
|
|
|
def send_response(self, code, message=None): # noqa: ANN001
|
|
self.status_code = code
|
|
|
|
def send_header(self, key, value): # noqa: ANN001
|
|
self.response_headers.append((key, value))
|
|
|
|
def end_headers(self):
|
|
return None
|
|
|
|
def payload(self) -> dict:
|
|
return json.loads(self.wfile.getvalue().decode("utf-8"))
|
|
|
|
|
|
class _GetCaptureHandler(JyotishAPIHandler):
|
|
def __init__(self, path: str) -> None:
|
|
self.headers = _FakeHeaders()
|
|
self.server = _FakeServer()
|
|
self.path = path
|
|
self.wfile = BytesIO()
|
|
self.status_code = None
|
|
self.response_headers = []
|
|
self.client_address = ("test-heavy-compute", 0)
|
|
|
|
def send_response(self, code, message=None): # noqa: ANN001
|
|
self.status_code = code
|
|
|
|
def send_header(self, key, value): # noqa: ANN001
|
|
self.response_headers.append((key, value))
|
|
|
|
def end_headers(self):
|
|
return None
|
|
|
|
def payload(self) -> dict:
|
|
return json.loads(self.wfile.getvalue().decode("utf-8"))
|
|
|
|
|
|
@pytest.fixture
|
|
def limit_one_gate(monkeypatch):
|
|
monkeypatch.setenv("JYOTISH_HEAVY_COMPUTE_CONCURRENCY", "1")
|
|
monkeypatch.setenv("JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS", "2")
|
|
monkeypatch.setenv("JYOTISH_API_RATE_LIMIT_PER_MINUTE", "0")
|
|
# 原值：reset() 使用默认「满了立刻拒绝」
|
|
# 新值：wait_seconds=0，本夹具里的 429 断言仍是立刻拒绝
|
|
# 原因：产品改为最多等 8 秒；排队成功、等待者满、超时另有测试
|
|
reset_heavy_compute_gate(wait_seconds=0)
|
|
yield
|
|
monkeypatch.delenv("JYOTISH_HEAVY_COMPUTE_CONCURRENCY", raising=False)
|
|
monkeypatch.delenv("JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS", raising=False)
|
|
reset_heavy_compute_gate()
|
|
|
|
|
|
def test_health_and_light_paths_are_not_gated() -> None:
|
|
assert not is_heavy_compute_path("/api/health")
|
|
assert not is_heavy_compute_path("/api/cities")
|
|
assert not is_heavy_compute_path("/api/location/resolve")
|
|
assert not is_heavy_compute_path("/api/chart/jobs/abc")
|
|
assert not is_heavy_compute_path("/api/high_rigor_workflow/jobs/abc")
|
|
assert not is_heavy_compute_path("/api/rectification/v5/versions")
|
|
assert is_heavy_compute_path("/api/high_rigor_workflow")
|
|
assert is_heavy_compute_path("/api/professional_report_reference")
|
|
assert is_heavy_compute_path("/api/rectification/sensitivity_scan")
|
|
assert is_heavy_compute_path("/api/vedastro_gateway/jobs/job1/run")
|
|
assert "/api/chart" not in HEAVY_COMPUTE_PATHS
|
|
assert is_heavy_compute_path("/api/qizheng")
|
|
assert is_heavy_compute_path("/api/western")
|
|
assert is_heavy_compute_path("/api/ephemeris_events")
|
|
assert DEFAULT_CONCURRENCY == 2
|
|
|
|
|
|
def test_acquire_fail_fast_then_succeeds_after_release() -> None:
|
|
# 原值：reset(limit=1, retry_after_seconds=3) 后第二路立刻 HeavyComputeBusy
|
|
# 新值：多传 wait_seconds=0
|
|
# 原因：默认改为短排队；本测试仍锁 Retry-After=3 与 ERR_COMPUTE_BUSY，不锁「禁止等待」
|
|
reset_heavy_compute_gate(limit=1, retry_after_seconds=3, wait_seconds=0)
|
|
first = acquire_heavy_compute_slot("/api/high_rigor_workflow")
|
|
assert first is not None
|
|
with pytest.raises(HeavyComputeBusy) as caught:
|
|
acquire_heavy_compute_slot("/api/rectification/sensitivity_scan")
|
|
assert caught.value.retry_after_seconds == 3
|
|
assert caught.value.error_code == "ERR_COMPUTE_BUSY"
|
|
light = acquire_heavy_compute_slot("/api/health")
|
|
assert light is None
|
|
release_heavy_compute_slot(first)
|
|
second = acquire_heavy_compute_slot("/api/consultation_workflow")
|
|
assert second is not None
|
|
release_heavy_compute_slot(second)
|
|
reset_heavy_compute_gate()
|
|
|
|
|
|
def test_saturated_request_returns_429_with_retry_after(limit_one_gate, monkeypatch) -> None:
|
|
started = threading.Event()
|
|
release_first = threading.Event()
|
|
first_status = {}
|
|
|
|
def _slow_compute(self, body): # noqa: ANN001
|
|
started.set()
|
|
assert release_first.wait(timeout=5)
|
|
return {"ok": True, "endpoint": "high_rigor_workflow"}
|
|
|
|
monkeypatch.setattr(JyotishAPIHandler, "_compute_high_rigor_workflow", _slow_compute)
|
|
|
|
def _run_first() -> None:
|
|
handler = _PostCaptureHandler("/api/high_rigor_workflow", {})
|
|
handler.do_POST()
|
|
first_status["code"] = handler.status_code
|
|
first_status["payload"] = handler.payload()
|
|
|
|
worker = threading.Thread(target=_run_first)
|
|
worker.start()
|
|
assert started.wait(timeout=5)
|
|
|
|
blocked = _PostCaptureHandler("/api/high_rigor_workflow", {})
|
|
blocked.do_POST()
|
|
assert blocked.status_code == 429
|
|
assert ("Retry-After", "2") in blocked.response_headers
|
|
payload = blocked.payload()
|
|
assert payload["success"] is False
|
|
assert payload["error_code"] == "ERR_COMPUTE_BUSY"
|
|
|
|
health = _GetCaptureHandler("/api/health")
|
|
health.do_GET()
|
|
assert health.status_code == 200
|
|
assert health.payload()["status"] == "ok"
|
|
|
|
light = _PostCaptureHandler("/api/location/resolve", {"city": "beijing"})
|
|
light.do_POST()
|
|
assert light.status_code == 200
|
|
assert light.payload()["status"] == "local_city_match"
|
|
|
|
release_first.set()
|
|
worker.join(timeout=5)
|
|
assert not worker.is_alive()
|
|
assert first_status["code"] == 200
|
|
|
|
recovered = _PostCaptureHandler("/api/high_rigor_workflow", {})
|
|
recovered.do_POST()
|
|
assert recovered.status_code == 200
|
|
assert recovered.payload()["ok"] is True
|
|
|
|
|
|
def test_warm_consultation_does_not_consume_the_user_slot(limit_one_gate, monkeypatch) -> None:
|
|
"""A warm in progress must leave the only heavy-compute slot for a real question."""
|
|
warm_entered = threading.Event()
|
|
release_warm = threading.Event()
|
|
user_entered = threading.Event()
|
|
release_user = threading.Event()
|
|
|
|
def _compute(self, body): # noqa: ANN001
|
|
if body.get("consult_card_warm_only") is True:
|
|
warm_entered.set()
|
|
assert release_warm.wait(timeout=5)
|
|
return {"success": True, "warmed": True}
|
|
user_entered.set()
|
|
assert release_user.wait(timeout=5)
|
|
return {"success": True, "endpoint": "consultation_workflow"}
|
|
|
|
monkeypatch.setattr(JyotishAPIHandler, "_compute_consultation_workflow", _compute)
|
|
|
|
warm_box: dict = {}
|
|
user_box: dict = {}
|
|
|
|
def _run_warm() -> None:
|
|
handler = _PostCaptureHandler(
|
|
"/api/consultation_workflow",
|
|
{"consult_card_warm_only": True},
|
|
)
|
|
handler.do_POST()
|
|
warm_box["code"] = handler.status_code
|
|
warm_box["payload"] = handler.payload()
|
|
|
|
def _run_user() -> None:
|
|
handler = _PostCaptureHandler("/api/consultation_workflow", {"question": "career"})
|
|
handler.do_POST()
|
|
user_box["code"] = handler.status_code
|
|
user_box["payload"] = handler.payload()
|
|
|
|
warm_thread = threading.Thread(target=_run_warm)
|
|
warm_thread.start()
|
|
assert warm_entered.wait(timeout=5)
|
|
|
|
user_thread = threading.Thread(target=_run_user)
|
|
user_thread.start()
|
|
assert user_entered.wait(timeout=5)
|
|
|
|
blocked = _PostCaptureHandler("/api/consultation_workflow", {"question": "again"})
|
|
blocked.do_POST()
|
|
assert blocked.status_code == 429
|
|
assert blocked.payload()["error_code"] == "ERR_COMPUTE_BUSY"
|
|
|
|
release_user.set()
|
|
release_warm.set()
|
|
user_thread.join(timeout=5)
|
|
warm_thread.join(timeout=5)
|
|
assert not user_thread.is_alive()
|
|
assert not warm_thread.is_alive()
|
|
assert user_box["code"] == 200
|
|
assert user_box["payload"]["endpoint"] == "consultation_workflow"
|
|
assert warm_box["code"] == 200
|
|
assert warm_box["payload"] == {"success": True, "warmed": True}
|
|
assert "packet" not in warm_box["payload"]
|
|
|
|
|
|
def _wait_until(predicate, timeout: float = 2.0) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.01)
|
|
return False
|
|
|
|
|
|
def test_third_request_gets_the_slot_after_one_holder_releases() -> None:
|
|
reset_heavy_compute_gate(limit=2, retry_after_seconds=2, wait_seconds=2, max_waiters=4)
|
|
first = acquire_heavy_compute_slot("/api/consultation_workflow")
|
|
second = acquire_heavy_compute_slot("/api/rectification/v5/score")
|
|
assert first is not None and second is not None
|
|
found: dict = {}
|
|
|
|
def _third() -> None:
|
|
found["slot"] = acquire_heavy_compute_slot("/api/consultation_workflow")
|
|
|
|
worker = threading.Thread(target=_third)
|
|
worker.start()
|
|
assert _wait_until(lambda: heavy_compute_waiter_count() == 1)
|
|
release_heavy_compute_slot(first)
|
|
worker.join(timeout=3)
|
|
assert not worker.is_alive()
|
|
assert found["slot"] is not None
|
|
release_heavy_compute_slot(second)
|
|
release_heavy_compute_slot(found["slot"])
|
|
reset_heavy_compute_gate()
|
|
|
|
|
|
def test_waiter_overflow_rejects_immediately() -> None:
|
|
reset_heavy_compute_gate(limit=1, retry_after_seconds=2, wait_seconds=2, max_waiters=1)
|
|
held = acquire_heavy_compute_slot("/api/high_rigor_workflow")
|
|
assert held is not None
|
|
parked: dict = {}
|
|
|
|
def _wait() -> None:
|
|
try:
|
|
parked["slot"] = acquire_heavy_compute_slot("/api/consultation_workflow")
|
|
except HeavyComputeBusy as exc:
|
|
parked["error"] = exc
|
|
|
|
waiter = threading.Thread(target=_wait)
|
|
waiter.start()
|
|
assert _wait_until(lambda: heavy_compute_waiter_count() == 1)
|
|
started = time.monotonic()
|
|
with pytest.raises(HeavyComputeBusy) as caught:
|
|
acquire_heavy_compute_slot("/api/rectification/v5/score")
|
|
assert time.monotonic() - started < 0.3
|
|
assert caught.value.retry_after_seconds == 2
|
|
assert caught.value.error_code == "ERR_COMPUTE_BUSY"
|
|
release_heavy_compute_slot(held)
|
|
waiter.join(timeout=3)
|
|
assert not waiter.is_alive()
|
|
assert parked.get("slot") is not None
|
|
release_heavy_compute_slot(parked["slot"])
|
|
reset_heavy_compute_gate()
|
|
|
|
|
|
def test_queue_timeout_still_returns_busy() -> None:
|
|
reset_heavy_compute_gate(limit=1, retry_after_seconds=2, wait_seconds=0.2, max_waiters=4)
|
|
held = acquire_heavy_compute_slot("/api/professional_report_reference")
|
|
assert held is not None
|
|
started = time.monotonic()
|
|
with pytest.raises(HeavyComputeBusy) as caught:
|
|
acquire_heavy_compute_slot("/api/consultation_workflow")
|
|
elapsed = time.monotonic() - started
|
|
assert 0.1 <= elapsed < 2.0
|
|
assert caught.value.retry_after_seconds == 2
|
|
assert caught.value.error_code == "ERR_COMPUTE_BUSY"
|
|
light = acquire_heavy_compute_slot("/api/health")
|
|
assert light is None
|
|
release_heavy_compute_slot(held)
|
|
reset_heavy_compute_gate()
|
|
|
|
|
|
def test_quality_gate_runs_heavy_compute_gate() -> None:
|
|
from scripts.run_quality_gate import CORE_PYTEST_TARGETS
|
|
|
|
assert "tests/test_api_heavy_compute_gate.py" in CORE_PYTEST_TARGETS
|
|
assert gate.DEFAULT_CONCURRENCY == 2
|
|
assert gate.DEFAULT_QUEUE_WAIT_SECONDS == 8
|
|
assert gate.DEFAULT_MAX_WAITERS == 4
|