Jyotisha/scripts/api_process_pool/pool.py

"""Run the heaviest compute endpoints in worker processes (issue #133).

The API server is one ``ThreadingHTTPServer`` process. Its heavy routes are
pure-Python number crunching, so two requests inside the gate's two slots
still shared one core under the GIL: one score alone 1.19 s, two at once
2.35-2.6 s each. Here the same compute methods run on a fresh
``JyotishComputeService`` in a spawn-started worker, one worker per heavy slot,
so two requests use two cores. The heavy-compute gate still decides who runs.

- Enabled only by the server's ``main()`` (and ``JYOTISH_HEAVY_COMPUTE_PROCESSES``
  not ``0``); tests that call compute methods in process are unchanged.
- Workers start with ``PYTHONHASHSEED=0`` (ERR-111: Shadbala fingerprints vary
  with the hash seed), and are started up front so the first user does not
  pay the import.
- A crashed pool is replaced; the request that hit it runs in process.
"""

from __future__ import annotations

import multiprocessing
import os
import sys
import threading
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures.process import BrokenProcessPool
from typing import Any

ENV_PROCESSES = "JYOTISH_HEAVY_COMPUTE_PROCESSES"
POOLED_METHODS = frozenset({
    "_compute_rectification_v5_score",
    "_compute_rectification_v5_block_scan",
    "professional_report_reference",
})

_guard = threading.Lock()
_executor: ProcessPoolExecutor | None = None
_workers = 0
_config: tuple[int, list[str]] | None = None
_service = None


def _worker_init(paths: list[str]) -> None:
    for path in reversed(paths):
        if path not in sys.path:
            sys.path.insert(0, path)


def _worker_call(method: str, body: dict[str, Any]) -> Any:
    global _service
    from scripts.jyotish_api_server import BadRequest, JyotishComputeService

    if _service is None:
        _service = JyotishComputeService()
    if method == "professional_report_reference":
        from scripts.professional_report_reference import (
            ProfessionalReportReferenceInputError,
            build_professional_report_reference,
        )
        try:
            return build_professional_report_reference(_service, body)
        except ProfessionalReportReferenceInputError as exc:
            raise BadRequest(str(exc)) from exc
    return getattr(_service, method)(body)


def _warm() -> int:
    from scripts import jyotish_api_server  # noqa: F401 - import the engine once per worker

    return os.getpid()


def enabled_by_env() -> bool:
    return os.environ.get(ENV_PROCESSES, "1").strip().lower() not in {"0", "false", "no", "off"}


def start(workers: int, paths: list[str]) -> bool:
    """Start the pool with ``workers`` processes; False when switched off."""
    global _executor, _workers, _config
    if workers < 1 or not enabled_by_env():
        return False
    _config = (workers, list(paths))
    with _guard:
        if _executor is not None:
            return True
        previous = os.environ.get("PYTHONHASHSEED")
        os.environ["PYTHONHASHSEED"] = "0"
        try:
            executor = ProcessPoolExecutor(
                max_workers=workers,
                mp_context=multiprocessing.get_context("spawn"),
                initializer=_worker_init,
                initargs=(paths,),
            )
            # Start every worker now, while the hash seed is in the environment.
            for future in [executor.submit(_warm) for _ in range(workers)]:
                future.result(timeout=120)
        finally:
            if previous is None:
                os.environ.pop("PYTHONHASHSEED", None)
            else:
                os.environ["PYTHONHASHSEED"] = previous
        _executor, _workers = executor, workers
        return True


def active() -> bool:
    return _executor is not None


def run(method: str, body: dict[str, Any], inline) -> Any:
    """``method`` in a worker when the pool is up, else ``inline(body)`` in this process."""
    executor = _executor
    if executor is None or method not in POOLED_METHODS:
        return inline(body)
    try:
        return executor.submit(_worker_call, method, body).result()
    except BrokenProcessPool:
        _replace_broken(executor)
        return inline(body)


def _replace_broken(broken: ProcessPoolExecutor) -> None:
    global _executor
    with _guard:
        if _executor is broken:
            _executor = None
    broken.shutdown(wait=False, cancel_futures=True)
    if _config is not None:
        threading.Thread(target=start, args=_config, name="heavy-pool-restart", daemon=True).start()


def stop() -> None:
    global _executor
    with _guard:
        executor, _executor = _executor, None
    if executor is not None:
        executor.shutdown(wait=True, cancel_futures=True)