"""Public S3 corpus build. Run --plan first, then --run on the batch machine.""" import argparse import asyncio import json from pathlib import Path from app.config import settings from app.ingest.__main__ import parser as ingest_parser from app.ingest.__main__ import run as ingest from app.ingest.fetch import atomic_json, fetch_pages from app.ingest.materialize import materialize from app.ingest.scope import CTGOV_QUERY, SCOPE_VERSION from app.ingest.snapshot import fingerprint, load_manifest, plan_manifest def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser(description=__doc__) action = result.add_mutually_exclusive_group() action.add_argument("--plan", action="store_true", help="Read only the manifest (the default)") action.add_argument("--run", action="store_true", help="Scan, normalize, embed, link and index") result.add_argument("--manifest", default=settings.openalex_snapshot_manifest) result.add_argument("--data-dir", type=Path, default=Path("data/research-s3")) result.add_argument("--max-files", type=int, help="Explicit partial snapshot for smoke tests") result.add_argument( "--largest-first", action="store_true", help="Order parts by descending size so a partial --max-files run covers more of the corpus", ) result.add_argument( "--skip-files", type=int, default=0, help="Skip leading manifest parts to continue an earlier --max-files run; use a new data dir", ) result.add_argument("--max-snapshot-bytes", type=int, default=5_000_000_000) result.add_argument( "--max-records", type=int, help="Explicit partial record cap; default has no cap" ) result.add_argument( "--registry-limit", type=int, help="Explicit partial registry cap; default exhausts cursor" ) result.add_argument("--skip-registry", action="store_true") result.add_argument("--skip-embeddings", action="store_true") result.add_argument( "--work-ids", type=Path, help="Reference backfill from S3; requires --skip-registry" ) result.add_argument("--batch-size", type=int, default=100) result.add_argument("--min-free-bytes", type=int, default=1_000_000_000) return result async def bootstrap(args) -> dict: if args.work_ids and not args.skip_registry: raise ValueError("Reference backfill requires --skip-registry") if args.registry_limit is not None and args.registry_limit <= 0: raise ValueError("registry-limit must be positive") if args.batch_size <= 0: raise ValueError("batch-size must be positive") if args.max_records is not None and args.max_records <= 0: raise ValueError("max-records must be positive") if args.max_snapshot_bytes <= 0 or args.min_free_bytes < 0: raise ValueError("Snapshot budget must be positive and free-space threshold nonnegative") manifest = await load_manifest(args.manifest) plan = plan_manifest( manifest, max_files=args.max_files, largest_first=args.largest_first, skip_files=args.skip_files, ) summary = {k: v for k, v in plan.items() if k != "files"} summary.update( source="public_s3", scope="reference_backfill" if args.work_ids else SCOPE_VERSION, registry_scope="all statuses, study types and dates for the hypertension/kidney condition union", target_index=settings.elastic_index, note="Complete means all matches under the declared scope in the pinned snapshot, not all research worldwide.", ) if not args.run: return summary if plan["physical_bytes_budgeted"] > args.max_snapshot_bytes: raise ValueError( f"Plan requires {plan['physical_bytes_budgeted']} physical bytes; set --max-snapshot-bytes after inspecting capacity on the batch host" ) from app.embeddings import get_embedder from app.repository import ElasticRepository args.data_dir.mkdir(parents=True, exist_ok=True) manifest_path = args.data_dir / "manifest.json" if manifest_path.exists() and json.loads(manifest_path.read_text()) != manifest: raise ValueError("Snapshot changed; use a new data directory for the new release") atomic_json(manifest_path, manifest) report_path = args.data_dir / "corpus-summary.json" summary.update(complete=False, stage="starting", stages={}) atomic_json(report_path, summary) repository = ElasticRepository() embedder = ( None if args.skip_embeddings else get_embedder(settings.embedding_model, settings.embedding_device) ) def progress(stage): def emit(state): summary["stage"] = stage summary["stages"][stage] = state atomic_json(report_path, summary) print(json.dumps({"stage": stage, "records": state["records"]}), flush=True) return emit try: await repository.ensure_index() registry = None if not args.skip_registry: registry_path = args.data_dir / "ctgov.jsonl" registry = await fetch_pages( "ctgov", registry_path, query=CTGOV_QUERY, filters="", limit=args.registry_limit, resume=True, ) summary["registry"] = registry await materialize( registry_path, repository, embedder=embedder, batch_size=args.batch_size, progress=progress("registry_index"), ) paper_path = args.data_dir / "openalex.jsonl" commands = [ "snapshot", "--manifest", str(manifest_path), "--output", str(paper_path), "--profile", "hypertension-kidney", "--resume", "--max-bytes", str(args.max_snapshot_bytes), "--min-free-bytes", str(args.min_free_bytes), ] for flag, value in ( ("--skip-files", args.skip_files or None), ("--max-files", args.max_files), ("--limit", args.max_records), ("--work-ids", args.work_ids), ): if value is not None: commands.extend([flag, str(value)]) summary["stage"] = "snapshot_scan" atomic_json(report_path, summary) snapshot = await ingest(ingest_parser().parse_args(commands)) summary["snapshot"] = snapshot if "://" in args.manifest and snapshot["selected_parts_complete"]: if fingerprint(await load_manifest(args.manifest)) != plan["manifest_sha256"]: snapshot["complete_snapshot_scope"] = False snapshot["manifest_changed_during_scan"] = True atomic_json(paper_path.with_suffix(".jsonl.checkpoint.json"), snapshot) raise RuntimeError( "Public snapshot changed during the scan; result is not a complete release" ) await materialize( paper_path, repository, embedder=embedder, batch_size=args.batch_size, progress=progress("snapshot_index"), ) summary.update( stage="finished", complete=bool( snapshot["complete_snapshot_scope"] and registry is not None and registry["done"] ), index=await repository.health(), ) if args.work_ids: summary["reference_scan_complete"] = snapshot["complete_snapshot_scope"] atomic_json(report_path, summary) return summary except Exception as exc: summary.update(stage="interrupted", complete=False, error=type(exc).__name__) atomic_json(report_path, summary) raise finally: await repository.close() def main(): args = parser().parse_args() try: print(json.dumps(asyncio.run(bootstrap(args)), indent=2)) except (ValueError, RuntimeError, FileExistsError, TypeError) as exc: raise SystemExit(str(exc)) from None if __name__ == "__main__": main()