"""Fail-fast bounded concurrency for heavy Jyotish API compute endpoints. Request threads that run rectification scans or high-rigor workflows share one process-wide semaphore sized for the 2 vCPU production host. Saturated requests return immediately; they are not queued. Health checks and other light routes must not call this gate. """ from __future__ import annotations import logging import os import threading import time DEFAULT_CONCURRENCY = 2 DEFAULT_RETRY_AFTER_SECONDS = 2 ENV_CONCURRENCY = "JYOTISH_HEAVY_COMPUTE_CONCURRENCY" ENV_RETRY_AFTER = "JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS" ENV_GATE_CHART_OVERVIEW = "JYOTISH_GATE_CHART_OVERVIEW" _LOG = logging.getLogger("jyotish.heavy_compute") HEAVY_COMPUTE_PATHS = frozenset( { "/api/rectification/sensitivity_scan", "/api/active_rectification_events", "/api/active_rectification_events_v4", "/api/rectification/v5/candidate-features", "/api/rectification/v5/score", "/api/rectification/v5/varga-segments", "/api/rectification/v5/diagnostics", "/api/rectification/v5/vedastro-validate", "/api/rectification/v5/range_reading", "/api/rectification/v5/block_scan", "/api/dynamic_rectification_opportunities", "/api/dynamic_rectification_score", "/api/high_rigor_workflow", "/api/consultation_workflow", "/api/professional_reading", "/api/professional_report_reference", "/api/vedastro/range_scan", "/api/vedastro_gateway/run", "/api/thematic_report", "/api/qizheng", "/api/western", "/api/ephemeris_events", } ) _state_lock = threading.Lock() _semaphore: threading.BoundedSemaphore | None = None _retry_after_seconds = DEFAULT_RETRY_AFTER_SECONDS _limit = DEFAULT_CONCURRENCY _holders: list[dict] = [] _holder_seq = 0 class _HeldSlot: """One acquired heavy-compute permit, paired with the path that holds it.""" def __init__(self, semaphore: threading.BoundedSemaphore, holder_id: int) -> None: self._semaphore = semaphore self._holder_id = holder_id self._released = False def release(self) -> None: with _state_lock: if self._released: return self._released = True _holders[:] = [item for item in _holders if item["id"] != self._holder_id] self._semaphore.release() class HeavyComputeBusy(RuntimeError): """No heavy-compute slot is free; callers must fail fast with HTTP 429.""" error_code = "ERR_COMPUTE_BUSY" def __init__(self, retry_after_seconds: int) -> None: super().__init__( "Heavy compute capacity is saturated; retry after the Retry-After delay." ) self.retry_after_seconds = max(int(retry_after_seconds), 1) def _parse_positive_int(raw: str | None, default: int) -> int: try: value = int(str(raw or "").strip()) except (TypeError, ValueError): return default return value if value >= 1 else default def is_heavy_compute_path(path: str) -> bool: if path in HEAVY_COMPUTE_PATHS: return True return path.startswith("/api/vedastro_gateway/jobs/") and path.endswith("/run") def reset_heavy_compute_gate( *, limit: int | None = None, retry_after_seconds: int | None = None, ) -> None: """Rebuild the process-wide semaphore. Tests must call this after env changes.""" global _semaphore, _retry_after_seconds, _limit, _holders resolved_limit = ( limit if limit is not None else _parse_positive_int(os.environ.get(ENV_CONCURRENCY), DEFAULT_CONCURRENCY) ) resolved_retry = ( retry_after_seconds if retry_after_seconds is not None else _parse_positive_int(os.environ.get(ENV_RETRY_AFTER), DEFAULT_RETRY_AFTER_SECONDS) ) with _state_lock: _retry_after_seconds = resolved_retry _limit = resolved_limit _holders = [] _semaphore = threading.BoundedSemaphore(resolved_limit) def _ensure_locked() -> tuple[threading.BoundedSemaphore, int]: global _semaphore, _retry_after_seconds, _limit if _semaphore is None: _limit = _parse_positive_int(os.environ.get(ENV_CONCURRENCY), DEFAULT_CONCURRENCY) _retry_after_seconds = _parse_positive_int( os.environ.get(ENV_RETRY_AFTER), DEFAULT_RETRY_AFTER_SECONDS ) _semaphore = threading.BoundedSemaphore(_limit) return _semaphore, _retry_after_seconds def _holder_view(now: float) -> list[dict]: return [ { "path": item["path"], "held_seconds": round(max(now - item["started"], 0.0), 3), } for item in _holders ] def _acquire_recorded(path: str) -> _HeldSlot: """Take one permit and remember which path holds it. Fail fast when full.""" global _holder_seq with _state_lock: semaphore, retry_after = _ensure_locked() if not semaphore.acquire(blocking=False): snapshot = _holder_view(time.monotonic()) failed = True slot = None else: _holder_seq += 1 holder_id = _holder_seq _holders.append({"id": holder_id, "path": path, "started": time.monotonic()}) failed = False snapshot = [] slot = _HeldSlot(semaphore, holder_id) if failed: _LOG.warning( "heavy_compute_busy rejected_path=%s holders=%s", path, snapshot, ) raise HeavyComputeBusy(retry_after) assert slot is not None return slot def acquire_heavy_compute_slot(path: str) -> _HeldSlot | None: """Acquire a slot for a gated path. Light paths return None. Fail-fast on saturation.""" if not is_heavy_compute_path(path): return None return _acquire_recorded(path) def _chart_overview_gate_enabled() -> bool: return os.environ.get(ENV_GATE_CHART_OVERVIEW, "").strip() == "1" def acquire_heavy_compute_slot_for_request(path: str, body: dict | None) -> _HeldSlot | None: """Path gate, plus the chart-overview switch and the warm-consultation exception. ``/api/chart`` stays out of ``HEAVY_COMPUTE_PATHS``. It takes a slot only when ``JYOTISH_GATE_CHART_OVERVIEW=1``, the body does not skip the VedAstro overview, and the body is not an async enqueue. Default is off. """ payload = body if isinstance(body, dict) else {} if path == "/api/consultation_workflow" and payload.get("consult_card_warm_only") is True: return None if path == "/api/chart" and _chart_overview_gate_enabled(): if payload.get("async") or payload.get("enqueue"): return None if payload.get("skip_vedastro_main_entry_overview"): return None return _acquire_recorded(path) return acquire_heavy_compute_slot(path) def heavy_compute_status() -> dict: """Process-local counts only. No request body and no user identity.""" with _state_lock: _ensure_locked() return {"limit": _limit, "in_use": len(_holders)} def release_heavy_compute_slot(slot: _HeldSlot | None) -> None: if slot is not None: slot.release()