"""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)