Jyotisha/scripts/research/api_chart_subprocess_ab.py

"""A/B for issue #64. A is this branch with the PyJHora short-circuit forced off
(the staging control flow: always spawn). B is the short-circuit on.

Chart skip/concurrency is measured on the real chart call. The gate schemes
O2/O3/O4 are simulated only and are not implemented.

Volatile fields dropped before the report byte compare:
generated_at, created_at, timestamp, elapsed_seconds, elapsed_ms, duration_ms, request_id.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import statistics
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "scripts"))

os.environ.setdefault("JYOTISH_API_CHART_CACHE_TTL_SECONDS", "0")
os.environ.setdefault("VEDASTRO_ENABLE_NETWORK", "0")
os.environ.setdefault("JYOTISH_GATE_CHART_OVERVIEW", "0")

import annual_pyjhora_replay as replay  # noqa: E402
import tajika as tajika_top  # noqa: E402
from scripts.jyotish_api_server import JyotishAPIHandler  # noqa: E402
from scripts.jyotish_engine import reset_oracle_progress_snapshot_cache  # noqa: E402
from scripts.professional_report_reference import build_professional_report_reference  # noqa: E402
from scripts.tajika import reset_external_pyjhora_python_cache  # noqa: E402

BIRTH = {
    "year": 1990,
    "month": 1,
    "day": 1,
    "hour": 12,
    "minute": 0,
    "second": 0,
    "lat": 28.6139,
    "lon": 77.2090,
    "tz": 5.5,
    "ayanamsa": "lahiri",
    "node_mode": "true",
}
PARTNER = {**BIRTH, "year": 1992, "day": 2}
VOLATILE = {"generated_at", "created_at", "timestamp", "elapsed_seconds", "elapsed_ms", "duration_ms", "request_id"}


class _Counter:
    def __init__(self) -> None:
        self.n = 0
        self._real = subprocess.run
        self._lock = threading.Lock()

    def install(self) -> None:
        def _run(*args, **kwargs):
            with self._lock:
                self.n += 1
            return self._real(*args, **kwargs)

        subprocess.run = _run  # type: ignore[assignment]
        for module in list(sys.modules.values()):
            bound = getattr(module, "subprocess", None)
            if getattr(bound, "run", None) is self._real:
                bound.run = _run  # type: ignore[union-attr]

    def reset(self) -> None:
        with self._lock:
            self.n = 0


def _median(values: list[float]) -> float:
    return statistics.median(values) if values else 0.0


def _strip(value):
    if isinstance(value, dict):
        return {key: _strip(item) for key, item in value.items() if key not in VOLATILE}
    if isinstance(value, list):
        return [_strip(item) for item in value]
    return value


def _canon(value) -> str:
    return json.dumps(_strip(value), sort_keys=True, ensure_ascii=True, default=str)


def _chart_summary(chart: dict) -> dict:
    planets = chart.get("planets") if isinstance(chart, dict) else None
    moon = planets.get("Moon") if isinstance(planets, dict) else None
    lon = None
    if isinstance(moon, dict):
        lon = moon.get("lon", moon.get("longitude"))
    return {
        "ascendant": chart.get("ascendant") if isinstance(chart, dict) else None,
        "planets": planets,
        "moon_longitude": lon,
    }


def _one_chart(body: dict) -> tuple[float, dict]:
    started = time.perf_counter()
    chart = object.__new__(JyotishAPIHandler)._compute_chart_sync(body)
    return time.perf_counter() - started, chart


def _measure_charts(counter: _Counter) -> dict:
    serial_times = []
    concurrent_times = []
    serial_subs = []
    concurrent_subs = []
    same = True
    for _ in range(3):
        counter.reset()
        started = time.perf_counter()
        _, self_a = _one_chart(BIRTH)
        _, partner_a = _one_chart(PARTNER)
        serial_times.append(time.perf_counter() - started)
        serial_subs.append(counter.n)
        counter.reset()
        started = time.perf_counter()
        with ThreadPoolExecutor(max_workers=2) as pool:
            self_b = pool.submit(_one_chart, {**BIRTH, "skip_vedastro_main_entry_overview": True})
            partner_b = pool.submit(_one_chart, {**PARTNER, "skip_vedastro_main_entry_overview": True})
            _, self_chart = self_b.result()
            _, partner_chart = partner_b.result()
        concurrent_times.append(time.perf_counter() - started)
        concurrent_subs.append(counter.n)
        same = same and _canon(_chart_summary(self_a)) == _canon(_chart_summary(self_chart))
        same = same and _canon(_chart_summary(partner_a)) == _canon(_chart_summary(partner_chart))
    return {
        "serial_seconds_median": round(_median(serial_times), 4),
        "concurrent_skip_seconds_median": round(_median(concurrent_times), 4),
        "serial_subprocesses_median": _median(serial_subs),
        "concurrent_skip_subprocesses_median": _median(concurrent_subs),
        "summary_bytes_match": same,
    }


def _measure_report(counter: _Counter, *, spawn: bool) -> dict:
    times = []
    subs = []
    payload = None
    body = {**BIRTH, "format": "markdown", "edition": "full_data", "languages": ["zh", "en"], "include_fact_tables": True}
    for _ in range(3):
        replay.reset_jhora_import_cache()
        if spawn:
            replay._jhora_spec_cache = True
        reset_oracle_progress_snapshot_cache()
        reset_external_pyjhora_python_cache()
        tajika_top.reset_external_pyjhora_python_cache()
        counter.reset()
        started = time.perf_counter()
        payload = build_professional_report_reference(
            object.__new__(JyotishAPIHandler),
            body,
        )
        times.append(time.perf_counter() - started)
        subs.append(counter.n)
    return {
        "seconds_median": round(_median(times), 4),
        "subprocesses_median": _median(subs),
        "payload": payload,
    }


def _simulate() -> list[dict]:
    """One fixed minute. O2 adds a second pool and is the only scheme that touches BUG-717."""
    rng = random.Random(64)
    arrivals = [("report", index * 30.0, 13.0) for index in range(2)]
    arrivals += [("consult", rng.random() * 60.0, 0.2) for _ in range(20)]
    arrivals += [("score", rng.random() * 60.0, 8.0) for _ in range(6)]
    arrivals.sort(key=lambda item: item[1])

    def run(mode: str) -> dict:
        heavy: list[tuple[float, str]] = []
        light: list[float] = []
        rejected = {"consult": 0, "report": 0, "score": 0}
        accepted = {"consult": 0, "report": 0, "score": 0}
        peak = 0

        def occupy(at: float) -> tuple[int, int]:
            running = [item for item in heavy if item[0] > at]
            lights = [item for item in light if item > at]
            return len(running), len(lights)

        for kind, arrival, duration in arrivals:
            heavy_n, light_n = occupy(arrival)
            heavy = [item for item in heavy if item[0] > arrival]
            light = [item for item in light if item > arrival]
            placed = False
            stamp = arrival
            if mode == "O2" and kind == "consult":
                if light_n < 2:
                    light.append(arrival + duration)
                    placed = True
            elif mode == "O3" and kind == "consult":
                start = arrival
                while start <= arrival + 1.5:
                    heavy_n, _light_n = occupy(start)
                    if heavy_n < 2:
                        heavy.append((start + duration, kind))
                        placed = True
                        stamp = start
                        break
                    nxt = min(item[0] for item in heavy)
                    if nxt > arrival + 1.5:
                        break
                    start = nxt
            elif mode == "O4" and kind != "consult":
                heavy_jobs = sum(1 for _end, slot in heavy if slot != "consult")
                if heavy_jobs < 1 and heavy_n < 2:
                    heavy.append((arrival + duration, kind))
                    placed = True
            else:
                if heavy_n < 2:
                    heavy.append((arrival + duration, kind))
                    placed = True
            if placed:
                accepted[kind] += 1
            else:
                rejected[kind] += 1
            running, lights = occupy(stamp)
            peak = max(peak, running + lights)

        def rate(kind: str) -> str:
            total = accepted[kind] + rejected[kind]
            return "0" if total == 0 else f"{rejected[kind] / total:.2f}"

        return {
            "scheme": mode,
            "consult_429": rate("consult"),
            "report_429": rate("report"),
            "score_429": rate("score"),
            "peak": peak,
            "touches_bug_717": mode == "O2",
        }

    return [run(mode) for mode in ("O1", "O2", "O3", "O4")]


def resolve_record_path(explicit: str | None) -> Path:
    """Write the A/B JSON only where asked. The committed note stays untouched."""
    if explicit:
        return Path(explicit)
    directory = Path(tempfile.mkdtemp(prefix="api-chart-subprocess-ab-"))
    return directory / "api-chart-subprocess-gate-ab-20261013.json"


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="A/B for issue #64 chart subprocess gate")
    parser.add_argument(
        "--out",
        default=None,
        help="Write the JSON record to this path. Omit to use a fresh temp directory, not the committed docs/testing file.",
    )
    args = parser.parse_args(argv)
    counter = _Counter()
    counter.install()
    charts = _measure_charts(counter)
    forced = _measure_report(counter, spawn=True)
    native = _measure_report(counter, spawn=False)
    report_match = _canon(forced["payload"]) == _canon(native["payload"])
    schemes = _simulate()
    out = {
        "charts": charts,
        "report_forced_subprocess": {key: value for key, value in forced.items() if key != "payload"},
        "report_short_circuit": {key: value for key, value in native.items() if key != "payload"},
        "report_bytes_match": report_match,
        "schemes": schemes,
        "jhora_installed": False,
    }
    text = json.dumps(out, ensure_ascii=False, indent=2)
    dest = resolve_record_path(args.out)
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(text + "\n", encoding="utf-8")
    print(dest, file=sys.stderr)
    print(text)


if __name__ == "__main__":
    main()