Whole-dataset validation

Code available on GitHub.

The per-batch validation overhead tutorial measures what it costs to validate inside a training loop, a cost every consumer of the data pays on every epoch. This one measures the alternative: validate the whole dataset once, where it’s produced, so the training jobs downstream can trust it without checking again. The question it answers is how wall-clock validation time grows with dataset size.

Datasets

Four synthetic datasets shaped like common training data, all valid by construction:

Dataset	Modeled on	Keys and checks	Bytes/row
continuous control	TorchRL / D4RL HalfCheetah replay buffer	`observation` (17) and `next_observation` (17) in [-100, 100], `action` (6) in [-1, 1], `reward`, `done`, `terminated`	166
Atari frames	DQN frame stacks	`pixels` uint8 (4×84×84) with no blank frames, `action` in the 18-action set, clipped `reward`, `done`	28,237
tabular classification	supervised tabular training set	`features` (32) finite, `category` (12) valid one-hot rows, `label` in [0, 9], `sample_weight` ≥ 0	188
token sequences	LM fine-tuning	`input_ids` (512) in the 32k vocab, `attention_mask`, `labels` that are token ids or the -100 ignore index	8,704

Setting up the environment

The benchmark runs on 4 CPUs with 12 GiB of memory and the CPU-only torch wheel. torch’s thread count is pinned to the CPU request, since os.cpu_count() in a pod reports the node’s cores.

main.py
                
image = (
    flyte.Image.from_debian_base(python_version=(3, 12), name="pandera-tensordict-dataset-bench")
    .with_pip_packages("torch==2.*", index_url="https://download.pytorch.org/whl/cpu")
    .with_pip_packages("pandera[torch]==0.34.0", "matplotlib")
    .with_apt_packages("fonts-inter")
)

# torch threads must match the CPU request: os.cpu_count() in a pod reports the
# node's cores, and oversubscribing the limit makes everything crawl.
CPUS = 4

env = flyte.TaskEnvironment(
    name="pandera_tensordict_dataset_bench",
    image=image,
    resources=flyte.Resources(cpu=CPUS, memory="12Gi"),
)

The schemas

Every key gets dtype and shape checks, and every non-boolean key gets a value check that scans every element. The one-hot, finite, blank-frame, and ignore-index checks are custom Checks, which is why the schemas use the object-based TensorDictSchema API: TensorDictModel fields don’t accept custom checks in pandera 0.34.

main.py
                
                    
                
            
                
            

                
                    
                
            
def build_schema(kind: str):
    """The schema for one dataset kind.

    These use the object-based TensorDictSchema API because several keys need
    custom checks, which TensorDictModel fields don't accept in pandera 0.34.
    """
    import torch
    import pandera.tensordict as pa
    from pandera import Check

    def t(dtype, shape, *checks):
        return pa.Tensor(dtype=dtype, shape=shape, checks=list(checks) or None)

    if kind == "continuous_control":
        # HalfCheetah: 17-dim observations, 6 continuous actions bounded to [-1, 1].
        keys = {
            "observation": t(torch.float32, (None, 17), Check.in_range(-100.0, 100.0)),
            "action": t(torch.float32, (None, 6), Check.in_range(-1.0, 1.0)),
            "reward": t(torch.float32, (None,), Check.in_range(-100.0, 100.0)),
            "next_observation": t(torch.float32, (None, 17), Check.in_range(-100.0, 100.0)),
            "done": t(torch.bool, (None,)),
            "terminated": t(torch.bool, (None,)),
        }
    elif kind == "atari_frames":
        # uint8 already bounds pixel values, so the pixel check catches blank frames
        # instead, a common symptom of a bad reset or a broken renderer.
        not_blank = Check(lambda x: x.flatten(2).amax(dim=-1).gt(0).all(dim=-1), name="no_blank_frames")
        keys = {
            "pixels": t(torch.uint8, (None, 4, 84, 84), not_blank),
            "action": t(torch.int64, (None,), Check.in_range(0, 17)),
            "reward": t(torch.float32, (None,), Check.in_range(-1.0, 1.0)),
            "done": t(torch.bool, (None,)),
        }
    elif kind == "tabular_classification":
        # A valid one-hot row is all zeros and ones and sums to exactly one.
        one_hot = Check(lambda x: ((x == 0) | (x == 1)).all(dim=-1) & (x.sum(dim=-1) == 1), name="one_hot")
        finite = Check(lambda x: torch.isfinite(x), name="finite")
        keys = {
            "features": t(torch.float32, (None, 32), finite),
            "category": t(torch.float32, (None, 12), one_hot),
            "label": t(torch.int64, (None,), Check.in_range(0, 9)),
            "sample_weight": t(torch.float32, (None,), Check.ge(0.0)),
        }
    elif kind == "token_sequences":
        vocab = 32_000
        # Labels are token ids, or -100 where the loss should ignore the position.
        label_ids = Check(lambda x: (x == -100) | ((x >= 0) & (x < vocab)), name="label_ids")
        keys = {
            "input_ids": t(torch.int64, (None, 512), Check.in_range(0, vocab - 1)),
            "attention_mask": t(torch.bool, (None, 512)),
            "labels": t(torch.int64, (None, 512), label_ids),
        }
    else:
        raise ValueError(f"unknown dataset kind: {kind}")
    return pa.TensorDictSchema(keys=keys, batch_size=(None,))

