Convert hybrid HDF5+Parquet data into wds

config/tol_hybrid_to_wds_example.yaml

@@ -0,0 +1,46 @@
+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_hybrid_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"

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.

docs/hybrid_to_webdataset_README.md

@@ -0,0 +1,63 @@
+# 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.
+
+## 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/<uuid>` 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. 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 .
+```
+
+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.
+- 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.
+
+## 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:
+
+- `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).
+
+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:
+
+```bash
+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).
+
+## Validating output
+
+[`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.

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

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

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

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

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

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_hybrid_to_wds,
+)
 

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)

src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md

@@ -0,0 +1,46 @@
+# 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.
+
+## 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`.
+
+## Configuration highlights
+
+Add a `tol_hybrid_to_wds` block to the standard toolbox config template:
+
+```yaml
+tol_hybrid_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_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.
+
+## 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.

src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py

@@ -0,0 +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.
+"""
+
+from TreeOfLife_toolbox.tol_hybrid_to_wds import classes  # noqa: F401