Conformance kit

The checks a generator, stage or component plugin runs in its own CI, and TrackGen runs over every built-in; see Testing a Plugin. track_gen.testing is a subpackage outside track_gen.__all__: import it explicitly. import track_gen does not import it, and importing it imports neither torch nor pytest.

Conformance kit for generator and pipeline-stage plugins (and TrackGen’s own).

A plugin author runs one call in their own CI and learns whether their generator (or relax / width / elevation / validity stage) honours the contract (here the built-ins polar and relax.xpbd; a plugin passes its own registered names):

from track_gen.testing import check_generator, check_stage

report = check_generator("polar", params={"polar_num_knots": 9})  # cpu (+ cuda:0)
report.raise_for_failures()
check_stage("relax", "xpbd").raise_for_failures()

or, in pytest, loads the opt-in plugin (pytest_plugins = ["track_gen.testing.pytest_plugin"]) and calls the track_gen_conformance fixture. TrackGen’s own suite runs the same checks over every registered generator. Among them, params.per_env sets every dynamic param of the unit’s stage per env and checks that env e equals a run with e’s values as scalars (a stage must read param[e], never param[0] or a host copy, and never write its param arrays).

Modules:

import track_gen does not import this package, and importing it does not import torch (only the torch capture modes do, inside the mode).

Conformance checks

Does a generator honour the contract? check_generator / check_gate_generator.

One call runs a registered generator through every check the generator contract implies (see The Stage contract and Hard rules) and returns a Report (here for the built-in polar; a plugin passes its own registered name):

from track_gen.testing import check_generator

report = check_generator("polar", params={"polar_num_knots": 9})  # cpu (+ cuda:0)
report.raise_for_failures()

Every check runs per device. A check is a named CheckResult; its name is stable, so a test can assert exactly which check a broken generator fails.

Components are checked by track_gen.testing.check_component() (track_gen.testing.components; CHECKS["component"]).

Checks (CHECKS), for a track generator:

construct

The generator resolves, validates its params, allocates, runs its warm-up pass and a first generate() (into a poisoned raw output: the batch the checks read).

output.shape

The raw output (the [E*R] vec2f centerline and the [E] int32 validity flags) and every Track field have the documented shape and dtype; flags are 0 or 1; the result’s layout is the Program’s (canonical device) and result.validate() passes.

output.finite

The raw centerline is finite for every env the generator marked valid, and the final Track is finite before count[e] for every valid env. The raw output is poisoned before the generation the checks read, so a point the generator never writes is reported as such (it would otherwise hold stale, possibly finite memory).

output.padding

Every per-point Track field is NaN past count[e] (every env).

output.count

0 <= count[e] <= N_max for every env, 3 <= count[e] for valid ones.

output.closed

The raw loop closes (the last-to-first gap is at most 3x the median step, or 1.5x the longest step between consecutive points) and the pipeline accepts it (valid fraction > 0).

output.view (every kind)

The unit’s pipeline built with out= the middle view big.view(1, E + 1) of a byte-poisoned result of E + 2 envs (with the kit’s explicit seeds) writes exactly the bytes of the standalone batch into the view, after construction and after a masked advance, and leaves every byte outside it untouched: a kernel that indexes past [0, E) of a result array (or reads a result buffer it did not write) fails here. check_stage builds the stage’s composed pipeline the same way.

attributes.written / attributes.carried

For a generator that writes per-point attributes (raw.z, raw.half_width, raw.<pkg>.<name>; a gate generator anchors.<pkg>.<name>; skipped “no attributes” otherwise): every raw slot of every attribute is written for every env the generator marks valid (a gate generator: every slot before anchors.count; the raw buffers are poisoned like the centerline), and every EXPORTED result attribute (Track.attributes / GateSequence.attributes) holds its padding (NaN / -1) beyond count[e] and, on the real points of every env the generator marked valid, finite values; a linear attribute lies within the env’s raw [min, max] per component (slack 4 * spacing(float32(max|v|))), a nearest / previous / int32 one (or one carried by index) in the env’s raw value set, an angle one in [-pi, pi].

determinism

Two fresh generators built with the same seeds produce bitwise identical batches (raw output and every result field, NaN padding included).

diversity

Different seeds produce different tracks: no env repeats under a disjoint set of seeds, and no two envs of one batch are identical (envs valid in both batches only: two rejected envs may legitimately be alike). Reported skipped for a generator that declares seeded = False (an ingestion generator: its output does not depend on the seed).

batch_independence

Env e’s track depends on its own seed only: it is identical alone (E = 1), inside a larger batch, and when the other envs’ seeds are permuted.

statelessness

generate() twice with unchanged seeds reproduces the first batch, a reseed reproduces a fresh generator with the new seeds, and reseeding back reproduces the first batch (no hidden state carried between calls).

params.per_env (every kind)

For the unit’s stage (the generator, or the stage under test) with dynamic fields that no randomizer draws (skipped “no dynamic parameters” otherwise): every such field is set to a per-env [E] array (float32: d + s * 0.25 * span * (e + 1) / E, clamped to the declared bounds, toward the farther bound, span = hi - lo or, with a bound missing, max(|d|, 0.1): away from a field’s one bound, alternating directions for fields without bounds; int32: d + (e % 3) - 1, bool: (d + e) % 2), one generation (eager, same seeds); then, for each env e, every field is set to env e’s value as a scalar and env e of that generation must equal env e of the per-env one, bitwise (raw output and every result field). The live arrays (and the stage’s latched copies) must read back what was set: a stage reads param[e], never param[0] or a host copy, and never writes its param arrays. The fields are restored afterwards.

bounded.absorbing

For a unit whose enqueue runs a ctx.bounded_loop with active= (seen while the unit runs, whatever it declares; skipped “not applicable” otherwise): with every active flag poisoned to 0 first, an eager generation with the host poll lowering (a), which stops at the first window whose check reads 0, writes the same bytes (raw output, result, the stage’s write channels) as one with every window unrolled (c), for batch A and batch B. Fails for a non-absorbing loop (an iteration after active[0] reads 0 changes an output) and for a body that does not write active before the first check. Skipped “not exercised” when no window was skipped for the kit’s seeds.

