Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
| """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" | |
| 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)) | |