Per-batch validation overhead

Code available on GitHub.

pandera 0.34 can validate PyTorch TensorDict batches: dtypes, shapes, and values on every key. The obvious place to do that is inside the training loop, right before each batch reaches the model. This benchmark measures what that costs in wall-clock time.

It trains a 3-layer MLP policy (REINFORCE-style loss, Adam) on TensorDict batches of RL transitions. Each batch has a 5% chance of corruption (corrupt_rate=0.05, the 5% in the chart footnotes); with the fixed seed, 8 of the 200 pooled batches (4%) are corrupted, in one of three ways: float64 observations, one action outside the action space, or one NaN reward. Every training length runs in two modes on the same pod:

Mode	How it skips corrupt batches
without pandera	looks up a precomputed `is_valid` mask, so skipping is free
with pandera	calls `Transition.validate(td)` and skips on `SchemaError`

Both modes train on exactly the same batches, and the benchmark asserts that pandera skips the same steps as the mask, so the difference in wall-clock time is the cost of validation and nothing else. The baseline is deliberately generous. A real pipeline has no mask, and without one the corrupt batches would crash the step (float64 input to a float32 layer, an out-of-range index into gather) or silently poison the weights (NaN). The cost of not validating tutorial measures that side.

Setting up the environment

The benchmark runs on CPU with the CPU-only torch wheel. Two settings matter for the measurement:

torch’s thread count is pinned to the CPU request. In a pod, os.cpu_count() reports the node’s cores, not the container’s limit, and an oversubscribed 4-CPU limit made a run that normally takes a few minutes still be going after 30.
The CPU limit sits above the thread count. With the limit equal to the thread count, the larger models trained 5 to 17% faster with validation than without, consistent with CFS throttling: a matmul-heavy step uses up the quota, and short validation calls between steps change where the throttling lands. Raising the limit removed the effect, and the step-size sweep records cgroup throttling per run to confirm it stays at zero.

main.py
                