Proving the checks scan everything

A schema that could pass without scanning every row would make validation look cheaper than it is. Before timing a dataset, the benchmark corrupts one value in the last row of a small copy and asserts that validation raises. If it doesn’t, the benchmark refuses to run. The check also serves as a warmup, so first-call costs don’t land on the smallest size.

main.py
                
                    
                
            
                
            

                
                    
                
            
def corrupt_last_row(kind: str, td) -> None:
    """Break one value in the last row, so a schema that passes didn't scan it all."""
    if kind == "continuous_control":
        td["action"][-1, 0] = 2.0
    elif kind == "atari_frames":
        td["pixels"][-1, 3] = 0
    elif kind == "tabular_classification":
        td["category"][-1] = 0.0
    elif kind == "token_sequences":
        td["labels"][-1, 300] = -5

def assert_checks_scan_everything(kind: str, schema, chunk_rows: int) -> None:
    import pandera.tensordict as pa

    td = make_dataset(kind, 3 * chunk_rows + 7)
    validate_streamed(schema, td, chunk_rows)
    corrupt_last_row(kind, td)
    try:
        validate_streamed(schema, td, chunk_rows)
    except pa.SchemaError:
        return
    raise AssertionError(f"{kind}: schema didn't catch a corrupted value in the last row")

Streaming validation in chunks

Each dataset is one in-memory TensorDict, validated by streaming it through validate(inplace=True) in 64 MB chunks. Slicing a TensorDict returns views, so chunking adds no copies, and inplace=True skips the defensive clone validate() makes by default.

main.py
                
def validate_streamed(schema, td, chunk_rows: int) -> float:
    """Validate ``td`` in chunks of ``chunk_rows`` rows and return wall-clock seconds.

    Slicing a TensorDict returns views, so chunking adds no copies, and
    ``inplace=True`` skips the defensive clone ``validate()`` makes by default.
    """
    n = td.batch_size[0]
    start = time.perf_counter()
    for i in range(0, n, chunk_rows):
        schema.validate(td[i : i + chunk_rows], inplace=True)
    return time.perf_counter() - start

Timing every dataset and size

Each dataset is validated at 16 MB, 64 MB, 256 MB, 1 GB, and 4 GB, with data generation outside the timer and three repeats per point. A second sweep validates the same 1 GB continuous-control dataset in chunks of 256, 4,096, 65,536, and 1,048,576 rows, plus the whole dataset at once, to show what chunk size does to the cost.

