File size: 9,319 Bytes
7c1aeef | 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 | """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)
|