pipeline/submit.py

"""Stage 1a — submit every corpus asset to GET /process/asset.

Two things about the endpoint drive this module's shape
(`services/info-endpoint/info_endpoint/src/api/processing.py`):

1. `refresh_metadata` defaults to **False**. With it off the call enqueues
   plugin jobs but never writes the base document, so the asset gets embeddings
   and a graph and remains absent from search. We always pass it and we assert
   `metadata_refreshed` in the response — a 200 that did not write is a failure,
   not a success.
2. The base write is synchronous inside the request (it calls `update_meta`
   directly), while plugin dispatch is a Redis enqueue. So submission returning
   cleanly means the document count is exactly N; only the plugin half needs
   polling afterwards (see status.py).
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Set

import httpx

from .config import BUILD_DIR, Config, active_plugins
from .enumerate_assets import Asset, load_assets
from .queues import connect, queue_depths
from .restore import assets_to_process

CHECKPOINT = BUILD_DIR / "submitted.jsonl"


@dataclass
class SubmitResult:
    url: str
    ok: bool
    reason: str = ""
    plugins: Optional[List[str]] = None

    def to_json(self) -> str:
        return json.dumps(
            {"url": self.url, "ok": self.ok, "reason": self.reason, "plugins": self.plugins or []}
        )


def load_checkpoint(path: Path = CHECKPOINT) -> Set[str]:
    """URLs already submitted successfully, so a resumed run skips them."""
    if not path.exists():
        return set()
    done = set()
    for line in path.read_text().splitlines():
        if not line.strip():
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if record.get("ok"):
            done.add(record["url"])
    return done


async def submit_one(
    client: httpx.AsyncClient,
    base_url: str,
    asset: Asset,
    max_retries: int,
    plugins: tuple,
) -> SubmitResult:
    params = [
        ("url", asset.url),
        # Both explicit. refresh_metadata in particular defaults to false.
        ("refresh_metadata", "true"),
        ("refresh_plugins", "true"),
        ("priority", "priority"),
    ]
    # Defense in depth alongside the info-endpoint's plugin.config.yaml gate:
    # the endpoint intersects this list with its own active plugins.
    params += [("plugins", p) for p in plugins]

    delay = 1.0
    last = "unknown"
    for attempt in range(max_retries + 1):
        try:
            response = await client.get(f"{base_url}/process/asset", params=params)
        except (httpx.TransportError, httpx.TimeoutException) as exc:
            last = f"transport: {type(exc).__name__}"
        else:
            if response.status_code == 404:
                # The corpus names an object the bucket does not have. Terminal
                # and interesting — a corpus bug, not a transient fault.
                return SubmitResult(asset.url, False, "not_found_on_storage")
            if response.status_code == 422:
                return SubmitResult(asset.url, False, f"invalid_url: {response.text[:200]}")
            if response.status_code < 400:
                body = response.json()
                if not body.get("metadata_refreshed"):
                    # 200 without a metadata write means no base document.
                    return SubmitResult(
                        asset.url, False, "metadata_not_refreshed", body.get("plugins")
                    )
                return SubmitResult(asset.url, True, "", body.get("plugins"))
            last = f"http {response.status_code}: {response.text[:200]}"

        if attempt < max_retries:
            await asyncio.sleep(delay)
            delay = min(delay * 2, 30.0)

    return SubmitResult(asset.url, False, f"retries_exhausted ({last})")


async def wait_for_queue_headroom(
    redis_client, ceiling: int, plugins: tuple, poll: float = 5.0
) -> None:
    """Block while the plugin queues are deeper than `ceiling`.

    Without this a large corpus pushes every job into Redis long before a worker
    drains one, which buys nothing and makes the queue the memory bottleneck.
    """
    if ceiling <= 0:
        return
    while True:
        depths = await queue_depths(redis_client, plugins)
        pending = sum(v for k, v in depths.items() if k != "results")
        if pending <= ceiling:
            return
        print(f"    backpressure: {pending} jobs queued (ceiling {ceiling}), waiting ...", flush=True)
        await asyncio.sleep(poll)


async def run_async(cfg: Config, resume: bool = True) -> int:
    plugins = active_plugins(cfg)
    corpus_assets = load_assets()
    assets = assets_to_process(corpus_assets)
    already = load_checkpoint() if resume else set()
    todo = [a for a in assets if a.url not in already]

    if len(assets) != len(corpus_assets):
        print(f"unchanged: {len(corpus_assets) - len(assets)} pinned asset(s) already in the baseline (skipping)")
    print(f"assets:    {len(assets)}")
    print(f"submitted: {len(already)} (skipping)" if already else "submitted: 0")
    print(f"to submit: {len(todo)}")
    if not todo:
        print("\nnothing to submit.")
        return 0

    concurrency = cfg.get_int("SUBMIT_CONCURRENCY", 16)
    max_retries = cfg.get_int("SUBMIT_MAX_RETRIES", 5)
    ceiling = cfg.get_int("QUEUE_DEPTH_CEILING", 5000)
    redis_url = f"redis://localhost:{cfg.get_int('REDIS_HOST_PORT', 6380)}"

    print(f"concurrency={concurrency} retries={max_retries} queue_ceiling={ceiling}\n")

    semaphore = asyncio.Semaphore(concurrency)
    results: List[SubmitResult] = []
    redis_client = connect(redis_url)

    limits = httpx.Limits(max_connections=concurrency * 2, max_keepalive_connections=concurrency)
    timeout = httpx.Timeout(connect=10.0, read=120.0, write=30.0, pool=60.0)

    if not resume and CHECKPOINT.exists():
        CHECKPOINT.unlink()
    BUILD_DIR.mkdir(parents=True, exist_ok=True)
    checkpoint_lock = asyncio.Lock()

    async with httpx.AsyncClient(limits=limits, timeout=timeout) as http_client:
        with CHECKPOINT.open("a") as fh:

            async def worker(asset: Asset) -> None:
                async with semaphore:
                    result = await submit_one(http_client, cfg.info_endpoint, asset, max_retries, plugins)
                results.append(result)
                async with checkpoint_lock:
                    fh.write(result.to_json() + "\n")
                    fh.flush()
                done = len(results)
                if done % 250 == 0 or done == len(todo):
                    failed = sum(1 for r in results if not r.ok)
                    print(f"  {done}/{len(todo)} submitted ({failed} failed)", flush=True)

            # Submit in chunks so backpressure is checked between batches
            # rather than only once at the start.
            chunk_size = max(concurrency * 10, 100)
            for start in range(0, len(todo), chunk_size):
                await wait_for_queue_headroom(redis_client, ceiling, plugins)
                chunk = todo[start:start + chunk_size]
                await asyncio.gather(*(worker(a) for a in chunk))

    await redis_client.aclose()

    failures = [r for r in results if not r.ok]
    by_reason: Dict[str, int] = {}
    for failure in failures:
        key = failure.reason.split(":")[0].split("(")[0].strip()
        by_reason[key] = by_reason.get(key, 0) + 1

    print(f"\nsubmitted ok: {len(results) - len(failures)}")
    print(f"failed:       {len(failures)}")
    for reason, count in sorted(by_reason.items(), key=lambda kv: -kv[1]):
        print(f"  {reason}: {count}")

    if failures:
        path = BUILD_DIR / "submit-failures.json"
        path.write_text(
            json.dumps([{"url": f.url, "reason": f.reason} for f in failures], indent=2) + "\n"
        )
        print(f"\nwrote {path.name}. Re-run `submit` to retry only these.")
        return 1
    return 0


def run(cfg: Config, resume: bool = True) -> int:
    return asyncio.run(run_async(cfg, resume=resume))