refresh_coverage.py

"""Inventory published annotation assemblies using completion markers and metadata.

Completed shards use assembly IDs from their packed sources, under the pipeline's
strict unpack/publication contract. Shards without markers read annotation IDs
directly. No sequences or probability columns are read. Cached results are tied
to source and publication content hashes, so reruns reuse unchanged shards.
"""
import argparse
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
import time

from huggingface_hub import HfApi, HfFileSystem, get_token
import pyarrow.parquet as pq

from remote_catalog import MeasuredFile, source_info

ROOT = Path(__file__).resolve().parent
ANNOTATIONS = "HuggingFaceBio/genbank-annotations"
INPUTS = "HuggingFaceBio/GENERanno-packed-inference-input"


def fingerprint(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest()


def read_accessions(fs, bucket, path):
    accessions = set()
    with MeasuredFile(fs, f"buckets/{bucket}/{path}") as stream:
        pf = pq.ParquetFile(stream)
        column = pf.schema.names.index("assembly_accession")
        for group in range(pf.num_row_groups):
            stats = pf.metadata.row_group(group).column(column).statistics
            if stats and stats.has_min_max and stats.min == stats.max and stats.null_count == 0:
                accessions.add(stats.min)
            else:
                accessions.update(pf.read_row_group(group, columns=["assembly_accession"], use_threads=False)["assembly_accession"].to_pylist())
        fetched = stream.bytes_read
    if any(not isinstance(a, str) or not re.fullmatch(r"GCA_\d+\.\d+", a) for a in accessions):
        raise ValueError(f"Invalid assembly accession in {path}")
    return accessions, fetched


def collect(cache_dir, output, workers):
    cache_dir.mkdir(parents=True, exist_ok=True)
    token = get_token()
    api = HfApi(token=token)
    started = datetime.now(timezone.utc).isoformat()
    shards = defaultdict(lambda: {"objects": [], "marker": None})
    print("Listing published annotation objects…", flush=True)
    for item in api.list_bucket_tree(ANNOTATIONS, "annotations", recursive=True):
        if item.type != "file":
            continue
        parts = item.path.split("/")
        if len(parts) == 3 and parts[-1].endswith(".parquet"):
            # Older publishers wrote one file per shard, without a shard folder.
            shards["/".join(parts[1:])]["objects"].append((item.path, item.xet_hash, item.size))
            continue
        if len(parts) != 4:
            if item.path.endswith(".parquet"):
                raise ValueError(f"Unrecognized annotation layout: {item.path}")
            continue
        shard = "/".join(parts[1:3])
        if parts[-1].endswith(".parquet"):
            shards[shard]["objects"].append((item.path, item.xet_hash, item.size))
        elif parts[-1] == "_SUCCESS.json":
            shards[shard]["marker"] = (item.path, item.xet_hash)
    shards = {key: val for key, val in shards.items() if val["objects"]}
    if not shards:
        raise ValueError("No annotation objects found")
    print(f"Found {len(shards):,} published shards, {sum(len(v['objects']) for v in shards.values()):,} objects", flush=True)

    def scan(shard, entry):
        fs, worker_api = HfFileSystem(token=token), HfApi(token=token)
        source = None
        if entry["marker"]:
            path = f"final_overlapped_pack/packed_blocks_by_division/{shard}/sources.parquet"
            info = source_info(worker_api, INPUTS, path)
            if info is None:
                raise ValueError(f"Missing packed source metadata for {shard}")
            source = (path, info.xet_hash)
        identity = {"objects": sorted(entry["objects"]), "marker": entry["marker"], "source": source, "version": 1}
        digest = fingerprint(identity)
        target = cache_dir / (shard.replace("/", "__") + ".json")
        if target.exists():
            previous = json.loads(target.read_text())
            if previous["fingerprint"] == digest:
                return previous
        accessions, fetched = set(), 0
        marker = None
        if source:
            marker = json.loads(fs.cat(f"buckets/{ANNOTATIONS}/{entry['marker'][0]}"))
            division, shard_id = shard.split("/")
            if (marker.get("division") != division or marker.get("shard") != shard_id
                    or int(marker.get("expected_parts", 0)) <= 0
                    or int(marker.get("pack_count", 0)) < int(marker["expected_parts"])
                    or int(marker.get("prob_count", 0)) < int(marker["expected_parts"])):
                raise ValueError(f"Invalid completion marker for {shard}")
            accessions, fetched = read_accessions(fs, INPUTS, source[0])
            after = source_info(worker_api, INPUTS, source[0])
            if after is None or after.xet_hash != source[1]:
                raise ValueError(f"Source changed while reading {shard}; retry")
            after_marker = source_info(worker_api, ANNOTATIONS, entry['marker'][0])
            if after_marker is None or after_marker.xet_hash != entry['marker'][1]:
                raise ValueError(f"Completion marker changed while reading {shard}; retry")
        else:
            for path, content_hash, _ in entry["objects"]:
                ids, size = read_accessions(fs, ANNOTATIONS, path)
                accessions.update(ids)
                fetched += size
                after = source_info(worker_api, ANNOTATIONS, path)
                if after is None or after.xet_hash != content_hash:
                    raise ValueError(f"Annotation changed while reading {path}; retry")
        result = {"shard": shard, "fingerprint": digest, "accessions": sorted(accessions),
                  "annotation_objects": len(entry["objects"]), "metadata_bytes_read": fetched,
                  "method": "strict_completion_marker_and_packed_sources" if source else "annotation_metadata",
                  "source": source, "marker": entry["marker"], "completion": marker}
        temporary = target.with_suffix(".tmp")
        temporary.write_text(json.dumps(result))
        temporary.replace(target)
        return result

    def retry_scan(shard, entry):
        for attempt in range(3):
            try:
                return scan(shard, entry)
            except Exception:
                if attempt == 2:
                    raise
                time.sleep(2 ** attempt)

    results, errors = [], []
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(retry_scan, shard, entry): shard for shard, entry in shards.items()}
        for future in as_completed(futures):
            try:
                results.append(future.result())
            except Exception as exc:
                errors.append(f"{futures[future]}: {type(exc).__name__}: {exc}")
                print(errors[-1], flush=True)
            if (len(results) + len(errors)) % 20 == 0:
                print(f"Inventoried {len(results):,}/{len(shards):,} shards; {len(errors)} errors", flush=True)
    if errors:
        raise RuntimeError(f"Coverage snapshot not replaced: {len(errors)} shards failed. Retry to reuse cached successes.\n" + "\n".join(errors))
    accessions = sorted({a for r in results for a in r["accessions"]})
    inventory = {"created_at": datetime.now(timezone.utc).isoformat(), "listing_started_at": started,
                 "annotation_bucket": ANNOTATIONS, "input_bucket": INPUTS,
                 "scope": "Assemblies with published annotations; includes partial assemblies. Exact accession versions; no base completeness claim.",
                 "method": "Completed shard assembly IDs from packed sources under strict unpack publication contract; unmarked shards read from annotation metadata. Source namespace is assumed to remain associated with its published completion marker.",
                 "annotation_object_count": sum(r["annotation_objects"] for r in results),
                 "shard_count": len(results), "metadata_bytes_read": sum(r["metadata_bytes_read"] for r in results),
                 "inventory_fingerprint": fingerprint(sorted(r["fingerprint"] for r in results)),
                 "accessions": accessions,
                 "shards": [{k: v for k, v in r.items() if k != "accessions"} for r in sorted(results, key=lambda r: r['shard'])]}
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(inventory))
    temporary.replace(output)
    print(f"Saved {len(accessions):,} distinct assembly versions to {output}", flush=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache-dir", type=Path, default=ROOT / ".cache/coverage/shards")
    parser.add_argument("--output", type=Path, default=ROOT / "data/coverage.json")
    parser.add_argument("--workers", type=int, default=12)
    args = parser.parse_args()
    if not 1 <= args.workers <= 32:
        parser.error("workers must be between 1 and 32")
    collect(args.cache_dir, args.output, args.workers)