main.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def run_dataset_benchmark(
    sizes_mb: list[int],
    repeats: int,
    chunk_mb: int,
    sweep_size_mb: int,
    sweep_chunk_rows: list[int],
) -> str:
    """Time whole-dataset validation for every kind and size on this one pod."""
    import gc
    import os
    import platform

    import torch

    torch.set_num_threads(min(CPUS, os.cpu_count() or 1))

    size_rows = []
    for kind in KINDS:
        schema = build_schema(kind)
        bpr = bytes_per_row(kind)
        chunk_rows = max(1, chunk_mb * MB // bpr)
        # Doubles as a warmup, so first-call costs don't land on the smallest size.
        assert_checks_scan_everything(kind, schema, min(chunk_rows, 4096))
        for size_mb in sizes_mb:
            n = max(1, size_mb * MB // bpr)
            td = make_dataset(kind, n)
            runs = [validate_streamed(schema, td, chunk_rows) for _ in range(repeats)]
            seconds = statistics.median(runs)
            size_rows.append({
                "kind": kind,
                "size_mb": size_mb,
                "rows": n,
                "bytes": n * bpr,
                "chunk_rows": chunk_rows,
                "seconds": seconds,
                "seconds_runs": runs,
                "gb_per_s": n * bpr / 1e9 / seconds,
                "ns_per_row": 1e9 * seconds / n,
            })
            print(f"{kind:>24} {size_mb:>6} MB  {n:>11,} rows  {seconds:8.3f} s", flush=True)
            del td
            gc.collect()

    # Chunk-size sweep on one dataset: the same 1 GB validated 256 rows at a time
    # (like a training batch) up to all at once.
    sweep_kind = "continuous_control"
    schema = build_schema(sweep_kind)
    bpr = bytes_per_row(sweep_kind)
    n = sweep_size_mb * MB // bpr
    td = make_dataset(sweep_kind, n)
    sweep_rows = []
    for chunk_rows in [*sweep_chunk_rows, n]:
        runs = [validate_streamed(schema, td, min(chunk_rows, n)) for _ in range(repeats)]
        seconds = statistics.median(runs)
        sweep_rows.append({
            "chunk_rows": min(chunk_rows, n),
            "whole_dataset": chunk_rows >= n,
            "seconds": seconds,
            "seconds_runs": runs,
            "ns_per_row": 1e9 * seconds / n,
        })
        print(f"sweep chunk={min(chunk_rows, n):>11,} rows  {seconds:8.3f} s", flush=True)

    return json.dumps({
        "config": {
            "sizes_mb": sizes_mb, "repeats": repeats, "chunk_mb": chunk_mb,
            "sweep_kind": sweep_kind, "sweep_size_mb": sweep_size_mb, "sweep_rows": n,
            "bytes_per_row": {k: bytes_per_row(k) for k in KINDS},
        },
        "environment": {
            "torch": torch.__version__,
            "pandera": __import__("pandera").__version__,
            "tensordict": __import__("tensordict").__version__,
            "python": platform.python_version(),
            "cpu_count": os.cpu_count(),
            "torch_threads": torch.get_num_threads(),
            "machine": platform.machine(),
        },
        "sizes": size_rows,
        "sweep": sweep_rows,
    }, indent=2)

Results

The run used 4 torch threads on x86_64, with torch 2.14.1+cpu, pandera 0.34.0, and tensordict 0.14.2.

Validation time by dataset size, and the same 1 GB validated at different chunk sizes

Validation time at each size (median of 3, streamed in 64 MB chunks):

Dataset	16 MB	256 MB	1 GB	4 GB	Throughput
continuous control	3 ms	42 ms	168 ms	658 ms	6.5 GB/s · 25 ns/row
Atari frames	0.5 ms	8 ms	29 ms	112 ms	38 GB/s · 734 ns/row
tabular classification	6 ms	115 ms	458 ms	1.77 s	2.4 GB/s · 77 ns/row
token sequences	1 ms	29 ms	114 ms	453 ms	9.5 GB/s · 917 ns/row

The same 1 GB of continuous-control data (6.5M rows) at different chunk sizes:

Chunk	Time	ns/row
256 rows	5.17 s	800
4,096 rows	634 ms	98
65,536 rows	186 ms	29
1,048,576 rows	165 ms	26
whole dataset at once	303 ms	47

How to read it:

Time is linear in dataset size for every dataset, so throughput is flat from 16 MB to 4 GB. Extrapolating at the slowest rate here (2.4 GB/s), a 40 GB dataset would take about 17 seconds on 4 cores; sizes past 4 GB weren’t measured.
Throughput depends on how much work the checks do per byte, not on how many bytes there are. The Atari frames are mostly uint8 pixels with one cheap reduction per frame, so they validate at 38 GB/s. Tabular data is slowest because the one-hot check does several passes (equality, sum) over a small float column.
Chunk size matters a lot at the small end. In 256-row chunks, the size of a training batch, the same gigabyte takes 31x longer than in 1M-row chunks, because every validate() call pays a fixed overhead. That’s the per-batch cost the companion benchmark measures, paid here on one pass instead of every epoch.
Validating everything in one call was slower than 1M-row chunks (303 ms vs 165 ms). The likely reason is that each check’s comparison builds a boolean mask the size of the whole tensor, which falls out of cache; that wasn’t profiled. Either way, moderate chunks were faster and use less peak memory.
For a rough comparison with the per-batch benchmark (its schema is a little different, and it used the default inplace=False): there, validating each 256-row batch inside the loop cost about 0.3 ms per step. Validating continuous-control data once, upstream, in large chunks costs about 26 ns per row, or about 7 µs per 256 rows.

Rendering the report

The plot task renders the chart in light and dark themes and writes it, with both tables, into the run’s report tab.

main.py
                
                    
                
            
                
            

                
                    
                
            
@env.task(report=True)
async def plot(results: str) -> None:
    """Render the chart in both themes into the run's report tab."""
    data = json.loads(results)
    pngs = {}
    for theme in ("light", "dark"):
        out = Path(f"dataset-validation-{theme}.png")
        render_plot(data, out, theme=theme)
        pngs[theme] = out.read_bytes()
    await flyte.report.replace.aio(report_html(data, pngs["light"], pngs["dark"]))
    await flyte.report.flush.aio()

@env.task
async def main(
    sizes_mb: list[int] = [16, 64, 256, 1024, 4096],
    repeats: int = 3,
    chunk_mb: int = 64,
    sweep_size_mb: int = 1024,
    sweep_chunk_rows: list[int] = [256, 4_096, 65_536, 1_048_576],
) -> str:
    results = await run_dataset_benchmark(sizes_mb, repeats, chunk_mb, sweep_size_mb, sweep_chunk_rows)
    await plot(results)
    return results

Run it

        
    
# From v2/tutorials/pandera_dataset_validation in your unionai-examples clone.
# On the cluster in your Flyte config; builds the image remotely, then writes results/ locally
uv run main.py

# Locally, with small sizes
uv run main.py --local --sizes-mb 8 32 128 --repeats 2 \
    --sweep-size-mb 64 --sweep-chunk-rows 256 4096 65536