services/pdf_utils.py

"""
Helpers for downloading LHAPDF grids from a HF dataset and for loading
a PDF set that has already been validated against the supported error
types (see services/error_types.py).
"""

import logging
import os
import re
import shutil
import threading
import time

import lhapdf
from huggingface_hub import snapshot_download

from services.config import COMPUTE_SEMAPHORE_LIMIT, LHAPDF_DATA_PATH
from services.error_types import validate_error_type

logger = logging.getLogger(__name__)

# LHAPDF's C++ core (via its Python bindings) is not thread-safe: concurrent
# calls into grid loading, mkPDFs(), or xfxQ() evaluation from different
# threads at the same time can corrupt its internal global state and crash
# the whole worker process — this is what shows up here as a bare
# "500 Internal Server Error" with no useful detail. FastAPI runs sync
# `def` endpoints in a thread pool, so a visitor rapidly switching between
# flavors/energies/datasets fires several concurrent requests that would
# otherwise all touch lhapdf at once. EVERY call into lhapdf anywhere in
# this app — here, and the pdf.xfxQ()/pdf_set.uncertainty() calls in
# services/plotting.py — must hold this lock (app.py wraps the
# figure-building calls with it; see compute_*_plot endpoints).
LHAPDF_LOCK = threading.RLock()

# Caps how many "not cached, need to compute" requests are allowed to
# actively hold a PDFSet/members and build a figure at the same time --
# independent of (and in addition to) LHAPDF_LOCK above. LHAPDF_LOCK
# already serializes the actual lhapdf calls to one at a time, but that
# alone doesn't stop dozens of requests from a burst (e.g. the website's
# Compare Datasets mode across several datasets/flavors/energies) from all
# piling in, each holding its own PDFSet/numpy allocations while it waits
# its turn for the lock. On a resource-limited Space container that's
# enough to get OOM-killed, which looks exactly like the Space "getting
# stuck" or restarting for no visible reason in the logs. app.py's four
# compute_*_plot endpoints acquire this around the whole "not cached"
# branch (PDFSet/members loading through figure building), so excess
# requests block here -- before consuming that memory -- rather than all
# proceeding at once. See services/config.py::COMPUTE_SEMAPHORE_LIMIT.
COMPUTE_SEMAPHORE = threading.Semaphore(COMPUTE_SEMAPHORE_LIMIT)

# (dataset_repo, grid_name) -> lhapdf.PDFSet, already validated.
_pdf_set_cache = {}
# (dataset_repo, grid_name) -> list of lhapdf.PDF members (pdf_set.mkPDFs()).
# This is the expensive step (parses every member's grid file from disk),
# so it's cached separately and reused across every plot request for the
# same PDF set instead of being rebuilt per-request.
_pdf_members_cache = {}


def ensure_grid_downloaded(dataset_repo: str, grid_name: str, hf_token: str, force: bool = False) -> None:
    """
    Download the LHAPDF grid folder for `grid_name` from `dataset_repo`
    into LHAPDF_DATA_PATH, unless it is already present locally.

    A directory that merely exists and is non-empty is not proof the
    download actually completed — snapshot_download() can be interrupted
    partway (network hiccup, HF rate limit, container recycled), which
    leaves a partial set of member grid files on disk. Since this Space's
    disk is ephemeral, that mainly bites right after a cold start. A
    partial grid isn't necessarily caught here or by lhapdf.getPDFSet()
    (which only reads the .info file); it can surface much later, deep
    inside PDFSet.uncertainty(), as an opaque `std::bad_array_new_length`
    when mkPDFs() silently returns fewer members than the .info file's
    NumMembers promises. get_pdf_members() below cross-checks that count
    and calls back in here with force=True to wipe and redownload if it's
    ever wrong, so this flag exists to let that self-heal actually replace
    a bad local copy instead of leaving it in place forever.
    """
    local_grid_dir = os.path.join(LHAPDF_DATA_PATH, grid_name)
    if not force and os.path.isdir(local_grid_dir) and os.listdir(local_grid_dir):
        logger.info("grid '%s' already present locally, skipping download", grid_name)
        return

    if force and os.path.isdir(local_grid_dir):
        logger.warning("forcing redownload of grid '%s' from %s -- wiping local copy at %s", grid_name, dataset_repo, local_grid_dir)
        shutil.rmtree(local_grid_dir)

    logger.info("downloading grid '%s' from dataset repo %s ...", grid_name, dataset_repo)
    t0 = time.monotonic()
    snapshot_download(
        repo_id=dataset_repo,
        repo_type="dataset",
        allow_patterns=f"{grid_name}/*",
        local_dir=LHAPDF_DATA_PATH,
        token=hf_token,
    )
    logger.info("finished downloading grid '%s' in %.2fs", grid_name, time.monotonic() - t0)

    if LHAPDF_DATA_PATH not in lhapdf.paths():
        lhapdf.pathsPrepend(LHAPDF_DATA_PATH)
        logger.info("prepended %s to lhapdf search paths", LHAPDF_DATA_PATH)


