Spaces:
Sleeping
Sleeping
File size: 8,227 Bytes
b92b744 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """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()
|