From 5009164850b35af3be1c8b0cb0ca1292ce125e8c Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Tue, 30 Jun 2026 14:32:00 -0400 Subject: [PATCH 01/10] Convert the hybrid HDF5+Parquet dataset into WebDataset training shards Reads WebP images from the HDF5 store joined with resolved-taxonomy and uuid->h5 lookup parquet, writing .tar shards whose samples carry the ten text-prompt sidecars used for BioCLIP-style training (scientific, taxonomic, and common-name variants). When no vernacular common name exists, the common name falls back to the scientific name without author citation. Co-Authored-By: Claude Opus 4.8 --- config/tol_hdf5_to_wds_example.yaml | 49 +++ generate_hdf5_lookup.py | 122 ++++++ pyproject.toml | 3 + scripts/copy_implicated_shards.py | 222 ++++++++++ scripts/find_empty_taxonomy.py | 67 +++ scripts/publish_scrubbed_shards.py | 174 ++++++++ scripts/remove_bad_samples.py | 89 ++++ scripts/remove_bad_samples_parallel.py | 180 ++++++++ scripts/run_incremental_scrub.sh | 57 +++ scripts/tools_filter.slurm | 17 +- scripts/tools_scheduler.slurm | 7 +- scripts/tools_submit.sh | 0 scripts/tools_verifier.slurm | 7 +- scripts/tools_worker.slurm | 7 +- src/TreeOfLife_toolbox/__init__.py | 7 +- .../tol_hdf5_to_wds/README.md | 58 +++ .../tol_hdf5_to_wds/__init__.py | 8 + .../tol_hdf5_to_wds/classes.py | 404 ++++++++++++++++++ .../tol_hdf5_to_wds/utils.py | 171 ++++++++ tol_hdf5_to_wds.py | 324 ++++++++++++++ validate_tars.py | 250 +++++++++++ wds_taxoncom_audit.py | 231 ++++++++++ 22 files changed, 2432 insertions(+), 22 deletions(-) create mode 100644 config/tol_hdf5_to_wds_example.yaml create mode 100644 generate_hdf5_lookup.py create mode 100644 scripts/copy_implicated_shards.py create mode 100644 scripts/find_empty_taxonomy.py create mode 100644 scripts/publish_scrubbed_shards.py create mode 100644 scripts/remove_bad_samples.py create mode 100644 scripts/remove_bad_samples_parallel.py create mode 100644 scripts/run_incremental_scrub.sh mode change 100644 => 100755 scripts/tools_submit.sh create mode 100644 src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md create mode 100644 src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py create mode 100644 src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py create mode 100644 src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py create mode 100644 tol_hdf5_to_wds.py create mode 100644 validate_tars.py create mode 100644 wds_taxoncom_audit.py diff --git a/config/tol_hdf5_to_wds_example.yaml b/config/tol_hdf5_to_wds_example.yaml new file mode 100644 index 0000000..8ef5f85 --- /dev/null +++ b/config/tol_hdf5_to_wds_example.yaml @@ -0,0 +1,49 @@ +account: "PAS2136" + +# Input directory should point to the root of the parquet_to_hdf5 outputs. +path_to_input: "/fs/ess/PAS2136/TreeOfLife_HDF5/converted" +# Working directory for toolbox artifacts (logs, checkpoints, etc.). +path_to_output_folder: "/fs/scratch/PAS2136/TreeOfLife_hdf5_to_wds" + +tools_parameters: + num_workers: 8 + max_nodes: 2 + workers_per_node: 4 + cpu_per_worker: 32 + threshold_size: 224 + new_resize_size: 720 + +scripts: + tools_submitter: "scripts/tools_submit.sh" + tools_filter_script: "scripts/tools_filter.slurm" + tools_scheduling_script: "scripts/tools_scheduler.slurm" + tools_worker_script: "scripts/tools_worker.slurm" + tools_verification_script: "scripts/tools_verifier.slurm" + +output_structure: + urls_folder: "urls" + logs_folder: "logs" + images_folder: "images" + schedules_folder: "schedules" + profiles_table: "profiles.csv" + ignored_table: "ignored.csv" + inner_checkpoint_file: "inner_checkpoint.yaml" + tools_folder: "tools" + +tol_hdf5_to_wds: + shard_size: 10000 + shard_limit: 0 + metadata_glob: "**/*_metadata.parquet" + include_sources: [] + include_servers: [] + tar_output_root: "/fs/ess/PAS2136/TreeOfLife_WDS/shards" + resize_size: 224 + shard_prefix: "shard" + runner_timeout_seconds: 3600 + # Lookup selecting which images to convert: a parquet/CSV with a `uuid` + # column. The `.h5` path is auto-derived as the sibling `*_images.h5` of each + # `*_metadata.parquet`, so no path column is required (an optional `hdf5_path` + # or `h5_file` column overrides it). + lookup_table_path: "/fs/ess/PAS2136/TreeOfLife/lookups/subset.parquet" + # Resolved-taxonomy parquet, joined by `uuid`; required for the text prompts. + taxa_glob: "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/source=*/*.parquet" diff --git a/generate_hdf5_lookup.py b/generate_hdf5_lookup.py new file mode 100644 index 0000000..c368d9b --- /dev/null +++ b/generate_hdf5_lookup.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python +""" +Generate a UUID-to-HDF5 lookup table from a converted Tree of Life dataset. + +Usage: + python generate_hdf5_lookup.py \ + --data-root /fs/scratch/PAS2136/TreeOfLife_test-wds/data \ + --output /fs/scratch/PAS2136/TreeOfLife_test-wds/lookup-tables/all/hdf5_lookup.parquet \ + --sample 100 +""" + +from __future__ import annotations + +import argparse +import glob +import os +from typing import Iterable, List + +import polars as pl + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Create a lookup Parquet of UUIDs to HDF5/metadata paths." + ) + parser.add_argument( + "--data-root", + required=True, + help="Root directory containing source=*/server=*/data_*_{metadata,images} files.", + ) + parser.add_argument( + "--output", + required=True, + help="Destination Parquet file (e.g. /.../lookup-tables/all/lookup.parquet).", + ) + parser.add_argument( + "--sample", + type=float, + default=100.0, + help="Percentage of UUIDs to keep (0 < sample <= 100). Default: 100.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed used when down-sampling.", + ) + return parser.parse_args() + + +def discover_metadata_files(data_root: str) -> List[str]: + pattern = os.path.join(data_root, "source=*", "server=*", "data_*_metadata.parquet") + return sorted(glob.glob(pattern)) + + +def build_lazy_frames(metadata_files: Iterable[str]) -> List[pl.LazyFrame]: + lazy_frames: List[pl.LazyFrame] = [] + for meta_path in metadata_files: + source = _extract_between(meta_path, "source=", "/server=") + server = _extract_between(meta_path, "server=", "/data_") + base_name = os.path.basename(meta_path).replace("_metadata.parquet", "") + hdf5_path = meta_path.replace("_metadata.parquet", "_images.h5") + + lazy = ( + pl.scan_parquet(meta_path) + .select("uuid") + .with_columns( + pl.lit(meta_path).alias("metadata_path"), + pl.lit(hdf5_path).alias("hdf5_path"), + pl.lit(source).alias("source"), + pl.lit(server).alias("server"), + pl.lit(base_name).alias("base_name"), + ) + ) + lazy_frames.append(lazy) + return lazy_frames + + +def _extract_between(value: str, start_token: str, end_token: str) -> str: + start_idx = value.find(start_token) + if start_idx == -1: + return "" + start_idx += len(start_token) + end_idx = value.find(end_token, start_idx) + if end_idx == -1: + return value[start_idx:] + return value[start_idx:end_idx] + + +def main() -> None: + args = parse_args() + data_root = os.path.abspath(args.data_root) + metadata_files = discover_metadata_files(data_root) + + if not metadata_files: + raise SystemExit( + f"No metadata parquet files found under {data_root}. " + "Expected pattern source=*/server=*/data_*_metadata.parquet" + ) + + lazy_frames = build_lazy_frames(metadata_files) + lookup_lazy = pl.concat(lazy_frames) + + if args.sample <= 0 or args.sample > 100: + raise SystemExit("--sample must be in the interval (0, 100].") + + if args.sample < 100: + frac = args.sample / 100.0 + lookup_lazy = lookup_lazy.sample( + fraction=frac, with_replacement=False, shuffle=True, seed=args.seed + ) + + lookup_df = lookup_lazy.collect(streaming=True) + + output_path = os.path.abspath(args.output) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + lookup_df.write_parquet(output_path, compression="zstd") + print(f"Wrote {lookup_df.height} records to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 02b2e82..792ccfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,11 +30,13 @@ dependencies = [ "cramjam", "cython", "fsspec", + "h5py", "inflate64", "mpi4py", "multivolumefile", "opencv-python", "pandas", + "polars", "pathspec", "pillow", "psutil", @@ -51,6 +53,7 @@ dependencies = [ "texttable", "trove-classifiers", "typing-extensions", + "webdataset", "wheel" ] diff --git a/scripts/copy_implicated_shards.py b/scripts/copy_implicated_shards.py new file mode 100644 index 0000000..5c57d61 --- /dev/null +++ b/scripts/copy_implicated_shards.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +""" +Copy shard tar files implicated by missing-taxa.csv into a staging directory, +but ONLY when the source shard appears complete. + +Completeness criteria (intrinsic): +- `tar tf` succeeds (tar structure readable) +- exactly 10,000 *.jpg entries +- exactly 10,000 entries for each required TXT sidecar suffix + +Staging/update semantics: +- If a scrubbed output already exists (in --scrubbed-dir), skip entirely. +- If destination shard exists AND same byte size as source, skip. +- Otherwise, (re)copy the source shard into destination (atomic temp + rename), + but only if the source passes completeness. + +Usage: + python scripts/copy_implicated_shards.py \ + --map missing-taxa.csv \ + --src-dir /path/to/shards \ + --dst-dir /path/to/shards_to_scrub \ + --scrubbed-dir /path/to/shards_scrubbed +""" + +from __future__ import annotations + +import argparse +import csv +import os +import shutil +import subprocess +from pathlib import Path +from typing import Dict, List, Set + +from tqdm import tqdm + + +SIDE_TXT_SUFFIXES: List[str] = [ + "com.txt", + "common_name.txt", + "sci.txt", + "sci_com.txt", + "scientific_name.txt", + "taxon.txt", + "taxonTag.txt", + "taxonTag_com.txt", + "taxon_com.txt", + "taxonomic_name.txt", +] + +EXPECTED_PER_TYPE = 10_000 + + +def shard_name(shard_id: int) -> str: + return f"shard-{shard_id:05d}.tar" + + +def count_data_rows(path: Path) -> int: + with path.open("rb") as f: + return max(0, sum(1 for _ in f) - 1) + + +def load_unique_shard_ids(csv_path: Path, total_rows: int | None = None) -> Set[int]: + shard_ids: Set[int] = set() + with csv_path.open(newline="") as fh: + reader = csv.DictReader(fh) + for row in tqdm(reader, total=total_rows, desc="Scanning missing-taxa.csv", unit="rows"): + shard_ids.add(int(row["shard_id"])) + return shard_ids + + +def tar_counts_via_tar_tf(tar_path: Path) -> Dict[str, int] | None: + counts: Dict[str, int] = {sfx: 0 for sfx in SIDE_TXT_SUFFIXES} + counts["jpg"] = 0 + counts["total"] = 0 + + proc = subprocess.Popen( + ["tar", "tf", str(tar_path)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + + assert proc.stdout is not None + try: + for line in proc.stdout: + name = line.strip() + if not name: + continue + counts["total"] += 1 + + if name.endswith(".jpg"): + counts["jpg"] += 1 + continue + + for sfx in SIDE_TXT_SUFFIXES: + if name.endswith("." + sfx): + counts[sfx] += 1 + break + finally: + proc.stdout.close() + rc = proc.wait() + + if rc != 0: + return None + return counts + + +def is_complete_shard(tar_path: Path) -> bool: + counts = tar_counts_via_tar_tf(tar_path) + if counts is None: + return False + + if counts["jpg"] != EXPECTED_PER_TYPE: + return False + + for sfx in SIDE_TXT_SUFFIXES: + if counts[sfx] != EXPECTED_PER_TYPE: + return False + + return True + + +def copy_atomic(src: Path, dst: Path) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.with_name(dst.name + ".tmp") + shutil.copy2(src, tmp) # preserve mtime for staging comparisons if you care later + os.replace(tmp, dst) + + +def same_size(src: Path, dst: Path) -> bool: + try: + return src.stat().st_size == dst.stat().st_size + except FileNotFoundError: + return False + + +def main() -> None: + p = argparse.ArgumentParser( + description="Copy implicated shard tars into a staging directory (only when complete)." + ) + p.add_argument("--map", required=True, type=Path, help="missing-taxa.csv with columns uuid,shard_id.") + p.add_argument("--src-dir", required=True, type=Path, help="Directory containing shard-XXXXX.tar files.") + p.add_argument("--dst-dir", required=True, type=Path, help="Destination directory for shards to scrub.") + p.add_argument( + "--scrubbed-dir", + required=True, + type=Path, + help="Directory containing scrubbed shard outputs; if shard exists here, it is skipped.", + ) + p.add_argument("--limit", type=int, default=0, help="Optional: stop after copying N shards (0 = no limit).") + args = p.parse_args() + + total_rows = count_data_rows(args.map) + shard_ids = load_unique_shard_ids(args.map, total_rows=total_rows) + shard_id_list = sorted(shard_ids) + + copied = 0 + updated = 0 + + skipped_missing = 0 + skipped_incomplete = 0 + skipped_same_size = 0 + skipped_already_scrubbed = 0 + + for sid in tqdm(shard_id_list, desc="Checking/copying shards", unit="shards"): + src = args.src_dir / shard_name(sid) + dst = args.dst_dir / src.name + scrubbed = args.scrubbed_dir / src.name + + # Robust: if already scrubbed, never stage again (prevents noise if scrubbed published back to src-dir) + if scrubbed.exists(): + skipped_already_scrubbed += 1 + continue + + if not src.exists(): + skipped_missing += 1 + continue + + # If destination exists and size matches, skip. + if dst.exists() and same_size(src, dst): + skipped_same_size += 1 + continue + + # Only copy/update if source is intrinsically complete. + if not is_complete_shard(src): + skipped_incomplete += 1 + continue + + was_update = dst.exists() + copy_atomic(src, dst) + if was_update: + updated += 1 + else: + copied += 1 + + if (copied + updated) % 25 == 0: + print( + f"[INFO] copied={copied} updated={updated} " + f"missing={skipped_missing} incomplete={skipped_incomplete} " + f"same_size_skip={skipped_same_size} already_scrubbed_skip={skipped_already_scrubbed}" + ) + + if args.limit and (copied + updated) >= args.limit: + break + + print( + "[DONE] " + f"unique_shards={len(shard_id_list)} " + f"copied={copied} " + f"updated={updated} " + f"skipped_missing={skipped_missing} " + f"skipped_incomplete={skipped_incomplete} " + f"skipped_same_size={skipped_same_size} " + f"skipped_already_scrubbed={skipped_already_scrubbed}" + ) + + +if __name__ == "__main__": + main() + diff --git a/scripts/find_empty_taxonomy.py b/scripts/find_empty_taxonomy.py new file mode 100644 index 0000000..dc8dd8e --- /dev/null +++ b/scripts/find_empty_taxonomy.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python +""" +Scan shard_metadata for UUIDs lacking any taxonomy fields and emit a CSV map. + +Usage: + python scripts/find_empty_taxonomy.py /path/to/shard_metadata output.csv +""" + +import argparse +import csv +from pathlib import Path + +import polars as pl + + +def gather_null_taxonomy(shard_metadata_root: Path) -> pl.DataFrame: + dataset = pl.scan_parquet( + str(shard_metadata_root / "shard_id=*" / "*.parquet"), hive_partitioning=True + ) + taxonomy_cols = [ + "scientific_name", + "common_name", + "provided_common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + filter_expr = None + for col in taxonomy_cols: + expr = pl.col(col).is_null() | (pl.col(col).str.strip_chars().eq("")) + filter_expr = expr if filter_expr is None else (filter_expr & expr) + + empty_taxa_df = ( + dataset.filter(filter_expr) + .select("uuid", "shard_id") + .collect() + .sort("shard_id") + ) + return empty_taxa_df + + +def main(): + parser = argparse.ArgumentParser( + description="Find shard metadata rows with empty taxonomy." + ) + parser.add_argument("shard_metadata_root", type=Path) + parser.add_argument("output_csv", type=Path) + args = parser.parse_args() + + df = gather_null_taxonomy(args.shard_metadata_root) + if df.is_empty(): + args.output_csv.write_text("uuid,shard_id\n") + return + + with args.output_csv.open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["uuid", "shard_id"]) + for row in df.iter_rows(): + writer.writerow(row) + + +if __name__ == "__main__": + main() diff --git a/scripts/publish_scrubbed_shards.py b/scripts/publish_scrubbed_shards.py new file mode 100644 index 0000000..f065dc4 --- /dev/null +++ b/scripts/publish_scrubbed_shards.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +""" +Publish scrubbed shards back into the primary shards directory. + +Design goals: +- Only overwrite shards/ with scrubbed versions from shards_scrubbed/. +- Do NOT touch shards_to_scrub/ (your backup copies). +- Avoid clobbering a shard that is still being generated: only publish onto a "complete" source shard. +- Publish via temp + atomic rename so Globus never sees a partial file. +- Do NOT preserve mtime (important for Globus sync level 2: mtime-based). + +Usage: + python scripts/publish_scrubbed_shards.py \ + --shards /path/to/shards \ + --shards-scrubbed /path/to/shards_scrubbed +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +from pathlib import Path +from typing import Dict, List, Optional + + +SIDE_TXT_SUFFIXES: List[str] = [ + "com.txt", + "common_name.txt", + "sci.txt", + "sci_com.txt", + "scientific_name.txt", + "taxon.txt", + "taxonTag.txt", + "taxonTag_com.txt", + "taxon_com.txt", + "taxonomic_name.txt", +] + +EXPECTED_PER_TYPE = 10_000 + + +def tar_counts_via_tar_tf(tar_path: Path) -> Optional[Dict[str, int]]: + """ + Return counts of jpg + each sidecar suffix for tar_path by streaming `tar tf`. + Returns None if `tar tf` fails. + """ + counts: Dict[str, int] = {sfx: 0 for sfx in SIDE_TXT_SUFFIXES} + counts["jpg"] = 0 + + proc = subprocess.Popen( + ["tar", "tf", str(tar_path)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + + assert proc.stdout is not None + try: + for line in proc.stdout: + name = line.strip() + if not name: + continue + + if name.endswith(".jpg"): + counts["jpg"] += 1 + continue + + for sfx in SIDE_TXT_SUFFIXES: + if name.endswith("." + sfx): + counts[sfx] += 1 + break + finally: + proc.stdout.close() + rc = proc.wait() + + return None if rc != 0 else counts + + +def is_complete_source_shard(tar_path: Path) -> bool: + """ + Only publish onto a shard in shards/ if it still looks like a complete original shard. + (Prevents racing the generator / avoids clobbering partials.) + """ + counts = tar_counts_via_tar_tf(tar_path) + if counts is None: + return False + + if counts["jpg"] != EXPECTED_PER_TYPE: + return False + + for sfx in SIDE_TXT_SUFFIXES: + if counts[sfx] != EXPECTED_PER_TYPE: + return False + + return True + + +def copy_atomic_no_mtime(src: Path, dst: Path) -> None: + """ + Copy src -> dst through a temp file in dst's directory, then atomic rename. + Does NOT preserve metadata/mtime (good for Globus sync level 2). + """ + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.with_name(dst.name + ".tmp") + shutil.copyfile(src, tmp) # intentionally not copy2() + os.replace(tmp, dst) + + +def main() -> None: + p = argparse.ArgumentParser(description="Publish scrubbed shard tars back into shards/ safely.") + p.add_argument("--shards", required=True, type=Path, help="Primary shards/ directory to overwrite.") + p.add_argument("--shards-scrubbed", required=True, type=Path, help="Directory containing scrubbed shard outputs.") + p.add_argument("--limit", type=int, default=0, help="Optional: stop after publishing N shards (0 = no limit).") + args = p.parse_args() + + shards_dir: Path = args.shards + scrubbed_dir: Path = args.shards_scrubbed + + published = 0 + skipped_missing_src = 0 + skipped_incomplete_src = 0 + skipped_no_change = 0 + + for scrubbed_tar in sorted(scrubbed_dir.glob("shard-*.tar")): + name = scrubbed_tar.name + dst_tar = shards_dir / name + + if not dst_tar.exists(): + # If the generator hasn't created it yet, don't create it here. + skipped_missing_src += 1 + continue + + # Never clobber a shard that doesn't look "complete" yet. + if not is_complete_source_shard(dst_tar): + skipped_incomplete_src += 1 + continue + + # If shards/ already matches scrubbed by size, assume already published. + # (Cheap heuristic; avoids unnecessary rewrite/mtime churn.) + try: + if dst_tar.stat().st_size == scrubbed_tar.stat().st_size: + skipped_no_change += 1 + continue + except FileNotFoundError: + # race: try again next loop + continue + + copy_atomic_no_mtime(scrubbed_tar, dst_tar) + published += 1 + + if published % 25 == 0: + print( + f"[INFO] published={published} no_change={skipped_no_change} " + f"missing_src={skipped_missing_src} incomplete_src={skipped_incomplete_src}" + ) + + if args.limit and published >= args.limit: + break + + print( + "[DONE] " + f"published={published} " + f"skipped_no_change={skipped_no_change} " + f"skipped_missing_src={skipped_missing_src} " + f"skipped_incomplete_src={skipped_incomplete_src}" + ) + + +if __name__ == "__main__": + main() + diff --git a/scripts/remove_bad_samples.py b/scripts/remove_bad_samples.py new file mode 100644 index 0000000..8bf48a6 --- /dev/null +++ b/scripts/remove_bad_samples.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python +""" +Remove specified UUIDs from shard tar files. + +Usage: + python scripts/remove_bad_samples.py \ + --shards /path/to/shards \ + --map missing-taxa.csv \ + [--backup-dir /path/to/backups] +""" + +from __future__ import annotations + +import argparse +import csv +import os +import tarfile +from collections import defaultdict +from pathlib import Path +from typing import Dict, Iterable, List + + +SIDE_EXTS = [ + "jpg", + "com.txt", + "common_name.txt", + "sci.txt", + "sci_com.txt", + "scientific_name.txt", + "taxon.txt", + "taxonTag.txt", + "taxonTag_com.txt", + "taxon_com.txt", + "taxonomic_name.txt", +] + + +def load_uuid_map(csv_path: Path) -> Dict[int, List[str]]: + shard_map: Dict[int, List[str]] = defaultdict(list) + with csv_path.open() as fh: + reader = csv.DictReader(fh) + for row in reader: + shard_map[int(row["shard_id"])].append(row["uuid"]) + return shard_map + + +def scrub_shard(shard_path: Path, uuids: Iterable[str], backup_dir: Path | None) -> None: + if not shard_path.exists(): + print(f"[WARN] Shard not found: {shard_path}") + return + + tmp_path = shard_path.with_suffix(".tmp") + drop_prefixes = set(uuids) + drop_members = {f"{u}.{ext}" for u in drop_prefixes for ext in SIDE_EXTS} + + with tarfile.open(shard_path, "r") as src, tarfile.open(tmp_path, "w") as dst: + for member in src: + base = member.name.split(".", 1)[0] + if base in drop_prefixes and member.name in drop_members: + continue + data = src.extractfile(member) if member.isfile() else None + dst.addfile(member, data) + + if backup_dir: + backup_dir.mkdir(parents=True, exist_ok=True) + shard_backup = backup_dir / shard_path.name + os.replace(shard_path, shard_backup) + else: + shard_path.unlink() + os.replace(tmp_path, shard_path) + + +def main(): + parser = argparse.ArgumentParser(description="Remove bad UUIDs from shard tar files.") + parser.add_argument("--shards", required=True, type=Path, help="Directory containing shard-XXXXX.tar files.") + parser.add_argument("--map", required=True, type=Path, help="CSV with columns uuid,shard_id.") + parser.add_argument("--backup-dir", type=Path, help="Optional dir to store original tar backups.") + args = parser.parse_args() + + shard_map = load_uuid_map(args.map) + for shard_id, uuids in shard_map.items(): + shard_name = f"shard-{shard_id:05d}.tar" + shard_path = args.shards / shard_name + print(f"[INFO] Scrubbing {len(uuids)} samples from {shard_name}") + scrub_shard(shard_path, uuids, args.backup_dir) + + +if __name__ == "__main__": + main() diff --git a/scripts/remove_bad_samples_parallel.py b/scripts/remove_bad_samples_parallel.py new file mode 100644 index 0000000..dcde917 --- /dev/null +++ b/scripts/remove_bad_samples_parallel.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +""" +Remove specified UUIDs (and sidecars) from shard tar files using parallel workers. + +Incremental behavior: +- Reads from --shards (input directory) +- Writes scrubbed shards to --shards-scrubbed (required output directory) +- If scrubbed output already exists for a shard, that shard is skipped +- Missing input shards are skipped + +Usage: + python scripts/remove_bad_samples_parallel.py \ + --shards /path/to/shards_to_scrub \ + --shards-scrubbed /path/to/shards_scrubbed \ + --map missing-taxa.csv \ + --workers 8 \ + --max-inflight 32 +""" + +from __future__ import annotations + +import argparse +import csv +import os +import tarfile +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path +from typing import Iterable, Iterator, List, Optional, Tuple + +SIDE_EXTS = { + "jpg", + "com.txt", + "common_name.txt", + "sci.txt", + "sci_com.txt", + "scientific_name.txt", + "taxon.txt", + "taxonTag.txt", + "taxonTag_com.txt", + "taxon_com.txt", + "taxonomic_name.txt", +} + + +def iter_shard_groups(csv_path: Path) -> Iterator[Tuple[int, List[str]]]: + """ + Stream (shard_id, [uuid...]) groups from a CSV sorted by shard_id. + """ + with csv_path.open(newline="") as fh: + reader = csv.DictReader(fh) + current_sid: Optional[int] = None + uuids: List[str] = [] + + for row in reader: + sid = int(row["shard_id"]) + u = row["uuid"] + + if current_sid is None: + current_sid = sid + + if sid != current_sid: + yield current_sid, uuids + current_sid = sid + uuids = [u] + else: + uuids.append(u) + + if current_sid is not None and uuids: + yield current_sid, uuids + + +def scrub_shard( + src_shard_path: Path, + dst_shard_path: Path, + uuids: Iterable[str], +) -> Tuple[str, int, int]: + """ + Rewrite shard tar, dropping members whose basename matches .. + Writes scrubbed output to dst_shard_path using temp + atomic rename. + Returns (shard_name, removed_count, kept_count). + """ + if not src_shard_path.exists(): + return (src_shard_path.name, 0, 0) + + dst_shard_path.parent.mkdir(parents=True, exist_ok=True) + tmp_out = dst_shard_path.with_name(dst_shard_path.name + ".tmp") + + drop_prefixes = set(uuids) + removed = 0 + kept = 0 + + with tarfile.open(src_shard_path, "r") as src, tarfile.open(tmp_out, "w") as dst: + for member in src: + filename = os.path.basename(member.name) + + dot = filename.find(".") + if dot != -1: + base = filename[:dot] + if base in drop_prefixes: + suffix = filename[dot + 1 :] + if suffix in SIDE_EXTS: + removed += 1 + continue + + data = src.extractfile(member) if member.isfile() else None + dst.addfile(member, data) + kept += 1 + + os.replace(tmp_out, dst_shard_path) + return (src_shard_path.name, removed, kept) + + +def _worker(args: Tuple[Path, Path, List[str]]) -> Tuple[str, int, int]: + src_shard_path, dst_shard_path, uuids = args + return scrub_shard(src_shard_path, dst_shard_path, uuids) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Remove bad UUIDs from shard tar files (parallel, incremental).") + parser.add_argument("--shards", required=True, type=Path, help="Input directory containing shard-XXXXX.tar files.") + parser.add_argument("--shards-scrubbed", required=True, type=Path, help="Output directory for scrubbed shards.") + parser.add_argument("--map", required=True, type=Path, help="CSV with columns uuid,shard_id (sorted by shard_id recommended).") + parser.add_argument("--workers", type=int, default=8, help="Parallel shard workers (I/O heavy; tune to cluster limits).") + parser.add_argument("--max-inflight", type=int, default=32, help="Bound queued shard jobs to limit FS pressure.") + args = parser.parse_args() + + shards_dir: Path = args.shards + scrubbed_dir: Path = args.shards_scrubbed + workers = max(1, args.workers) + max_inflight = max(workers, args.max_inflight) + + scrubbed_dir.mkdir(parents=True, exist_ok=True) + + inflight = [] + total_removed = 0 + shards_processed = 0 + shards_skipped_missing = 0 + shards_skipped_already = 0 + + with ProcessPoolExecutor(max_workers=workers) as executor: + for shard_id, uuids in iter_shard_groups(args.map): + shard_name = f"shard-{shard_id:05d}.tar" + src_shard_path = shards_dir / shard_name + dst_shard_path = scrubbed_dir / shard_name + + if not src_shard_path.exists(): + shards_skipped_missing += 1 + continue + + # Incremental: if output already exists, skip + if dst_shard_path.exists(): + shards_skipped_already += 1 + continue + + fut = executor.submit(_worker, (src_shard_path, dst_shard_path, uuids)) + inflight.append(fut) + + if len(inflight) >= max_inflight: + name, removed, kept = inflight.pop(0).result() + total_removed += removed + shards_processed += 1 + print(f"[INFO] {name}: removed={removed} kept={kept}") + + for fut in as_completed(inflight): + name, removed, kept = fut.result() + total_removed += removed + shards_processed += 1 + print(f"[INFO] {name}: removed={removed} kept={kept}") + + print( + "[DONE] " + f"shards_processed={shards_processed} " + f"shards_skipped_missing={shards_skipped_missing} " + f"shards_skipped_already={shards_skipped_already} " + f"members_removed={total_removed}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_incremental_scrub.sh b/scripts/run_incremental_scrub.sh new file mode 100644 index 0000000..881513e --- /dev/null +++ b/scripts/run_incremental_scrub.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +set -euo pipefail + +START_TS=$(date +%s) +MAX_SECONDS=$((16 * 60 * 60)) # 16 hours +SLEEP_SECONDS=60 # 1 minute between cycles + +ts() { date +"%Y-%m-%d %H:%M:%S"; } + +echo "[$(ts)] [INFO] Starting incremental copy+scrub loop" +echo "[$(ts)] [INFO] Will stop after $((MAX_SECONDS/3600)) hours" + +while true; do + NOW_TS=$(date +%s) + ELAPSED=$((NOW_TS - START_TS)) + + if [ "$ELAPSED" -ge "$MAX_SECONDS" ]; then + echo "[$(ts)] [INFO] Reached $((MAX_SECONDS/3600))-hour limit; exiting." + break + fi + + echo + echo "[$(ts)] [INFO] ===== elapsed=$((ELAPSED/60)) min =====" + + echo "[$(ts)] [INFO] Running copy_implicated_shards.py" + python scripts/copy_implicated_shards.py \ + --map missing-taxa.csv \ + --src-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ + --dst-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_to_scrub \ + --scrubbed-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed + + echo "[$(ts)] [INFO] Running remove_bad_samples_parallel.py" + python scripts/remove_bad_samples_parallel.py \ + --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_to_scrub \ + --shards-scrubbed /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed \ + --map missing-taxa.csv \ + --workers 2 \ + --max-inflight 2 + + echo "[$(ts)] [INFO] Publishing scrubbed shards back into primary shards/" + python scripts/publish_scrubbed_shards.py \ + --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ + --shards-scrubbed /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed + + echo "[$(ts)] [INFO] Updating sizes.json" + python scripts/sizes_incremental_fast.py \ + --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ + --map missing-taxa.csv \ + --min-age-seconds 180 \ + --progress-every 50 \ + --heartbeat-seconds 30 + + echo "[$(ts)] [INFO] Sleeping ${SLEEP_SECONDS}s" + sleep "$SLEEP_SECONDS" +done + +echo "[$(ts)] [DONE] Incremental loop finished" diff --git a/scripts/tools_filter.slurm b/scripts/tools_filter.slurm index 6aee3f6..f6d8c69 100644 --- a/scripts/tools_filter.slurm +++ b/scripts/tools_filter.slurm @@ -1,7 +1,8 @@ #!/bin/bash #SBATCH --job-name tool_filter -#SBATCH --mem=0 -#SBATCH --time=01:00:00 +#SBATCH --partition=cpu +#SBATCH --mem=200G +#SBATCH --time=2:00:00 if [ "$#" -eq 0 ]; then echo "Usage: $0 tool_name" @@ -10,17 +11,21 @@ fi tool_name=$1 -logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" +logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER:-${SCRATCH:-$PWD}/tools_logs}" mkdir -p "$logs_dir" driver_memory="110G" executor_memory="64G" -module load spark/3.4.1 -module load miniconda3/23.3.1-py310 +module load spark/3.5.1 +module load python/3.12 +if [ -z "$REPO_ROOT" ]; then + echo "REPO_ROOT is not set. Please export REPO_ROOT to the repository path." >&2 + exit 1 +fi source "${REPO_ROOT}/.venv/bin/activate" -pbs-spark-submit \ +spark-submit \ --driver-memory $driver_memory \ --executor-memory $executor_memory \ "${TOOLBOX_PATH}/main/filter.py" \ diff --git a/scripts/tools_scheduler.slurm b/scripts/tools_scheduler.slurm index ea35a32..8b159e4 100644 --- a/scripts/tools_scheduler.slurm +++ b/scripts/tools_scheduler.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_scheduler -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=00:05:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes=1 \ --ntasks-per-node=1 \ --cpus-per-task=1 \ - --mem=0 \ --output="${logs_dir}/tool_scheduler.log" \ python "${TOOLBOX_PATH}/main/scheduler.py" "${tool_name}" diff --git a/scripts/tools_submit.sh b/scripts/tools_submit.sh old mode 100644 new mode 100755 diff --git a/scripts/tools_verifier.slurm b/scripts/tools_verifier.slurm index 6a3b75e..ac332f5 100644 --- a/scripts/tools_verifier.slurm +++ b/scripts/tools_verifier.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_verifier -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=00:05:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes=1 \ --ntasks-per-node=1 \ --cpus-per-task=1 \ - --mem=0 \ --output="${logs_dir}/tool_verifier.log" \ python "${TOOLBOX_PATH}/main/verification.py" "${tool_name}" diff --git a/scripts/tools_worker.slurm b/scripts/tools_worker.slurm index 4856e62..601bd8d 100644 --- a/scripts/tools_worker.slurm +++ b/scripts/tools_worker.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_worker -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=03:00:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes="$TOOLS_MAX_NODES" \ --ntasks-per-node="$TOOLS_WORKERS_PER_NODE" \ --cpus-per-task="$TOOLS_CPU_PER_WORKER" \ - --mem=0 \ --output="${logs_dir}/tool_worker-%2t.log" \ python "${TOOLBOX_PATH}/main/runner.py" "${tool_name}" diff --git a/src/TreeOfLife_toolbox/__init__.py b/src/TreeOfLife_toolbox/__init__.py index c1a2a93..8d14536 100644 --- a/src/TreeOfLife_toolbox/__init__.py +++ b/src/TreeOfLife_toolbox/__init__.py @@ -1,2 +1,7 @@ -from TreeOfLife_toolbox import tol200m_fathom_net_crop, tol200m_bioscan_data_transfer, data_transfer +from TreeOfLife_toolbox import ( + tol200m_fathom_net_crop, + tol200m_bioscan_data_transfer, + data_transfer, + tol_hdf5_to_wds, +) diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md new file mode 100644 index 0000000..e523652 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md @@ -0,0 +1,58 @@ +# tol_hdf5_to_wds + +Converts the hybrid Tree of Life dataset (HDF5 images + Parquet metadata) into +[WebDataset](https://github.com/webdataset/webdataset) `.tar` shards that mirror +the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. + +## Pipeline steps + +1. **Filter**: Scans the converted dataset for `*_metadata.parquet` files, + enriches each row with the matching `*_images.h5` path, and writes shard-sized + parquet partitions under `tools/tol_hdf5_to_wds/shard_metadata`. +2. **Scheduler**: Assigns shard identifiers to MPI ranks and writes + `schedule.csv` for the toolbox runner scripts. +3. **Runner**: For every scheduled shard, loads the shard metadata, reads + WebP-compressed bytes from the referenced HDF5 file, converts them to JPEG + (with an optional resize), generates ten taxonomy/common-name text prompts, + and emits `shard-XXXXX.tar` archives via `webdataset.TarWriter`. + +## Configuration highlights + +Add a `tol_hdf5_to_wds` block to the standard toolbox config template: + +```yaml +tol_hdf5_to_wds: + shard_size: 10000 # rows per shard + shard_limit: 0 # optional cap, 0 disables + metadata_glob: "**/*_metadata.parquet" + tar_output_root: "/fs/.../wds" # defaults to tools folder if omitted + resize_size: 224 # square resize in pixels (0 keeps original) + lookup_table_path: "/fs/.../subset.parquet" # optional UUID list (CSV or Parquet) + lookup_path_column: "path" # only needed for CSVs with raw paths + taxa_glob: "/fs/.../annotations/source=*/*.parquet" # external taxonomy parquet + # If the lookup includes an `hdf5_path` column it will override the auto-resolved path. +``` + +Set `path_to_input` to the root of the converted dataset (where the +`*_metadata.parquet` and `*_images.h5` files live) and `path_to_output_folder` +to the working directory for toolbox artifacts. + +If `lookup_table_path` is provided, it can be either CSV or Parquet and must +contain at least a `uuid` column. Include a column named `hdf5_path` to override +the derived file path, or a `path` column (configurable via +`lookup_path_column`) when providing CSVs that still reference the original +`data_*.parquet` names. Only UUIDs listed in the lookup are converted. + +Run the tool via the standard CLI: + +```bash +CONFIG_PATH=config/tol_hdf5_to_wds_example.yaml \ +tree_of_life_toolbox config/tol_hdf5_to_wds_example.yaml tol_hdf5_to_wds +``` + +Each shard sample contains the JPEG plus ten text prompts: +`scientific_name.txt`, `taxonomic_name.txt`, `sci.txt`, `taxon.txt`, +`taxonTag.txt`, `common_name.txt`, `com.txt`, `sci_com.txt`, +`taxon_com.txt`, and `taxonTag_com.txt`. Files ending in `_com` append +the resolved common name (or `Unknown` when no vernacular exists) to +their prompt text. diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py new file mode 100644 index 0000000..9bbb162 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py @@ -0,0 +1,8 @@ +""" +Tooling for converting HDF5-backed Tree of Life data into WebDataset shards. + +Importing this module registers filter/scheduler/runner implementations with +the toolbox registry so they can be scheduled through the generic CLI. +""" + +from TreeOfLife_toolbox.tol_hdf5_to_wds import classes # noqa: F401 diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py new file mode 100644 index 0000000..4f84c88 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py @@ -0,0 +1,404 @@ +""" +Toolbox integration for converting HDF5-backed TreeOfLife data into WebDatasets. +""" + +from __future__ import annotations + +import glob +import os +from pathlib import Path +from typing import Dict, List + +import h5py +import pandas as pd +import polars as pl +import pyspark.sql.functions as func +from pyspark.sql.window import Window + +from TreeOfLife_toolbox.main.config import Config +from TreeOfLife_toolbox.main.filters import FilterRegister, SparkFilterToolBase +from TreeOfLife_toolbox.main.runners import MPIRunnerTool, RunnerRegister +from TreeOfLife_toolbox.main.schedulers import DefaultScheduler, SchedulerRegister +from TreeOfLife_toolbox.tol_hdf5_to_wds.utils import ( + convert_webp_to_jpeg, + generate_text_files, + init_shard_writer, +) + + +def _sanitize_path_column(column): + return func.regexp_replace(column, "^file:", "") + + +@FilterRegister("tol_hdf5_to_wds") +class TolHDF5ToWDSFilter(SparkFilterToolBase): + """ + Builds shard-level metadata derived from the parquet_to_hdf5 outputs. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hdf5_to_wds" + tool_cfg = self.config.get("tol_hdf5_to_wds", {}) or {} + self.shard_size = int(tool_cfg.get("shard_size", 10000)) + self.shard_limit = int(tool_cfg.get("shard_limit", 0)) + self.metadata_glob = tool_cfg.get("metadata_glob", "**/*_metadata.parquet") + self.include_sources = tool_cfg.get("include_sources") or [] + self.include_servers = tool_cfg.get("include_servers") or [] + self.taxa_glob = tool_cfg.get("taxa_glob") + self.lookup_table_path = tool_cfg.get("lookup_table_path") + self.lookup_path_column = tool_cfg.get("lookup_path_column", "path") + self.shard_metadata_dir = os.path.join( + self.tools_path, self.filter_name, "shard_metadata" + ) + + if self.shard_size <= 0: + raise ValueError("tol_hdf5_to_wds.shard_size must be greater than zero") + + os.makedirs(self.shard_metadata_dir, exist_ok=True) + + def _discover_metadata_files(self) -> List[str]: + dataset_root = self.config.get("path_to_input") + if not dataset_root or not os.path.exists(dataset_root): + raise ValueError("path_to_input must point to converted HDF5 dataset") + + pattern = os.path.join(dataset_root, self.metadata_glob) + files = glob.glob(pattern, recursive=True) + return sorted(files) + + def run(self): + metadata_files = self._discover_metadata_files() + if not metadata_files: + raise ValueError( + f"No metadata parquet files found using pattern: {self.metadata_glob}" + ) + + self.logger.info("Discovered %d metadata files", len(metadata_files)) + + # Read only the UUID column to avoid schema conflicts from unused fields + metadata_df = self.spark.read.parquet(*metadata_files).select("uuid") + metadata_df = metadata_df.withColumn( + "metadata_file", _sanitize_path_column(func.input_file_name()) + ) + metadata_df = metadata_df.withColumn( + "hdf5_path", + func.regexp_replace("metadata_file", "_metadata\\.parquet$", "_images.h5"), + ) + metadata_df = metadata_df.withColumn( + "base_name", + func.regexp_extract("metadata_file", r"([^/]+)_metadata\.parquet$", 1), + ) + + if "source" not in metadata_df.columns: + metadata_df = metadata_df.withColumn( + "source", func.regexp_extract("metadata_file", r"source=([^/]+)", 1) + ) + if "server" not in metadata_df.columns: + metadata_df = metadata_df.withColumn( + "server", func.regexp_extract("metadata_file", r"server=([^/]+)", 1) + ) + + if self.taxa_glob: + taxa_files = glob.glob(self.taxa_glob) + if not taxa_files: + raise ValueError( + f"No taxonomy parquet files found for pattern: {self.taxa_glob}" + ) + taxa_df = self.spark.read.parquet(*taxa_files) + taxa_columns = [ + col_name + for col_name in [ + "uuid", + "scientific_name", + "provided_common_name", + "common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + if col_name in taxa_df.columns + ] + taxa_df = taxa_df.select(*taxa_columns) + metadata_df = metadata_df.join(taxa_df, on="uuid", how="left") + + if self.include_sources: + metadata_df = metadata_df.filter(func.col("source").isin(self.include_sources)) + if self.include_servers: + metadata_df = metadata_df.filter(func.col("server").isin(self.include_servers)) + + if self.lookup_table_path: + lookup_path_lower = self.lookup_table_path.lower() + if lookup_path_lower.endswith(".parquet"): + lookup_df = self.spark.read.parquet(self.lookup_table_path) + else: + lookup_df = ( + self.spark.read.option("header", True).csv(self.lookup_table_path) + ) + + if "uuid" not in lookup_df.columns: + raise ValueError("lookup table must contain a 'uuid' column") + + has_lookup_path = ( + self.lookup_path_column in lookup_df.columns + if self.lookup_path_column + else False + ) + has_hdf5_override = "hdf5_path" in lookup_df.columns + if not has_hdf5_override and "h5_file" in lookup_df.columns: + lookup_df = lookup_df.withColumnRenamed("h5_file", "hdf5_path") + has_hdf5_override = True + + if has_lookup_path: + lookup_df = lookup_df.withColumn( + "lookup_base_name", + func.regexp_replace( + func.regexp_extract( + func.col(self.lookup_path_column), r"([^/]+)$", 1 + ), + r"\.parquet$", + "", + ), + ) + else: + lookup_df = lookup_df.withColumn("lookup_base_name", func.lit(None)) + + if has_hdf5_override: + lookup_df = lookup_df.withColumnRenamed("hdf5_path", "lookup_hdf5_path") + + select_cols = ["uuid", "lookup_base_name"] + if has_hdf5_override: + select_cols.append("lookup_hdf5_path") + + lookup_df = lookup_df.select(*select_cols).dropDuplicates(["uuid"]) + metadata_df = metadata_df.join(lookup_df, on="uuid", how="inner") + + if has_lookup_path: + metadata_df = metadata_df.filter( + func.col("lookup_base_name").isNull() + | (func.col("lookup_base_name") == func.col("base_name")) + ) + + if has_hdf5_override: + if "hdf5_path" in metadata_df.columns: + metadata_df = metadata_df.drop("hdf5_path") + metadata_df = metadata_df.withColumnRenamed( + "lookup_hdf5_path", "hdf5_path" + ) + + metadata_df = metadata_df.drop("lookup_base_name") + + # Drop duplicate UUIDs to avoid double writes + metadata_df = metadata_df.dropDuplicates(["uuid"]) + + required_columns = [ + "uuid", + "hdf5_path", + "source", + "server", + "scientific_name", + "common_name", + "provided_common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + + string_columns = set(required_columns) + + for col_name in required_columns: + if col_name not in metadata_df.columns: + default_value = func.lit(None) + if col_name in string_columns: + default_value = default_value.cast("string") + metadata_df = metadata_df.withColumn(col_name, default_value) + + ordering_window = Window.orderBy("metadata_file", "uuid") + metadata_df = ( + metadata_df.withColumn( + "row_number", func.row_number().over(ordering_window) + ) + .withColumn( + "shard_id", + func.floor((func.col("row_number") - 1) / self.shard_size).cast("int"), + ) + .drop("row_number") + ) + + if self.shard_limit > 0: + metadata_df = metadata_df.filter(func.col("shard_id") < self.shard_limit) + + if metadata_df.count() == 0: + raise ValueError("No metadata rows left after filtering/shard limiting") + + selected_columns = required_columns + ["shard_id"] + sharded_df = metadata_df.select(*selected_columns) + + ( + sharded_df.repartition("shard_id") + .write.partitionBy("shard_id") + .mode("overwrite") + .parquet(self.shard_metadata_dir) + ) + self.logger.info("Wrote shard metadata to %s", self.shard_metadata_dir) + + schedule_df = ( + sharded_df.select("shard_id") + .dropDuplicates() + .withColumn( + "metadata_path", + func.concat( + func.lit(self.shard_metadata_dir), + func.lit("/shard_id="), + func.col("shard_id"), + ), + ) + .withColumn("server_name", func.lit("tol_hdf5_to_wds")) + .withColumn("partition_id", func.col("shard_id")) + ) + + self.save_filter(schedule_df) + total_shards = schedule_df.count() + self.logger.info("Prepared %d shards", total_shards) + + +@SchedulerRegister("tol_hdf5_to_wds") +class TolHDF5ToWDSScheduler(DefaultScheduler): + """ + Scheduler that treats each metadata shard as an independent work unit. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hdf5_to_wds" + # Scheduling keys only. The data columns (metadata_path, shard_id) ride along in the filter table and are picked up by the runner's data_scheme; including them here duplicates them into schedule.csv and collides on the runner's merge. + self.scheme = ["server_name", "partition_id"] + + +@RunnerRegister("tol_hdf5_to_wds") +class TolHDF5ToWDSRunner(MPIRunnerTool): + """ + Converts shard metadata into WebDataset tar archives. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hdf5_to_wds" + self.data_scheme = ["metadata_path", "shard_id", "server_name", "partition_id"] + self.verification_scheme = ["server_name", "partition_id"] + + params = self.config.get("tol_hdf5_to_wds", {}) or {} + self.tar_output_root = params.get("tar_output_root") or os.path.join( + self.tools_path, self.filter_name, "tar_dataset" + ) + self.resize_size = int(params.get("resize_size", 224)) + self.shard_prefix = params.get("shard_prefix", "shard") + self.total_time = int(params.get("runner_timeout_seconds", 3600)) + + os.makedirs(self.tar_output_root, exist_ok=True) + + def _load_shard_metadata(self, metadata_path: str) -> pl.DataFrame: + metadata_dir = Path(metadata_path) + if not metadata_dir.exists(): + raise FileNotFoundError(f"Shard metadata directory missing: {metadata_path}") + + files = sorted(metadata_dir.glob("*.parquet")) + if not files: + raise FileNotFoundError(f"No Parquet files found under {metadata_path}") + + return pl.read_parquet([str(path) for path in files]) + + def _iter_records(self, shard_df: pl.DataFrame): + for row in shard_df.iter_rows(named=True): + yield row + + def _open_hdf5(self, path: str, cache: Dict[str, h5py.File]) -> h5py.File: + if path not in cache: + if not os.path.exists(path): + raise FileNotFoundError(f"HDF5 file missing: {path}") + cache[path] = h5py.File(path, "r") + return cache[path] + + def apply_filter( + self, filtering_df: pd.DataFrame, server_name: str, partition_id: int + ) -> int: + shard_id = int(filtering_df.iloc[0]["shard_id"]) + metadata_path = filtering_df.iloc[0]["metadata_path"] + return self.process_shard(shard_id, metadata_path) + + def process_shard(self, shard_id: int, metadata_path: str) -> int: + self.is_enough_time() + shard_df = self._load_shard_metadata(metadata_path) + if shard_df.height == 0: + self.logger.warning("Shard %s is empty, skipping", shard_id) + return 0 + + tar_writer, tar_path = init_shard_writer( + self.tar_output_root, shard_id, self.shard_prefix + ) + hdf5_cache: Dict[str, h5py.File] = {} + written = 0 + failed = 0 + + try: + for row in self._iter_records(shard_df): + uuid = row.get("uuid") + hdf5_path = row.get("hdf5_path") + if not uuid or not hdf5_path: + failed += 1 + continue + + try: + h5_file = self._open_hdf5(hdf5_path, hdf5_cache) + images_group = h5_file.get("images") + if images_group is None or uuid not in images_group: + failed += 1 + continue + + dataset = images_group[uuid][:] + webp_bytes = dataset.tobytes() + jpeg_bytes = convert_webp_to_jpeg(webp_bytes, self.resize_size) + + taxon_dict = { + "scientific_name": row.get("scientific_name"), + "common_name": row.get("common_name"), + "kingdom": row.get("kingdom"), + "phylum": row.get("phylum"), + "class": row.get("class"), + "order": row.get("order"), + "family": row.get("family"), + "genus": row.get("genus"), + "species": row.get("species"), + } + + sample = {"__key__": uuid, "jpg": jpeg_bytes} + for ext, content in generate_text_files(taxon_dict).items(): + sample[ext] = content.encode("utf-8") + + tar_writer.write(sample) + written += 1 + except Exception as exc: # noqa: BLE001 + failed += 1 + self.logger.error( + "Failed to process uuid %s in shard %s: %s", uuid, shard_id, exc + ) + finally: + tar_writer.close() + for handle in hdf5_cache.values(): + handle.close() + + self.logger.info( + "Shard %s written to %s (%s records, %s failures)", + shard_id, + tar_path, + written, + failed, + ) + return written diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py new file mode 100644 index 0000000..14c04b8 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py @@ -0,0 +1,171 @@ +""" +Utility helpers for converting HDF5-backed metadata rows into WebDataset shards. +""" + +from __future__ import annotations + +import io +import logging +import os +from typing import Dict, Tuple + +import webdataset as wds +from PIL import Image + + +def convert_webp_to_jpeg(webp_bytes: bytes, resize_size: int = 0) -> bytes: + """ + Decode WebP bytes, optionally resize to a square, and re-encode as JPEG. + """ + image = Image.open(io.BytesIO(webp_bytes)).convert("RGB") + if resize_size and resize_size > 0: + resample = getattr(Image, "Resampling", Image).LANCZOS # Pillow>=9 compat + image = image.resize((resize_size, resize_size), resample) + + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=95) + return buffer.getvalue() + + +def determine_most_specific_known_rank(taxon_dict: Dict[str, str | None]) -> str | None: + """ + Return the most specific rank (closest to species) whose value is not unknown. + """ + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + for rank in reversed(taxonomic_ranks): + value = taxon_dict.get(rank) + if value and value.strip().lower() != "unknown": + return rank + return None + + +def _clean_species_name(species_value: str | None) -> str | None: + if not species_value: + return None + parts = species_value.split() + if len(parts) > 1: + return parts[-1] + return species_value + + +def create_taxon_tag_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Compose a descriptive sentence enumerating the known taxonomy ranks. + """ + if not most_specific_rank: + return "a photo of unknown taxonomy." + + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + text_parts = [] + for rank in taxonomic_ranks: + value = taxon_dict.get(rank) + if not value or value.lower() == "unknown": + if rank == most_specific_rank: + break + continue + + if rank == "species": + cleaned = _clean_species_name(value) + if cleaned: + text_parts.append(f"{rank} {cleaned}") + else: + text_parts.append(f"{rank} {value}") + + if rank == most_specific_rank: + break + + return "a photo of " + " ".join(text_parts) + "." if text_parts else "a photo of unknown taxonomy." + + +def create_taxonomic_name_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Build a space-delimited taxonomy string up to the most specific known rank. + """ + if not most_specific_rank: + return "" + + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + parts = [] + for rank in taxonomic_ranks: + value = taxon_dict.get(rank) + if not value or value.lower() == "unknown": + if rank == most_specific_rank: + break + continue + + if rank == "species": + cleaned = _clean_species_name(value) + if cleaned: + parts.append(cleaned) + else: + parts.append(value) + + if rank == most_specific_rank: + break + + return " ".join(parts) + + +def create_scientific_name_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Build the scientific name from the most specific known rank. + """ + if not most_specific_rank: + return "Unknown" + + value = taxon_dict.get(most_specific_rank) + if not value or value.lower() == "unknown": + return "Unknown" + return value.rstrip(".") + + +def _get_common_name(taxon_dict: Dict[str, str | None]) -> str: + value = taxon_dict.get("common_name") + if value and value.strip(): + return value.strip().rstrip(".") + scientific = taxon_dict.get("scientific_name") + if scientific and scientific.strip(): + return scientific.strip().rstrip(".") + return "Unknown" + + +def generate_text_files(taxon_dict: Dict[str, str | None]) -> Dict[str, str]: + """ + Generate the five auxiliary text files used by downstream training pipelines. + """ + most_specific_rank = determine_most_specific_known_rank(taxon_dict) + taxonomic_name = create_taxonomic_name_text(taxon_dict, most_specific_rank) + scientific_name = create_scientific_name_text(taxon_dict, most_specific_rank) + taxon_for_common = dict(taxon_dict) + taxon_for_common["scientific_name"] = scientific_name + common_name = _get_common_name(taxon_for_common) + taxon_tag_text = create_taxon_tag_text(taxon_dict, most_specific_rank) + + files = { + "taxonomic_name.txt": taxonomic_name, + "scientific_name.txt": scientific_name, + "sci.txt": f"a photo of {scientific_name}.", + "taxon.txt": f"a photo of {taxonomic_name}.", + "taxonTag.txt": taxon_tag_text, + "common_name.txt": common_name, + "com.txt": f"a photo of {common_name}.", + "sci_com.txt": f"a photo of {scientific_name} with common name {common_name}.", + "taxon_com.txt": f"a photo of {taxonomic_name} with common name {common_name}.", + "taxonTag_com.txt": f"{taxon_tag_text.rstrip('.')} with common name {common_name}.", + } + + return files + + +def init_shard_writer(output_dir: str, shard_id: int, prefix: str = "shard") -> Tuple[wds.TarWriter, str]: + """ + Create (or truncate) the shard tarball and return the TarWriter plus path. + """ + os.makedirs(output_dir, exist_ok=True) + shard_name = f"{prefix}-{shard_id:05d}.tar" + shard_path = os.path.join(output_dir, shard_name) + if os.path.exists(shard_path): + os.remove(shard_path) + + logging.info("Creating shard %s at %s", shard_id, shard_path) + return wds.TarWriter(shard_path), shard_path diff --git a/tol_hdf5_to_wds.py b/tol_hdf5_to_wds.py new file mode 100644 index 0000000..0308f33 --- /dev/null +++ b/tol_hdf5_to_wds.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python +""" +Single-node converter from the hybrid Tree of Life HDF5/metadata dataset to WebDataset shards. + +Example: + python tol_hdf5_to_wds.py \ + --input-root /fs/ess/PAS2136/TreeOfLife/data \ + --lookup /fs/scratch/PAS2136/TreeOfLife_test-wds/lookup-tables/all/hdf5_lookup.parquet \ + --taxa-glob "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/*/*.parquet" \ + --output-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/manual \ + --metadata-glob "**/*_metadata.parquet" +""" + +from __future__ import annotations + +import argparse +import contextlib +import glob +import logging +import os +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + +import h5py +import polars as pl +import webdataset as wds + +try: + from TreeOfLife_toolbox.tol_hdf5_to_wds.utils import ( + convert_webp_to_jpeg, + generate_text_files, + ) +except ImportError as exc: # pragma: no cover - script usage only + raise SystemExit( + "Unable to import TreeOfLife_toolbox.tol_hdf5_to_wds.utils. " + "Please install the toolbox package (pip install -e .) before running this script." + ) from exc + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Create WebDataset shards from HDF5 dataset") + parser.add_argument("--input-root", required=True, help="Root folder containing *_metadata.parquet + *_images.h5") + parser.add_argument( + "--metadata-glob", + default="**/*_metadata.parquet", + help="Glob pattern under --input-root to find metadata files (default: %(default)s)", + ) + parser.add_argument( + "--lookup", + required=True, + help="Lookup table (CSV or Parquet) listing UUIDs to include. " + "Must contain a 'uuid' column; optional 'hdf5_path' overrides file locations.", + ) + parser.add_argument( + "--taxa-glob", + help="Optional glob pointing to taxonomy parquet files (e.g. /path/source=*/*.parquet).", + ) + parser.add_argument( + "--output-dir", + required=True, + help="Directory where shard-XXXXX.tar files will be written.", + ) + parser.add_argument( + "--shard-size", + type=int, + default=10000, + help="Number of samples per shard (default: %(default)s)", + ) + parser.add_argument( + "--resize", + type=int, + default=224, + help="Output JPEG square size in pixels (0 keeps original WebP dimensions).", + ) + parser.add_argument( + "--shard-prefix", + default="shard", + help="Prefix for tar files (default: %(default)s)", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Logging verbosity (default: %(default)s)", + ) + return parser.parse_args() + + +def setup_logger(level: str) -> logging.Logger: + logging.basicConfig( + level=getattr(logging, level.upper()), + format="%(asctime)s - %(levelname)s - %(message)s", + ) + return logging.getLogger("tol_hdf5_to_wds") + + +def discover_metadata_files(root: str, pattern: str) -> List[str]: + search_pattern = os.path.join(root, pattern) + files = sorted(glob.glob(search_pattern, recursive=True)) + if not files: + raise FileNotFoundError(f"No metadata parquet files found using pattern: {search_pattern}") + return files + + +def build_metadata_lazy(metadata_files: List[str]) -> pl.LazyFrame: + # Include file paths to derive HDF5 locations. + metadata_lazy = pl.scan_parquet(metadata_files, include_file_paths="metadata_file") + metadata_lazy = metadata_lazy.with_columns( + pl.col("metadata_file") + .str.replace(r"^file:", "", literal=True) + .alias("metadata_file_clean") + ) + metadata_lazy = metadata_lazy.with_columns( + pl.col("metadata_file_clean") + .str.replace("_metadata.parquet$", "_images.h5") + .alias("hdf5_path"), + pl.col("metadata_file_clean") + .str.extract(r"([^/]+)_metadata\.parquet$", group_index=1) + .alias("base_name"), + ) + schema_names = metadata_lazy.collect_schema().names() + drop_cols = [c for c in ("source", "server") if c in schema_names] + if drop_cols: + metadata_lazy = metadata_lazy.drop(drop_cols) + metadata_lazy = metadata_lazy.with_columns( + pl.col("metadata_file_clean").str.extract(r"source=([^/]+)", group_index=1).alias("source"), + pl.col("metadata_file_clean").str.extract(r"server=([^/]+)", group_index=1).alias("server"), + ) + desired_cols = [ + "uuid", + "hdf5_path", + "base_name", + "source", + "server", + "scientific_name", + "provided_common_name", + "common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + cols_present = [c for c in desired_cols if c in metadata_lazy.collect_schema().names()] + metadata_lazy = metadata_lazy.select(cols_present + ["metadata_file_clean"]) + return metadata_lazy + + +def load_lookup_table(path: str) -> pl.LazyFrame: + ext = Path(path).suffix.lower() + if ext == ".parquet": + lookup = pl.scan_parquet(path) + else: + lookup = pl.scan_csv(path) + schema = lookup.collect_schema() + schema_names = schema.names() + if "uuid" not in schema: + raise ValueError(f"Lookup table {path} must contain a 'uuid' column") + keep_cols = ["uuid"] + if "hdf5_path" in schema_names: + keep_cols.append("hdf5_path") + if "path" in schema_names: + keep_cols.append("path") + return lookup.select(keep_cols) + + +def load_taxa_lazy(glob_pattern: str) -> pl.LazyFrame: + taxa_files = glob.glob(glob_pattern) + if not taxa_files: + raise FileNotFoundError(f"No taxonomy files match {glob_pattern}") + taxa_lazy = pl.scan_parquet(taxa_files, extra_columns="ignore") + schema_names = taxa_lazy.collect_schema().names() + cols = [col for col in [ + "uuid", + "common_name", + "provided_common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] if col in schema_names] + return taxa_lazy.select(cols) + + +def merge_metadata( + metadata_lazy: pl.LazyFrame, + lookup_lazy: pl.LazyFrame, + taxa_lazy: pl.LazyFrame | None, + logger: logging.Logger, +) -> pl.DataFrame: + joined = metadata_lazy.join(lookup_lazy, on="uuid", how="inner") + # Ensure column types align before sorting/collecting + schema_names = joined.collect_schema().names() + cast_exprs = [] + if "common_name" in schema_names: + cast_exprs.append(pl.col("common_name").cast(pl.Utf8, strict=False)) + if "provided_common_name" in schema_names: + cast_exprs.append(pl.col("provided_common_name").cast(pl.Utf8, strict=False)) + if "source" in schema_names: + cast_exprs.append(pl.col("source").cast(pl.Utf8, strict=False)) + if "server" in schema_names: + cast_exprs.append(pl.col("server").cast(pl.Utf8, strict=False)) + if cast_exprs: + joined = joined.with_columns(cast_exprs) + joined_schema = joined.collect_schema().names() + if "hdf5_path_right" in joined_schema: + joined = joined.with_columns( + pl.when(pl.col("hdf5_path_right").is_not_null()) + .then(pl.col("hdf5_path_right")) + .otherwise(pl.col("hdf5_path")) + .alias("hdf5_path") + ).drop("hdf5_path_right") + if taxa_lazy is not None: + joined = joined.join(taxa_lazy, on="uuid", how="left") + joined = joined.unique(subset=["uuid"]) + result = joined.sort(["hdf5_path", "uuid"]).collect(streaming=True) + logger.info("Collected %s metadata rows", result.height) + return result + + +def chunk_records(df: pl.DataFrame, shard_size: int) -> Iterable[Tuple[int, List[dict]]]: + buffer: List[dict] = [] + shard_id = 0 + for row in df.iter_rows(named=True): + buffer.append(row) + if len(buffer) >= shard_size: + yield shard_id, buffer + buffer = [] + shard_id += 1 + if buffer: + yield shard_id, buffer + + +def process_shard( + shard_id: int, + records: List[dict], + output_dir: Path, + resize: int, + shard_prefix: str, + logger: logging.Logger, +) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + shard_path = output_dir / f"{shard_prefix}-{shard_id:05d}.tar" + writer = wds.TarWriter(str(shard_path)) + hdf5_cache: Dict[str, h5py.File] = {} + written = failed = 0 + try: + for record in records: + uuid = record.get("uuid") + hdf5_path = record.get("hdf5_path") + if not uuid or not hdf5_path: + failed += 1 + continue + try: + h5_file = _open_hdf5(hdf5_path, hdf5_cache) + dataset = h5_file["images"][uuid][:] + webp_bytes = dataset.tobytes() + jpeg_bytes = convert_webp_to_jpeg(webp_bytes, resize) + + taxon_dict = { + "scientific_name": record.get("scientific_name"), + "common_name": record.get("common_name"), + "kingdom": record.get("kingdom"), + "phylum": record.get("phylum"), + "class": record.get("class"), + "order": record.get("order"), + "family": record.get("family"), + "genus": record.get("genus"), + "species": record.get("species"), + } + sample = {"__key__": uuid, "jpg": jpeg_bytes} + for ext, content in generate_text_files(taxon_dict).items(): + sample[ext] = content.encode("utf-8") + writer.write(sample) + written += 1 + except Exception as exc: # pragma: no cover + failed += 1 + logger.error("Shard %s: failed uuid %s (%s)", shard_id, uuid, exc) + logger.info("Shard %s written with %s samples (%s failures)", shard_id, written, failed) + finally: + writer.close() + for handle in hdf5_cache.values(): + with contextlib.suppress(Exception): + handle.close() + + +def _open_hdf5(path: str, cache: Dict[str, h5py.File]) -> h5py.File: + if path not in cache: + cache[path] = h5py.File(path, "r") + return cache[path] + + +def main(): + args = parse_args() + logger = setup_logger(args.log_level) + + metadata_files = discover_metadata_files(args.input_root, args.metadata_glob) + logger.info("Found %s metadata parquet files", len(metadata_files)) + + metadata_lazy = build_metadata_lazy(metadata_files) + lookup_lazy = load_lookup_table(args.lookup) + taxa_lazy = load_taxa_lazy(args.taxa_glob) if args.taxa_glob else None + + metadata = merge_metadata(metadata_lazy, lookup_lazy, taxa_lazy, logger) + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + last_shard = -1 + for shard_id, records in chunk_records(metadata, args.shard_size): + last_shard = shard_id + process_shard(shard_id, records, output_dir, args.resize, args.shard_prefix, logger) + + total_shards = (last_shard + 1) if metadata.height > 0 else 0 + logger.info("Completed WebDataset creation: %s shards written", total_shards) + + +if __name__ == "__main__": + main() diff --git a/validate_tars.py b/validate_tars.py new file mode 100644 index 0000000..9d2e262 --- /dev/null +++ b/validate_tars.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +import os +import json +import argparse +from pathlib import Path +from concurrent.futures import ProcessPoolExecutor +from time import time +import tarfile +from collections import defaultdict + +# The 10 required sidecar suffixes +REQUIRED_SIDECARS = [ + 'com.txt', + 'common_name.txt', + 'sci.txt', + 'sci_com.txt', + 'scientific_name.txt', + 'taxon.txt', + 'taxonTag.txt', + 'taxonTag_com.txt', + 'taxon_com.txt', + 'taxonomic_name.txt', +] + +def validate_tar(tar_path: str) -> tuple[str, dict]: + """ + Validate that each JPG has all 10 required text sidecars AND they contain text. + Returns (tar_name, validation_result) + """ + tar_name = os.path.basename(tar_path) + + # Track all files by UUID + files_by_uuid = defaultdict(dict) # uuid -> {extension: has_content} + jpg_uuids = set() + + try: + with tarfile.open(tar_path, 'r|') as tar: + for member in tar: + if not member.isfile(): + continue + + name = member.name + basename = os.path.basename(name) + + if '.' not in basename: + continue + + # Split on first dot to get UUID + parts = basename.split('.', 1) + if len(parts) != 2: + continue + + uuid = parts[0] + extension = parts[1] + + # Track JPGs + if extension.lower() == 'jpg': + jpg_uuids.add(uuid) + files_by_uuid[uuid][extension] = True # JPGs don't need content check + + # For text files, check if they have content + elif extension in REQUIRED_SIDECARS: + has_content = False + try: + f = tar.extractfile(member) + if f: + content = f.read() + # Check if file has any non-whitespace content + has_content = len(content.strip()) > 0 + f.close() + except Exception: + has_content = False + + files_by_uuid[uuid][extension] = has_content + + # Now validate: check each JPG has all sidecars with content + issues = [] + for uuid in sorted(jpg_uuids): + extensions = files_by_uuid[uuid] + missing = [] + empty = [] + + for required in REQUIRED_SIDECARS: + if required not in extensions: + missing.append(required) + elif not extensions[required]: + empty.append(required) + + if missing or empty: + issue = {'uuid': uuid} + if missing: + issue['missing'] = missing + if empty: + issue['empty'] = empty + issues.append(issue) + + return (tar_name, { + 'valid': len(issues) == 0, + 'jpg_count': len(jpg_uuids), + 'issues': issues + }) + + except Exception as e: + return (tar_name, { + 'valid': False, + 'error': str(e), + 'jpg_count': 0, + 'issues': [] + }) + +def process_chunk(tar_paths: list[str]) -> list[tuple[str, dict]]: + """Process multiple tars in one worker.""" + return [validate_tar(p) for p in tar_paths] + +def chunks(lst, n): + """Yield successive n-sized chunks.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + +def validate_all_tars(source_dir: str, workers: int, chunk_size: int = 1) -> None: + src = Path(source_dir) + tar_files = sorted([str(p) for p in src.glob("*.tar")]) + n = len(tar_files) + + if n == 0: + raise SystemExit(f"No .tar files in {src}") + + print(f"Validating {n} shards with {workers} workers...", flush=True) + print("Checking: file existence + non-empty content", flush=True) + + results = {} + all_issues = [] + total_jpgs = 0 + valid_tars = 0 + total_missing = 0 + total_empty = 0 + + t0 = time() + processed = 0 + + # Split work into chunks + tar_chunks = list(chunks(tar_files, chunk_size)) + + with ProcessPoolExecutor(max_workers=workers) as ex: + for result_chunk in ex.map(process_chunk, tar_chunks): + for tar_name, validation in result_chunk: + results[tar_name] = validation + + jpg_count = validation.get('jpg_count', 0) + total_jpgs += jpg_count + + if validation['valid']: + valid_tars += 1 + else: + # Record issues for this tar + for issue in validation.get('issues', []): + issue_record = { + 'tar': tar_name, + 'uuid': issue['uuid'] + } + if 'missing' in issue: + issue_record['missing'] = issue['missing'] + total_missing += len(issue['missing']) + if 'empty' in issue: + issue_record['empty'] = issue['empty'] + total_empty += len(issue['empty']) + all_issues.append(issue_record) + + # Also note if there was an error + if 'error' in validation: + all_issues.append({ + 'tar': tar_name, + 'error': validation['error'] + }) + + processed += 1 + if processed % 500 == 0 or processed == n: + elapsed = time() - t0 + rate = processed / elapsed + eta = (n - processed) / rate if rate > 0 else 0 + print(f"[{processed}/{n}] {rate:.1f}/s | Valid: {valid_tars}/{processed} | Issues: {len(all_issues)} (missing: {total_missing}, empty: {total_empty}) | ETA: {eta:.0f}s", flush=True) + + # Write detailed results + out_path = src / "tar_validation_results.json" + summary = { + 'summary': { + 'total_tars': n, + 'valid_tars': valid_tars, + 'invalid_tars': n - valid_tars, + 'total_jpgs': total_jpgs, + 'total_issues': len(all_issues), + 'total_missing_files': total_missing, + 'total_empty_files': total_empty + }, + 'per_tar': results + } + + with open(out_path, "w") as f: + json.dump(summary, f, indent=2) + + # Write issue log (more human-readable) + if all_issues: + issues_path = src / "tar_validation_issues.log" + with open(issues_path, "w") as f: + f.write(f"Validation Issues - {len(all_issues)} total\n") + f.write(f"Missing files: {total_missing}, Empty files: {total_empty}\n") + f.write("=" * 80 + "\n\n") + + # Group by tar file + by_tar = defaultdict(list) + for issue in all_issues: + by_tar[issue['tar']].append(issue) + + for tar_name in sorted(by_tar.keys()): + issues = by_tar[tar_name] + f.write(f"TAR: {tar_name} ({len(issues)} issues)\n") + f.write("-" * 80 + "\n") + + for issue in issues: + if 'error' in issue: + f.write(f" ERROR: {issue['error']}\n") + else: + f.write(f" UUID: {issue['uuid']}\n") + if 'missing' in issue: + f.write(f" Missing: {', '.join(issue['missing'])}\n") + if 'empty' in issue: + f.write(f" Empty: {', '.join(issue['empty'])}\n") + + f.write("\n") + + print(f"\n⚠️ Issues found: {len(all_issues)}") + print(f" Missing files: {total_missing}") + print(f" Empty files: {total_empty}") + print(f" Wrote details to: {issues_path}") + + elapsed = time() - t0 + print(f"\n✓ Done in {elapsed:.1f}s ({n/elapsed:.1f} shards/s)") + print(f"✓ Valid tars: {valid_tars}/{n} ({100*valid_tars/n:.1f}%)") + print(f"✓ Total JPGs: {total_jpgs:,}") + print(f"✓ Wrote summary to: {out_path}") + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Validate JPG sidecars in tar shards (checks content)") + ap.add_argument("--source_dir", required=True, help="Directory containing .tar files") + ap.add_argument("--workers", type=int, default=96, help="Number of parallel workers") + ap.add_argument("--chunk-size", type=int, default=1, help="Tars per worker batch") + args = ap.parse_args() + + validate_all_tars(args.source_dir, args.workers, args.chunk_size) + diff --git a/wds_taxoncom_audit.py b/wds_taxoncom_audit.py new file mode 100644 index 0000000..2a79eb7 --- /dev/null +++ b/wds_taxoncom_audit.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +""" +wds_taxoncom_audit.py + +One-node audit / training-pipeline analog to pinpoint where samples are missing `taxon_com.txt`. + +Logs: + - wds_audit_logs/missing_taxon_com.workerXYZ.jsonl (samples missing taxon_com.txt after grouping) + - wds_audit_logs/bad_records.workerXYZ.jsonl (non-empty dicts missing fname/data) + - wds_audit_logs/stats.workerXYZ.json (summary per worker) + +Usage: + python wds_taxoncom_audit.py \ + --shards "/path/to/shards/shard-{00000..99999}.tar" \ + --workers 96 \ + --out-dir wds_audit_logs +""" + +import argparse +import json +import logging +import os +from datetime import datetime + +import webdataset as wds +from webdataset.tariterators import base_plus_ext, url_opener, tar_file_expander, valid_sample + + +# ---------------------------- +# Logging / utilities +# ---------------------------- + +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + + +def now_ts(): + return datetime.now().isoformat(timespec="seconds") + + +def worker_id(): + """ + Best-effort worker id detection. + + Note: torch dataloader workers don't reliably export a standard env var. + We'll use these if present, else -1. + """ + for k in ("WDS_WORKER", "WORKER", "LOCAL_RANK", "RANK"): + v = os.environ.get(k) + if v is not None and str(v).lstrip("-").isdigit(): + return int(v) + return -1 + + +def append_jsonl(path, rec): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "a") as f: + f.write(json.dumps(rec) + "\n") + + +# ---------------------------- +# WDS "nothrow" tar parsing +# ---------------------------- + +def log_and_continue(exn): + logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.") + return True + + +# These globals get set in main() so that nested functions can log without threading args everywhere. +OUT_DIR = "wds_audit_logs" + + +def log_bad_record(filesample, reason="missing_fname_or_data"): + wid = worker_id() + path = os.path.join(OUT_DIR, f"bad_records.worker{wid:03d}.jsonl") + rec = { + "ts": now_ts(), + "reason": reason, + "url": filesample.get("__url__", filesample.get("url")), + "keys": sorted(list(filesample.keys())), + } + append_jsonl(path, rec) + + +def log_missing_taxon_com(sample, reason="missing taxon_com.txt (pre-rename)"): + wid = worker_id() + path = os.path.join(OUT_DIR, f"missing_taxon_com.worker{wid:03d}.jsonl") + rec = { + "ts": now_ts(), + "reason": reason, + "url": sample.get("__url__", sample.get("url")), + "key": sample.get("__key__"), + "present_suffixes": sorted([k for k in sample.keys() if not k.startswith("__")]), + } + append_jsonl(path, rec) + + +def group_by_keys_nothrow(data, keys=base_plus_ext, lcase=True, suffixes=None, handler=None): + """ + Robust grouping that tolerates: + - empty dicts (often sentinel-ish): skip silently + - malformed dicts missing fname/data: log (if non-empty) and skip + + This is meant to be a faithful-ish version of the modified training pipeline, + while remaining diagnostic. + """ + current_sample = None + + for filesample in data: + if not isinstance(filesample, dict): + logging.warning(f"Skipping non-dict filesample: {type(filesample)}") + continue + + # IMPORTANT: empty dicts can appear due to iterator plumbing; skip quietly. + if not filesample: + continue + + # Defensive: sometimes expander yields dicts without fname/data + if "fname" not in filesample or "data" not in filesample: + # Log only non-empty dicts (otherwise it's useless noise). + log_bad_record(filesample, reason="malformed_tar_record_missing_fname_or_data") + continue + + fname, value = filesample["fname"], filesample["data"] + prefix, suffix = keys(fname) + if prefix is None: + continue + if lcase: + suffix = suffix.lower() + + # Same collision handling rationale as training data.py + if current_sample is None or prefix != current_sample["__key__"] or suffix in current_sample: + if valid_sample(current_sample): + yield current_sample + current_sample = dict(__key__=prefix, __url__=filesample.get("__url__")) + + if suffixes is None or suffix in suffixes: + current_sample[suffix] = value + + if valid_sample(current_sample): + yield current_sample + + +def tarfile_to_samples_nothrow(src, handler=log_and_continue): + streams = url_opener(src, handler=handler) + files = tar_file_expander(streams, handler=handler) + return group_by_keys_nothrow(files, handler=handler) + + +# ---------------------------- +# Filters for audit +# ---------------------------- + +def has_image(sample): + return ("png" in sample or "jpg" in sample or "jpeg" in sample or "webp" in sample) + + +def has_any_txt(sample): + # training's filter_no_caption_or_no_image checks "any('txt' in key for key in sample)" + return any("txt" in k for k in sample.keys()) + + +def require_taxon_com_txt(sample): + """ + Audit hook: record and drop samples that lack taxon_com.txt at the grouped-sample stage. + + This is the strongest check for grouping/boundary/corruption-induced misses. + """ + if "taxon_com.txt" not in sample: + log_missing_taxon_com(sample, reason="missing taxon_com.txt (pre-rename)") + return False + return True + + +# ---------------------------- +# Main audit +# ---------------------------- + +def audit(shards_pattern: str, workers: int, log_every: int, limit_kept: int): + shards = list(wds.shardlists.expand_urls(shards_pattern)) + logging.info(f"Expanded to {len(shards)} shard URLs") + + # Build a training-ish pipeline (single node; split_by_node not needed here) + dataset = wds.DataPipeline( + wds.SimpleShardList(shards), + wds.split_by_worker, + tarfile_to_samples_nothrow, + wds.select(lambda s: has_any_txt(s) and has_image(s)), + wds.select(lambda s: require_taxon_com_txt(s)), + # We intentionally avoid decode/tokenizer here: this audit is about tar member presence. + ) + + loader = wds.WebLoader( + dataset, + batch_size=None, + shuffle=False, + num_workers=workers, + persistent_workers=True, + ) + + kept = 0 + for sample in loader: + kept += 1 + if log_every and kept % log_every == 0: + logging.info(f"Kept {kept} samples (after filters)") + + if limit_kept and kept >= limit_kept: + break + + logging.info(f"Done. Kept {kept} samples (after filters). Logs in {OUT_DIR}/") + + +def main(): + global OUT_DIR + + ap = argparse.ArgumentParser() + ap.add_argument("--shards", required=True, help='Braceexpand pattern like ".../shard-{00000..99999}.tar"') + ap.add_argument("--workers", type=int, default=96) + ap.add_argument("--out-dir", default="wds_audit_logs") + ap.add_argument("--log-every", type=int, default=200000) + ap.add_argument("--limit", type=int, default=0, help="Stop after this many kept samples (0 = no limit)") + args = ap.parse_args() + + OUT_DIR = args.out_dir + os.makedirs(OUT_DIR, exist_ok=True) + + audit(args.shards, args.workers, args.log_every, args.limit) + + +if __name__ == "__main__": + main() From b6ce6f32405070e969b24c85ab921ae8aa818a87 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Tue, 30 Jun 2026 14:55:58 -0400 Subject: [PATCH 02/10] Document WebDataset creation from the hybrid HDF5+Parquet format Co-Authored-By: Claude Opus 4.8 --- docs/README.md | 6 ++ docs/hybrid_to_webdataset_README.md | 104 ++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 docs/hybrid_to_webdataset_README.md diff --git a/docs/README.md b/docs/README.md index f066e3c..4923c34 100644 --- a/docs/README.md +++ b/docs/README.md @@ -24,3 +24,9 @@ The dataset is divided into the following components, each with its own download 4. Monitor progress using the provided utilities. Each download component has its own README with detailed instructions. Click on the links above to access specific guidance for accessing images from each data provider. + +# Create webdataset + +Instructions for packaging the dataset into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards for model training, by storage format: + +- [Hybrid format (HDF5 + Parquet)](hybrid_to_webdataset_README.md) - create WebDataset shards from the hybrid HDF5-image + Parquet-metadata storage format. diff --git a/docs/hybrid_to_webdataset_README.md b/docs/hybrid_to_webdataset_README.md new file mode 100644 index 0000000..7df8182 --- /dev/null +++ b/docs/hybrid_to_webdataset_README.md @@ -0,0 +1,104 @@ +# Create a WebDataset from the hybrid HDF5 + Parquet format + +This guide covers packaging the Tree of Life dataset, stored in the hybrid +format of WebP images in HDF5 (`*_images.h5`) alongside Parquet metadata +(`*_metadata.parquet`), into [WebDataset](https://github.com/webdataset/webdataset) +`.tar` shards for model training. It is implemented by the +[`tol_hdf5_to_wds`](../src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md) tool. + +## What it produces + +One `shard-XXXXX.tar` per `shard_size` samples. Each sample is keyed by `uuid` +and contains a JPEG plus ten text-prompt sidecars used for BioCLIP-style +training: + +| | | +|---|---| +| `scientific_name.txt`, `common_name.txt`, `taxonomic_name.txt` | raw label values | +| `sci.txt`, `com.txt`, `taxon.txt`, `taxonTag.txt` | `a photo of {value}.` captions | +| `sci_com.txt`, `taxon_com.txt`, `taxonTag_com.txt` | the caption with ` with common name {common}.` appended | + +When no vernacular common name is available, the common name falls back to the +scientific name (without author citation). + +## Inputs + +All three are joined by `uuid`: + +1. **Hybrid dataset**: the root holding `*_images.h5` (WebP bytes under an + `/images/` group) and the paired `*_metadata.parquet`. +2. **Resolved taxonomy**: Parquet providing `scientific_name`, `common_name`, + and the `kingdom` through `species` ranks per `uuid` (the text-label source). +3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed + UUIDs are converted. The `.h5` path is auto-derived as the sibling + `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an + optional `hdf5_path` (or `h5_file`) column overrides it. + +## Prerequisites + +Install the toolbox (see the top-level [installation instructions](../README.md#installation-instructions)): + +```bash +pip install -e . +``` + +Option A (single node) needs only the Python dependencies that installs. Option B (distributed) runs on a cluster, and its Slurm scripts encode site-specific choices you must adapt: + +- An **MPI** implementation with `mpi4py` installed against it. The workers run under MPI; if `mpi4py` is built against a different MPI than the one loaded at runtime, they fail to load it. +- A **Spark** runtime for the filter stage. +- The **partition** and **module** names in `scripts/tools_*.slurm`, and the virtual environment those scripts activate (by default the repo-root `.venv`, so install the toolbox there or point the scripts at your environment). + +Edit `scripts/tools_*.slurm` to load your cluster's MPI and Spark modules, target your partition, and activate your environment. As committed they target one cluster (OSC Cardinal) as a worked example to replace with your own. + +## Option A: single node (quick / small subsets) + +Best for testing or modest UUID lists. No Spark/MPI required: + +```bash +python tol_hdf5_to_wds.py \ + --input-root /path/to/hybrid_dataset \ + --lookup /path/to/uuid_lookup.parquet \ + --taxa-glob "/path/to/resolved_taxa/source=*/*.parquet" \ + --output-dir /path/to/wds \ + --resize 224 --shard-size 10000 +``` + +(Generate a `uuid`-to-`h5` lookup for a UUID subset with +[`generate_hdf5_lookup.py`](../generate_hdf5_lookup.py).) + +## Option B: distributed (full dataset, Slurm) + +The toolbox runner reads one config and submits the whole pipeline (Spark +filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. +Start from [`config/tol_hdf5_to_wds_example.yaml`](../config/tol_hdf5_to_wds_example.yaml) +and set, for your run: + +- `account` and `tools_parameters` (nodes, workers, CPUs) for your allocation; +- `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the + toolbox working directory); +- in the `tol_hdf5_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` + (required for the taxonomy/common-name prompts), and `lookup_table_path` + (a parquet/CSV with a `uuid` column selecting which images to convert). + +See the [tool README](../src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md) for +the full key reference. + +Then run the pipeline with a single command: + +```bash +tree_of_life_toolbox config/tol_hdf5_to_wds_my_run.yaml tol_hdf5_to_wds +``` + +The runner reads the config, exports the environment it implies (account, +node/worker counts, output folders), and submits the four chained stages. The +**filter** discovers the `*_metadata.parquet` files under `path_to_input`, +derives each sibling `*_images.h5`, joins the taxonomy, and restricts to the +lookup UUIDs; the **scheduler** assigns shards to MPI ranks; the **workers** +read the WebP bytes, resize to JPEG, write the ten prompts, and emit the tars; +the **verifier** marks completion. The underlying Slurm scripts are described in +[scripts_README](scripts_README.md). + +## Validating output + +`validate_tars.py` checks shard integrity, and `wds_taxoncom_audit.py` audits +the taxonomy/common-name prompts across shards. From a0dfec5d71432dc1ccba3d198a980fa70b36a105 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Tue, 30 Jun 2026 17:50:40 -0400 Subject: [PATCH 03/10] Exit the Spark filter process once filtering completes Co-Authored-By: Claude Opus 4.8 --- src/TreeOfLife_toolbox/main/filter.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/TreeOfLife_toolbox/main/filter.py b/src/TreeOfLife_toolbox/main/filter.py index ed526c5..b8c91a2 100644 --- a/src/TreeOfLife_toolbox/main/filter.py +++ b/src/TreeOfLife_toolbox/main/filter.py @@ -46,3 +46,10 @@ checkpoint["filtering_completed"] = True else: logger.info("Filtering was already completed") + + # Spark's driver JVM does not reliably exit after the context is stopped, which leaves the batch job lingering until its walltime even though the work is done. All output is already written and checkpointed, so stop the session and terminate immediately. + try: + tool_filter.spark.stop() + except Exception: + pass + os._exit(0) From e3ee9d1b065a5fdb1b6ab79395b094764d43fb73 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Tue, 30 Jun 2026 19:15:19 -0400 Subject: [PATCH 04/10] Rename the tool to tol_hybrid_to_wds and drop empty-taxonomy entries during the build Co-Authored-By: Claude Opus 4.8 --- ...le.yaml => tol_hybrid_to_wds_example.yaml} | 2 +- docs/hybrid_to_webdataset_README.md | 12 +++---- src/TreeOfLife_toolbox/__init__.py | 2 +- .../README.md | 12 +++---- .../__init__.py | 2 +- .../classes.py | 31 ++++++++++++------- .../utils.py | 0 tol_hdf5_to_wds.py => tol_hybrid_to_wds.py | 18 ++++++++--- 8 files changed, 49 insertions(+), 30 deletions(-) rename config/{tol_hdf5_to_wds_example.yaml => tol_hybrid_to_wds_example.yaml} (98%) rename src/TreeOfLife_toolbox/{tol_hdf5_to_wds => tol_hybrid_to_wds}/README.md (88%) rename src/TreeOfLife_toolbox/{tol_hdf5_to_wds => tol_hybrid_to_wds}/__init__.py (76%) rename src/TreeOfLife_toolbox/{tol_hdf5_to_wds => tol_hybrid_to_wds}/classes.py (92%) rename src/TreeOfLife_toolbox/{tol_hdf5_to_wds => tol_hybrid_to_wds}/utils.py (100%) rename tol_hdf5_to_wds.py => tol_hybrid_to_wds.py (92%) diff --git a/config/tol_hdf5_to_wds_example.yaml b/config/tol_hybrid_to_wds_example.yaml similarity index 98% rename from config/tol_hdf5_to_wds_example.yaml rename to config/tol_hybrid_to_wds_example.yaml index 8ef5f85..652a9de 100644 --- a/config/tol_hdf5_to_wds_example.yaml +++ b/config/tol_hybrid_to_wds_example.yaml @@ -30,7 +30,7 @@ output_structure: inner_checkpoint_file: "inner_checkpoint.yaml" tools_folder: "tools" -tol_hdf5_to_wds: +tol_hybrid_to_wds: shard_size: 10000 shard_limit: 0 metadata_glob: "**/*_metadata.parquet" diff --git a/docs/hybrid_to_webdataset_README.md b/docs/hybrid_to_webdataset_README.md index 7df8182..2c6160f 100644 --- a/docs/hybrid_to_webdataset_README.md +++ b/docs/hybrid_to_webdataset_README.md @@ -4,7 +4,7 @@ This guide covers packaging the Tree of Life dataset, stored in the hybrid format of WebP images in HDF5 (`*_images.h5`) alongside Parquet metadata (`*_metadata.parquet`), into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards for model training. It is implemented by the -[`tol_hdf5_to_wds`](../src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md) tool. +[`tol_hybrid_to_wds`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) tool. ## What it produces @@ -55,7 +55,7 @@ Edit `scripts/tools_*.slurm` to load your cluster's MPI and Spark modules, targe Best for testing or modest UUID lists. No Spark/MPI required: ```bash -python tol_hdf5_to_wds.py \ +python tol_hybrid_to_wds.py \ --input-root /path/to/hybrid_dataset \ --lookup /path/to/uuid_lookup.parquet \ --taxa-glob "/path/to/resolved_taxa/source=*/*.parquet" \ @@ -70,23 +70,23 @@ python tol_hdf5_to_wds.py \ The toolbox runner reads one config and submits the whole pipeline (Spark filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. -Start from [`config/tol_hdf5_to_wds_example.yaml`](../config/tol_hdf5_to_wds_example.yaml) +Start from [`config/tol_hybrid_to_wds_example.yaml`](../config/tol_hybrid_to_wds_example.yaml) and set, for your run: - `account` and `tools_parameters` (nodes, workers, CPUs) for your allocation; - `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the toolbox working directory); -- in the `tol_hdf5_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` +- in the `tol_hybrid_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` (required for the taxonomy/common-name prompts), and `lookup_table_path` (a parquet/CSV with a `uuid` column selecting which images to convert). -See the [tool README](../src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md) for +See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for the full key reference. Then run the pipeline with a single command: ```bash -tree_of_life_toolbox config/tol_hdf5_to_wds_my_run.yaml tol_hdf5_to_wds +tree_of_life_toolbox config/tol_hybrid_to_wds_my_run.yaml tol_hybrid_to_wds ``` The runner reads the config, exports the environment it implies (account, diff --git a/src/TreeOfLife_toolbox/__init__.py b/src/TreeOfLife_toolbox/__init__.py index 8d14536..b1dd744 100644 --- a/src/TreeOfLife_toolbox/__init__.py +++ b/src/TreeOfLife_toolbox/__init__.py @@ -2,6 +2,6 @@ tol200m_fathom_net_crop, tol200m_bioscan_data_transfer, data_transfer, - tol_hdf5_to_wds, + tol_hybrid_to_wds, ) diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md similarity index 88% rename from src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md index e523652..82f7b8d 100644 --- a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/README.md +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md @@ -1,4 +1,4 @@ -# tol_hdf5_to_wds +# tol_hybrid_to_wds Converts the hybrid Tree of Life dataset (HDF5 images + Parquet metadata) into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards that mirror @@ -8,7 +8,7 @@ the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. 1. **Filter**: Scans the converted dataset for `*_metadata.parquet` files, enriches each row with the matching `*_images.h5` path, and writes shard-sized - parquet partitions under `tools/tol_hdf5_to_wds/shard_metadata`. + parquet partitions under `tools/tol_hybrid_to_wds/shard_metadata`. 2. **Scheduler**: Assigns shard identifiers to MPI ranks and writes `schedule.csv` for the toolbox runner scripts. 3. **Runner**: For every scheduled shard, loads the shard metadata, reads @@ -18,10 +18,10 @@ the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. ## Configuration highlights -Add a `tol_hdf5_to_wds` block to the standard toolbox config template: +Add a `tol_hybrid_to_wds` block to the standard toolbox config template: ```yaml -tol_hdf5_to_wds: +tol_hybrid_to_wds: shard_size: 10000 # rows per shard shard_limit: 0 # optional cap, 0 disables metadata_glob: "**/*_metadata.parquet" @@ -46,8 +46,8 @@ the derived file path, or a `path` column (configurable via Run the tool via the standard CLI: ```bash -CONFIG_PATH=config/tol_hdf5_to_wds_example.yaml \ -tree_of_life_toolbox config/tol_hdf5_to_wds_example.yaml tol_hdf5_to_wds +CONFIG_PATH=config/tol_hybrid_to_wds_example.yaml \ +tree_of_life_toolbox config/tol_hybrid_to_wds_example.yaml tol_hybrid_to_wds ``` Each shard sample contains the JPEG plus ten text prompts: diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py similarity index 76% rename from src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py index 9bbb162..2acf7d8 100644 --- a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/__init__.py +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py @@ -5,4 +5,4 @@ the toolbox registry so they can be scheduled through the generic CLI. """ -from TreeOfLife_toolbox.tol_hdf5_to_wds import classes # noqa: F401 +from TreeOfLife_toolbox.tol_hybrid_to_wds import classes # noqa: F401 diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py similarity index 92% rename from src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py index 4f84c88..700080c 100644 --- a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/classes.py +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py @@ -19,7 +19,7 @@ from TreeOfLife_toolbox.main.filters import FilterRegister, SparkFilterToolBase from TreeOfLife_toolbox.main.runners import MPIRunnerTool, RunnerRegister from TreeOfLife_toolbox.main.schedulers import DefaultScheduler, SchedulerRegister -from TreeOfLife_toolbox.tol_hdf5_to_wds.utils import ( +from TreeOfLife_toolbox.tol_hybrid_to_wds.utils import ( convert_webp_to_jpeg, generate_text_files, init_shard_writer, @@ -30,7 +30,7 @@ def _sanitize_path_column(column): return func.regexp_replace(column, "^file:", "") -@FilterRegister("tol_hdf5_to_wds") +@FilterRegister("tol_hybrid_to_wds") class TolHDF5ToWDSFilter(SparkFilterToolBase): """ Builds shard-level metadata derived from the parquet_to_hdf5 outputs. @@ -38,8 +38,8 @@ class TolHDF5ToWDSFilter(SparkFilterToolBase): def __init__(self, cfg: Config): super().__init__(cfg) - self.filter_name = "tol_hdf5_to_wds" - tool_cfg = self.config.get("tol_hdf5_to_wds", {}) or {} + self.filter_name = "tol_hybrid_to_wds" + tool_cfg = self.config.get("tol_hybrid_to_wds", {}) or {} self.shard_size = int(tool_cfg.get("shard_size", 10000)) self.shard_limit = int(tool_cfg.get("shard_limit", 0)) self.metadata_glob = tool_cfg.get("metadata_glob", "**/*_metadata.parquet") @@ -53,7 +53,7 @@ def __init__(self, cfg: Config): ) if self.shard_size <= 0: - raise ValueError("tol_hdf5_to_wds.shard_size must be greater than zero") + raise ValueError("tol_hybrid_to_wds.shard_size must be greater than zero") os.makedirs(self.shard_metadata_dir, exist_ok=True) @@ -125,6 +125,15 @@ def run(self): taxa_df = taxa_df.select(*taxa_columns) metadata_df = metadata_df.join(taxa_df, on="uuid", how="left") + # Drop a row only when every taxonomy field is null or blank, so no degenerate text sidecars are ever written. Robust to incomplete upstream taxonomy and a no-op on inputs that are already complete. + taxonomy_fields = [c for c in taxa_columns if c != "uuid"] + all_empty = None + for col_name in taxonomy_fields: + blank = func.col(col_name).isNull() | (func.trim(func.col(col_name)) == "") + all_empty = blank if all_empty is None else (all_empty & blank) + if all_empty is not None: + metadata_df = metadata_df.filter(~all_empty) + if self.include_sources: metadata_df = metadata_df.filter(func.col("source").isin(self.include_sources)) if self.include_servers: @@ -260,7 +269,7 @@ def run(self): func.col("shard_id"), ), ) - .withColumn("server_name", func.lit("tol_hdf5_to_wds")) + .withColumn("server_name", func.lit("tol_hybrid_to_wds")) .withColumn("partition_id", func.col("shard_id")) ) @@ -269,7 +278,7 @@ def run(self): self.logger.info("Prepared %d shards", total_shards) -@SchedulerRegister("tol_hdf5_to_wds") +@SchedulerRegister("tol_hybrid_to_wds") class TolHDF5ToWDSScheduler(DefaultScheduler): """ Scheduler that treats each metadata shard as an independent work unit. @@ -277,12 +286,12 @@ class TolHDF5ToWDSScheduler(DefaultScheduler): def __init__(self, cfg: Config): super().__init__(cfg) - self.filter_name = "tol_hdf5_to_wds" + self.filter_name = "tol_hybrid_to_wds" # Scheduling keys only. The data columns (metadata_path, shard_id) ride along in the filter table and are picked up by the runner's data_scheme; including them here duplicates them into schedule.csv and collides on the runner's merge. self.scheme = ["server_name", "partition_id"] -@RunnerRegister("tol_hdf5_to_wds") +@RunnerRegister("tol_hybrid_to_wds") class TolHDF5ToWDSRunner(MPIRunnerTool): """ Converts shard metadata into WebDataset tar archives. @@ -290,11 +299,11 @@ class TolHDF5ToWDSRunner(MPIRunnerTool): def __init__(self, cfg: Config): super().__init__(cfg) - self.filter_name = "tol_hdf5_to_wds" + self.filter_name = "tol_hybrid_to_wds" self.data_scheme = ["metadata_path", "shard_id", "server_name", "partition_id"] self.verification_scheme = ["server_name", "partition_id"] - params = self.config.get("tol_hdf5_to_wds", {}) or {} + params = self.config.get("tol_hybrid_to_wds", {}) or {} self.tar_output_root = params.get("tar_output_root") or os.path.join( self.tools_path, self.filter_name, "tar_dataset" ) diff --git a/src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/utils.py similarity index 100% rename from src/TreeOfLife_toolbox/tol_hdf5_to_wds/utils.py rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/utils.py diff --git a/tol_hdf5_to_wds.py b/tol_hybrid_to_wds.py similarity index 92% rename from tol_hdf5_to_wds.py rename to tol_hybrid_to_wds.py index 0308f33..cca5a79 100644 --- a/tol_hdf5_to_wds.py +++ b/tol_hybrid_to_wds.py @@ -3,7 +3,7 @@ Single-node converter from the hybrid Tree of Life HDF5/metadata dataset to WebDataset shards. Example: - python tol_hdf5_to_wds.py \ + python tol_hybrid_to_wds.py \ --input-root /fs/ess/PAS2136/TreeOfLife/data \ --lookup /fs/scratch/PAS2136/TreeOfLife_test-wds/lookup-tables/all/hdf5_lookup.parquet \ --taxa-glob "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/*/*.parquet" \ @@ -26,13 +26,13 @@ import webdataset as wds try: - from TreeOfLife_toolbox.tol_hdf5_to_wds.utils import ( + from TreeOfLife_toolbox.tol_hybrid_to_wds.utils import ( convert_webp_to_jpeg, generate_text_files, ) except ImportError as exc: # pragma: no cover - script usage only raise SystemExit( - "Unable to import TreeOfLife_toolbox.tol_hdf5_to_wds.utils. " + "Unable to import TreeOfLife_toolbox.tol_hybrid_to_wds.utils. " "Please install the toolbox package (pip install -e .) before running this script." ) from exc @@ -91,7 +91,7 @@ def setup_logger(level: str) -> logging.Logger: level=getattr(logging, level.upper()), format="%(asctime)s - %(levelname)s - %(message)s", ) - return logging.getLogger("tol_hdf5_to_wds") + return logging.getLogger("tol_hybrid_to_wds") def discover_metadata_files(root: str, pattern: str) -> List[str]: @@ -217,6 +217,16 @@ def merge_metadata( ).drop("hdf5_path_right") if taxa_lazy is not None: joined = joined.join(taxa_lazy, on="uuid", how="left") + # Drop a row only when every taxonomy field is null or blank, so no degenerate text sidecars are ever written. Robust to incomplete upstream taxonomy and a no-op on inputs that are already complete. + tax_cols = [c for c in [ + "scientific_name", "common_name", "provided_common_name", "kingdom", + "phylum", "class", "order", "family", "genus", "species", + ] if c in joined.collect_schema().names()] + if tax_cols: + all_blank = pl.all_horizontal( + [pl.col(c).is_null() | (pl.col(c).str.strip_chars() == "") for c in tax_cols] + ) + joined = joined.filter(~all_blank) joined = joined.unique(subset=["uuid"]) result = joined.sort(["hdf5_path", "uuid"]).collect(streaming=True) logger.info("Collected %s metadata rows", result.height) From 0ac8c2eb67469bfb7ebbeb42095f2a7dcce6a376 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Tue, 30 Jun 2026 19:47:29 -0400 Subject: [PATCH 05/10] Relocate the one-off taxonomy-scrub scripts out of scripts/ Co-Authored-By: Claude Opus 4.8 --- scripts/copy_implicated_shards.py | 222 ------------------------- scripts/find_empty_taxonomy.py | 67 -------- scripts/publish_scrubbed_shards.py | 174 ------------------- scripts/remove_bad_samples.py | 89 ---------- scripts/remove_bad_samples_parallel.py | 180 -------------------- scripts/run_incremental_scrub.sh | 57 ------- 6 files changed, 789 deletions(-) delete mode 100644 scripts/copy_implicated_shards.py delete mode 100644 scripts/find_empty_taxonomy.py delete mode 100644 scripts/publish_scrubbed_shards.py delete mode 100644 scripts/remove_bad_samples.py delete mode 100644 scripts/remove_bad_samples_parallel.py delete mode 100644 scripts/run_incremental_scrub.sh diff --git a/scripts/copy_implicated_shards.py b/scripts/copy_implicated_shards.py deleted file mode 100644 index 5c57d61..0000000 --- a/scripts/copy_implicated_shards.py +++ /dev/null @@ -1,222 +0,0 @@ -#!/usr/bin/env python3 -""" -Copy shard tar files implicated by missing-taxa.csv into a staging directory, -but ONLY when the source shard appears complete. - -Completeness criteria (intrinsic): -- `tar tf` succeeds (tar structure readable) -- exactly 10,000 *.jpg entries -- exactly 10,000 entries for each required TXT sidecar suffix - -Staging/update semantics: -- If a scrubbed output already exists (in --scrubbed-dir), skip entirely. -- If destination shard exists AND same byte size as source, skip. -- Otherwise, (re)copy the source shard into destination (atomic temp + rename), - but only if the source passes completeness. - -Usage: - python scripts/copy_implicated_shards.py \ - --map missing-taxa.csv \ - --src-dir /path/to/shards \ - --dst-dir /path/to/shards_to_scrub \ - --scrubbed-dir /path/to/shards_scrubbed -""" - -from __future__ import annotations - -import argparse -import csv -import os -import shutil -import subprocess -from pathlib import Path -from typing import Dict, List, Set - -from tqdm import tqdm - - -SIDE_TXT_SUFFIXES: List[str] = [ - "com.txt", - "common_name.txt", - "sci.txt", - "sci_com.txt", - "scientific_name.txt", - "taxon.txt", - "taxonTag.txt", - "taxonTag_com.txt", - "taxon_com.txt", - "taxonomic_name.txt", -] - -EXPECTED_PER_TYPE = 10_000 - - -def shard_name(shard_id: int) -> str: - return f"shard-{shard_id:05d}.tar" - - -def count_data_rows(path: Path) -> int: - with path.open("rb") as f: - return max(0, sum(1 for _ in f) - 1) - - -def load_unique_shard_ids(csv_path: Path, total_rows: int | None = None) -> Set[int]: - shard_ids: Set[int] = set() - with csv_path.open(newline="") as fh: - reader = csv.DictReader(fh) - for row in tqdm(reader, total=total_rows, desc="Scanning missing-taxa.csv", unit="rows"): - shard_ids.add(int(row["shard_id"])) - return shard_ids - - -def tar_counts_via_tar_tf(tar_path: Path) -> Dict[str, int] | None: - counts: Dict[str, int] = {sfx: 0 for sfx in SIDE_TXT_SUFFIXES} - counts["jpg"] = 0 - counts["total"] = 0 - - proc = subprocess.Popen( - ["tar", "tf", str(tar_path)], - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - bufsize=1, - ) - - assert proc.stdout is not None - try: - for line in proc.stdout: - name = line.strip() - if not name: - continue - counts["total"] += 1 - - if name.endswith(".jpg"): - counts["jpg"] += 1 - continue - - for sfx in SIDE_TXT_SUFFIXES: - if name.endswith("." + sfx): - counts[sfx] += 1 - break - finally: - proc.stdout.close() - rc = proc.wait() - - if rc != 0: - return None - return counts - - -def is_complete_shard(tar_path: Path) -> bool: - counts = tar_counts_via_tar_tf(tar_path) - if counts is None: - return False - - if counts["jpg"] != EXPECTED_PER_TYPE: - return False - - for sfx in SIDE_TXT_SUFFIXES: - if counts[sfx] != EXPECTED_PER_TYPE: - return False - - return True - - -def copy_atomic(src: Path, dst: Path) -> None: - dst.parent.mkdir(parents=True, exist_ok=True) - tmp = dst.with_name(dst.name + ".tmp") - shutil.copy2(src, tmp) # preserve mtime for staging comparisons if you care later - os.replace(tmp, dst) - - -def same_size(src: Path, dst: Path) -> bool: - try: - return src.stat().st_size == dst.stat().st_size - except FileNotFoundError: - return False - - -def main() -> None: - p = argparse.ArgumentParser( - description="Copy implicated shard tars into a staging directory (only when complete)." - ) - p.add_argument("--map", required=True, type=Path, help="missing-taxa.csv with columns uuid,shard_id.") - p.add_argument("--src-dir", required=True, type=Path, help="Directory containing shard-XXXXX.tar files.") - p.add_argument("--dst-dir", required=True, type=Path, help="Destination directory for shards to scrub.") - p.add_argument( - "--scrubbed-dir", - required=True, - type=Path, - help="Directory containing scrubbed shard outputs; if shard exists here, it is skipped.", - ) - p.add_argument("--limit", type=int, default=0, help="Optional: stop after copying N shards (0 = no limit).") - args = p.parse_args() - - total_rows = count_data_rows(args.map) - shard_ids = load_unique_shard_ids(args.map, total_rows=total_rows) - shard_id_list = sorted(shard_ids) - - copied = 0 - updated = 0 - - skipped_missing = 0 - skipped_incomplete = 0 - skipped_same_size = 0 - skipped_already_scrubbed = 0 - - for sid in tqdm(shard_id_list, desc="Checking/copying shards", unit="shards"): - src = args.src_dir / shard_name(sid) - dst = args.dst_dir / src.name - scrubbed = args.scrubbed_dir / src.name - - # Robust: if already scrubbed, never stage again (prevents noise if scrubbed published back to src-dir) - if scrubbed.exists(): - skipped_already_scrubbed += 1 - continue - - if not src.exists(): - skipped_missing += 1 - continue - - # If destination exists and size matches, skip. - if dst.exists() and same_size(src, dst): - skipped_same_size += 1 - continue - - # Only copy/update if source is intrinsically complete. - if not is_complete_shard(src): - skipped_incomplete += 1 - continue - - was_update = dst.exists() - copy_atomic(src, dst) - if was_update: - updated += 1 - else: - copied += 1 - - if (copied + updated) % 25 == 0: - print( - f"[INFO] copied={copied} updated={updated} " - f"missing={skipped_missing} incomplete={skipped_incomplete} " - f"same_size_skip={skipped_same_size} already_scrubbed_skip={skipped_already_scrubbed}" - ) - - if args.limit and (copied + updated) >= args.limit: - break - - print( - "[DONE] " - f"unique_shards={len(shard_id_list)} " - f"copied={copied} " - f"updated={updated} " - f"skipped_missing={skipped_missing} " - f"skipped_incomplete={skipped_incomplete} " - f"skipped_same_size={skipped_same_size} " - f"skipped_already_scrubbed={skipped_already_scrubbed}" - ) - - -if __name__ == "__main__": - main() - diff --git a/scripts/find_empty_taxonomy.py b/scripts/find_empty_taxonomy.py deleted file mode 100644 index dc8dd8e..0000000 --- a/scripts/find_empty_taxonomy.py +++ /dev/null @@ -1,67 +0,0 @@ -#!/usr/bin/env python -""" -Scan shard_metadata for UUIDs lacking any taxonomy fields and emit a CSV map. - -Usage: - python scripts/find_empty_taxonomy.py /path/to/shard_metadata output.csv -""" - -import argparse -import csv -from pathlib import Path - -import polars as pl - - -def gather_null_taxonomy(shard_metadata_root: Path) -> pl.DataFrame: - dataset = pl.scan_parquet( - str(shard_metadata_root / "shard_id=*" / "*.parquet"), hive_partitioning=True - ) - taxonomy_cols = [ - "scientific_name", - "common_name", - "provided_common_name", - "kingdom", - "phylum", - "class", - "order", - "family", - "genus", - "species", - ] - filter_expr = None - for col in taxonomy_cols: - expr = pl.col(col).is_null() | (pl.col(col).str.strip_chars().eq("")) - filter_expr = expr if filter_expr is None else (filter_expr & expr) - - empty_taxa_df = ( - dataset.filter(filter_expr) - .select("uuid", "shard_id") - .collect() - .sort("shard_id") - ) - return empty_taxa_df - - -def main(): - parser = argparse.ArgumentParser( - description="Find shard metadata rows with empty taxonomy." - ) - parser.add_argument("shard_metadata_root", type=Path) - parser.add_argument("output_csv", type=Path) - args = parser.parse_args() - - df = gather_null_taxonomy(args.shard_metadata_root) - if df.is_empty(): - args.output_csv.write_text("uuid,shard_id\n") - return - - with args.output_csv.open("w", newline="") as f: - writer = csv.writer(f) - writer.writerow(["uuid", "shard_id"]) - for row in df.iter_rows(): - writer.writerow(row) - - -if __name__ == "__main__": - main() diff --git a/scripts/publish_scrubbed_shards.py b/scripts/publish_scrubbed_shards.py deleted file mode 100644 index f065dc4..0000000 --- a/scripts/publish_scrubbed_shards.py +++ /dev/null @@ -1,174 +0,0 @@ -#!/usr/bin/env python3 -""" -Publish scrubbed shards back into the primary shards directory. - -Design goals: -- Only overwrite shards/ with scrubbed versions from shards_scrubbed/. -- Do NOT touch shards_to_scrub/ (your backup copies). -- Avoid clobbering a shard that is still being generated: only publish onto a "complete" source shard. -- Publish via temp + atomic rename so Globus never sees a partial file. -- Do NOT preserve mtime (important for Globus sync level 2: mtime-based). - -Usage: - python scripts/publish_scrubbed_shards.py \ - --shards /path/to/shards \ - --shards-scrubbed /path/to/shards_scrubbed -""" - -from __future__ import annotations - -import argparse -import os -import shutil -import subprocess -from pathlib import Path -from typing import Dict, List, Optional - - -SIDE_TXT_SUFFIXES: List[str] = [ - "com.txt", - "common_name.txt", - "sci.txt", - "sci_com.txt", - "scientific_name.txt", - "taxon.txt", - "taxonTag.txt", - "taxonTag_com.txt", - "taxon_com.txt", - "taxonomic_name.txt", -] - -EXPECTED_PER_TYPE = 10_000 - - -def tar_counts_via_tar_tf(tar_path: Path) -> Optional[Dict[str, int]]: - """ - Return counts of jpg + each sidecar suffix for tar_path by streaming `tar tf`. - Returns None if `tar tf` fails. - """ - counts: Dict[str, int] = {sfx: 0 for sfx in SIDE_TXT_SUFFIXES} - counts["jpg"] = 0 - - proc = subprocess.Popen( - ["tar", "tf", str(tar_path)], - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - bufsize=1, - ) - - assert proc.stdout is not None - try: - for line in proc.stdout: - name = line.strip() - if not name: - continue - - if name.endswith(".jpg"): - counts["jpg"] += 1 - continue - - for sfx in SIDE_TXT_SUFFIXES: - if name.endswith("." + sfx): - counts[sfx] += 1 - break - finally: - proc.stdout.close() - rc = proc.wait() - - return None if rc != 0 else counts - - -def is_complete_source_shard(tar_path: Path) -> bool: - """ - Only publish onto a shard in shards/ if it still looks like a complete original shard. - (Prevents racing the generator / avoids clobbering partials.) - """ - counts = tar_counts_via_tar_tf(tar_path) - if counts is None: - return False - - if counts["jpg"] != EXPECTED_PER_TYPE: - return False - - for sfx in SIDE_TXT_SUFFIXES: - if counts[sfx] != EXPECTED_PER_TYPE: - return False - - return True - - -def copy_atomic_no_mtime(src: Path, dst: Path) -> None: - """ - Copy src -> dst through a temp file in dst's directory, then atomic rename. - Does NOT preserve metadata/mtime (good for Globus sync level 2). - """ - dst.parent.mkdir(parents=True, exist_ok=True) - tmp = dst.with_name(dst.name + ".tmp") - shutil.copyfile(src, tmp) # intentionally not copy2() - os.replace(tmp, dst) - - -def main() -> None: - p = argparse.ArgumentParser(description="Publish scrubbed shard tars back into shards/ safely.") - p.add_argument("--shards", required=True, type=Path, help="Primary shards/ directory to overwrite.") - p.add_argument("--shards-scrubbed", required=True, type=Path, help="Directory containing scrubbed shard outputs.") - p.add_argument("--limit", type=int, default=0, help="Optional: stop after publishing N shards (0 = no limit).") - args = p.parse_args() - - shards_dir: Path = args.shards - scrubbed_dir: Path = args.shards_scrubbed - - published = 0 - skipped_missing_src = 0 - skipped_incomplete_src = 0 - skipped_no_change = 0 - - for scrubbed_tar in sorted(scrubbed_dir.glob("shard-*.tar")): - name = scrubbed_tar.name - dst_tar = shards_dir / name - - if not dst_tar.exists(): - # If the generator hasn't created it yet, don't create it here. - skipped_missing_src += 1 - continue - - # Never clobber a shard that doesn't look "complete" yet. - if not is_complete_source_shard(dst_tar): - skipped_incomplete_src += 1 - continue - - # If shards/ already matches scrubbed by size, assume already published. - # (Cheap heuristic; avoids unnecessary rewrite/mtime churn.) - try: - if dst_tar.stat().st_size == scrubbed_tar.stat().st_size: - skipped_no_change += 1 - continue - except FileNotFoundError: - # race: try again next loop - continue - - copy_atomic_no_mtime(scrubbed_tar, dst_tar) - published += 1 - - if published % 25 == 0: - print( - f"[INFO] published={published} no_change={skipped_no_change} " - f"missing_src={skipped_missing_src} incomplete_src={skipped_incomplete_src}" - ) - - if args.limit and published >= args.limit: - break - - print( - "[DONE] " - f"published={published} " - f"skipped_no_change={skipped_no_change} " - f"skipped_missing_src={skipped_missing_src} " - f"skipped_incomplete_src={skipped_incomplete_src}" - ) - - -if __name__ == "__main__": - main() - diff --git a/scripts/remove_bad_samples.py b/scripts/remove_bad_samples.py deleted file mode 100644 index 8bf48a6..0000000 --- a/scripts/remove_bad_samples.py +++ /dev/null @@ -1,89 +0,0 @@ -#!/usr/bin/env python -""" -Remove specified UUIDs from shard tar files. - -Usage: - python scripts/remove_bad_samples.py \ - --shards /path/to/shards \ - --map missing-taxa.csv \ - [--backup-dir /path/to/backups] -""" - -from __future__ import annotations - -import argparse -import csv -import os -import tarfile -from collections import defaultdict -from pathlib import Path -from typing import Dict, Iterable, List - - -SIDE_EXTS = [ - "jpg", - "com.txt", - "common_name.txt", - "sci.txt", - "sci_com.txt", - "scientific_name.txt", - "taxon.txt", - "taxonTag.txt", - "taxonTag_com.txt", - "taxon_com.txt", - "taxonomic_name.txt", -] - - -def load_uuid_map(csv_path: Path) -> Dict[int, List[str]]: - shard_map: Dict[int, List[str]] = defaultdict(list) - with csv_path.open() as fh: - reader = csv.DictReader(fh) - for row in reader: - shard_map[int(row["shard_id"])].append(row["uuid"]) - return shard_map - - -def scrub_shard(shard_path: Path, uuids: Iterable[str], backup_dir: Path | None) -> None: - if not shard_path.exists(): - print(f"[WARN] Shard not found: {shard_path}") - return - - tmp_path = shard_path.with_suffix(".tmp") - drop_prefixes = set(uuids) - drop_members = {f"{u}.{ext}" for u in drop_prefixes for ext in SIDE_EXTS} - - with tarfile.open(shard_path, "r") as src, tarfile.open(tmp_path, "w") as dst: - for member in src: - base = member.name.split(".", 1)[0] - if base in drop_prefixes and member.name in drop_members: - continue - data = src.extractfile(member) if member.isfile() else None - dst.addfile(member, data) - - if backup_dir: - backup_dir.mkdir(parents=True, exist_ok=True) - shard_backup = backup_dir / shard_path.name - os.replace(shard_path, shard_backup) - else: - shard_path.unlink() - os.replace(tmp_path, shard_path) - - -def main(): - parser = argparse.ArgumentParser(description="Remove bad UUIDs from shard tar files.") - parser.add_argument("--shards", required=True, type=Path, help="Directory containing shard-XXXXX.tar files.") - parser.add_argument("--map", required=True, type=Path, help="CSV with columns uuid,shard_id.") - parser.add_argument("--backup-dir", type=Path, help="Optional dir to store original tar backups.") - args = parser.parse_args() - - shard_map = load_uuid_map(args.map) - for shard_id, uuids in shard_map.items(): - shard_name = f"shard-{shard_id:05d}.tar" - shard_path = args.shards / shard_name - print(f"[INFO] Scrubbing {len(uuids)} samples from {shard_name}") - scrub_shard(shard_path, uuids, args.backup_dir) - - -if __name__ == "__main__": - main() diff --git a/scripts/remove_bad_samples_parallel.py b/scripts/remove_bad_samples_parallel.py deleted file mode 100644 index dcde917..0000000 --- a/scripts/remove_bad_samples_parallel.py +++ /dev/null @@ -1,180 +0,0 @@ -#!/usr/bin/env python3 -""" -Remove specified UUIDs (and sidecars) from shard tar files using parallel workers. - -Incremental behavior: -- Reads from --shards (input directory) -- Writes scrubbed shards to --shards-scrubbed (required output directory) -- If scrubbed output already exists for a shard, that shard is skipped -- Missing input shards are skipped - -Usage: - python scripts/remove_bad_samples_parallel.py \ - --shards /path/to/shards_to_scrub \ - --shards-scrubbed /path/to/shards_scrubbed \ - --map missing-taxa.csv \ - --workers 8 \ - --max-inflight 32 -""" - -from __future__ import annotations - -import argparse -import csv -import os -import tarfile -from concurrent.futures import ProcessPoolExecutor, as_completed -from pathlib import Path -from typing import Iterable, Iterator, List, Optional, Tuple - -SIDE_EXTS = { - "jpg", - "com.txt", - "common_name.txt", - "sci.txt", - "sci_com.txt", - "scientific_name.txt", - "taxon.txt", - "taxonTag.txt", - "taxonTag_com.txt", - "taxon_com.txt", - "taxonomic_name.txt", -} - - -def iter_shard_groups(csv_path: Path) -> Iterator[Tuple[int, List[str]]]: - """ - Stream (shard_id, [uuid...]) groups from a CSV sorted by shard_id. - """ - with csv_path.open(newline="") as fh: - reader = csv.DictReader(fh) - current_sid: Optional[int] = None - uuids: List[str] = [] - - for row in reader: - sid = int(row["shard_id"]) - u = row["uuid"] - - if current_sid is None: - current_sid = sid - - if sid != current_sid: - yield current_sid, uuids - current_sid = sid - uuids = [u] - else: - uuids.append(u) - - if current_sid is not None and uuids: - yield current_sid, uuids - - -def scrub_shard( - src_shard_path: Path, - dst_shard_path: Path, - uuids: Iterable[str], -) -> Tuple[str, int, int]: - """ - Rewrite shard tar, dropping members whose basename matches .. - Writes scrubbed output to dst_shard_path using temp + atomic rename. - Returns (shard_name, removed_count, kept_count). - """ - if not src_shard_path.exists(): - return (src_shard_path.name, 0, 0) - - dst_shard_path.parent.mkdir(parents=True, exist_ok=True) - tmp_out = dst_shard_path.with_name(dst_shard_path.name + ".tmp") - - drop_prefixes = set(uuids) - removed = 0 - kept = 0 - - with tarfile.open(src_shard_path, "r") as src, tarfile.open(tmp_out, "w") as dst: - for member in src: - filename = os.path.basename(member.name) - - dot = filename.find(".") - if dot != -1: - base = filename[:dot] - if base in drop_prefixes: - suffix = filename[dot + 1 :] - if suffix in SIDE_EXTS: - removed += 1 - continue - - data = src.extractfile(member) if member.isfile() else None - dst.addfile(member, data) - kept += 1 - - os.replace(tmp_out, dst_shard_path) - return (src_shard_path.name, removed, kept) - - -def _worker(args: Tuple[Path, Path, List[str]]) -> Tuple[str, int, int]: - src_shard_path, dst_shard_path, uuids = args - return scrub_shard(src_shard_path, dst_shard_path, uuids) - - -def main() -> None: - parser = argparse.ArgumentParser(description="Remove bad UUIDs from shard tar files (parallel, incremental).") - parser.add_argument("--shards", required=True, type=Path, help="Input directory containing shard-XXXXX.tar files.") - parser.add_argument("--shards-scrubbed", required=True, type=Path, help="Output directory for scrubbed shards.") - parser.add_argument("--map", required=True, type=Path, help="CSV with columns uuid,shard_id (sorted by shard_id recommended).") - parser.add_argument("--workers", type=int, default=8, help="Parallel shard workers (I/O heavy; tune to cluster limits).") - parser.add_argument("--max-inflight", type=int, default=32, help="Bound queued shard jobs to limit FS pressure.") - args = parser.parse_args() - - shards_dir: Path = args.shards - scrubbed_dir: Path = args.shards_scrubbed - workers = max(1, args.workers) - max_inflight = max(workers, args.max_inflight) - - scrubbed_dir.mkdir(parents=True, exist_ok=True) - - inflight = [] - total_removed = 0 - shards_processed = 0 - shards_skipped_missing = 0 - shards_skipped_already = 0 - - with ProcessPoolExecutor(max_workers=workers) as executor: - for shard_id, uuids in iter_shard_groups(args.map): - shard_name = f"shard-{shard_id:05d}.tar" - src_shard_path = shards_dir / shard_name - dst_shard_path = scrubbed_dir / shard_name - - if not src_shard_path.exists(): - shards_skipped_missing += 1 - continue - - # Incremental: if output already exists, skip - if dst_shard_path.exists(): - shards_skipped_already += 1 - continue - - fut = executor.submit(_worker, (src_shard_path, dst_shard_path, uuids)) - inflight.append(fut) - - if len(inflight) >= max_inflight: - name, removed, kept = inflight.pop(0).result() - total_removed += removed - shards_processed += 1 - print(f"[INFO] {name}: removed={removed} kept={kept}") - - for fut in as_completed(inflight): - name, removed, kept = fut.result() - total_removed += removed - shards_processed += 1 - print(f"[INFO] {name}: removed={removed} kept={kept}") - - print( - "[DONE] " - f"shards_processed={shards_processed} " - f"shards_skipped_missing={shards_skipped_missing} " - f"shards_skipped_already={shards_skipped_already} " - f"members_removed={total_removed}" - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/run_incremental_scrub.sh b/scripts/run_incremental_scrub.sh deleted file mode 100644 index 881513e..0000000 --- a/scripts/run_incremental_scrub.sh +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -START_TS=$(date +%s) -MAX_SECONDS=$((16 * 60 * 60)) # 16 hours -SLEEP_SECONDS=60 # 1 minute between cycles - -ts() { date +"%Y-%m-%d %H:%M:%S"; } - -echo "[$(ts)] [INFO] Starting incremental copy+scrub loop" -echo "[$(ts)] [INFO] Will stop after $((MAX_SECONDS/3600)) hours" - -while true; do - NOW_TS=$(date +%s) - ELAPSED=$((NOW_TS - START_TS)) - - if [ "$ELAPSED" -ge "$MAX_SECONDS" ]; then - echo "[$(ts)] [INFO] Reached $((MAX_SECONDS/3600))-hour limit; exiting." - break - fi - - echo - echo "[$(ts)] [INFO] ===== elapsed=$((ELAPSED/60)) min =====" - - echo "[$(ts)] [INFO] Running copy_implicated_shards.py" - python scripts/copy_implicated_shards.py \ - --map missing-taxa.csv \ - --src-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ - --dst-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_to_scrub \ - --scrubbed-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed - - echo "[$(ts)] [INFO] Running remove_bad_samples_parallel.py" - python scripts/remove_bad_samples_parallel.py \ - --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_to_scrub \ - --shards-scrubbed /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed \ - --map missing-taxa.csv \ - --workers 2 \ - --max-inflight 2 - - echo "[$(ts)] [INFO] Publishing scrubbed shards back into primary shards/" - python scripts/publish_scrubbed_shards.py \ - --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ - --shards-scrubbed /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards_scrubbed - - echo "[$(ts)] [INFO] Updating sizes.json" - python scripts/sizes_incremental_fast.py \ - --shards /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/2025-12_common-name-fix/shards \ - --map missing-taxa.csv \ - --min-age-seconds 180 \ - --progress-every 50 \ - --heartbeat-seconds 30 - - echo "[$(ts)] [INFO] Sleeping ${SLEEP_SECONDS}s" - sleep "$SLEEP_SECONDS" -done - -echo "[$(ts)] [DONE] Incremental loop finished" From a32dd0305a0b8f1308ce7c97eff40238aff432b5 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Wed, 1 Jul 2026 13:38:09 -0400 Subject: [PATCH 06/10] Un-force-wrap comments and markdown. --- config/tol_hybrid_to_wds_example.yaml | 5 +- docs/hybrid_to_webdataset_README.md | 54 +++++-------------- .../tol_hybrid_to_wds/README.md | 33 +++--------- .../tol_hybrid_to_wds/__init__.py | 3 +- 4 files changed, 22 insertions(+), 73 deletions(-) diff --git a/config/tol_hybrid_to_wds_example.yaml b/config/tol_hybrid_to_wds_example.yaml index 652a9de..45e023d 100644 --- a/config/tol_hybrid_to_wds_example.yaml +++ b/config/tol_hybrid_to_wds_example.yaml @@ -40,10 +40,7 @@ tol_hybrid_to_wds: resize_size: 224 shard_prefix: "shard" runner_timeout_seconds: 3600 - # Lookup selecting which images to convert: a parquet/CSV with a `uuid` - # column. The `.h5` path is auto-derived as the sibling `*_images.h5` of each - # `*_metadata.parquet`, so no path column is required (an optional `hdf5_path` - # or `h5_file` column overrides it). + # Lookup selecting which images to convert: a parquet/CSV with a `uuid` column. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is required (an optional `hdf5_path` or `h5_file` column overrides it). lookup_table_path: "/fs/ess/PAS2136/TreeOfLife/lookups/subset.parquet" # Resolved-taxonomy parquet, joined by `uuid`; required for the text prompts. taxa_glob: "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/source=*/*.parquet" diff --git a/docs/hybrid_to_webdataset_README.md b/docs/hybrid_to_webdataset_README.md index 2c6160f..eca5ba3 100644 --- a/docs/hybrid_to_webdataset_README.md +++ b/docs/hybrid_to_webdataset_README.md @@ -1,16 +1,10 @@ # Create a WebDataset from the hybrid HDF5 + Parquet format -This guide covers packaging the Tree of Life dataset, stored in the hybrid -format of WebP images in HDF5 (`*_images.h5`) alongside Parquet metadata -(`*_metadata.parquet`), into [WebDataset](https://github.com/webdataset/webdataset) -`.tar` shards for model training. It is implemented by the -[`tol_hybrid_to_wds`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) tool. +This guide covers packaging the Tree of Life dataset, stored in the hybrid format of WebP images in HDF5 (`*_images.h5`) alongside Parquet metadata (`*_metadata.parquet`), into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards for model training. It is implemented by the [`tol_hybrid_to_wds`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) tool. ## What it produces -One `shard-XXXXX.tar` per `shard_size` samples. Each sample is keyed by `uuid` -and contains a JPEG plus ten text-prompt sidecars used for BioCLIP-style -training: +One `shard-XXXXX.tar` per `shard_size` samples. Each sample is keyed by `uuid` and contains a JPEG plus ten text-prompt sidecars used for BioCLIP-style training: | | | |---|---| @@ -18,21 +12,15 @@ training: | `sci.txt`, `com.txt`, `taxon.txt`, `taxonTag.txt` | `a photo of {value}.` captions | | `sci_com.txt`, `taxon_com.txt`, `taxonTag_com.txt` | the caption with ` with common name {common}.` appended | -When no vernacular common name is available, the common name falls back to the -scientific name (without author citation). +When no vernacular common name is available, the common name falls back to the scientific name (without author citation). ## Inputs All three are joined by `uuid`: -1. **Hybrid dataset**: the root holding `*_images.h5` (WebP bytes under an - `/images/` group) and the paired `*_metadata.parquet`. -2. **Resolved taxonomy**: Parquet providing `scientific_name`, `common_name`, - and the `kingdom` through `species` ranks per `uuid` (the text-label source). -3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed - UUIDs are converted. The `.h5` path is auto-derived as the sibling - `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an - optional `hdf5_path` (or `h5_file`) column overrides it. +1. **Hybrid dataset**: the root holding `*_images.h5` (WebP bytes under an `/images/` group) and the paired `*_metadata.parquet`. +2. **Resolved taxonomy**: Parquet providing `scientific_name`, `common_name`, and the `kingdom` through `species` ranks per `uuid` (the text-label source). +3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed UUIDs are converted. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an optional `hdf5_path` (or `h5_file`) column overrides it. ## Prerequisites @@ -63,25 +51,17 @@ python tol_hybrid_to_wds.py \ --resize 224 --shard-size 10000 ``` -(Generate a `uuid`-to-`h5` lookup for a UUID subset with -[`generate_hdf5_lookup.py`](../generate_hdf5_lookup.py).) +(Generate a `uuid`-to-`h5` lookup for a UUID subset with [`generate_hdf5_lookup.py`](../generate_hdf5_lookup.py).) ## Option B: distributed (full dataset, Slurm) -The toolbox runner reads one config and submits the whole pipeline (Spark -filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. -Start from [`config/tol_hybrid_to_wds_example.yaml`](../config/tol_hybrid_to_wds_example.yaml) -and set, for your run: +The toolbox runner reads one config and submits the whole pipeline (Spark filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. Start from [`config/tol_hybrid_to_wds_example.yaml`](../config/tol_hybrid_to_wds_example.yaml) and set, for your run: - `account` and `tools_parameters` (nodes, workers, CPUs) for your allocation; -- `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the - toolbox working directory); -- in the `tol_hybrid_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` - (required for the taxonomy/common-name prompts), and `lookup_table_path` - (a parquet/CSV with a `uuid` column selecting which images to convert). +- `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the toolbox working directory); +- in the `tol_hybrid_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` (required for the taxonomy/common-name prompts), and `lookup_table_path` (a parquet/CSV with a `uuid` column selecting which images to convert). -See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for -the full key reference. +See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for the full key reference. Then run the pipeline with a single command: @@ -89,16 +69,8 @@ Then run the pipeline with a single command: tree_of_life_toolbox config/tol_hybrid_to_wds_my_run.yaml tol_hybrid_to_wds ``` -The runner reads the config, exports the environment it implies (account, -node/worker counts, output folders), and submits the four chained stages. The -**filter** discovers the `*_metadata.parquet` files under `path_to_input`, -derives each sibling `*_images.h5`, joins the taxonomy, and restricts to the -lookup UUIDs; the **scheduler** assigns shards to MPI ranks; the **workers** -read the WebP bytes, resize to JPEG, write the ten prompts, and emit the tars; -the **verifier** marks completion. The underlying Slurm scripts are described in -[scripts_README](scripts_README.md). +The runner reads the config, exports the environment it implies (account, node/worker counts, output folders), and submits the four chained stages. The **filter** discovers the `*_metadata.parquet` files under `path_to_input`, derives each sibling `*_images.h5`, joins the taxonomy, and restricts to the lookup UUIDs; the **scheduler** assigns shards to MPI ranks; the **workers** read the WebP bytes, resize to JPEG, write the ten prompts, and emit the tars; the **verifier** marks completion. The underlying Slurm scripts are described in [scripts_README](scripts_README.md). ## Validating output -`validate_tars.py` checks shard integrity, and `wds_taxoncom_audit.py` audits -the taxonomy/common-name prompts across shards. +`validate_tars.py` checks shard integrity, and `wds_taxoncom_audit.py` audits the taxonomy/common-name prompts across shards. diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md index 82f7b8d..8cba429 100644 --- a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md @@ -1,20 +1,12 @@ # tol_hybrid_to_wds -Converts the hybrid Tree of Life dataset (HDF5 images + Parquet metadata) into -[WebDataset](https://github.com/webdataset/webdataset) `.tar` shards that mirror -the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. +Converts the hybrid Tree of Life dataset (HDF5 images + Parquet metadata) into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards that mirror the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. ## Pipeline steps -1. **Filter**: Scans the converted dataset for `*_metadata.parquet` files, - enriches each row with the matching `*_images.h5` path, and writes shard-sized - parquet partitions under `tools/tol_hybrid_to_wds/shard_metadata`. -2. **Scheduler**: Assigns shard identifiers to MPI ranks and writes - `schedule.csv` for the toolbox runner scripts. -3. **Runner**: For every scheduled shard, loads the shard metadata, reads - WebP-compressed bytes from the referenced HDF5 file, converts them to JPEG - (with an optional resize), generates ten taxonomy/common-name text prompts, - and emits `shard-XXXXX.tar` archives via `webdataset.TarWriter`. +1. **Filter**: Scans the converted dataset for `*_metadata.parquet` files, enriches each row with the matching `*_images.h5` path, and writes shard-sized parquet partitions under `tools/tol_hybrid_to_wds/shard_metadata`. +2. **Scheduler**: Assigns shard identifiers to MPI ranks and writes `schedule.csv` for the toolbox runner scripts. +3. **Runner**: For every scheduled shard, loads the shard metadata, reads WebP-compressed bytes from the referenced HDF5 file, converts them to JPEG (with an optional resize), generates ten taxonomy/common-name text prompts, and emits `shard-XXXXX.tar` archives via `webdataset.TarWriter`. ## Configuration highlights @@ -33,15 +25,9 @@ tol_hybrid_to_wds: # If the lookup includes an `hdf5_path` column it will override the auto-resolved path. ``` -Set `path_to_input` to the root of the converted dataset (where the -`*_metadata.parquet` and `*_images.h5` files live) and `path_to_output_folder` -to the working directory for toolbox artifacts. +Set `path_to_input` to the root of the converted dataset (where the `*_metadata.parquet` and `*_images.h5` files live) and `path_to_output_folder` to the working directory for toolbox artifacts. -If `lookup_table_path` is provided, it can be either CSV or Parquet and must -contain at least a `uuid` column. Include a column named `hdf5_path` to override -the derived file path, or a `path` column (configurable via -`lookup_path_column`) when providing CSVs that still reference the original -`data_*.parquet` names. Only UUIDs listed in the lookup are converted. +If `lookup_table_path` is provided, it can be either CSV or Parquet and must contain at least a `uuid` column. Include a column named `hdf5_path` to override the derived file path, or a `path` column (configurable via `lookup_path_column`) when providing CSVs that still reference the original `data_*.parquet` names. Only UUIDs listed in the lookup are converted. Run the tool via the standard CLI: @@ -50,9 +36,4 @@ CONFIG_PATH=config/tol_hybrid_to_wds_example.yaml \ tree_of_life_toolbox config/tol_hybrid_to_wds_example.yaml tol_hybrid_to_wds ``` -Each shard sample contains the JPEG plus ten text prompts: -`scientific_name.txt`, `taxonomic_name.txt`, `sci.txt`, `taxon.txt`, -`taxonTag.txt`, `common_name.txt`, `com.txt`, `sci_com.txt`, -`taxon_com.txt`, and `taxonTag_com.txt`. Files ending in `_com` append -the resolved common name (or `Unknown` when no vernacular exists) to -their prompt text. +Each shard sample contains the JPEG plus ten text prompts: `scientific_name.txt`, `taxonomic_name.txt`, `sci.txt`, `taxon.txt`, `taxonTag.txt`, `common_name.txt`, `com.txt`, `sci_com.txt`, `taxon_com.txt`, and `taxonTag_com.txt`. Files ending in `_com` append the resolved common name (or `Unknown` when no vernacular exists) to their prompt text. diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py index 2acf7d8..636ae60 100644 --- a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py @@ -1,8 +1,7 @@ """ Tooling for converting HDF5-backed Tree of Life data into WebDataset shards. -Importing this module registers filter/scheduler/runner implementations with -the toolbox registry so they can be scheduled through the generic CLI. +Importing this module registers filter/scheduler/runner implementations with the toolbox registry so they can be scheduled through the generic CLI. """ from TreeOfLife_toolbox.tol_hybrid_to_wds import classes # noqa: F401 From e9adbec57ad80903c952b3fe48a6d40d2c0e6129 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Wed, 1 Jul 2026 13:57:01 -0400 Subject: [PATCH 07/10] Move the WDS companion scripts under the tool --- src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md | 7 +++++++ .../tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py | 0 .../tol_hybrid_to_wds/scripts/validate_tars.py | 0 3 files changed, 7 insertions(+) rename generate_hdf5_lookup.py => src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py (100%) rename validate_tars.py => src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py (100%) diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md index 8cba429..756b15a 100644 --- a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md @@ -37,3 +37,10 @@ tree_of_life_toolbox config/tol_hybrid_to_wds_example.yaml tol_hybrid_to_wds ``` Each shard sample contains the JPEG plus ten text prompts: `scientific_name.txt`, `taxonomic_name.txt`, `sci.txt`, `taxon.txt`, `taxonTag.txt`, `common_name.txt`, `com.txt`, `sci_com.txt`, `taxon_com.txt`, and `taxonTag_com.txt`. Files ending in `_com` append the resolved common name (or `Unknown` when no vernacular exists) to their prompt text. + +## Companion scripts + +Standalone helpers live in [`scripts/`](scripts/) beside this tool: + +- `generate_hdf5_lookup.py` builds a `uuid`-to-`h5` lookup table (Parquet) from a converted dataset, optionally sampling a subset, for use as `lookup_table_path`. +- `validate_tars.py` checks output shard integrity: every tar is readable and each sample carries the JPEG and all ten text sidecars. diff --git a/generate_hdf5_lookup.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py similarity index 100% rename from generate_hdf5_lookup.py rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py diff --git a/validate_tars.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py similarity index 100% rename from validate_tars.py rename to src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py From 45e7447cc0cd6c84faf637381f033722f6788864 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Wed, 1 Jul 2026 13:57:08 -0400 Subject: [PATCH 08/10] Relocate the standalone single-node converter out of the repo root --- tol_hybrid_to_wds.py | 334 ------------------------------------------- 1 file changed, 334 deletions(-) delete mode 100644 tol_hybrid_to_wds.py diff --git a/tol_hybrid_to_wds.py b/tol_hybrid_to_wds.py deleted file mode 100644 index cca5a79..0000000 --- a/tol_hybrid_to_wds.py +++ /dev/null @@ -1,334 +0,0 @@ -#!/usr/bin/env python -""" -Single-node converter from the hybrid Tree of Life HDF5/metadata dataset to WebDataset shards. - -Example: - python tol_hybrid_to_wds.py \ - --input-root /fs/ess/PAS2136/TreeOfLife/data \ - --lookup /fs/scratch/PAS2136/TreeOfLife_test-wds/lookup-tables/all/hdf5_lookup.parquet \ - --taxa-glob "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/*/*.parquet" \ - --output-dir /fs/scratch/PAS2136/TreeOfLife_test-wds/wds/manual \ - --metadata-glob "**/*_metadata.parquet" -""" - -from __future__ import annotations - -import argparse -import contextlib -import glob -import logging -import os -from pathlib import Path -from typing import Dict, Iterable, List, Tuple - -import h5py -import polars as pl -import webdataset as wds - -try: - from TreeOfLife_toolbox.tol_hybrid_to_wds.utils import ( - convert_webp_to_jpeg, - generate_text_files, - ) -except ImportError as exc: # pragma: no cover - script usage only - raise SystemExit( - "Unable to import TreeOfLife_toolbox.tol_hybrid_to_wds.utils. " - "Please install the toolbox package (pip install -e .) before running this script." - ) from exc - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Create WebDataset shards from HDF5 dataset") - parser.add_argument("--input-root", required=True, help="Root folder containing *_metadata.parquet + *_images.h5") - parser.add_argument( - "--metadata-glob", - default="**/*_metadata.parquet", - help="Glob pattern under --input-root to find metadata files (default: %(default)s)", - ) - parser.add_argument( - "--lookup", - required=True, - help="Lookup table (CSV or Parquet) listing UUIDs to include. " - "Must contain a 'uuid' column; optional 'hdf5_path' overrides file locations.", - ) - parser.add_argument( - "--taxa-glob", - help="Optional glob pointing to taxonomy parquet files (e.g. /path/source=*/*.parquet).", - ) - parser.add_argument( - "--output-dir", - required=True, - help="Directory where shard-XXXXX.tar files will be written.", - ) - parser.add_argument( - "--shard-size", - type=int, - default=10000, - help="Number of samples per shard (default: %(default)s)", - ) - parser.add_argument( - "--resize", - type=int, - default=224, - help="Output JPEG square size in pixels (0 keeps original WebP dimensions).", - ) - parser.add_argument( - "--shard-prefix", - default="shard", - help="Prefix for tar files (default: %(default)s)", - ) - parser.add_argument( - "--log-level", - default="INFO", - choices=["DEBUG", "INFO", "WARNING", "ERROR"], - help="Logging verbosity (default: %(default)s)", - ) - return parser.parse_args() - - -def setup_logger(level: str) -> logging.Logger: - logging.basicConfig( - level=getattr(logging, level.upper()), - format="%(asctime)s - %(levelname)s - %(message)s", - ) - return logging.getLogger("tol_hybrid_to_wds") - - -def discover_metadata_files(root: str, pattern: str) -> List[str]: - search_pattern = os.path.join(root, pattern) - files = sorted(glob.glob(search_pattern, recursive=True)) - if not files: - raise FileNotFoundError(f"No metadata parquet files found using pattern: {search_pattern}") - return files - - -def build_metadata_lazy(metadata_files: List[str]) -> pl.LazyFrame: - # Include file paths to derive HDF5 locations. - metadata_lazy = pl.scan_parquet(metadata_files, include_file_paths="metadata_file") - metadata_lazy = metadata_lazy.with_columns( - pl.col("metadata_file") - .str.replace(r"^file:", "", literal=True) - .alias("metadata_file_clean") - ) - metadata_lazy = metadata_lazy.with_columns( - pl.col("metadata_file_clean") - .str.replace("_metadata.parquet$", "_images.h5") - .alias("hdf5_path"), - pl.col("metadata_file_clean") - .str.extract(r"([^/]+)_metadata\.parquet$", group_index=1) - .alias("base_name"), - ) - schema_names = metadata_lazy.collect_schema().names() - drop_cols = [c for c in ("source", "server") if c in schema_names] - if drop_cols: - metadata_lazy = metadata_lazy.drop(drop_cols) - metadata_lazy = metadata_lazy.with_columns( - pl.col("metadata_file_clean").str.extract(r"source=([^/]+)", group_index=1).alias("source"), - pl.col("metadata_file_clean").str.extract(r"server=([^/]+)", group_index=1).alias("server"), - ) - desired_cols = [ - "uuid", - "hdf5_path", - "base_name", - "source", - "server", - "scientific_name", - "provided_common_name", - "common_name", - "kingdom", - "phylum", - "class", - "order", - "family", - "genus", - "species", - ] - cols_present = [c for c in desired_cols if c in metadata_lazy.collect_schema().names()] - metadata_lazy = metadata_lazy.select(cols_present + ["metadata_file_clean"]) - return metadata_lazy - - -def load_lookup_table(path: str) -> pl.LazyFrame: - ext = Path(path).suffix.lower() - if ext == ".parquet": - lookup = pl.scan_parquet(path) - else: - lookup = pl.scan_csv(path) - schema = lookup.collect_schema() - schema_names = schema.names() - if "uuid" not in schema: - raise ValueError(f"Lookup table {path} must contain a 'uuid' column") - keep_cols = ["uuid"] - if "hdf5_path" in schema_names: - keep_cols.append("hdf5_path") - if "path" in schema_names: - keep_cols.append("path") - return lookup.select(keep_cols) - - -def load_taxa_lazy(glob_pattern: str) -> pl.LazyFrame: - taxa_files = glob.glob(glob_pattern) - if not taxa_files: - raise FileNotFoundError(f"No taxonomy files match {glob_pattern}") - taxa_lazy = pl.scan_parquet(taxa_files, extra_columns="ignore") - schema_names = taxa_lazy.collect_schema().names() - cols = [col for col in [ - "uuid", - "common_name", - "provided_common_name", - "kingdom", - "phylum", - "class", - "order", - "family", - "genus", - "species", - ] if col in schema_names] - return taxa_lazy.select(cols) - - -def merge_metadata( - metadata_lazy: pl.LazyFrame, - lookup_lazy: pl.LazyFrame, - taxa_lazy: pl.LazyFrame | None, - logger: logging.Logger, -) -> pl.DataFrame: - joined = metadata_lazy.join(lookup_lazy, on="uuid", how="inner") - # Ensure column types align before sorting/collecting - schema_names = joined.collect_schema().names() - cast_exprs = [] - if "common_name" in schema_names: - cast_exprs.append(pl.col("common_name").cast(pl.Utf8, strict=False)) - if "provided_common_name" in schema_names: - cast_exprs.append(pl.col("provided_common_name").cast(pl.Utf8, strict=False)) - if "source" in schema_names: - cast_exprs.append(pl.col("source").cast(pl.Utf8, strict=False)) - if "server" in schema_names: - cast_exprs.append(pl.col("server").cast(pl.Utf8, strict=False)) - if cast_exprs: - joined = joined.with_columns(cast_exprs) - joined_schema = joined.collect_schema().names() - if "hdf5_path_right" in joined_schema: - joined = joined.with_columns( - pl.when(pl.col("hdf5_path_right").is_not_null()) - .then(pl.col("hdf5_path_right")) - .otherwise(pl.col("hdf5_path")) - .alias("hdf5_path") - ).drop("hdf5_path_right") - if taxa_lazy is not None: - joined = joined.join(taxa_lazy, on="uuid", how="left") - # Drop a row only when every taxonomy field is null or blank, so no degenerate text sidecars are ever written. Robust to incomplete upstream taxonomy and a no-op on inputs that are already complete. - tax_cols = [c for c in [ - "scientific_name", "common_name", "provided_common_name", "kingdom", - "phylum", "class", "order", "family", "genus", "species", - ] if c in joined.collect_schema().names()] - if tax_cols: - all_blank = pl.all_horizontal( - [pl.col(c).is_null() | (pl.col(c).str.strip_chars() == "") for c in tax_cols] - ) - joined = joined.filter(~all_blank) - joined = joined.unique(subset=["uuid"]) - result = joined.sort(["hdf5_path", "uuid"]).collect(streaming=True) - logger.info("Collected %s metadata rows", result.height) - return result - - -def chunk_records(df: pl.DataFrame, shard_size: int) -> Iterable[Tuple[int, List[dict]]]: - buffer: List[dict] = [] - shard_id = 0 - for row in df.iter_rows(named=True): - buffer.append(row) - if len(buffer) >= shard_size: - yield shard_id, buffer - buffer = [] - shard_id += 1 - if buffer: - yield shard_id, buffer - - -def process_shard( - shard_id: int, - records: List[dict], - output_dir: Path, - resize: int, - shard_prefix: str, - logger: logging.Logger, -) -> None: - output_dir.mkdir(parents=True, exist_ok=True) - shard_path = output_dir / f"{shard_prefix}-{shard_id:05d}.tar" - writer = wds.TarWriter(str(shard_path)) - hdf5_cache: Dict[str, h5py.File] = {} - written = failed = 0 - try: - for record in records: - uuid = record.get("uuid") - hdf5_path = record.get("hdf5_path") - if not uuid or not hdf5_path: - failed += 1 - continue - try: - h5_file = _open_hdf5(hdf5_path, hdf5_cache) - dataset = h5_file["images"][uuid][:] - webp_bytes = dataset.tobytes() - jpeg_bytes = convert_webp_to_jpeg(webp_bytes, resize) - - taxon_dict = { - "scientific_name": record.get("scientific_name"), - "common_name": record.get("common_name"), - "kingdom": record.get("kingdom"), - "phylum": record.get("phylum"), - "class": record.get("class"), - "order": record.get("order"), - "family": record.get("family"), - "genus": record.get("genus"), - "species": record.get("species"), - } - sample = {"__key__": uuid, "jpg": jpeg_bytes} - for ext, content in generate_text_files(taxon_dict).items(): - sample[ext] = content.encode("utf-8") - writer.write(sample) - written += 1 - except Exception as exc: # pragma: no cover - failed += 1 - logger.error("Shard %s: failed uuid %s (%s)", shard_id, uuid, exc) - logger.info("Shard %s written with %s samples (%s failures)", shard_id, written, failed) - finally: - writer.close() - for handle in hdf5_cache.values(): - with contextlib.suppress(Exception): - handle.close() - - -def _open_hdf5(path: str, cache: Dict[str, h5py.File]) -> h5py.File: - if path not in cache: - cache[path] = h5py.File(path, "r") - return cache[path] - - -def main(): - args = parse_args() - logger = setup_logger(args.log_level) - - metadata_files = discover_metadata_files(args.input_root, args.metadata_glob) - logger.info("Found %s metadata parquet files", len(metadata_files)) - - metadata_lazy = build_metadata_lazy(metadata_files) - lookup_lazy = load_lookup_table(args.lookup) - taxa_lazy = load_taxa_lazy(args.taxa_glob) if args.taxa_glob else None - - metadata = merge_metadata(metadata_lazy, lookup_lazy, taxa_lazy, logger) - - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - last_shard = -1 - for shard_id, records in chunk_records(metadata, args.shard_size): - last_shard = shard_id - process_shard(shard_id, records, output_dir, args.resize, args.shard_prefix, logger) - - total_shards = (last_shard + 1) if metadata.height > 0 else 0 - logger.info("Completed WebDataset creation: %s shards written", total_shards) - - -if __name__ == "__main__": - main() From cf994068f37afcd9b9adb215c1d0b3e9e262dfa6 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Wed, 1 Jul 2026 13:57:08 -0400 Subject: [PATCH 09/10] Relocate the one-off taxon_com audit out of the repo root --- wds_taxoncom_audit.py | 231 ------------------------------------------ 1 file changed, 231 deletions(-) delete mode 100644 wds_taxoncom_audit.py diff --git a/wds_taxoncom_audit.py b/wds_taxoncom_audit.py deleted file mode 100644 index 2a79eb7..0000000 --- a/wds_taxoncom_audit.py +++ /dev/null @@ -1,231 +0,0 @@ -#!/usr/bin/env python3 -""" -wds_taxoncom_audit.py - -One-node audit / training-pipeline analog to pinpoint where samples are missing `taxon_com.txt`. - -Logs: - - wds_audit_logs/missing_taxon_com.workerXYZ.jsonl (samples missing taxon_com.txt after grouping) - - wds_audit_logs/bad_records.workerXYZ.jsonl (non-empty dicts missing fname/data) - - wds_audit_logs/stats.workerXYZ.json (summary per worker) - -Usage: - python wds_taxoncom_audit.py \ - --shards "/path/to/shards/shard-{00000..99999}.tar" \ - --workers 96 \ - --out-dir wds_audit_logs -""" - -import argparse -import json -import logging -import os -from datetime import datetime - -import webdataset as wds -from webdataset.tariterators import base_plus_ext, url_opener, tar_file_expander, valid_sample - - -# ---------------------------- -# Logging / utilities -# ---------------------------- - -logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") - - -def now_ts(): - return datetime.now().isoformat(timespec="seconds") - - -def worker_id(): - """ - Best-effort worker id detection. - - Note: torch dataloader workers don't reliably export a standard env var. - We'll use these if present, else -1. - """ - for k in ("WDS_WORKER", "WORKER", "LOCAL_RANK", "RANK"): - v = os.environ.get(k) - if v is not None and str(v).lstrip("-").isdigit(): - return int(v) - return -1 - - -def append_jsonl(path, rec): - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "a") as f: - f.write(json.dumps(rec) + "\n") - - -# ---------------------------- -# WDS "nothrow" tar parsing -# ---------------------------- - -def log_and_continue(exn): - logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.") - return True - - -# These globals get set in main() so that nested functions can log without threading args everywhere. -OUT_DIR = "wds_audit_logs" - - -def log_bad_record(filesample, reason="missing_fname_or_data"): - wid = worker_id() - path = os.path.join(OUT_DIR, f"bad_records.worker{wid:03d}.jsonl") - rec = { - "ts": now_ts(), - "reason": reason, - "url": filesample.get("__url__", filesample.get("url")), - "keys": sorted(list(filesample.keys())), - } - append_jsonl(path, rec) - - -def log_missing_taxon_com(sample, reason="missing taxon_com.txt (pre-rename)"): - wid = worker_id() - path = os.path.join(OUT_DIR, f"missing_taxon_com.worker{wid:03d}.jsonl") - rec = { - "ts": now_ts(), - "reason": reason, - "url": sample.get("__url__", sample.get("url")), - "key": sample.get("__key__"), - "present_suffixes": sorted([k for k in sample.keys() if not k.startswith("__")]), - } - append_jsonl(path, rec) - - -def group_by_keys_nothrow(data, keys=base_plus_ext, lcase=True, suffixes=None, handler=None): - """ - Robust grouping that tolerates: - - empty dicts (often sentinel-ish): skip silently - - malformed dicts missing fname/data: log (if non-empty) and skip - - This is meant to be a faithful-ish version of the modified training pipeline, - while remaining diagnostic. - """ - current_sample = None - - for filesample in data: - if not isinstance(filesample, dict): - logging.warning(f"Skipping non-dict filesample: {type(filesample)}") - continue - - # IMPORTANT: empty dicts can appear due to iterator plumbing; skip quietly. - if not filesample: - continue - - # Defensive: sometimes expander yields dicts without fname/data - if "fname" not in filesample or "data" not in filesample: - # Log only non-empty dicts (otherwise it's useless noise). - log_bad_record(filesample, reason="malformed_tar_record_missing_fname_or_data") - continue - - fname, value = filesample["fname"], filesample["data"] - prefix, suffix = keys(fname) - if prefix is None: - continue - if lcase: - suffix = suffix.lower() - - # Same collision handling rationale as training data.py - if current_sample is None or prefix != current_sample["__key__"] or suffix in current_sample: - if valid_sample(current_sample): - yield current_sample - current_sample = dict(__key__=prefix, __url__=filesample.get("__url__")) - - if suffixes is None or suffix in suffixes: - current_sample[suffix] = value - - if valid_sample(current_sample): - yield current_sample - - -def tarfile_to_samples_nothrow(src, handler=log_and_continue): - streams = url_opener(src, handler=handler) - files = tar_file_expander(streams, handler=handler) - return group_by_keys_nothrow(files, handler=handler) - - -# ---------------------------- -# Filters for audit -# ---------------------------- - -def has_image(sample): - return ("png" in sample or "jpg" in sample or "jpeg" in sample or "webp" in sample) - - -def has_any_txt(sample): - # training's filter_no_caption_or_no_image checks "any('txt' in key for key in sample)" - return any("txt" in k for k in sample.keys()) - - -def require_taxon_com_txt(sample): - """ - Audit hook: record and drop samples that lack taxon_com.txt at the grouped-sample stage. - - This is the strongest check for grouping/boundary/corruption-induced misses. - """ - if "taxon_com.txt" not in sample: - log_missing_taxon_com(sample, reason="missing taxon_com.txt (pre-rename)") - return False - return True - - -# ---------------------------- -# Main audit -# ---------------------------- - -def audit(shards_pattern: str, workers: int, log_every: int, limit_kept: int): - shards = list(wds.shardlists.expand_urls(shards_pattern)) - logging.info(f"Expanded to {len(shards)} shard URLs") - - # Build a training-ish pipeline (single node; split_by_node not needed here) - dataset = wds.DataPipeline( - wds.SimpleShardList(shards), - wds.split_by_worker, - tarfile_to_samples_nothrow, - wds.select(lambda s: has_any_txt(s) and has_image(s)), - wds.select(lambda s: require_taxon_com_txt(s)), - # We intentionally avoid decode/tokenizer here: this audit is about tar member presence. - ) - - loader = wds.WebLoader( - dataset, - batch_size=None, - shuffle=False, - num_workers=workers, - persistent_workers=True, - ) - - kept = 0 - for sample in loader: - kept += 1 - if log_every and kept % log_every == 0: - logging.info(f"Kept {kept} samples (after filters)") - - if limit_kept and kept >= limit_kept: - break - - logging.info(f"Done. Kept {kept} samples (after filters). Logs in {OUT_DIR}/") - - -def main(): - global OUT_DIR - - ap = argparse.ArgumentParser() - ap.add_argument("--shards", required=True, help='Braceexpand pattern like ".../shard-{00000..99999}.tar"') - ap.add_argument("--workers", type=int, default=96) - ap.add_argument("--out-dir", default="wds_audit_logs") - ap.add_argument("--log-every", type=int, default=200000) - ap.add_argument("--limit", type=int, default=0, help="Stop after this many kept samples (0 = no limit)") - args = ap.parse_args() - - OUT_DIR = args.out_dir - os.makedirs(OUT_DIR, exist_ok=True) - - audit(args.shards, args.workers, args.log_every, args.limit) - - -if __name__ == "__main__": - main() From f5154fbbad1f92bb9e29f34514a028dcefd15a18 Mon Sep 17 00:00:00 2001 From: Matthew Thompson Date: Wed, 1 Jul 2026 13:57:08 -0400 Subject: [PATCH 10/10] Update the WDS docs for the relocated scripts --- docs/hybrid_to_webdataset_README.md | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/docs/hybrid_to_webdataset_README.md b/docs/hybrid_to_webdataset_README.md index eca5ba3..4417d89 100644 --- a/docs/hybrid_to_webdataset_README.md +++ b/docs/hybrid_to_webdataset_README.md @@ -20,17 +20,19 @@ All three are joined by `uuid`: 1. **Hybrid dataset**: the root holding `*_images.h5` (WebP bytes under an `/images/` group) and the paired `*_metadata.parquet`. 2. **Resolved taxonomy**: Parquet providing `scientific_name`, `common_name`, and the `kingdom` through `species` ranks per `uuid` (the text-label source). -3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed UUIDs are converted. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an optional `hdf5_path` (or `h5_file`) column overrides it. +3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed UUIDs are converted. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an optional `hdf5_path` (or `h5_file`) column overrides it. To build a lookup for a UUID subset from a converted dataset, use [`generate_hdf5_lookup.py`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py). ## Prerequisites +The pipeline runs on a cluster through Slurm, using Spark for the filter stage and MPI for the workers. This holds for any run size: for a small or low-resource job you scale the pipeline down through the config (see below) rather than switching to a different code path, so Spark and MPI are always required. + Install the toolbox (see the top-level [installation instructions](../README.md#installation-instructions)): ```bash pip install -e . ``` -Option A (single node) needs only the Python dependencies that installs. Option B (distributed) runs on a cluster, and its Slurm scripts encode site-specific choices you must adapt: +The Slurm scripts encode site-specific choices you must adapt: - An **MPI** implementation with `mpi4py` installed against it. The workers run under MPI; if `mpi4py` is built against a different MPI than the one loaded at runtime, they fail to load it. - A **Spark** runtime for the filter stage. @@ -38,22 +40,7 @@ Option A (single node) needs only the Python dependencies that installs. Option Edit `scripts/tools_*.slurm` to load your cluster's MPI and Spark modules, target your partition, and activate your environment. As committed they target one cluster (OSC Cardinal) as a worked example to replace with your own. -## Option A: single node (quick / small subsets) - -Best for testing or modest UUID lists. No Spark/MPI required: - -```bash -python tol_hybrid_to_wds.py \ - --input-root /path/to/hybrid_dataset \ - --lookup /path/to/uuid_lookup.parquet \ - --taxa-glob "/path/to/resolved_taxa/source=*/*.parquet" \ - --output-dir /path/to/wds \ - --resize 224 --shard-size 10000 -``` - -(Generate a `uuid`-to-`h5` lookup for a UUID subset with [`generate_hdf5_lookup.py`](../generate_hdf5_lookup.py).) - -## Option B: distributed (full dataset, Slurm) +## Running the pipeline The toolbox runner reads one config and submits the whole pipeline (Spark filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. Start from [`config/tol_hybrid_to_wds_example.yaml`](../config/tol_hybrid_to_wds_example.yaml) and set, for your run: @@ -61,7 +48,7 @@ The toolbox runner reads one config and submits the whole pipeline (Spark filter - `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the toolbox working directory); - in the `tol_hybrid_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` (required for the taxonomy/common-name prompts), and `lookup_table_path` (a parquet/CSV with a `uuid` column selecting which images to convert). -See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for the full key reference. +For a small or low-resource run, scale `tools_parameters` down (for example `num_workers: 1`, `max_nodes: 1`, `workers_per_node: 1`); the execution path is identical. See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for the full key reference. Then run the pipeline with a single command: @@ -73,4 +60,4 @@ The runner reads the config, exports the environment it implies (account, node/w ## Validating output -`validate_tars.py` checks shard integrity, and `wds_taxoncom_audit.py` audits the taxonomy/common-name prompts across shards. +[`validate_tars.py`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py) checks shard integrity: that every tar is readable and each sample carries the JPEG and all ten text sidecars.