bounded.lowering (CUDA)

For the same units: the eager default (c), the forced host poll (a), a Warp capture with conditional windows (b; skipped without conditional-node support) and the same capture unrolled (c), replayed into poisoned outputs, before and after a reseed, are bitwise equal.

host.require_caller_capturable (CPU)

For a pipeline with HOST stages (skipped “not applicable” otherwise): its spec with require_caller_capturable=True is refused at build with a PlanError naming every HOST stage.

hot_path.host_sync / hot_path.allocation / hot_path.stream / hot_path.capture (CPU)

Two generate() calls under hot_path_guard(), every ctx.bounded_loop unrolled (the superset of what any lowering emits): no host sync or readback, no allocation, every launch / copy from track_gen or the plugin’s package passes stream= (no assign() / zero_() / fill_()), no nested capture or graph launch. On CPU the stream check sees that stream= is passed, not which stream (stream=None is what a CPU LaunchContext passes): a launch on the wrong stream is caught by capture.caller on CUDA. A unit with a bounded loop also runs once with the default lowering: its host polls are the sanctioned readback (details["sanctioned"]), any other one fails hot_path.host_sync. A HOST stage’s ctx.host() operations are sanctioned too (details["host_ops"]); its own syncs, readbacks and allocations fail.

hot_path.topology (CPU)

The two calls emit the same launch sequence (LaunchTrace): kernels, dimensions, buffers and scalar arguments.

hot_path.buffers (CPU)

Every pointer launched on belongs to the generator (its result, scratch, seed and parameter buffers; see buffers_of()), and a second generator built alike launches on none of the first one’s buffers.

capture.caller (CUDA)

assert_caller_capturable(): recorded into a caller’s capture, replays equal the eager batch, before and after a reseed. A pipeline with HOST stages must instead refuse: generate(stream=s, advance=True) inside a caller’s Warp capture raises NotCapturableError naming every HOST stage before any launch, and the caller’s graph replays its own work only (outputs, seeds and epochs untouched).

capture.self_graph (CUDA)

generate() captures the generator’s own graph (a pipeline with HOST stages: one graph per non-HOST segment, none at the HOST ones, graph stays None); its replays (into poisoned outputs, before and after a reseed) equal the eager batch.

A gate generator runs the same checks on its GateSequence (raw output: the vec2f anchors before count[e]: anchors.xy, [E*A], for a Stage generator; the [E*G] gates a legacy spec writes), with output.accepted (valid fraction > 0) in place of output.closed and min_gates <= count[e] <= max_gates for valid envs; a Stage generator’s anchors.count must also lie within its capacity() (min_gates <= anchors.count[e] <= gates) for every env.

A pipeline with HOST stages (run eagerly, between the graph launches of the rest of the pipeline on CUDA) waives hot_path.topology and hot_path.buffers (reported as skipped: its launches may depend on host data); hot_path.host_sync and hot_path.allocation still run, with the HOST stages’ ctx.host() operations sanctioned. A legacy capturable=False generator spec (which cannot use ctx.host()) also waives hot_path.host_sync and hot_path.allocation. capture.caller then checks that the pipeline refuses a caller’s capture and capture.self_graph the per-segment CUDA graphs it builds of its own outside a capture.

check_stage() runs a relax / width / elevation / validity stage inside the default track pipeline (track_gen.pipeline.presets.track(), the stage in its slot; a validity check after validity.default), and an order / stations stage (or a validity check that reads gate channels) inside the default gate pipeline (track_gen.pipeline.presets.gates(): an order on the stations of the bezier track generator, stations on its raw loop, relaxed_stations for a stage that reads cl.*; a gate check after validity.gate_default; with pipeline="gates", an elevation stage on the bezier gate generator’s gates, its ports rewired to gates.*), through the same machinery, with these checks (no diversity / output.closed: they are the generator’s business):

output.accepted (relax / width / elevation / order / stations)

The pipeline accepts some env (valid fraction > 0). When it accepts none, the same pipeline with the slot’s neutral built-in (relax.none, width.constant, elevation.flat, order.raw, stations.raw) is run: if it accepts none either, the check is skipped (the stage did not reject them).

stage.isolation

After a generation, re-running ONLY the stage’s enqueue on the same inputs leaves the versions it read unchanged (it never writes a read port) and reproduces the versions it wrote (no stale scratch, no dependency on its own previous output).

validity.bits (validity)

The stage owns at least one fail bit (StageBuilder.fail_bit), sets no bit it does not own (the fail_mask of the pipeline without it, OR its own bits), and Track.valid == (fail_mask == 0).

width.values (width)

The stage’s version of cl.half_width (tap("cl.half_width", after=<id>)) is finite and >= 0 on every real point of every env the generator marks valid (floor violations are validity’s job, bit 4). A width stage that reads cl.half_width (a rewrite) runs right after width.constant; a required read port ("cl.*") is wired with ports=.

determinism, batch_independence, statelessness, hot_path.* and capture.* are the generator kit’s, on the whole Program (capturable: the built Program’s).

check_generator(name, params=None, *, devices=None, num_envs=8, seeds=(11, 1009), config=None, capture_modes=('warp', 'external', 'torch'), skip=(), modules=())[source]

Run every conformance check (CHECKS ["generator"]) on the registered track generator name.

