From ab8991d5ace6e0aac0cfffe8436ad67a4c2292ea Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Mon, 31 Aug 2026 10:29:52 -0400 Subject: [PATCH 1/3] Add TOL-200M batch embedding configs and SLURM jobs template (BioCLIP 2 & BioCLIP 2.5) --- .../configs/embed_bioclip_2.yaml | 27 +++++++ .../configs/embed_bioclip_2_5.yaml | 30 ++++++++ .../scripts/build_chunks.py | 65 +++++++++++++++++ .../scripts/count_images.py | 70 +++++++++++++++++++ .../scripts/embed.slurm | 65 +++++++++++++++++ .../scripts/embed_array.slurm | 66 +++++++++++++++++ 6 files changed, 323 insertions(+) create mode 100644 processing/tol200m_image_embeddings/configs/embed_bioclip_2.yaml create mode 100644 processing/tol200m_image_embeddings/configs/embed_bioclip_2_5.yaml create mode 100755 processing/tol200m_image_embeddings/scripts/build_chunks.py create mode 100755 processing/tol200m_image_embeddings/scripts/count_images.py create mode 100644 processing/tol200m_image_embeddings/scripts/embed.slurm create mode 100644 processing/tol200m_image_embeddings/scripts/embed_array.slurm diff --git a/processing/tol200m_image_embeddings/configs/embed_bioclip_2.yaml b/processing/tol200m_image_embeddings/configs/embed_bioclip_2.yaml new file mode 100644 index 0000000..efb0d8e --- /dev/null +++ b/processing/tol200m_image_embeddings/configs/embed_bioclip_2.yaml @@ -0,0 +1,27 @@ +# BioCLIP 2 (ViT-L/14, 768-dim) - HDF5 batch embedding config +# for hpc-inference open_clip_embed (--input_type hdf5). +# +# Same pipeline as embed_bioclip_2_5.yaml, different weights. ViT-L is lighter, +# so a larger batch and full worker count keep the GPU fed (On H100, this model tends to +# be data-loader bound rather than compute bound). +# +# Outputs RAW float32 embeddings. + +models: + bioclip_2: + name: hf-hub:imageomics/bioclip-2 + pretrained: null + +batch_size: 128 +num_workers: 24 +prefetch_factor: 16 + +group_name: images # HDF5 group holding uuid -> image bytes +color_mode: RGB +validate_images: false + +evenly_distribute: true # LPT file assignment across ranks (by file size) +stagger: false + +max_rows_per_file: 500000 +out_prefix: bioclip_2 # output column will be emb_bioclip_2 diff --git a/processing/tol200m_image_embeddings/configs/embed_bioclip_2_5.yaml b/processing/tol200m_image_embeddings/configs/embed_bioclip_2_5.yaml new file mode 100644 index 0000000..2d7dd76 --- /dev/null +++ b/processing/tol200m_image_embeddings/configs/embed_bioclip_2_5.yaml @@ -0,0 +1,30 @@ +# BioCLIP 2.5 Huge (ViT-H/14, 1024-dim) - HDF5 batch embedding config +# for hpc-inference open_clip_embed (--input_type hdf5). +# +# Compute-bound model: the GPU forward pass dominates, so worker/prefetch tuning +# matters less than for smaller models. Sized for 1x H100 94GB w/ 26 CPU cores. +# +# NOTE: the original production run additionally used bf16 autocast with +# in-pipeline L2-normalization and float16 storage (features not yet in the +# released hpc-inference package). This config runs the released fp32 path and +# outputs RAW float32 embeddings; normalization + float16 casting happen in the +# standardization step (see ../README.md). bf16-vs-fp32 cosine: 0.9998 mean. + +models: + bioclip_2_5: + name: hf-hub:imageomics/bioclip-2.5-vith14 + pretrained: null + +batch_size: 32 +num_workers: 24 +prefetch_factor: 8 + +group_name: images # HDF5 group holding uuid -> image bytes +color_mode: RGB +validate_images: false + +evenly_distribute: true # LPT file assignment across ranks (by file size) +stagger: false + +max_rows_per_file: 500000 +out_prefix: bioclip_2_5 # output column will be emb_bioclip_2_5 diff --git a/processing/tol200m_image_embeddings/scripts/build_chunks.py b/processing/tol200m_image_embeddings/scripts/build_chunks.py new file mode 100755 index 0000000..d54153f --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/build_chunks.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python +"""Build size-balanced array chunks from a per-file image-count parquet. + +Greedy-packs the h5 files into N chunks each holding ~target_images, so every array +task does ~the same amount of work (embedding time is per-image) regardless of the +wide per-file variance in gbif. Deterministic: same counts + target -> same chunks, +so resume keeps the same task<->file mapping. Idempotent: refuses to overwrite an +existing chunk set (delete chunks_manifest.parquet to rebuild). + +Usage: python build_chunks.py +Writes: /task_/filelist.txt and /chunks_manifest.parquet +Prints: NTASKS (number of chunks) to stdout. +""" +import os +import sys + +import pyarrow as pa +import pyarrow.parquet as pq + + +def main() -> None: + counts_parquet, out_root, target = sys.argv[1], sys.argv[2], int(sys.argv[3]) + manifest = os.path.join(out_root, "chunks_manifest.parquet") + + if os.path.exists(manifest): # already built -> reuse (keeps task<->file mapping stable) + print(pq.read_metadata(manifest).num_rows) + return + + t = pq.read_table(counts_parquet).to_pydict() + files = [(p, n) for p, n in zip(t["h5_path"], t["n_images"]) if n > 0] + total = sum(n for _, n in files) + ntasks = max(1, round(total / target)) + per = total / ntasks # even target per chunk + + chunks, cur, acc = [[]], 0, 0 + for path, n in files: + chunks[cur].append((path, n)) + acc += n + if acc >= per and cur < ntasks - 1: + chunks.append([]) + cur += 1 + acc = 0 + chunks = [c for c in chunks if c] # drop any trailing empty + + rows_id, rows_nf, rows_ni = [], [], [] + for i, ch in enumerate(chunks): + d = os.path.join(out_root, f"task_{i}") + os.makedirs(d, exist_ok=True) + with open(os.path.join(d, "filelist.txt"), "w") as fh: + fh.write("\n".join(p for p, _ in ch) + "\n") + rows_id.append(i) + rows_nf.append(len(ch)) + rows_ni.append(sum(n for _, n in ch)) + + table = pa.table({"task_id": rows_id, "n_files": rows_nf, "n_images": rows_ni}) + pq.write_table(table, manifest, compression="zstd") + sys.stderr.write( + f"Built {len(chunks)} chunks from {len(files)} files ({total:,} images); " + f"per-chunk images min={min(rows_ni):,} max={max(rows_ni):,} target={int(per):,}\n" + ) + print(len(chunks)) + + +if __name__ == "__main__": + main() diff --git a/processing/tol200m_image_embeddings/scripts/count_images.py b/processing/tol200m_image_embeddings/scripts/count_images.py new file mode 100755 index 0000000..8782dba --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/count_images.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python +"""Count images per h5 file for a source, in parallel, and write a parquet table. + +Work (embedding time) is per-image, so this per-file image count is what the array +driver balances chunks on. Reads the paired *_metadata.parquet footer (fast, no data +read); falls back to the h5 'images' group size if the parquet is missing. + +Usage: python count_images.py [num_procs] +Output: with columns [h5_path: str, n_images: int64], sorted by h5_path. +Run via a SLURM cpu job. +""" +import glob +import os +import sys +from concurrent.futures import ProcessPoolExecutor + +import pyarrow as pa +import pyarrow.parquet as pq + + +def count_one(h5_path: str) -> tuple: + parquet = h5_path.replace("_images.h5", "_metadata.parquet") + try: + return h5_path, pq.read_metadata(parquet).num_rows + except Exception: + try: + import h5py + with h5py.File(h5_path, "r") as f: + return h5_path, len(f["images"].keys()) + except Exception as e: + sys.stderr.write(f"WARN could not count {h5_path}: {e}\n") + return h5_path, -1 + + +def main() -> None: + target = sys.argv[1] + out_parquet = sys.argv[2] + nprocs = int(sys.argv[3]) if len(sys.argv) > 3 else (os.cpu_count() or 8) + + h5s = sorted( + glob.glob(os.path.join(target, "**", "*.h5"), recursive=True) + + glob.glob(os.path.join(target, "**", "*.hdf5"), recursive=True) + ) + sys.stderr.write(f"Counting {len(h5s)} h5 files with {nprocs} procs...\n") + + with ProcessPoolExecutor(max_workers=nprocs) as ex: + rows = list(ex.map(count_one, h5s, chunksize=16)) # kept in sorted(h5s) order + + paths = [h for h, _ in rows] + counts = [n for _, n in rows] + bad = sum(1 for n in counts if n < 0) + total = sum(n for n in counts if n >= 0) + + table = pa.table( + {"h5_path": pa.array(paths, pa.string()), + "n_images": pa.array(counts, pa.int64())} + ) + tmp = out_parquet + ".tmp" + pq.write_table(table, tmp, compression="zstd") + os.replace(tmp, out_parquet) + sys.stderr.write( + f"Wrote {out_parquet}: {len(rows)} files, {total:,} images" + + (f", {bad} FAILED\n" if bad else "\n") + ) + if bad: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/processing/tol200m_image_embeddings/scripts/embed.slurm b/processing/tol200m_image_embeddings/scripts/embed.slurm new file mode 100644 index 0000000..a0d4ec2 --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/embed.slurm @@ -0,0 +1,65 @@ +#!/bin/bash +# ============================================================================= +# TreeOfLife HDF5 batch embedding - multi-rank GPU job (one task = one GPU rank). +# For sources with a moderate file count (bioscan, eol, fathomnet). For gbif-scale +# sources use count_images.py + build_chunks.py + embed_array.slurm instead. +# +# Submit with the RUN_* variables exported, e.g.: +# export RUN_CONFIG=configs/embed_bioclip_2_5.yaml +# export RUN_TARGET_DIR=/source=bioscan +# export RUN_OUTPUT_DIR=/model=bioclip_2_5/source=bioscan +# sbatch --nodes=4 --time=02:00:00 embed.slurm # 4 nodes x 1 GPU = 4 ranks +# +# Rank sizing: ranks split whole files, and each rank's dataloader workers also +# split whole files, so keep files_per_rank >= num_workers (else workers idle) +# and world_size <= num_files (else a rank gets nothing). +# ============================================================================= +#SBATCH --job-name=tol_embed +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 # 1 task == 1 rank == 1 GPU +#SBATCH --gpus-per-task=1 +#SBATCH --cpus-per-task=26 # match the cluster's per-GPU CPU share +#SBATCH --mem=234G # match the cluster's per-GPU memory share +#SBATCH --partition=gpu +#SBATCH --time=02:00:00 +#SBATCH --account=YOUR_ACCOUNT +#SBATCH --mail-type=END,FAIL +#SBATCH --output=embed_%x_%j.out +#SBATCH --error=embed_%x_%j.err + +set -euo pipefail + +: "${RUN_CONFIG:?set RUN_CONFIG (path to embed_*.yaml)}" +: "${RUN_TARGET_DIR:?set RUN_TARGET_DIR (source=... dir with .h5 files)}" +: "${RUN_OUTPUT_DIR:?set RUN_OUTPUT_DIR (embedding output dir)}" +RUN_FILE_LIST="${RUN_FILE_LIST:-}" # optional: restrict to a subset of files + +module load cuda # match your venv's torch CUDA major +source /path/to/venv/bin/activate + +# Cache compile artifacts + HF weights on scratch (avoids recompiling / +# re-downloading the model every job). +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-/path/to/scratch/cache/torchinductor}" +export HF_HOME="${HF_HOME:-/path/to/scratch/cache/huggingface}" +export OMP_NUM_THREADS=1 # keep dataloader workers from oversubscribing cores +export HDF5_USE_FILE_LOCKING=FALSE # read-only access: safe, skips flock on GPFS +export SRUN_CPUS_PER_TASK="${SLURM_CPUS_PER_TASK}" # Slurm 23.11+ srun guard +mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${HF_HOME}" "${RUN_OUTPUT_DIR}" + +echo "Job ${SLURM_JOB_ID}: ntasks=${SLURM_NTASKS} config=${RUN_CONFIG}" +echo " target=${RUN_TARGET_DIR}" +echo " output=${RUN_OUTPUT_DIR}" +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true + +EXTRA_ARGS=() +[[ -n "${RUN_FILE_LIST}" ]] && EXTRA_ARGS+=(--file_list "${RUN_FILE_LIST}") + +srun python -m hpc_inference.inference.embed.open_clip_embed \ + "${RUN_TARGET_DIR}" \ + "${RUN_OUTPUT_DIR}" \ + --input_type hdf5 \ + --config "${RUN_CONFIG}" \ + "${EXTRA_ARGS[@]}" + +echo "Done. Embeddings + profile_results under: ${RUN_OUTPUT_DIR}" +exit 0 diff --git a/processing/tol200m_image_embeddings/scripts/embed_array.slurm b/processing/tol200m_image_embeddings/scripts/embed_array.slurm new file mode 100644 index 0000000..3ac8a9c --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/embed_array.slurm @@ -0,0 +1,66 @@ +#!/bin/bash +# ============================================================================= +# TreeOfLife HDF5 batch embedding - SLURM ARRAY task (1 GPU per task), for +# sources with many files of highly variable size (gbif: ~24.6K files, 1-20K +# images each). +# +# Prerequisite (one CPU job / login-safe for small sources): +# python count_images.py /source=gbif file_counts.parquet +# NTASKS=$(python build_chunks.py file_counts.parquet /model=/source=gbif 7000000) +# +# Then submit (cap concurrency to be a good citizen on a shared account): +# export RUN_CONFIG=configs/embed_bioclip_2_5.yaml +# export RUN_TARGET_DIR=/source=gbif +# export RUN_OUT_ROOT=/model=bioclip_2_5/source=gbif +# sbatch --array=0-$((NTASKS-1))%6 --time=06:00:00 embed_array.slurm +# +# Each task N reads its precomputed image-count-balanced chunk +# (task_/filelist.txt) and writes task-scoped output under task_/ so +# tasks never collide. Budget --time for the WORST chunk: chunks with many +# small files run slower per image (per-file HDF5 open overhead). +# ============================================================================= +#SBATCH --job-name=tol_embed_array +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 # 1 task == world_size 1 == rank 0 +#SBATCH --gpus-per-task=1 +#SBATCH --cpus-per-task=26 +#SBATCH --mem=234G +#SBATCH --partition=gpu +#SBATCH --account=YOUR_ACCOUNT +#SBATCH --mail-type=FAIL +#SBATCH --output=embed_array_%A_%a.out +#SBATCH --error=embed_array_%A_%a.err + +set -euo pipefail + +: "${RUN_CONFIG:?set RUN_CONFIG}" +: "${RUN_TARGET_DIR:?set RUN_TARGET_DIR}" +: "${RUN_OUT_ROOT:?set RUN_OUT_ROOT}" +: "${SLURM_ARRAY_TASK_ID:?this script must run as a job array}" + +module load cuda +source /path/to/venv/bin/activate + +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-/path/to/scratch/cache/torchinductor}" +export HF_HOME="${HF_HOME:-/path/to/scratch/cache/huggingface}" +export OMP_NUM_THREADS=1 +export HDF5_USE_FILE_LOCKING=FALSE +export SRUN_CPUS_PER_TASK="${SLURM_CPUS_PER_TASK}" +mkdir -p "${TORCHINDUCTOR_CACHE_DIR}" "${HF_HOME}" + +N="${SLURM_ARRAY_TASK_ID}" +TASK_DIR="${RUN_OUT_ROOT}/task_${N}" +CHUNK_LIST="${TASK_DIR}/filelist.txt" +[[ -f "${CHUNK_LIST}" ]] || { echo "ERROR: no chunk list at ${CHUNK_LIST} (run build_chunks.py first)"; exit 1; } + +echo "Array task ${N}: $(wc -l < "${CHUNK_LIST}") files -> ${TASK_DIR}" + +srun python -m hpc_inference.inference.embed.open_clip_embed \ + "${RUN_TARGET_DIR}" \ + "${TASK_DIR}" \ + --input_type hdf5 \ + --config "${RUN_CONFIG}" \ + --file_list "${CHUNK_LIST}" + +echo "Task ${N} done -> ${TASK_DIR}" +exit 0 From 23dcc0913b9b385961e1219d677c59fcb5f972da Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Mon, 31 Aug 2026 10:30:57 -0400 Subject: [PATCH 2/3] Add embedding standardization and validation scripts (TOL catalog join, taxa sorting, optional normalization, float16 contract) --- .../scripts/build_spine.py | 67 ++++++++ .../scripts/standardize_config.py | 158 ++++++++++++++++++ .../scripts/validate_config.py | 122 ++++++++++++++ 3 files changed, 347 insertions(+) create mode 100644 processing/tol200m_image_embeddings/scripts/build_spine.py create mode 100644 processing/tol200m_image_embeddings/scripts/standardize_config.py create mode 100644 processing/tol200m_image_embeddings/scripts/validate_config.py diff --git a/processing/tol200m_image_embeddings/scripts/build_spine.py b/processing/tol200m_image_embeddings/scripts/build_spine.py new file mode 100644 index 0000000..20555a1 --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/build_spine.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python +"""Build the metadata SPINE: the published bioclip-2_float16 config minus `emb`, +plus a global `_pos` row index. + +The spine defines the row set, row order, and 15 metadata columns shared by ALL +embedding configs (it equals the TreeOfLife-200M catalog restricted to its +233,055,986 rows, in the published global taxonomic sort). Deriving it from the +published files (instead of re-sorting the catalog) preserves row order exactly: +the sort keys have large tie groups, so a fresh sort would not reproduce the +published order within ties. + +One output file per published file, physical row order preserved (pyarrow +read_table/write_table are order-preserving). `_pos` is the global row index +computed from cumulative per-file row counts. + +Usage: python build_spine.py [num_procs] + PUBLISHED_DIR: local copy of the published bioclip-2_float16/ files + SPINE_DIR: output dir (train-NNNNN-of-NNNNN.parquet, uuid + 15 meta + _pos) +Run via a SLURM cpu job. +""" +import glob +import os +import sys +from concurrent.futures import ProcessPoolExecutor + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +EXPECT_ROWS = 233_055_986 + + +def build_one(args): + src, dst, offset = args + t = pq.read_table(src).drop_columns(["emb"]) + pos = pa.array(np.arange(offset, offset + t.num_rows, dtype=np.int64)) + t = t.append_column("_pos", pos) + pq.write_table(t, dst, compression="zstd", row_group_size=50_000) + return t.num_rows + + +def main(): + published_dir, spine_dir = sys.argv[1], sys.argv[2] + nprocs = int(sys.argv[3]) if len(sys.argv) > 3 else int(os.environ.get("SLURM_CPUS_PER_TASK", 16)) + os.makedirs(spine_dir, exist_ok=True) + + srcs = sorted(glob.glob(os.path.join(published_dir, "train-*.parquet"))) + if not srcs: + sys.exit(f"no train-*.parquet under {published_dir}") + + # cumulative offsets from actual per-file row counts (footer reads, fast) + counts = [pq.read_metadata(f).num_rows for f in srcs] + offsets = np.concatenate(([0], np.cumsum(counts)[:-1])) + jobs = [(src, os.path.join(spine_dir, os.path.basename(src)), int(off)) + for src, off in zip(srcs, offsets)] + + with ProcessPoolExecutor(max_workers=nprocs) as ex: + total = sum(ex.map(build_one, jobs)) + + status = "OK" if total == EXPECT_ROWS else "MISMATCH" + print(f"spine rows: {total:,} across {len(srcs)} files (expect {EXPECT_ROWS:,}) -> {status}") + if total != EXPECT_ROWS: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/processing/tol200m_image_embeddings/scripts/standardize_config.py b/processing/tol200m_image_embeddings/scripts/standardize_config.py new file mode 100644 index 0000000..51b599b --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/standardize_config.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python +"""Standardize raw embeddings into one published config: spine metadata + one +model's embeddings, in spine order, in the published parquet contract. + +Pipeline: DuckDB inner-joins the spine (uuid + 15 metadata + _pos) to the raw +per-rank embedding parquets by uuid (raw rows not in the catalog are dropped), +orders by _pos, and streams 50K-row batches. A thread pool of writers then, per +output file: optionally L2-normalizes each embedding in float32 (--normalize; +used for bioclip-2.5, NOT for bioclip-2), casts to fixed_size_list[dim] +(lossless for normalized values), and writes the contract parquet: +ZSTD-3, 50K-row row groups, statistics + page indexes, 9 sorting_columns +(nulls_first=False), train-NNNNN-of-MMMMM.parquet with rows_per_file rows. + +rows_per_file targets ~500 MB/file: 500_000 for 512/768-dim, 250_000 for 1024-dim. + +Usage: + python standardize_config.py --dim D \ + --spine "/train-*.parquet" \ + --emb-glob "/model=/source=*/embeddings/rank_*/*.parquet" \ + [--emb-glob "/model=/source=gbif/task_*/embeddings/rank_*/*.parquet"] \ + --out [--normalize] [--threads T] [--write-threads W] [--mem-limit G] + +Run via a SLURM cpu job; the global join+sort spills through DuckDB temp space +(~470 GB peak for the 1024-dim config -> hugemem-class node; ~340 GB for 512-dim). +""" +import argparse +import os +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +import duckdb +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +BATCH = 50_000 # == published row group size + +SORT_KEYS = ["source_dataset", "kingdom", "phylum", "class", "order", "family", + "genus", "species", "common_name"] +META = ["source_dataset", "source_id", "kingdom", "phylum", "class", "order", "family", + "genus", "species", "scientific_name", "common_name", "publisher", + "basisOfRecord", "identifier", "img_type"] + + +def transform_emb(chunked, dim, normalize, emb_type): + """list column -> (optionally L2-normalized) fixed_size_list[dim].""" + arr = chunked.combine_chunks() if isinstance(chunked, pa.ChunkedArray) else chunked + fsl = arr.cast(pa.list_(pa.float32(), dim)) + mat = np.reshape(fsl.values.to_numpy(zero_copy_only=False), (len(arr), dim)) + if normalize: + norms = np.linalg.norm(mat, axis=1, keepdims=True) + norms[norms == 0] = 1.0 # keep all-zero vectors as-is + mat = mat / norms + flat = pa.array(mat.astype(np.float16).ravel()) + return pa.FixedSizeListArray.from_arrays(flat, dim).cast(emb_type) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("model_col", help="raw embedding column, e.g. emb_bioclip_2_5") + ap.add_argument("config_name", help="published config dir name") + ap.add_argument("rows_per_file", type=int) + ap.add_argument("--dim", type=int, required=True) + ap.add_argument("--spine", required=True, help="spine glob (from build_spine.py)") + ap.add_argument("--emb-glob", action="append", required=True, + help="raw embedding parquet glob(s); repeatable") + ap.add_argument("--out", required=True, help="staging root; writes //") + ap.add_argument("--normalize", action="store_true", + help="L2-normalize embeddings in float32 before the float16 cast") + ap.add_argument("--threads", type=int, default=int(os.environ.get("SLURM_CPUS_PER_TASK", 32))) + ap.add_argument("--write-threads", type=int, default=0, help="0 -> max(4, threads-4)") + ap.add_argument("--mem-limit", default="400GB") + ap.add_argument("--tmp", default=os.environ.get("TMPDIR", "/tmp"), help="DuckDB spill dir") + args = ap.parse_args() + + assert args.rows_per_file % BATCH == 0, f"rows_per_file must be a multiple of {BATCH}" + batches_per_file = args.rows_per_file // BATCH + write_threads = args.write_threads or max(4, args.threads - 4) + out_dir = os.path.join(args.out, args.config_name) + os.makedirs(out_dir, exist_ok=True) + + emb_type = pa.list_(pa.field("element", pa.float16()), args.dim) # published emb type + + con = duckdb.connect() + con.execute(f"SET memory_limit='{args.mem_limit}'; SET threads={args.threads}; " + f"SET temp_directory='{args.tmp}'; SET preserve_insertion_order=true;") + + glist = "['" + "','".join(args.emb_glob) + "']" + meta_sel = ", ".join(f's."{c}"' for c in META) + base = f""" + WITH sp AS (SELECT * FROM read_parquet('{args.spine}')), + emb AS (SELECT uuid, "{args.model_col}" AS emb FROM read_parquet({glist})) + """ + total_expected = con.execute( + base + "SELECT COUNT(*) FROM sp s JOIN emb e USING (uuid)").fetchone()[0] + n_files = (total_expected + args.rows_per_file - 1) // args.rows_per_file + print(f"expected {total_expected:,} rows -> {n_files} files @ {args.rows_per_file:,}; " + f"normalize={args.normalize}; duckdb threads={args.threads}, " + f"write threads={write_threads}", flush=True) + + reader = con.execute(base + f""" + SELECT s.uuid, e.emb AS emb, {meta_sel} + FROM sp s JOIN emb e USING (uuid) ORDER BY s._pos + """).fetch_record_batch(rows_per_batch=BATCH) + + raw_schema = None + sc = None + + def write_file(idx, raw_table): + # transform runs here so normalization + cast parallelize across the pool + ei = raw_table.schema.get_field_index("emb") + emb = transform_emb(raw_table.column(ei), args.dim, args.normalize, emb_type) + table = raw_table.set_column(ei, pa.field("emb", emb_type), emb) + path = os.path.join(out_dir, f"train-{idx:05d}-of-{n_files:05d}.parquet") + pq.write_table(table, path, compression="zstd", compression_level=3, + write_statistics=True, write_page_index=True, + sorting_columns=sc, row_group_size=BATCH) + + pool = ThreadPoolExecutor(max_workers=write_threads) + pending = deque() # in-flight writes (each holds a ~500 MB table) + max_inflight = write_threads + 4 + buf = [] + fi = total = 0 + + def flush(): + nonlocal buf, fi + if not buf: + return + table = pa.Table.from_batches(buf, schema=raw_schema) + while len(pending) >= max_inflight: # backpressure vs the DuckDB stream + pending.popleft().result() + pending.append(pool.submit(write_file, fi, table)) + fi += 1 + buf = [] + + for batch in reader: + if raw_schema is None: + raw_schema = batch.schema + ei = raw_schema.get_field_index("emb") + out_schema = raw_schema.set(ei, pa.field("emb", emb_type)) + sc = pq.SortingColumn.from_ordering(out_schema, [(k, "ascending") for k in SORT_KEYS]) + buf.append(batch) + total += batch.num_rows + if len(buf) == batches_per_file: + flush() + flush() # trailing partial file + + for f in pending: + f.result() + pool.shutdown() + + print(f"wrote {total:,} rows across {fi} files -> {out_dir}", flush=True) + assert total == total_expected, f"ROW MISMATCH {total} != {total_expected}" + assert fi == n_files, f"FILE COUNT MISMATCH {fi} != {n_files}" + + +if __name__ == "__main__": + main() diff --git a/processing/tol200m_image_embeddings/scripts/validate_config.py b/processing/tol200m_image_embeddings/scripts/validate_config.py new file mode 100644 index 0000000..f4e291a --- /dev/null +++ b/processing/tol200m_image_embeddings/scripts/validate_config.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python +"""Validate a standardized config before upload. + +Checks: + 1. file / row / distinct-uuid counts (rows == distinct == spine rows, 0 dups) + 2. parquet contract on first/middle/last file: emb type fixed_size_list[dim], + 50K row groups, 9 sorting_columns, page indexes, file size + 3. FULL uuid-at-position order match against the spine + 4. embedding recompute on a random row-group sample: normalize (if --normalize) + + float16-cast the RAW embeddings for the sampled uuids and byte-compare + against the stored values + +Usage: + python validate_config.py --dim D \ + --spine "/train-*.parquet" \ + --emb-glob "/.../rank_*/*.parquet" [--emb-glob ...] \ + [--normalize] [--sample-files 3] [--threads T] [--mem-limit G] + +Run via a SLURM cpu job (step 3 sorts 233M rows twice; step 4 joins a sample +against the raw embeddings). +""" +import argparse +import glob +import os +import random + +import duckdb +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("config_dir") + ap.add_argument("model_col") + ap.add_argument("--dim", type=int, required=True) + ap.add_argument("--spine", required=True) + ap.add_argument("--emb-glob", action="append", required=True) + ap.add_argument("--normalize", action="store_true") + ap.add_argument("--sample-files", type=int, default=3) + ap.add_argument("--threads", type=int, default=int(os.environ.get("SLURM_CPUS_PER_TASK", 32))) + ap.add_argument("--mem-limit", default="200GB") + ap.add_argument("--tmp", default=os.environ.get("TMPDIR", "/tmp")) + args = ap.parse_args() + + files = sorted(glob.glob(os.path.join(args.config_dir, "*.parquet"))) + con = duckdb.connect() + con.execute(f"SET memory_limit='{args.mem_limit}'; SET threads={args.threads}; " + f"SET temp_directory='{args.tmp}';") + ok = True + + # 1. counts + n = con.sql(f"SELECT COUNT(*) FROM read_parquet({files})").fetchone()[0] + d = con.sql(f"SELECT COUNT(DISTINCT uuid) FROM read_parquet({files})").fetchone()[0] + sp_n = con.sql(f"SELECT COUNT(*) FROM read_parquet('{args.spine}')").fetchone()[0] + print(f"1. files={len(files)} rows={n:,} distinct_uuid={d:,} dups={n-d:,} spine={sp_n:,}", + flush=True) + ok &= (n == d == sp_n) + + # 2. contract spot-check + for f in (files[0], files[len(files) // 2], files[-1]): + m = pq.ParquetFile(f).metadata + sch = pq.ParquetFile(f).schema_arrow + rg = m.row_group(0) + emb_t = sch.field("emb").type + print(f"2. {os.path.basename(f)}: emb={emb_t} rows={m.num_rows:,} " + f"n_rg={m.num_row_groups} rg0={rg.num_rows} sort={len(rg.sorting_columns)} " + f"pageidx={rg.column(0).has_offset_index} {os.path.getsize(f)/1e6:.0f}MB", + flush=True) + ok &= (str(emb_t) == f"fixed_size_list[{args.dim}]" + and rg.num_rows == 50_000 and len(rg.sorting_columns) == 9) + + # 3. full order vs spine + mism = con.sql(f""" + WITH o AS (SELECT row_number() OVER () rn, uuid FROM read_parquet({files})), + s AS (SELECT row_number() OVER (ORDER BY _pos) rn, uuid + FROM read_parquet('{args.spine}')) + SELECT COUNT(*) FROM o JOIN s USING(rn) WHERE o.uuid <> s.uuid""").fetchone()[0] + print(f"3. FULL order: uuid-at-position mismatches vs spine = {mism:,}", flush=True) + ok &= (mism == 0) + + # 4. recompute check on sampled files (first row group of each) + glist = "['" + "','".join(args.emb_glob) + "']" + rng = random.Random(0) + checked = bad = 0 + for f in rng.sample(files, min(args.sample_files, len(files))): + stored = pq.ParquetFile(f).read_row_group(0, columns=["uuid", "emb"]) + uuids = stored.column("uuid").to_pylist() + con.execute("CREATE OR REPLACE TEMP TABLE want(uuid VARCHAR)") + con.executemany("INSERT INTO want VALUES (?)", [(u,) for u in uuids]) + raw = con.sql(f""" + SELECT r.uuid, r."{args.model_col}" AS emb + FROM read_parquet({glist}) r JOIN want USING (uuid)""").arrow() + raw_map = dict(zip(raw.column("uuid").to_pylist(), + np.reshape(raw.column("emb").combine_chunks() + .cast(pa.list_(pa.float32(), args.dim)) + .values.to_numpy(zero_copy_only=False), + (raw.num_rows, args.dim)))) + got = np.reshape(stored.column("emb").combine_chunks() + .cast(pa.list_(pa.float16(), args.dim)) + .values.to_numpy(zero_copy_only=False), + (stored.num_rows, args.dim)) + for i, u in enumerate(uuids): + if u not in raw_map: + bad += 1 + continue + v = raw_map[u] + if args.normalize: + nrm = np.linalg.norm(v) + v = v / nrm if nrm != 0 else v + bad += not np.array_equal(v.astype(np.float16), got[i]) + checked += 1 + print(f"4. recompute check: {checked:,} rows sampled, mismatches={bad:,}", flush=True) + ok &= (bad == 0) + + print("\nVERDICT:", "PASS" if ok else "FAIL") + raise SystemExit(0 if ok else 1) + + +if __name__ == "__main__": + main() From c73c238db99aa3c098ba94572247c99f9c6ee6c1 Mon Sep 17 00:00:00 2001 From: Net Zhang Date: Mon, 31 Aug 2026 10:33:02 -0400 Subject: [PATCH 3/3] Document TOL-200M image embeddings provenance pipeline --- processing/README.md | 4 + .../requirements_tol200m_image_embeddings.txt | 10 ++ processing/tol200m_image_embeddings/README.md | 159 ++++++++++++++++++ 3 files changed, 173 insertions(+) create mode 100644 processing/docs/requirements_tol200m_image_embeddings.txt create mode 100644 processing/tol200m_image_embeddings/README.md diff --git a/processing/README.md b/processing/README.md index 11a049d..bb37a0e 100644 --- a/processing/README.md +++ b/processing/README.md @@ -25,6 +25,10 @@ Be sure to set appropriate `BASE_DIR` variables in `scripts/mongo/` and `scripts [BioCLIP 2 text embeddings of TreeOfLife-200M](https://huggingface.co/datasets/imageomics/TreeOfLife-200M/blob/main/embeddings/txt_emb_species.npy) were generated with [`make_txt_embedding.py`](scripts/make_txt_embedding.py), using the [`txt_emb_species.json`](https://huggingface.co/datasets/imageomics/TreeOfLife-200M/blob/main/embeddings/txt_emb_species.json) to provide the species names. More information about the `JSON` is provided in the [TreeOfLife-200M `embeddings/README`](https://huggingface.co/datasets/imageomics/TreeOfLife-200M/blob/main/embeddings/README.md). +## TreeOfLife-200M Image Embeddings + +The [`tol200m_image_embeddings/`](tol200m_image_embeddings) directory contains the SLURM jobs, configs, and standardization scripts used to produce the image-embedding configs of [TreeOfLife-200M-Embeddings](https://huggingface.co/datasets/imageomics/TreeOfLife-200M-Embeddings) (BioCLIP 2 and BioCLIP 2.5 Huge): batch embedding via the [hpc-inference](https://github.com/Imageomics/hpc-inference) package, then metadata join, taxonomic ordering, normalization (where applicable), float16 casting, and contract parquet writing. See its `README.md` for the pipeline and placeholder conventions. + ## Webdataset Construction The requirements and config files for taking the TreeOfLife structured dataset and putting it in webdataset format are `requirements_tol2webdataset.txt` and `tol2webdataset_full_224.yaml`, respectively. The code for this is run through `scripts/t2w_submit.sh` using the `tol2webdataset` scripts and modules. diff --git a/processing/docs/requirements_tol200m_image_embeddings.txt b/processing/docs/requirements_tol200m_image_embeddings.txt new file mode 100644 index 0000000..662a6c9 --- /dev/null +++ b/processing/docs/requirements_tol200m_image_embeddings.txt @@ -0,0 +1,10 @@ +# Requirements for processing/tol200m_image_embeddings (see its README.md). +# hpc-inference is not on PyPI; install from GitHub: +# uv pip install "hpc-inference[openclip] @ git+https://github.com/Imageomics/hpc-inference.git" +# It brings torch / open-clip-torch / pyarrow / h5py and the rest of the +# embedding-stage stack. The extras below cover profiling and standardization. +nvidia-ml-py +duckdb +pyarrow +numpy +h5py diff --git a/processing/tol200m_image_embeddings/README.md b/processing/tol200m_image_embeddings/README.md new file mode 100644 index 0000000..1cd1d86 --- /dev/null +++ b/processing/tol200m_image_embeddings/README.md @@ -0,0 +1,159 @@ +# TreeOfLife-200M Image Embeddings + +Provenance and reproduction scripts for the image-embedding configs published in +[imageomics/TreeOfLife-200M-Embeddings](https://huggingface.co/datasets/imageomics/TreeOfLife-200M-Embeddings): + +| Config | Model | Dim | Stored precision | Normalized | +|--------|-------|-----|------------------|------------| +| `bioclip-2_float16` | [BioCLIP 2](https://huggingface.co/imageomics/bioclip-2) (ViT-L/14) | 768 | float16 | No | +| `bioclip-2.5-vith14_float16` | [BioCLIP 2.5 Huge](https://huggingface.co/imageomics/bioclip-2.5-vith14) (ViT-H/14) | 1024 | float16 | Yes (L2) | + +Both models run through the **same pipeline**, differing only in model weights (and the +normalization decision at standardization time). The pipeline has two stages: + +``` +TreeOfLife-200M images (HDF5 shards) + | 1. batch embedding (hpc-inference, GPU SLURM jobs) + v +raw per-rank parquet (uuid, emb_ float32) + | 2. standardization (CPU SLURM job: catalog metadata join, taxonomic + | sort, L2-normalization if applicable, float16 cast, contract write) + v +train-NNNNN-of-MMMMM.parquet (published config layout) +``` + +Dataset layout, column semantics, sort order, and the parquet contract are documented in +the [dataset card](https://huggingface.co/datasets/imageomics/TreeOfLife-200M-Embeddings); +this README stays operational and does not repeat them. + +## Environment + +```bash +uv pip install "hpc-inference[openclip] @ git+https://github.com/Imageomics/hpc-inference.git" +uv pip install duckdb pyarrow h5py nvidia-ml-py +``` + +## Stage 1: batch embedding + +Input is the TreeOfLife-200M image collection materialized as HDF5 shards +(`/source=/server=/*_images.h5`, one `uuid -> image bytes` entry per +image, with a paired `*_metadata.parquet` per shard). Embedding uses the +[hpc-inference](https://github.com/Imageomics/hpc-inference) package's +`open_clip_embed` entry point (`--input_type hdf5`); each SLURM task is one GPU rank, +and ranks split the shard list automatically. + +For sources with a moderate file count (bioscan, eol, fathomnet: 14-654 files), +submit one multi-rank job. Keep `world_size <= num_files` and ideally +`files_per_rank >= num_workers`: + +```bash +export RUN_CONFIG=configs/embed_bioclip_2_5.yaml # or embed_bioclip_2.yaml +export RUN_TARGET_DIR=/source=bioscan +export RUN_OUTPUT_DIR=/model=bioclip_2_5/source=bioscan +sbatch --nodes=4 --time=02:00:00 scripts/embed.slurm # 4 nodes x 1 GPU = 4 ranks +``` + +For sources with many files of highly variable size (gbif: ~24.6K files, 1-20K +images each), first build image-count-balanced chunks, then submit a SLURM array +of single-GPU tasks. Balancing chunks by image count (not file count) is what +keeps task walltimes even: + +```bash +python scripts/count_images.py /source=gbif file_counts.parquet +NTASKS=$(python scripts/build_chunks.py file_counts.parquet \ + /model=bioclip_2_5/source=gbif 7000000) + +export RUN_CONFIG=configs/embed_bioclip_2_5.yaml +export RUN_TARGET_DIR=/source=gbif +export RUN_OUT_ROOT=/model=bioclip_2_5/source=gbif +sbatch --array=0-$((NTASKS-1))%6 --time=06:00:00 scripts/embed_array.slurm +``` + +Both paths write `.../embeddings/rank_/_rank__.parquet` with columns +`uuid` + `emb_` (raw float32, un-normalized). Budget the array `--time` +for the worst chunk: chunks packing many small files run slower per image because +of per-file HDF5 open overhead. + +> **Precision note.** The published `bioclip-2.5-vith14_float16` production run used +> bf16 autocast (plus in-pipeline normalization / float16 storage) for +> throughput gain; those pipeline features are not yet in the released `hpc-inference` +> package, so these scripts run the released fp32 path and defer normalization and +> casting to Stage 2. The result is equivalent up to compute precision: measured +> cosine similarity between bf16- and fp32-computed embeddings is 0.9998 mean +> (min 0.993), and after L2 normalization the stored float16 cast is lossless. + +## Stage 2: standardization + +Turns raw per-rank embeddings into a published config. Metadata is **not** re-derived: +every config shares the exact row set, row order, and 15 metadata columns of the +already-published `bioclip-2_float16` config, which is itself the +[TreeOfLife-200M catalog](https://huggingface.co/datasets/imageomics/TreeOfLife-200M) +(revision `94bbc0b`) restricted to its 233,055,986 rows and globally sorted by +``` +source_dataset > +kingdom > phylum > class > order > family > genus > species > common_name +``` +Reusing the published row order (rather than re-sorting the catalog) keeps all configs +row-aligned for free: the sort keys have large tie groups, so a fresh sort would not +reproduce the published order within ties. + +First build the spine (`uuid` + metadata + global `_pos` index, stripped from the +published config files, order preserved). Run once; every config build reuses it: + +```bash +python scripts/build_spine.py +``` + +Then build the config. DuckDB inner-joins the spine to the raw embeddings by `uuid` +(dropping raw rows not in the TOL catalog), orders by `_pos`, and a parallel writer pool +applies L2 normalization in float32 (`--normalize`, used for bioclip-2.5, not for +bioclip-2), casts to `fixed_size_list[dim]`, and writes the contract +parquet files: + +```bash +python scripts/standardize_config.py emb_bioclip_2_5 bioclip-2.5-vith14_float16 250000 \ + --dim 1024 --normalize \ + --spine "/train-*.parquet" \ + --emb-glob "/model=bioclip_2_5/source=*/embeddings/rank_*/*.parquet" \ + --emb-glob "/model=bioclip_2_5/source=gbif/task_*/embeddings/rank_*/*.parquet" \ + --out + +# bioclip-2 equivalent (no --normalize; 768-dim; 500K rows/file): +# python scripts/standardize_config.py emb_bioclip_2 bioclip-2_float16 500000 \ +# --dim 768 --spine ... --emb-glob ... --out +``` + +`rows_per_file` targets ~500 MB files: 500000 rows per file for 512/768-dim, 250000 for 1024-dim. +Run as a SLURM cpu job; the global join+sort spills through DuckDB temp space +(~470 GB peak for the 1024-dim config, so use a hugemem-class node; ~340 GB for +512-dim). + +## Stage 3: validation + +Checks a finished config: file / row / distinct-uuid counts, parquet contract fields, +FULL uuid-at-position order match against the spine, and a sample recompute of the +normalize+cast transform against the raw embeddings: + +```bash +python scripts/validate_config.py /bioclip-2.5-vith14_float16 emb_bioclip_2_5 \ + --dim 1024 --normalize \ + --spine "/train-*.parquet" \ + --emb-glob "/model=bioclip_2_5/source=*/embeddings/rank_*/*.parquet" \ + --emb-glob "/model=bioclip_2_5/source=gbif/task_*/embeddings/rank_*/*.parquet" +``` + +Exit code 0 and `VERDICT: PASS` mean the config is ready to upload. + +## Placeholders + +All `scripts/configs` use placeholders; set them before running: + +| Placeholder | Meaning | +|---|---| +| `YOUR_ACCOUNT` | SLURM account | +| `/path/to/venv` | Python venv with the environment above | +| `` | TreeOfLife-200M HDF5 shard root (`source=/...`) | +| `` | raw embedding output root (scratch) | +| `` | local copy of the published `bioclip-2_float16/` files | +| `` | spine output dir (intermediate, reusable) | +| `` | standardized config output root (upload staging) |