def _check_num_members_present(grid_name: str) -> None:
    """
    PDFSet.uncertainty() reads `NumMembers` directly out of the set's own
    `.info` file (a separate lookup from `pdf_set.size`, which is why
    get_pdf_members()'s member-count cross-check above doesn't catch this)
    to size an internal C++ array. If that field is missing, LHAPDF prints
    "Metadata for key: NumMembers not found." straight to stdout and then
    computes a garbage size, which surfaces as an opaque
    `std::bad_array_new_length` (MemoryError in Python) deep inside
    uncertainty() -- with every request for this grid failing the exact
    same way, since the problem is the .info file itself, not any one
    request's parameters. Reading the file directly here (rather than
    going through some lhapdf accessor whose exact behavior/API isn't
    verifiable without the compiled bindings installed) lets us fail once,
    clearly, before ever reaching uncertainty().
    """
    info_path = os.path.join(LHAPDF_DATA_PATH, grid_name, f"{grid_name}.info")
    try:
        with open(info_path, "r") as f:
            info_text = f.read()
    except OSError:
        logger.warning("could not read %s to check for NumMembers -- letting lhapdf's own validation run instead", info_path)
        return

    if not re.search(r"(?m)^\s*NumMembers\s*:", info_text):
        logger.error("%s has no 'NumMembers' field", info_path)
        raise RuntimeError(
            f"PDF set '{grid_name}' is missing the 'NumMembers' field in its .info "
            "file. This is a data problem in the dataset repo's .info file itself "
            "(not a download/caching issue on this Space) -- every uncertainty "
            "plot for this grid will fail until a correct 'NumMembers: <count>' "
            f"line is added to {grid_name}/{grid_name}.info in the dataset repo."
        )


def load_validated_pdf_set(dataset_repo: str, grid_name: str, hf_token: str):
    """
    Download the grid (if needed), load its PDFSet, and validate that its
    ErrorType is one this app can compute uncertainties for. Cached per
    (dataset_repo, grid_name) so repeat requests for the same PDF set
    don't re-validate every time.

    Raises services.error_types.UnsupportedErrorTypeError if the set's
    error type isn't supported yet.

    Returns
    -------
    lhapdf.PDFSet
    """
    key = (dataset_repo, grid_name)
    with LHAPDF_LOCK:
        cached = _pdf_set_cache.get(key)
        if cached is not None:
            logger.info("PDFSet cache hit for %s", key)
            return cached

        logger.info("PDFSet cache miss for %s, loading...", key)
        ensure_grid_downloaded(dataset_repo, grid_name, hf_token)
        _check_num_members_present(grid_name)
        pdf_set = lhapdf.getPDFSet(grid_name)
        validate_error_type(pdf_set, grid_name)
        logger.info("loaded PDFSet '%s': ErrorType=%s, size=%s", grid_name, getattr(pdf_set, "errorType", "?"), pdf_set.size)

        _pdf_set_cache[key] = pdf_set
        return pdf_set