Parameters:

  • name (str) – the registered generator name (registry.available("generator")).

  • params – its generator_params (a Params instance or a mapping of its fields); None: the defaults.

  • devices – Warp devices to run on (default: "cpu", plus "cuda:0" when Warp sees a CUDA device). The hot_path.* checks run on CPU devices, the capture.* checks on CUDA devices.

  • num_envs (int) – the batch size E (>= 2).

  • seeds (Sequence[int]) – two base seeds; batch A seeds env e with seeds[0] + e, batch B (the “different seeds” batch) with seeds[1] + e. They must not overlap.

  • config (Mapping | None) – extra TrackGenConfig fields (e.g. {"num_points": 64} to keep a slow generator fast); generator, device and num_envs are set by the kit.

  • capture_modes (Iterable[str]) – the modes of capture.caller (see assert_caller_capturable()).

  • skip (Iterable[str]) – check names (or a prefix such as "hot_path") to report as skipped.

  • modules (Iterable[str]) – more dotted module prefixes for hot_path.stream to probe, besides track_gen and the top-level package of the module defining the generator (registry.describe(kind, name).origin), e.g. a helper package of yours with another top-level name.

Returns:

A Report; call .raise_for_failures() to fail on any failed check.

Raises:

  • ValueError – invalid arguments (unknown check names, overlapping seeds, …), or an unknown generator name.

  • track_gen.plugin.PluginLoadError – the name belongs to a quarantined plugin (chained from the plugin’s load error).

Return type:

Report

check_gate_generator(name, params=None, *, devices=None, num_envs=8, seeds=(11, 1009), config=None, capture_modes=('warp', 'external', 'torch'), skip=(), modules=())[source]

As check_generator(), for the registered gate generator name (a gate_generator Stage or a legacy GateGeneratorSpec; CHECKS ["gate_generator"]). config holds extra GateGenConfig fields; without a gate_ordering the kit uses the generator’s preferred["order"] when it supports it, else "ccw" when supported, else its first supported ordering.

Parameters:
Return type:

Report

check_stage(kind, name, params=None, *, upstream=None, ports=None, pipeline=None, devices=None, num_envs=8, seeds=(11, 1009), capture_modes=('warp', 'external', 'torch'), skip=(), modules=())[source]

Run every conformance check (CHECKS [kind]) on the registered relax / width / elevation / validity / order / stations stage name, inside the default track pipeline (track_gen.pipeline.presets.track()) or, for an order / stations stage and a validity stage that reads gate channels, the default gate pipeline (track_gen.pipeline.presets.gates()). pipeline="gates" runs an elevation stage on gates (presets.gates( elevation=...), its ports rewired to gates.*; the default anchors generator bezier unless upstream names another): call it once per pipeline to check both.

The stage takes its slot (relax=, width=, elevation=, order=, stations=); a validity check runs after the default checks (upstream["checks"], default ("default",) / ("gate_default",)) with id name, and a validity stage that initialises valid.fail (writes it without reading it) replaces them. An order runs on the stations of the bezier TRACK generator (source="stations": no ordering whitelist), a stations stage on its raw loop (source="relaxed_stations" when it reads cl.*).

Parameters:

  • kind (str) – "relax", "width", "elevation", "validity", "order" or "stations".

  • name (str) – the registered stage name (registry.available(kind)).

  • params – its params (a Params instance or a mapping of its fields); None: the defaults.

  • ports (Mapping[str, str] | None) – the stage’s StageRef.ports (wires a required read port, e.g. {"src": "cl.acme.hw"} for width.from_channel, whose generator then goes in upstream).

  • upstream (Mapping | None) – more track() / gates() keyword arguments (generator, generator_params, num_points, n_max, half_width, the other kinds’ slots, checks (the checks BEFORE a validity check under test), on_invalid, source). num_envs, device and the slot under test are set by the kit.

  • pipeline (str | None) – None (default: the gate pipeline for an order / stations stage and a validity stage that reads gate channels, else the track pipeline), "track" or "gates". "gates" admits elevation, validity, order and stations stages, "track" relax, width, elevation and validity stages.

  • devices – as check_generator().

  • num_envs (int) – as check_generator().

  • seeds (Sequence[int]) – as check_generator().

  • capture_modes (Iterable[str]) – as check_generator().

  • skip (Iterable[str]) – as check_generator().

  • modules (Iterable[str]) – as check_generator().

Returns:

A Report; call .raise_for_failures() to fail on any failed check.

Raises:

  • ValueError – invalid arguments (an unknown kind, a pipeline the kind does not run in, a reserved upstream key, unknown check names, overlapping seeds, …), or an unknown stage name.

  • track_gen.plugin.PluginLoadError – the name belongs to a quarantined plugin (chained from the plugin’s load error).

Return type:

Report

CHECKS: Mapping[str, tuple[str, ...]]

The check names of each registry kind ("generator", "gate_generator", "relax", "width", "elevation", "validity", "order", "stations") and of "component" (check_component()), in run order. A read-only mapping.

CHECKS

The same mapping as track_gen.testing.CHECKS.

class Report[source]

Bases: object

Every CheckResult of one check_generator() / check_gate_generator() run (iterable, len()).

kind and name identify the generator or stage; results is the list of CheckResult, one per check and device.

kind: str
name: str
results: list
property ok: bool

Every check passed or was skipped.

property failures: list

The failed CheckResult s, in run order.

property failed_names: list

The names of the failed checks (each once, sorted).

result(name, device=None)[source]

The result of check name (on device, default: the first device).

Parameters:

  • name (str)

  • device (str | None)

Return type:

CheckResult

raise_for_failures()[source]

Raise one AssertionError listing every failed check; return self when all passed.

Return type:

Report

__init__(kind, name, results)
Parameters:
Return type:

None

class CheckResult[source]

Bases: object

The outcome of one check on one device.

name

the check name (see CHECKS).

Type:

str

device

the Warp device alias ("cpu", "cuda:0").

Type:

str

passed

True when the check passed or was skipped.

Type:

bool

message

why it failed (or was skipped); a short note otherwise.

Type:

str

skipped

the check did not run (not applicable, waived, or skip=).

Type:

bool

details

measurements, e.g. {"valid_fraction": 0.875} for output.closed.

Type:

Mapping[str, Any]

property status: str

"PASS", "FAIL" or "SKIP" (a skipped check also has passed).

__init__(name, device, passed, message='', skipped=False, details=<factory>)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)

