| |
| """Hydrate one non-redistributable LocateAnything media pool locally. |
| |
| The public Parquet relation is the only accepted source-file inventory. Every |
| upstream file is resolved relative to an explicitly supplied root and verified |
| before its bytes become part of a canonical, uncompressed TAR shard. This |
| program never downloads upstream data and never uploads hydrated output. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path, PurePosixPath |
| import re |
| import shutil |
| import subprocess |
| import sys |
| import tarfile |
| from typing import Any, Iterator, Sequence |
|
|
|
|
| SCHEMA = "locate-anything.local-hydration-receipt/v1" |
| PUBLIC_MAP_SCHEMA = "locate-anything.public-media-source-map/v1" |
| ENERGON_VERSION = "7.4.0" |
| DEFAULT_TARGET_BYTES = 4 * 1024 * 1024 * 1024 |
| MEMBER_RE = re.compile(r"^m/[0-9]{12}\.[A-Za-z0-9]+$") |
| RESTRICTED_POOLS = { |
| "crowdhuman": "locany--crowdhuman--train--source_image--v000001", |
| "deepfashion2": "locany--deepfashion2--train--source_image--v000001", |
| "flickr30k": "locany--flickr30k--train--source_image--v000001", |
| "partimagenet": "locany--imagenet--train--source_image--v000001", |
| "object365": "locany--object365--train--source_image--v000001", |
| "sku110k": "locany--sku110k--train--source_image--v000001", |
| "unsplash": "locany--unsplash--train--source_image--v000001", |
| } |
| PARQUET_COLUMNS = ( |
| "schema", |
| "pool_id", |
| "source_id", |
| "source_relative_path", |
| "source_bytes", |
| "source_sha256", |
| "member_name", |
| "content_sha256", |
| "encoded_bytes", |
| "media_kind", |
| "availability", |
| ) |
|
|
|
|
| class HydrationError(RuntimeError): |
| """The public mapping or local upstream media failed a strict check.""" |
|
|
|
|
| def _canonical_json(value: Any) -> bytes: |
| return ( |
| json.dumps( |
| value, |
| ensure_ascii=False, |
| sort_keys=True, |
| separators=(",", ":"), |
| allow_nan=False, |
| ) |
| + "\n" |
| ).encode("utf-8") |
|
|
|
|
| def _safe_relative(value: Any) -> PurePosixPath: |
| if not isinstance(value, str) or not value or "\\" in value or "\x00" in value: |
| raise HydrationError(f"unsafe source_relative_path: {value!r}") |
| relative = PurePosixPath(value) |
| if relative.is_absolute() or any( |
| part in {"", ".", ".."} for part in relative.parts |
| ): |
| raise HydrationError(f"non-canonical source_relative_path: {value!r}") |
| return relative |
|
|
|
|
| def _resolve_inside(root: Path, relative: PurePosixPath) -> Path: |
| candidate = root.joinpath(*relative.parts) |
| try: |
| resolved = candidate.resolve(strict=True) |
| except OSError as error: |
| raise HydrationError(f"missing upstream file: {relative.as_posix()}") from error |
| try: |
| resolved.relative_to(root) |
| except ValueError as error: |
| raise HydrationError( |
| f"upstream path escapes --source-root: {relative.as_posix()}" |
| ) from error |
| if not resolved.is_file(): |
| raise HydrationError( |
| f"upstream path is not a regular file: {relative.as_posix()}" |
| ) |
| return resolved |
|
|
|
|
| def _sha256_file(path: Path, chunk_bytes: int = 8 << 20) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| while block := handle.read(chunk_bytes): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| class _HashingReader: |
| def __init__(self, handle: Any) -> None: |
| self._handle = handle |
| self.digest = hashlib.sha256() |
| self.bytes_read = 0 |
|
|
| def read(self, size: int = -1) -> bytes: |
| block = self._handle.read(size) |
| self.digest.update(block) |
| self.bytes_read += len(block) |
| return block |
|
|
|
|
| class _ShardWriter: |
| def __init__(self, root: Path, target_bytes: int) -> None: |
| self.root = root |
| self.target_bytes = target_bytes |
| self.handle: tarfile.TarFile | None = None |
| self.current_bytes = 0 |
| self.shard_count = 0 |
| self.payload_count = 0 |
| self.payload_bytes = 0 |
|
|
| @staticmethod |
| def _contribution(size: int) -> int: |
| return 512 + ((size + 511) // 512) * 512 |
|
|
| def _open(self) -> None: |
| path = self.root / f"image-{self.shard_count:08d}.tar" |
| self.handle = tarfile.open(path, mode="w", format=tarfile.USTAR_FORMAT) |
| self.current_bytes = 0 |
| self.shard_count += 1 |
|
|
| def add( |
| self, |
| *, |
| source: Path, |
| member_name: str, |
| expected_size: int, |
| expected_sha256: str, |
| ) -> None: |
| contribution = self._contribution(expected_size) |
| if self.handle is None: |
| self._open() |
| elif ( |
| self.current_bytes |
| and self.current_bytes + contribution + 1024 > self.target_bytes |
| ): |
| self.handle.close() |
| self._open() |
| assert self.handle is not None |
| info = tarfile.TarInfo(member_name) |
| info.size = expected_size |
| info.mode = 0o644 |
| info.uid = 0 |
| info.gid = 0 |
| info.uname = "" |
| info.gname = "" |
| info.mtime = 0 |
| with source.open("rb") as raw: |
| reader = _HashingReader(raw) |
| self.handle.addfile(info, reader) |
| observed_sha256 = reader.digest.hexdigest() |
| if reader.bytes_read != expected_size: |
| raise HydrationError( |
| f"{source}: read {reader.bytes_read} bytes, expected {expected_size}" |
| ) |
| if observed_sha256 != expected_sha256: |
| raise HydrationError( |
| f"{source}: SHA-256 {observed_sha256} != {expected_sha256}" |
| ) |
| self.current_bytes += contribution |
| self.payload_count += 1 |
| self.payload_bytes += expected_size |
|
|
| def close(self) -> None: |
| if self.handle is not None: |
| self.handle.close() |
| self.handle = None |
|
|
|
|
| def _mapping_rows(mapping_root: Path) -> Iterator[dict[str, Any]]: |
| try: |
| import pyarrow.dataset as pads |
| except ImportError as error: |
| raise HydrationError( |
| "pyarrow is required to read the public Parquet mapping" |
| ) from error |
| parts = sorted(mapping_root.glob("part-*.parquet")) |
| if not parts: |
| raise HydrationError(f"no Parquet mapping parts under {mapping_root}") |
| dataset = pads.dataset([str(path) for path in parts], format="parquet") |
| missing = set(PARQUET_COLUMNS) - set(dataset.schema.names) |
| if missing: |
| raise HydrationError( |
| f"{mapping_root}: missing mapping columns {sorted(missing)!r}" |
| ) |
| scanner = dataset.scanner(columns=list(PARQUET_COLUMNS), batch_size=65_536) |
| for batch in scanner.to_batches(): |
| columns = batch.to_pydict() |
| for index in range(batch.num_rows): |
| yield {name: columns[name][index] for name in PARQUET_COLUMNS} |
|
|
|
|
| def _manifest_pool(repo_root: Path, dataset_id: str) -> str: |
| path = repo_root / "manifests" / "restricted-media.json" |
| if path.is_file(): |
| value = json.loads(path.read_text(encoding="utf-8")) |
| candidates: Any |
| if isinstance(value, list): |
| candidates = value |
| elif isinstance(value, dict): |
| candidates = value.get("pools") |
| if candidates is None: |
| candidates = value.get("external_media", {}).get("pools", []) |
| else: |
| candidates = [] |
| for row in candidates or []: |
| if isinstance(row, dict) and row.get("dataset_id") == dataset_id: |
| pool_id = row.get("pool_id") |
| if isinstance(pool_id, str) and pool_id: |
| return pool_id |
| try: |
| return RESTRICTED_POOLS[dataset_id] |
| except KeyError as error: |
| raise HydrationError( |
| f"{dataset_id!r} is not one of the seven external-media datasets" |
| ) from error |
|
|
|
|
| def _run_energon_prepare(pool_root: Path, python: Path, workers: int) -> None: |
| probe = subprocess.run( |
| [ |
| str(python), |
| "-c", |
| "import importlib.metadata as m; print(m.version('megatron-energon'))", |
| ], |
| check=False, |
| capture_output=True, |
| text=True, |
| ) |
| if probe.returncode != 0: |
| raise HydrationError( |
| "cannot import megatron-energon from --energon-python: " |
| + (probe.stderr or probe.stdout)[-2000:] |
| ) |
| version = probe.stdout.strip() |
| if version != ENERGON_VERSION: |
| raise HydrationError( |
| f"megatron-energon must be {ENERGON_VERSION}, got {version!r}" |
| ) |
| command = [ |
| str(python), |
| "-m", |
| "megatron.energon.tools.prepare", |
| str(pool_root), |
| "--no-progress", |
| "--num-workers", |
| str(max(1, workers)), |
| "--non-interactive", |
| "--split-parts", |
| r"train:^shards/image-[0-9]{8}\.tar$", |
| "--sample-type", |
| "CrudeWebdataset", |
| ] |
| result = subprocess.run(command, check=False, capture_output=True, text=True) |
| if result.returncode != 0: |
| raise HydrationError( |
| "Energon prepare failed: " + (result.stderr or result.stdout)[-4000:] |
| ) |
| required = ( |
| ".nv-meta/.info.json", |
| ".nv-meta/dataset.yaml", |
| ".nv-meta/index.sqlite", |
| ".nv-meta/index.uuid", |
| ".nv-meta/split.yaml", |
| ) |
| missing = [name for name in required if not (pool_root / name).is_file()] |
| if missing: |
| raise HydrationError(f"Energon omitted required files: {missing!r}") |
|
|
|
|
| def _verify_view_references( |
| repo_root: Path, dataset_id: str, pool_id: str, members: set[str] |
| ) -> int: |
| dataset_root = repo_root / "datasets" / dataset_id |
| if not dataset_root.is_dir(): |
| raise HydrationError(f"dataset directory is missing: {dataset_root}") |
| references = 0 |
| for records_path in sorted(dataset_root.glob("views/*/records.jsonl")): |
| with records_path.open("rb") as handle: |
| for line_number, raw in enumerate(handle, 1): |
| try: |
| record = json.loads(raw) |
| image = record["image"] |
| member_name = image["path"] |
| except (KeyError, TypeError, json.JSONDecodeError) as error: |
| raise HydrationError( |
| f"{records_path}:{line_number}: malformed image reference" |
| ) from error |
| if member_name not in members: |
| raise HydrationError( |
| f"{records_path}:{line_number}: missing member {member_name!r}" |
| ) |
| references += 1 |
| if references == 0: |
| raise HydrationError(f"{dataset_id}: no annotation image references found") |
| return references |
|
|
|
|
| def hydrate(args: argparse.Namespace) -> dict[str, Any]: |
| repo_root = args.repo_root.resolve(strict=True) |
| source_root = args.source_root.resolve(strict=True) |
| if not source_root.is_dir(): |
| raise HydrationError(f"--source-root is not a directory: {source_root}") |
| dataset_id = args.dataset |
| pool_id = args.pool_id or _manifest_pool(repo_root, dataset_id) |
| expected_pool = RESTRICTED_POOLS.get(dataset_id) |
| if expected_pool is not None and pool_id != expected_pool: |
| raise HydrationError( |
| f"{dataset_id}: pool {pool_id!r} != public contract {expected_pool!r}" |
| ) |
| mapping_root = repo_root / "mappings" / "media" / pool_id |
| final_root = repo_root / "media" / pool_id |
| stage = repo_root / "media" / f".{pool_id}.hydrating" |
| if stage.exists(): |
| if not args.force: |
| raise HydrationError( |
| f"stale hydration directory exists; inspect it or rerun with --force: {stage}" |
| ) |
| shutil.rmtree(stage) |
| if final_root.exists(): |
| unexpected = [ |
| path.name for path in final_root.iterdir() if path.name != "availability.json" |
| ] |
| if unexpected: |
| raise HydrationError( |
| f"{final_root} is already hydrated or contains unexpected files: " |
| f"{sorted(unexpected)!r}" |
| ) |
| (stage / "shards").mkdir(parents=True) |
| writer = _ShardWriter(stage / "shards", args.target_tar_bytes) |
| seen_members: set[str] = set() |
| source_rows = 0 |
| verified_source_bytes = 0 |
| try: |
| for row in _mapping_rows(mapping_root): |
| source_rows += 1 |
| if row["schema"] != PUBLIC_MAP_SCHEMA: |
| raise HydrationError( |
| f"mapping row {source_rows}: unsupported schema {row['schema']!r}" |
| ) |
| if row["pool_id"] != pool_id: |
| raise HydrationError( |
| f"mapping row {source_rows}: wrong pool {row['pool_id']!r}" |
| ) |
| if row["availability"] != "upstream_download_required": |
| raise HydrationError( |
| f"mapping row {source_rows}: media is not marked external" |
| ) |
| relative = _safe_relative(row["source_relative_path"]) |
| source = _resolve_inside(source_root, relative) |
| expected_bytes = int(row["source_bytes"]) |
| expected_sha256 = str(row["source_sha256"]) |
| member_name = str(row["member_name"]) |
| content_sha256 = str(row["content_sha256"]) |
| encoded_bytes = int(row["encoded_bytes"]) |
| if not MEMBER_RE.fullmatch(member_name): |
| raise HydrationError( |
| f"mapping row {source_rows}: invalid member_name {member_name!r}" |
| ) |
| if expected_bytes != encoded_bytes: |
| raise HydrationError( |
| f"mapping row {source_rows}: source/encoded byte mismatch" |
| ) |
| if expected_sha256 != content_sha256: |
| raise HydrationError( |
| f"mapping row {source_rows}: source/content SHA-256 mismatch" |
| ) |
| if source.stat().st_size != expected_bytes: |
| raise HydrationError( |
| f"{relative.as_posix()}: size {source.stat().st_size} " |
| f"!= {expected_bytes}" |
| ) |
| if member_name in seen_members: |
| observed = _sha256_file(source) |
| if observed != expected_sha256: |
| raise HydrationError( |
| f"{relative.as_posix()}: SHA-256 {observed} " |
| f"!= {expected_sha256}" |
| ) |
| else: |
| writer.add( |
| source=source, |
| member_name=member_name, |
| expected_size=expected_bytes, |
| expected_sha256=expected_sha256, |
| ) |
| seen_members.add(member_name) |
| verified_source_bytes += expected_bytes |
| finally: |
| writer.close() |
| if source_rows == 0 or not seen_members: |
| raise HydrationError(f"{mapping_root}: empty mapping") |
| _run_energon_prepare(stage, args.energon_python, args.workers) |
| references = _verify_view_references( |
| repo_root, dataset_id, pool_id, seen_members |
| ) |
| availability = { |
| "schema": "locate-anything.media-availability/v1", |
| "availability": "locally_hydrated", |
| "dataset_id": dataset_id, |
| "pool_id": pool_id, |
| "source_rows": source_rows, |
| "physical_media_count": len(seen_members), |
| "tar_count": writer.shard_count, |
| "payload_bytes": writer.payload_bytes, |
| } |
| (stage / "availability.json").write_bytes(_canonical_json(availability)) |
| if final_root.exists(): |
| shutil.rmtree(final_root) |
| os.replace(stage, final_root) |
| receipt = { |
| "schema": SCHEMA, |
| "status": "pass", |
| "dataset_id": dataset_id, |
| "pool_id": pool_id, |
| "source_rows": source_rows, |
| "verified_source_bytes": verified_source_bytes, |
| "physical_media_count": len(seen_members), |
| "payload_bytes": writer.payload_bytes, |
| "tar_count": writer.shard_count, |
| "annotation_reference_count": references, |
| "energon_version": ENERGON_VERSION, |
| } |
| receipt_root = repo_root / ".local" / "hydration-receipts" |
| receipt_root.mkdir(parents=True, exist_ok=True) |
| (receipt_root / f"{pool_id}.json").write_bytes(_canonical_json(receipt)) |
| return receipt |
|
|
|
|
| def _parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--repo-root", type=Path, required=True) |
| parser.add_argument("--dataset", choices=sorted(RESTRICTED_POOLS), required=True) |
| parser.add_argument("--pool-id") |
| parser.add_argument("--source-root", type=Path, required=True) |
| parser.add_argument( |
| "--energon-python", |
| type=Path, |
| default=Path(sys.executable), |
| help=f"Python environment containing megatron-energon {ENERGON_VERSION}.", |
| ) |
| parser.add_argument("--workers", type=int, default=max(1, os.cpu_count() or 1)) |
| parser.add_argument("--target-tar-bytes", type=int, default=DEFAULT_TARGET_BYTES) |
| parser.add_argument( |
| "--force", |
| action="store_true", |
| help="Remove only a stale hidden hydration directory before rebuilding.", |
| ) |
| return parser |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| args = _parser().parse_args(argv) |
| try: |
| result = hydrate(args) |
| except (HydrationError, OSError, ValueError) as error: |
| print(f"ERROR: {error}", file=sys.stderr) |
| return 2 |
| print(json.dumps(result, ensure_ascii=False, sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