def get_pdf_members(dataset_repo: str, grid_name: str, hf_token: str):
    """
    Return the PDF set's member objects (`pdf_set.mkPDFs()`), cached per
    (dataset_repo, grid_name). mkPDFs() parses every member's grid file
    off disk — for a set with hundreds of Hessian/replica members this is
    the actual expensive step, not loading the PDFSet object itself, so
    it's the one worth reusing across every plot request for the same set
    (standard, ratio, and correlation plots at any Q/flavor all share it).
    """
    key = (dataset_repo, grid_name)
    with LHAPDF_LOCK:
        cached = _pdf_members_cache.get(key)
        if cached is not None:
            logger.info("mkPDFs() cache hit for %s (%d members)", key, len(cached))
            return cached

        logger.info("mkPDFs() cache miss for %s, parsing member grid files from disk...", key)
        pdf_set = load_validated_pdf_set(dataset_repo, grid_name, hf_token)
        t0 = time.monotonic()
        pdfs = pdf_set.mkPDFs()
        logger.info("mkPDFs() for %s loaded %d member(s) in %.2fs (expected %s)", key, len(pdfs), time.monotonic() - t0, pdf_set.size)

        if len(pdfs) != pdf_set.size:
            # mkPDFs() loaded fewer (or more) members than the set's own
            # .info declares -- almost always a partial/stale local grid
            # download (see ensure_grid_downloaded's docstring), which
            # would otherwise crash PDFSet.uncertainty() with an opaque
            # std::bad_array_new_length on every request for this grid
            # until the container restarts. Wipe the cached PDFSet and the
            # on-disk grid, and try once more with a forced fresh download.
            logger.warning(
                "member count mismatch for %s: mkPDFs() returned %d but .info declares %d -- "
                "forcing a fresh redownload and retrying once",
                key, len(pdfs), pdf_set.size,
            )
            _pdf_set_cache.pop(key, None)
            ensure_grid_downloaded(dataset_repo, grid_name, hf_token, force=True)
            _check_num_members_present(grid_name)
            pdf_set = lhapdf.getPDFSet(grid_name)
            validate_error_type(pdf_set, grid_name)
            _pdf_set_cache[key] = pdf_set
            pdfs = pdf_set.mkPDFs()
            logger.info("retry mkPDFs() for %s loaded %d member(s) (expected %s)", key, len(pdfs), pdf_set.size)

            if len(pdfs) != pdf_set.size:
                logger.error("member count mismatch for %s persisted after redownload -- giving up", key)
                raise RuntimeError(
                    f"PDF set '{grid_name}' from {dataset_repo} has {len(pdfs)} "
                    f"loadable member(s) but its .info declares {pdf_set.size} -- "
                    "this persisted after a fresh redownload, so the grid files in "
                    "the dataset repo itself look incomplete or inconsistent with "
                    "its .info, not just a local caching issue."
                )

        _pdf_members_cache[key] = pdfs
        return pdfs


def invalidate_cache(dataset_repo: str, grid_name: str) -> None:
    """
    Evict the cached PDFSet and PDF members for (dataset_repo, grid_name).

    LHAPDF's C++ objects are not guaranteed safe against a computation that
    throws partway through -- e.g. PDFSet.uncertainty() failing (as with a
    missing 'NumMembers' field) after the loop has already called
    pdf.xfxQ() on several cached PDF member objects can plausibly leave
    those objects' internal interpolation/lookup state inconsistent. Since
    both caches are reused across every later request for the same grid,
    one failed request could otherwise poison every subsequent one until
    the container restarts -- which looks exactly like the Space "getting
    stuck", especially under concurrent/rapid-switch usage. app.py calls
    this from every compute_*_plot endpoint's outer except-block so a
    failure only ever costs the request that hit it: the next request for
    this grid starts from a guaranteed-fresh PDFSet + mkPDFs() load.
    """
    key = (dataset_repo, grid_name)
    with LHAPDF_LOCK:
        had_set = _pdf_set_cache.pop(key, None) is not None
        had_members = _pdf_members_cache.pop(key, None) is not None
    if had_set or had_members:
        logger.warning("invalidated cached PDFSet/members for %s after a failure, to avoid poisoning later requests", key)