Component checks

Does a component honour the contract? check_component.

One call runs a Component (a built-in adapter or your own) through every check the component contract implies (see Writing a Component) and returns a Report:

from track_gen import tools
from track_gen.testing import check_component

report = check_component(lambda cc: tools.Progress(),
                         upstream=[lambda cc: tools.Checkpoints(0.5)])
report.raise_for_failures()

The kit builds the producer (source: "track" = presets.track, "gates" = presets.gates(gate_width=0.1) (passable gates), or a PipelineSpec / TrackGenConfig / GateGenConfig) with explicit seeds and on_invalid="keep", attaches the upstream components (factories, in order) and then the component under test (factory(cc): the kit passes the ComponentContext, so a component that owns device arrays builds them for the kit’s device and E) through ComponentContext.for_producer. Every check and every capture drives the whole chain in order, upstream first (on_generate of each, then reset of each, then step of each), and compares only the component under test.

Owned vs aliased. The kit finds every wp.array reachable from the component’s result (a bare array; dataclass fields, mapping values, named-tuple, tuple and list items, __slots__ and instance attributes, nested ones named with dots, e.g. contact.oob); a result that exposes none fails construct. Each one is classified by pointer: aliased when it is an array of the read target, of an upstream result or of the producer’s result (e.g. GateCheckpoints: every array), owned otherwise. Only owned arrays are compared by the byte checks, and poisoned by the capture checks.

Inputs. Every input name the chain declares gets one [E * slots] array per run. drive(step, inputs, source) writes them before each step; the default drive knows position ([E] vec3f: walking along the source’s centerline, or between consecutive gates, with lateral offsets that leave the road, or the gate opening, now and then), box_position (that point plus per-slot offsets, every fourth slot inactive (NaN) on some steps), orientation (the yaw of the motion) and half_extents ((0.03, 0.015)). Any other input needs drive= (construct fails naming it).

Checks (CHECKS ["component"]), per device:

construct

The producer and the chain build (cc.add checks the declaration: name, source, requires / optional protocols, topology, inputs, after); layout and result are set (with at least one wp.array); the source has a valid env (else every check would be vacuous; details["valid_envs"]); the inputs are allocated and bound; a warm-up on_generate(ctx, None) and reset(ctx, all); then the attach contract: a chain attached now (no warm-up) steps to the same owned bytes as the warmed one (attach ends with on_generate(cc.ctx, None) and initialises the per-env state as reset does).

requires.complete

On the read target: a proxy exposing only the members of requires + optional (plus layout) attaches and gives the same owned bytes over a trajectory; for each required member, a proxy without it makes cc.add raise TypeError; for each optional member, a proxy without it attaches. A component that declares no protocol and reads the producer (no requires, optional or source) is attached to an empty read target (only a layout): it must attach and give the same owned bytes, then the check reports “reads no source” (skipped); reading an undeclared member fails. One that reads another component’s result (source) without declaring a protocol of it is skipped (nothing to check against).

topology

A proxy whose layout is a closed=False copy: "closed" -> cc.add raises ValueError, "any" -> it attaches. Skipped without requires / optional (the check is the protocols’).

result.stable

step() returns component.result (identity); the result’s arrays keep their pointers after on_generate, reset and step; each aliased array IS an array of the read target / upstream / producer.

determinism

Two producers and chains with the same seeds and drive: identical owned bytes every step.

statelessness (stateful = False)

After a history of steps, step equals a FRESH component’s (attached now to the same upstream instances: its attach ends with on_generate(ctx, None)) on the same inputs. Skipped for a stateful component.

reset.exact (stateful = True)

reset(ctx, m) mid-trajectory: from then on the masked envs equal a fresh component’s, the unmasked ones a never-reset twin’s, every step.

regenerate.masked

The producer advanced with mask m (generate(advance=True, mask=m)), the chain’s on_generate(ctx, m) and reset(ctx, m): from then on the unmasked envs equal a never-advanced twin, the masked envs a fresh chain on a twin producer advanced with m before attaching; details["masked_rows_changed"] reports non-vacuity.

regenerate.self_reset (stateful, reads a source that carries generation)

The same run WITHOUT reset gives the same bytes as with it (the generation-aware self-reset). Not applicable otherwise (skipped).

invalid.sentinel

K: the producer’s result with half the envs failed (fail_mask |= 1) by the commit’s own kernel at poison=0 (geometry kept, valid = 0); P: the same at poison=1 (count = 0, NaN geometry). The same chain on each, driven from K’s geometry: byte-identical owned outputs over the trajectory (an invalid env reads exactly like a poisoned one). Skipped when the result has aliased arrays and no owned one (its consumers are checked with it upstream).

capacity.clamp (a component that reads a cc.capacity key)

Attached with the key scaled down (by 4, or more until some env truncates when the result or the utility exposes truncated): each env’s owned rows equal those of the same chain attached to that env’s one-env view (no write past its own rows), and some env reports truncation. Skipped when no key is read.

source.view

The chain attached to big.view(1, E + 1) of a byte-poisoned E + 2 result holding the batch: equal owned bytes to the standalone run over a trajectory, and every byte of big unchanged (a component never writes its source, nor outside its env range). Skipped without requires / optional.

hot_path.host_sync / .allocation / .stream / .capture (CPU)