image = (
    flyte.Image.from_debian_base(python_version=(3, 12), name="pandera-tensordict-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 a 4-CPU limit makes every step crawl.
CPUS = 4

# The limit sits above the thread count on purpose. With limit == threads, a
# matmul-heavy step uses up the CFS quota and gets throttled, and the throttling
# pattern shifts when validation runs between steps, which skews the comparison.
env = flyte.TaskEnvironment(
    name="pandera_tensordict_bench",
    image=image,
    resources=flyte.Resources(cpu=(CPUS, 2 * CPUS), memory="8Gi"),
)

The schema

The Transition model checks dtype and shape on all four keys, plus value checks on observations, actions, and rewards:

main.py
                
def build_schema(obs_dim: int):
    import torch
    import pandera.tensordict as pa

    class Transition(pa.TensorDictModel):
        observation: torch.float32 = pa.Field(shape=(None, obs_dim), ge=-OBS_BOUND, le=OBS_BOUND)
        action: torch.int64 = pa.Field(shape=(None,), isin=list(range(N_ACTIONS)))
        reward: torch.float32 = pa.Field(shape=(None,), ge=-1.0, le=1.0)
        done: torch.bool = pa.Field(shape=(None,))

        class Config:
            batch_size = (None,)

    return Transition

Corrupt batches

Batches come from a pre-generated pool of 200 that the loop cycles through, so data generation stays outside the timer and memory stays flat at long lengths. Each batch carries its is_valid flag, which is the mask the baseline mode uses:

main.py
                
                    
                
            
                
            

                
                    
                
            
def make_pool(pool_size: int, batch_size: int, obs_dim: int, corrupt_rate: float, seed: int):
    """A fixed pool of batches the training loop cycles through.

    Generating batches up front keeps data generation out of the timed loop, and
    a pool (rather than one batch per step) keeps memory flat at long lengths.
    """
    import torch
    from tensordict import TensorDict

    gen = torch.Generator().manual_seed(seed)
    rng = random.Random(seed)
    pool = []
    for _ in range(pool_size):
        observation = torch.randn(batch_size, obs_dim, generator=gen).clamp(-OBS_BOUND, OBS_BOUND)
        action = torch.randint(0, N_ACTIONS, (batch_size,), generator=gen)
        reward = torch.rand(batch_size, generator=gen) * 2 - 1
        done = torch.rand(batch_size, generator=gen) < 0.01

        corruption = rng.choice(CORRUPTIONS) if rng.random() < corrupt_rate else None
        if corruption == "float64_observation":
            observation = observation.double()
        elif corruption == "action_out_of_range":
            action[rng.randrange(batch_size)] = N_ACTIONS
        elif corruption == "nan_reward":
            reward[rng.randrange(batch_size)] = float("nan")

        td = TensorDict(
            {"observation": observation, "action": action, "reward": reward, "done": done},
            batch_size=[batch_size],
        )
        pool.append((td, corruption is None))
    return pool

The training loop

When a schema is passed, the loop times each validate() call on its own, as well as the whole run, so the benchmark can report validation milliseconds per batch directly:

main.py
                
                    
                
            
                
            

                
                    
                
            
def train(n_steps: int, pool: list, obs_dim: int, hidden: int, schema=None, seed: int = 0) -> dict:
    """Run ``n_steps`` of a REINFORCE-style update, skipping invalid batches."""
    import torch
    import pandera.tensordict as pa

    torch.manual_seed(seed)
    model = torch.nn.Sequential(
        torch.nn.Linear(obs_dim, hidden),
        torch.nn.ReLU(),
        torch.nn.Linear(hidden, hidden),
        torch.nn.ReLU(),
        torch.nn.Linear(hidden, N_ACTIONS),
    )
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    trained, skipped, skipped_steps, validate_s = 0, 0, [], 0.0
    start = time.perf_counter()
    for step in range(n_steps):
        td, is_valid = pool[step % len(pool)]

        if schema is not None:
            t0 = time.perf_counter()
            try:
                schema.validate(td)
                is_valid = True
            except pa.SchemaError:
                is_valid = False
            validate_s += time.perf_counter() - t0

        if not is_valid:
            skipped += 1
            skipped_steps.append(step)
            continue

        logits = model(td["observation"])
        log_prob = torch.log_softmax(logits, dim=-1).gather(1, td["action"].unsqueeze(1)).squeeze(1)
        loss = -(log_prob * td["reward"]).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        trained += 1

    return {
        "seconds": time.perf_counter() - start,
        "validate_seconds": validate_s,
        "trained": trained,
        "skipped": skipped,
        "skipped_steps": skipped_steps,
    }

Timing both modes

Both paths are warmed up for 50 steps before any timing. Each length is repeated five times, alternating which mode runs first so drift on the node doesn’t favor one of them. The results come back as an inline JSON string rather than a File, so the local driver can read them from run.outputs() without credentials for the cluster’s object store.

main.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def run_benchmark(
    lengths: list[int],
    repeats: int,
    batch_size: int,
    obs_dim: int,
    hidden: int,
    corrupt_rate: float,
    pool_size: int,
) -> str:
    """Time both modes at every training length on this one pod."""
    import os
    import platform

    import torch

    torch.set_num_threads(min(CPUS, os.cpu_count() or 1))
    schema = build_schema(obs_dim)
    pool = make_pool(pool_size, batch_size, obs_dim, corrupt_rate, seed=0)

    # Warm up both paths so one-time costs (allocator, autograd, pandera's schema
    # build) don't land on the shortest run.
    train(50, pool, obs_dim, hidden)
    train(50, pool, obs_dim, hidden, schema=schema)

    rows = []
    for n_steps in lengths:
        base_runs, pandera_runs = [], []
        for r in range(repeats):
            # Alternate which mode goes first so drift on the node doesn't
            # systematically favor one of them.
            order = [("base", None), ("pandera", schema)]
            if r % 2:
                order.reverse()
            for mode, s in order:
                result = train(n_steps, pool, obs_dim, hidden, schema=s, seed=r)
                (base_runs if mode == "base" else pandera_runs).append(result)

        for b, p in zip(base_runs, pandera_runs):
            assert b["skipped_steps"] == p["skipped_steps"], "pandera skipped different batches than the mask"

        base_s = [x["seconds"] for x in base_runs]
        pandera_s = [x["seconds"] for x in pandera_runs]
        overhead = [100 * (p - b) / b for b, p in zip(base_s, pandera_s)]
        validated = n_steps * repeats
        row = {
            "steps": n_steps,
            "trained_batches": base_runs[0]["trained"],
            "skipped_batches": base_runs[0]["skipped"],
            "without_pandera_s": statistics.median(base_s),
            "without_pandera_s_runs": base_s,
            "with_pandera_s": statistics.median(pandera_s),
            "with_pandera_s_runs": pandera_s,
            "overhead_pct": statistics.median(overhead),
            "overhead_pct_runs": overhead,
            "validate_ms_per_batch": 1000 * sum(x["validate_seconds"] for x in pandera_runs) / validated,
            "train_ms_per_batch": 1000 * sum(base_s) / (base_runs[0]["trained"] * repeats),
        }
        rows.append(row)
        print(
            f"{n_steps:>6} steps  without={row['without_pandera_s']:.2f}s  "
            f"with={row['with_pandera_s']:.2f}s  overhead={row['overhead_pct']:.1f}%",
            flush=True,
        )

    results = {
        "config": {
            "lengths": lengths,
            "repeats": repeats,
            "batch_size": batch_size,
            "obs_dim": obs_dim,
            "hidden": hidden,
            "corrupt_rate": corrupt_rate,
            "pool_size": pool_size,
            "pool_corrupt_batches": sum(not ok for _, ok in pool),
        },
        "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(),
        },
        "rows": rows,
    }
    # Returned inline as a string rather than a File, so the driver can read it
    # from run.outputs() without credentials for the cluster's bucket.
    return json.dumps(results, 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.

Wall-clock time with and without pandera, and added runtime per training length

Steps	Without pandera (s)	With pandera (s)	Overhead
100	0.14	0.17	24.3%
300	0.42	0.52	24.3%
1,000	1.52	1.85	22.4%
3,000	4.41	5.37	21.2%
10,000	14.06	17.31	23.2%

Validation cost about 0.30 ms per batch against a training step of about 1.5 ms. An earlier run of the same code on a slower node measured 0.28 ms per batch and overhead between 16% and 20%: the percentage is a ratio, so a cheaper step makes the same validation cost look bigger.

How to read it:

The percentage stays flat as training gets longer, because the cost is paid per batch and there’s nothing to amortize.
The percentage is the ratio of a fixed validation cost to the per-step compute, and this setup makes the step cheap on purpose (a 3-layer MLP on 4 CPU cores). The step-size sweep below measures how the share falls as the step gets more expensive.
Most of the cost comes from the value checks (observation range, action set, reward range), which scan every element; dtype, shape, and key checks only read metadata. In a separate laptop measurement, the full schema took 0.064 ms per batch and a structural-only schema took 0.021 ms, so the value checks were about two-thirds of the cost.
validate() was called with its default inplace=False, which clones the TensorDict before checking it. Passing inplace=True skips that copy; in a laptop measurement it cut per-batch validation time by about 10%.
Every run skipped exactly the corrupt batches (2, 10, 40, 120, 400 across the lengths), matching the mask.
If the per-step cost matters, validate where data enters the system, such as on insertion into a replay buffer or once per data-loader worker batch, instead of on every sampled training batch. The whole-dataset validation tutorial measures validating once, upstream.

Step-size sweep

run_step_size_sweep keeps everything about validation fixed (the same 256-row batches, the same schema, the same 8 corrupt batches out of 200) and grows the MLP’s hidden width from 64 to 4,096, so the forward and backward pass gets more expensive while pandera’s work per step stays the same. Growing the batch instead would grow validation along with it. Each width runs enough steps to take about 15 seconds per run, five runs per mode, alternating which mode goes first.

main.py
                
                    
                
            
                
            

                
                    
                
            
@env.task
async def run_step_size_sweep(
    widths: list[int],
    repeats: int,
    batch_size: int,
    obs_dim: int,
    corrupt_rate: float,
    pool_size: int,
    target_seconds: float,
) -> str:
    """Hold the batch (and so the validation work) fixed and grow the model.

    Every width validates the same 256-row batches against the same schema, so
    pandera's cost per step stays constant while the forward/backward pass gets
    more expensive. Each width runs enough steps to take about ``target_seconds``.
    """
    import os
    import platform

    import torch

    torch.set_num_threads(min(CPUS, os.cpu_count() or 1))
    schema = build_schema(obs_dim)
    pool = make_pool(pool_size, batch_size, obs_dim, corrupt_rate, seed=0)
    train(50, pool, obs_dim, widths[0])
    train(50, pool, obs_dim, widths[0], schema=schema)

    rows = []
    for hidden in widths:
        probe = train(20, pool, obs_dim, hidden)
        step_s = probe["seconds"] / max(probe["trained"], 1)
        n_steps = int(min(3_000, max(50, target_seconds / step_s)))

        base_runs, pandera_runs = [], []
        for r in range(repeats):
            order = [("base", None), ("pandera", schema)]
            if r % 2:
                order.reverse()
            for mode, sch in order:
                before = cpu_throttled_ms()
                result = train(n_steps, pool, obs_dim, hidden, schema=sch, seed=r)
                after = cpu_throttled_ms()
                result["throttled_ms"] = after - before if before is not None and after is not None else None
                (base_runs if mode == "base" else pandera_runs).append(result)
        for b, p in zip(base_runs, pandera_runs):
            assert b["skipped_steps"] == p["skipped_steps"], "pandera skipped different batches than the mask"

        base_s = [x["seconds"] for x in base_runs]
        pandera_s = [x["seconds"] for x in pandera_runs]
        overhead = [100 * (p - b) / b for b, p in zip(base_s, pandera_s)]
        trained = base_runs[0]["trained"]
        row = {
            "hidden": hidden,
            "params": mlp_params(obs_dim, hidden),
            "steps": n_steps,
            "without_pandera_s": statistics.median(base_s),
            "without_pandera_s_runs": base_s,
            "with_pandera_s": statistics.median(pandera_s),
            "with_pandera_s_runs": pandera_s,
            "overhead_pct": statistics.median(overhead),
            "overhead_pct_runs": overhead,
            "validate_ms_per_batch": 1000 * sum(x["validate_seconds"] for x in pandera_runs) / (n_steps * repeats),
            "train_ms_per_step": 1000 * statistics.median(base_s) / trained,
            "throttled_ms_without": [x["throttled_ms"] for x in base_runs],
            "throttled_ms_with": [x["throttled_ms"] for x in pandera_runs],
        }
        rows.append(row)
        print(
            f"hidden={hidden:>5}  params={row['params']:>11,}  steps={n_steps:>5}  "
            f"train={row['train_ms_per_step']:8.2f} ms/step  validate={row['validate_ms_per_batch']:.3f} ms  "
            f"overhead={row['overhead_pct']:.1f}%",
            flush=True,
        )

    return json.dumps({
        "config": {
            "widths": widths, "repeats": repeats, "batch_size": batch_size, "obs_dim": obs_dim,
            "corrupt_rate": corrupt_rate, "pool_size": pool_size, "target_seconds": target_seconds,
        },
        "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(),
        },
        "rows": rows,
    }, indent=2)

Validation time as a share of the training step, by model size

The chart plots validation’s share of the step: milliseconds per validate() call, timed around each call, divided by the median training milliseconds per step without pandera. Both halves are timed directly, so the ratio isn’t exposed to run-to-run noise the way the end-to-end difference is.

Hidden width	Params	Train ms/step	Validate ms/batch	Share of step	End-to-end added runtime
64	8.6K	1.05	0.30	28.9%	32.7%
256	83K	1.67	0.28	16.7%	16.7%
1,024	1.1M	9.25	0.37	4.0%	4.3%
2,048	4.3M	29.3	0.41	1.4%	−3.5%
4,096	17.1M	142	0.46	0.3%	0.0%

How to read it:

The share falls roughly in proportion to step time: a 135x more expensive step took validation from 29% of the step to 0.3%.
Validation’s own cost crept up with model size (0.28 to 0.46 ms per batch) even though the batch didn’t change. The likely cause is cache: a larger model’s weights evict the batch between steps. It grows far slower than the step does, so the share still falls.
Where validation is a big share of the step, the end-to-end difference between the two modes agrees with it. From 1M parameters up, the end-to-end difference is smaller than the run-to-run noise on a shared node, and single pairs of runs land on both sides of zero. That’s why the chart plots the directly timed share instead.

Rendering the report

The plot task renders both charts in light and dark themes and writes them, with tables, into the run’s report tab. main runs the two benchmarks one after the other on separate pods and returns both JSON strings.

main.py
                
                    
                
            
                
            

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

@env.task
async def main(
    lengths: list[int] = [100, 300, 1_000, 3_000, 10_000],
    repeats: int = 5,
    batch_size: int = 256,
    obs_dim: int = 64,
    hidden: int = 256,
    corrupt_rate: float = 0.05,
    pool_size: int = 200,
    widths: list[int] = [64, 256, 1_024, 2_048, 4_096],
    sweep_repeats: int = 5,
    sweep_target_seconds: float = 15.0,
) -> tuple[str, str]:
    results = await run_benchmark(lengths, repeats, batch_size, obs_dim, hidden, corrupt_rate, pool_size)
    sweep = await run_step_size_sweep(
        widths, sweep_repeats, batch_size, obs_dim, corrupt_rate, pool_size, sweep_target_seconds
    )
    await plot(results, sweep)
    return results, sweep

Run it

        
    
# From v2/tutorials/pandera_validation_overhead 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 a short configuration
uv run main.py --local --lengths 100 300 --repeats 2 \
    --widths 64 256 --sweep-repeats 2 --sweep-target-seconds 1