on_generate(ctx, None), on_generate(ctx, m), reset(ctx, m), step(ctx) twice under hot_path_guard() (modules=("track_gen", <the component's package>)), on a chain whose hot methods never ran before (a throwaway chain ran them first, so module loads stay outside the guard): work done only on a first hot call (a lazy allocation, a one-time readback) is caught.

hot_path.topology (CPU)

Two calls of each hot method emit the same LaunchTrace; also step with the rebind flag of the wrapped utility armed (a LaunchContext reporting capturing) vs unarmed: arming it changes no emission.

hot_path.buffers (CPU)

Every pointer the hot calls launch on belongs to the component, its upstream, the source, the inputs or the masks (buffers_of()); and on_generate(ctx, None), on_generate(ctx, m) and reset(ctx, m) launch on no pointer of a bound input (a Course runs them in its regeneration body, whose own CUDA graph is keyed on (advance, mask) only and would bake the input: read inputs in step; a respawn snapshot of an input belongs in step, with reset arming a per-env flag that step consumes). The upstream components are bound to the same input arrays, so an input reached through an upstream result that aliases it is caught too.

capture.caller (CUDA)

Per mode of modes ("warp", "external", "torch"): on a chain whose hot methods never ran eagerly, the chain a Course records: [any_set(done) -> flag; ctx.skip_unless(flag, [producer.enqueue(advance=True, mask=done) -> chain on_generate(ctx, done) -> chain reset(ctx, done)]) -> chain step(ctx)], recorded on a side stream (mempool disabled, wp.capture_launch forbidden) and replayed steps times with host-written inputs and a deterministic done schedule whose every fourth mask is all zero (the skipped branch replays); the masked part is a CUDA conditional-node body where the capture allows one ("warp" with CUDA >= 12.4; unconditional otherwise), so on_generate / reset must emit only kernels, fills and device-to-device copies on ctx’s stream; an eager twin runs the same chain, skipping the masked part on an all-zero done exactly when the capture recorded it under a conditional node (the replay skips it there), so both runs skip or run the same on_generate / reset; before each replay / eager iteration the owned arrays of the component under test are poisoned in BOTH runs (a skipped on_generate leaves the poison in both); equal every step. details[mode]: "ok" or "skipped: <reason>" (torch modes without torch).

capture.self_reset (CUDA, as regenerate.self_reset)

The same graph without the chain’s reset: equal to the eager twin WITH it.

A component that declares no read (no requires, optional or source: a raw-mode DiscCollision) and attaches to an empty read target reports requires.complete, topology and source.view as skipped (“reads no source”); regenerate.masked and invalid.sentinel run for every component (they pass trivially for one that reads nothing), and regenerate.self_reset for every stateful one that reads a source whose read target carries generation (optional=(HasGeneration,) alone qualifies).

check_component(factory, *, source='track', upstream=(), devices=None, num_envs=8, steps=32, seeds=None, drive=None, modes=('warp', 'external', 'torch'))[source]

Run a component through the component contract’s checks (see the module docstring) and return the Report.

Parameters:

  • factory (Callable) – factory(cc) -> Component: a fresh, unattached component (cc is the ComponentContext it will be attached through: its layout gives the device and E).

  • source – "track" (presets.track), "gates" (presets.gates with gate_width=0.1), or a PipelineSpec / TrackGenConfig / GateGenConfig (its num_envs / device replaced; a spec’s on_invalid forced to "keep").

  • upstream (Sequence[Callable]) – Factories of the components it reads (after, source), attached first, in order.

  • devices – Devices to run on (default: "cpu", plus "cuda:0" when available).

  • num_envs (int) – E (>= 2).

  • steps (int) – Trajectory length of the byte checks (>= 6).

  • seeds – None (7 * e + 3), an int base seed (base + e) or E values.

  • drive (Callable | None) – drive(step, inputs, source) writing the bound inputs (name -> wp.array) before each step; default: see the module docstring.

  • modes – The capture modes of capture.* (CUDA).

Returns:

The Report (kind="component"; the checks of CHECKS["component"] per device).

Return type:

Report

Hot-path guards

Hot-path contract guards: no host sync, no allocation, stream-explicit work.

A capturable generator’s hot path (its generate / a Stage’s enqueue, and every library call it makes) must be recordable into a caller’s CUDA graph. On a GPU that is checked by capturing it (track_gen.testing.capture). Without one, the guards here check the call shape instead, on any device:

  • no_host_sync(): no wp.synchronize*, no host readback (wp.array.numpy() / wp.array.list(), np.asarray of a CPU array) and no wp.copy between two devices;

  • no_allocation(): no Warp array allocation (wp.empty / zeros / full / ones / clone / array(data) / empty_like / zeros_like / …, seen at Warp’s allocation chokepoint, so every path is covered) and no native Warp object (wp.HashGrid and HashGrid.reserve, wp.Mesh, wp.Volume, wp.Bvh);

  • stream_explicit(): every wp.launch / wp.launch_tiled / wp.Launch.launch / wp.copy issued from the probed modules passes stream= (also when it reaches Warp through a Warp utility the probed code called), and none of them calls wp.array.assign() / zero_() / fill_() (their copy or memset targets Warp’s current stream) or enters wp.ScopedStream. The check sees that a stream= is passed, not which stream: stream=None (what a CPU LaunchContext passes) and a launch through a context other than the call’s are only caught by recording into a real capture on CUDA (track_gen.testing.capture);

  • capture (in hot_path_guard()): no nested wp.ScopedCapture / capture_begin / capture_end / capture_launch / capture_if / capture_while (a graph launch cannot be recorded into another capture).

hot_path_guard() combines all four. Each returns a HotPathGuard, a context manager that yields itself:

from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.testing.contract import hot_path_guard

gen = TrackGenerator(presets.track(num_envs=64, device="cpu"))  # may allocate, sync
with hot_path_guard(modules=("track_gen", "acme_tracks")):      # + your package
    gen.generate()              # raises HotPathViolation on the first breach

With strict=True (the default) a violating call raises HotPathViolation before it runs, and the block also raises on exit if a violation was recorded (so a broad except in the code under test cannot hide one). With strict=False the guard only records: every call goes through and HotPathGuard.violations lists them.

The guards patch Warp’s entry points (wherever a module binds them, so from warp import launch is seen too) for as long as one is active, and dispatch on a context variable: calls made from other threads, or from a context without a guard, go straight to Warp and are not seen. While no_host_sync() is active the TRACKGEN_DEBUG_SYNC debug sync is switched off (it is a deliberate, opt-in host sync); the switch is the environment variable, so it is off for the whole process until the guard exits.

Guard a call that runs eagerly: on CPU every generation does; on CUDA use enqueue() (a CUDA generate() replays the generator’s own graph, which the capture check reports as a wp.capture_launch).

The host poll of LaunchContext.bounded_loop (its lowering (a)) is the one sanctioned eager readback of a capturable path: it runs only on CPU outside every capture (or under the library’s private forced-lowering seam). The guard lets it through, and lists it in HotPathGuard.sanctioned, only when it comes from the library’s poll function called by the loop’s window driver. A HOST stage’s host operations (HostOps, ctx.host()) are sanctioned the same way: the syncs and readbacks the HostOps methods make themselves pass and are listed in sanctioned. Any other sync or readback, a plugin’s own included (also one made inside a loop body or a HOST stage), is a violation.

hot_path_guard(*, host_sync=True, allocation=True, stream=True, capture=True, modules=('track_gen',), allow=(), strict=True)[source]

Every hot-path check at once (see HotPathGuard for the arguments).

Parameters:
Return type:

HotPathGuard

no_host_sync(*, allow=(), strict=True)[source]

No wp.synchronize* and no wp.array.numpy() / list() in the block.

allow: (api, module, function) keys let through (and recorded in HotPathGuard.allowed), e.g. a known readback scheduled for removal.

Parameters:
Return type:

HotPathGuard

no_allocation(*, strict=True)[source]

No Warp array allocation in the block.

Parameters:

strict (bool)

Return type:

HotPathGuard

stream_explicit(modules=('track_gen',), *, strict=True)[source]

Every wp.launch / wp.copy from modules passes stream=, and modules call no wp.array.zero_() / fill_() and enter no wp.ScopedStream.

Parameters:
Return type:

HotPathGuard

The four functions above are the entry points; the class below is the type they return.

class HotPathGuard[source]

Bases: object

Records (and, when strict, refuses) hot-path contract breaches in a block.

The type hot_path_guard() / no_host_sync() / no_allocation() / stream_explicit() return (build one with them); use it as a context manager (it yields itself). A guard can be entered again after it exited; its record is reset on every entry.

Parameters:

  • host_sync – check host syncs and readbacks.

  • allocation – check Warp array allocations.

  • stream – check that work from modules is stream-explicit.

  • capture – check for nested graph captures and graph launches.

  • modules – dotted module prefixes whose launches, copies, memsets and stream scopes the stream check covers ("acme" covers acme and acme.kernels). Host syncs, allocations and captures are checked from every caller.

  • allow – (api, module, function) keys of host syncs / readbacks to let through and record in allowed instead (api without the call parentheses, e.g. ("wp.array.numpy", "acme.gen", "generate"); the function is the direct caller’s qualified name, "Class.method" for a method on Python >= 3.11).

  • strict – raise HotPathViolation at the first breach (default), or only record.

__init__(*, host_sync=True, allocation=True, stream=True, capture=True, modules=('track_gen',), allow=(), strict=True)[source]
Parameters:
Return type:

None

violations: list[Violation]

Every breach recorded in the last (or current) block, in call order.

allowed: list[tuple]

The allow= keys that were hit (let through), in call order.

sanctioned: list[tuple]

The sanctioned host syncs / readbacks that were let through (their (api, module, function) keys), in call order: bounded_loop host polls and a HOST stage’s HostOps operations.

property ok: bool

No breach was recorded.

of(category)[source]

The recorded breaches of one category.

Parameters:

category (str)

Return type:

list[Violation]

allocated()[source]

The arrays allocated inside the block (kept alive by the guard).

Return type:

list

report()[source]

Every recorded breach, grouped by category, one per line.

Return type:

str

raise_for_violations()[source]

Raise HotPathViolation listing every recorded breach, if any.

Return type:

None

class Violation[source]

Bases: object

One contract breach recorded by a HotPathGuard.

category

"host_sync", "allocation", "stream" or "capture".

Type:

str

api

the Warp call, e.g. "wp.array.numpy()", "wp.empty()", "wp.launch without stream=".

Type:

str

site

where it was called from (the first frame outside Warp for a call Warp made on the caller’s behalf, e.g. an allocation).

Type:

track_gen.testing._interpose.Site

key

(api, module, function) of the direct caller, the form no_host_sync()’s allow= takes.

Type:

tuple

array

for an allocation, the wp.array it created (kept alive by the guard, see HotPathGuard.allocated()); None otherwise.

Type:

object | None

__init__(category, api, site, key=(), array=None)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
class Site[source]

Bases: object

A source location: module.function:line.

module

the module’s __name__.

Type:

str

function

the function’s qualified name ("Class.method" for a method on Python >= 3.11).

Type:

str

line

the line number of the call.

Type:

int

__init__(module, function, line)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
CATEGORIES: tuple[str, ...] = ("host_sync", "allocation", "stream", "capture")

The Violation categories, in report order (the argument of HotPathGuard.of()).

CATEGORIES

The same tuple as track_gen.testing.CATEGORIES.

exception HotPathViolation[source]

Bases: AssertionError

A guarded hot path broke the contract (raised by a strict guard).

Launch traces

Record the device work a call emits, and compare it across calls.

trace_launches() records every wp.launch / wp.launch_tiled / wp.Launch.launch and wp.copy / wp.array.assign issued in a block (by any code outside the warp package, in the current context) as a LaunchRecord: the kernel, the launch dimensions, the stream handle passed, the pointer of every array argument and the value of every other argument (the fields of a wp.struct argument included). Two properties of a capturable hot path follow from that record:

  • static topology (diff(), assert_same_topology()): a CUDA graph replays the launches it recorded, so two calls with the same inputs must emit the same launches, in the same order, with the same dimensions, buffers and scalar arguments; a launch sequence that depends on a Python counter, on data read back or on an allocation is not captured faithfully;

  • owned buffers (foreign_pointers(), assert_pointers_within()): every pointer must fall inside a buffer the producer owns (its result, scratch, seed and parameter buffers, collected with buffers_of()): a buffer allocated per call, or shared through a module global, breaks a recorded graph or two producers running side by side. Ownership is also checked from the other side by the conformance kit: two producers built alike must launch on disjoint buffers.

Example (on CPU, where every generation runs eagerly):

from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.testing import (assert_pointers_within, assert_same_topology,
                               buffers_of, trace_launches)

gen = TrackGenerator(presets.track(num_envs=64, device="cpu"))
with trace_launches() as first:
    gen.generate()
with trace_launches() as second:
    gen.generate()
owned = buffers_of(gen, root="gen")
assert_same_topology(first, second, buffers=owned)
assert_pointers_within(first, owned)
trace_launches(*, include_warp=False)[source]

A LaunchTrace to use as a context manager (it yields itself). Calls Warp makes internally, such as the copy inside wp.clone, are left out unless include_warp=True:

from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.testing import trace_launches

gen = TrackGenerator(presets.track(num_envs=64, device="cpu"))
with trace_launches() as trace:
    gen.generate()
print(len(trace.records), "launches and copies")
Parameters:

include_warp (bool)

Return type:

LaunchTrace

The type trace_launches() returns:

class LaunchTrace[source]

Bases: object

The launches and copies recorded in a trace_launches() block.

A sequence of LaunchRecord (len, iteration, indexing). Calls Warp makes internally (e.g. the copy inside wp.clone) are left out unless include_warp=True.

__init__(*, include_warp=False)[source]
Parameters:

include_warp (bool)

Return type:

None

kernels()[source]

The kernel (or "wp.copy") of every record, in order.

Return type:

list[str]

topology(buffers=None, *, scalars=True, streams=False, ignore=())[source]

One comparable key per record: (op, kernel, dim, arrays[, scalars] [, stream]). With buffers, a pointer inside one of them becomes (label, offset) and any other pointer "<foreign>"; pointers inside the ranges of ignore (an iterable of arrays, e.g. buffers allocated during the call, already reported elsewhere) become "<allocated>".

Parameters:
Return type:

list[tuple]

diff(other, *, buffers=None, scalars=True, streams=False, ignore=(), limit=5)[source]

Human-readable differences between this trace and other ([] when their topologies are identical; see topology()).

Parameters:
Return type:

list[str]

same_topology(other, **kwargs)[source]

Whether diff() finds no difference.

Parameters:

other (LaunchTrace)

Return type:

bool

foreign_pointers(buffers, *, ignore=())[source]

(record, argument index, pointer) of every non-null pointer outside buffers (and outside the arrays in ignore).

Parameters:
Return type:

list[tuple[LaunchRecord, int, int]]

class LaunchRecord[source]

Bases: object

One recorded launch or copy.

op

"launch" or "copy".

Type:

str

kernel

"<warp module>.<kernel key>" for a launch, "wp.copy" for a copy ("wp.array.assign" for an assign).

Type:

str

dim

the launch dimensions (a copy: (count or dest.size,)).

Type:

tuple

stream

the CUDA handle of the stream= argument; None when no stream (or None) was passed, e.g. on CPU.

Type:

int | None

arrays

(argument index, pointer) of every array argument (inputs then outputs; a copy: dest is 0, src is 1); a null array has pointer 0. A field of a wp.struct argument is keyed "<index>.<field>" ("3.staging").

Type:

tuple

scalars

(argument index, value) of every other argument (a copy: the dest_offset, src_offset and count), keyed the same way.

Type:

tuple

site

the code that emitted it (the caller of LaunchContext.launch / copy / fill when it went through one, else the caller).

Type:

track_gen.testing._interpose.Site

__init__(op, kernel, dim, stream, arrays, scalars, site)
Parameters:
Return type:

None

classmethod __new__(*args, **kwargs)
buffers_of(*objects, root=None, max_depth=10)[source]

Every wp.array reachable from objects: through attributes (__dict__ and __slots__), dataclass fields, wp.struct fields, mappings, sequences, the variables a function closes over and the object a bound method is bound to, up to max_depth levels. Warp objects other than arrays and structs (devices, streams, graphs, kernels), modules, types, a function’s globals, numpy data and registry specs (GeneratorSpec / GateGeneratorSpec, shared by every generator built from them) are not entered.

buffers_of(gen) collects a TrackGenerator’s / GateGenerator’s result, scratch, seed, parameter and input buffers.

Parameters:

  • root (str | Sequence[str] | None) – the label prefix of each object (default "obj0", "obj1", …).

  • max_depth (int)

Return type:

BufferSet

class BufferSet[source]

Bases: object

A set of device buffers (address ranges with labels) a pointer can be checked against. Build one with buffers_of(), or add() arrays yourself.

__init__()[source]
Return type:

None

add(array, label='')[source]

Add array’s [ptr, ptr + capacity) range (a null array is ignored).

Parameters:

label (str)

Return type:

BufferSet

locate(ptr)[source]

(label, byte offset) of the buffer holding ptr (the innermost one when buffers overlap), or None.

Parameters:

ptr (int)

Return type:

tuple[str, int] | None

labels()[source]

The label of every buffer in the set, in address order.

Return type:

list[str]

assert_same_topology(a, b, *, buffers=None, scalars=True, streams=False)[source]

Raise AssertionError when two traces differ (see LaunchTrace.diff()).

Parameters:
Return type:

None

assert_pointers_within(trace, buffers)[source]

Raise AssertionError when a launch or copy touches a buffer outside buffers.

Parameters:
Return type:

None

Caller capture (CUDA)

Record a producer into a caller’s CUDA graph and check the replays (CUDA only).

assert_caller_capturable() is the GPU half of the hot-path contract: it records a producer’s generation (producer.generate(stream=s) by default) into a capture the CALLER owns, then replays that graph and compares every output with an eager twin. The recording uses the discipline TrackGen’s own capture tests use:

  • a side stream that is not Warp’s current stream (so work that ignores the call’s stream escapes the capture and is caught);

  • the Warp mempool disabled on the device (Warp’s default CUDA allocator cannot allocate while a capture is active, so an allocation fails loudly);

  • wp.capture_launch replaced by a tripwire (a graph launch cannot be recorded into another capture, CUDA error 900);

  • outputs poisoned with FINITE, dtype-aware sentinels before each replay (outputs are NaN-padded and NaN compares equal in the checks, so a NaN poison would hide a missed write);

  • a reseed between replays: the graph must read the live seeds, so the replay after a reseed must equal a fresh twin built with the new seeds.

Capture modes (modes=):

  • "warp": a Warp wp.ScopedCapture on the side stream;

  • "external": a torch.cuda.graph capture registered with Warp through track_gen.interop.torch.registered() (ScopedCapture(external=True));

  • "torch": a plain torch.cuda.graph capture; the producer is called without stream= and must adopt torch’s capture stream.

The torch modes import torch inside the mode only, and are reported as skipped when torch (with CUDA) is not available. The helpers (side_stream(), mempool_disabled(), capture_launch_forbidden(), recording(), poison(), poison_all(), result_arrays()) are public building blocks for custom capture tests.

assert_caller_capturable(producer, *, modes=('warp', 'external', 'torch'), record=None, twin=None, reseed=None, seeded=True)[source]

Record producer into a caller’s capture, replay it and compare with an eager twin; reseed between replays and compare again.

Parameters:

  • producer – a CUDA TrackGenerator / GateGenerator (anything with generate(stream=), seeds (a SeedState) and config), whose epochs are all 0 (never advanced).

  • modes (Iterable[str]) – the capture modes to run (see the module docstring).

  • record (Callable | None) – record(stream) -> result recording one generation (default producer.generate(stream=stream); stream is None in the "torch" mode, where the producer must adopt torch’s capture stream). result is a dataclass of arrays (or a mapping name -> array).

  • twin (Callable | None) – twin(seeds) -> {name: numpy array}, the eager reference for the [E] int32 seeds (default: a fresh producer of the same type built from producer.config with those seeds, run with enqueue()).

  • reseed (ndarray | None) – the [E] seeds to reseed to between replays (default: the current seeds + 7919).

  • seeded (bool) – False for a producer whose output does not depend on its seeds (an ingestion generator, Stage.seeded = False): the reseed precondition is not required, and the second replay checks that a replay after a reseed still equals the (unchanged) eager batch.

Returns:

{mode: "ok" | "skipped: <reason>"}.

Raises:

  • AssertionError – a recording failed, a replay differs from the twin, or the recording self-captured or launched a graph.

  • ValueError – producer is not on a CUDA device, its epochs are not 0, an unknown mode, or (seeded=True) a reseed that does not change the batch (the reseed comparison would be vacuous).

Return type:

dict

CAPTURE_MODES: tuple[str, ...] = ("warp", "external", "torch")

Every capture mode assert_caller_capturable() knows, in run order.

CAPTURE_MODES

The same tuple as track_gen.testing.CAPTURE_MODES.

The building blocks of assert_caller_capturable(), for custom capture tests (import them from track_gen.testing.capture):

side_stream(device)[source]

A new wp.Stream on device that is not Warp’s current stream.

Return type:

Stream

mempool_disabled(device)[source]

Disable the Warp mempool on device for the block (an allocation inside a capture then fails instead of being served from the pool).

Return type:

Iterator[None]

capture_launch_forbidden()[source]

Make wp.capture_launch raise for the block (also the reference track_gen.exec.runner holds). Yields the list of attempted calls; on a normal exit the block also fails if one was made (an implementation that swallowed the AssertionError is caught too).

Return type:

Iterator[list]

recording(device, stream)[source]

The caller’s Warp capture on stream: mempool disabled, capture_launch forbidden. Yields the wp.ScopedCapture (its .graph after the block).

Context managers exit in reverse order: the tripwire is lifted before capture_end, and the mempool is restored after it.

Parameters:

stream (Stream)

Return type:

Iterator[object]

poison(arr)[source]

Fill arr with a finite, dtype-aware sentinel (-7 for signed and 7 for unsigned integers, -12345.0 for floats and every float component of a vector / quaternion), on Warp’s current stream. Never NaN: outputs are NaN-padded.

Parameters:

arr (array)

Return type:

None

poison_all(*results)[source]

poison() every array of every result, then synchronize the device(s).

Return type:

None

result_arrays(obj)[source]

name -> wp.array for every array field of a result dataclass (Track, GateSequence, …); None fields are skipped, and every array of a mapping field (Track.attributes) is listed as "<field>/<key>". A plain mapping of arrays is returned filtered the same way.

Return type:

dict

pytest plugin

track_gen.testing.pytest_plugin is opt-in (there is no pytest11 entry point): load it with pytest_plugins = ["track_gen.testing.pytest_plugin"] in a top-level conftest.py, or -p track_gen.testing.pytest_plugin. It imports pytest, so this reference documents it by hand rather than importing it.

track_gen_conformance(name, params=None, *, kind='generator', **kwargs)

A pytest fixture whose value is this factory. It runs check_generator() (with kind="gate_generator", check_gate_generator(); with kind one of "relax", "width", "elevation", "validity", "order", "stations", check_stage()) on name with params and kwargs, fails the test listing every failed check unless all pass, and returns the Report.

track_gen_check_component(factory, **kwargs)

A pytest fixture whose value is this factory. It runs check_component() on factory with kwargs, fails the test listing every failed check unless all pass, and returns the Report.

generator_names(kind='generator', *, builtin=None, exclude=()) → list[str]

The names registered under kind ("generator", "gate_generator" or one of the six stage kinds), sorted, without those in exclude. builtin=True keeps only the built-ins, builtin=False only the third-party ones.

parametrize_generators(kind='generator', *, argname='generator', builtin=None, exclude=())

pytest.mark.parametrize(argname, generator_names(kind, builtin=builtin, exclude=exclude)). It reads the registry at collection time, so the plugins its entry points name load then.