Testing a Plugin¶
Before you start
What you write: one conformance test per plugin name (a
check_*call, or the pytest fixture), plus tests of your plugin’s own properties.What you need to know: the contract of what you test: a track generator (Writing a Generator), a gate generator (Writing a Gate Generator), a pipeline stage (Writing a Stage) or a component (Writing a Component).
What to run: a CPU job on every pull request and, if you can, a CUDA job (A CI recipe), both with
TRACK_GEN_PLUGINS=strict.
track_gen.testing is the conformance kit for generator, stage and
component plugins. One call runs a registered generator through every check the
generator contract implies: output shapes and NaN
padding, determinism, batch independence, no hidden state, no host sync or allocation on
the hot path, a static launch sequence, the freeze of an iterative (BOUNDED) loop
and, on CUDA, recording into a caller’s graph (a capture). check_gate_generator
does the same for a gate generator, check_stage for a relax, width, elevation,
validity, order or stations stage (see Pipeline stages: check_stage), and check_component
for a runtime component (see Runtime components: check_component). TrackGen’s own test suite runs
the same checks over every built-in generator and stage, so a plugin is held to the
standard the built-ins meet. The kit’s functions are listed in Conformance kit.
The kit checks the contract, not the quality of the tracks: yield, shape variety and
speed are measured by benchmarks/compare_generators.py (see
How a generator is judged), and your generator’s own properties need tests of
their own (see Testing your plugin’s own properties).
track_gen.testing is a subpackage outside track_gen.__all__: import it
explicitly. import track_gen does not load it, and importing it imports neither torch
(only the torch capture modes do, and they are skipped when torch is missing) nor pytest
(only the opt-in pytest plugin does).
Quick start¶
With the Stage quick start’s acme.wobble registered (Writing a Generator):
from track_gen.testing import check_generator
report = check_generator("acme.wobble", params={"lobes": 5})
report.raise_for_failures()
check_generator looks the generator up in the registry, so the plugin must be
installed (its entry point registers it) or registered by hand first, as in the
generator guide. It runs every check on cpu, and on
cuda:0 too when Warp sees a CUDA device, and returns a Report.
raise_for_failures() raises one AssertionError listing every failed check, and
print(report) prints the whole table (here on a machine without CUDA; a skipped check
says why):
generator 'acme.wobble': 18/18 checks passed, 4 skipped
cpu construct PASS
cpu output.shape PASS
...
cpu attributes.written SKIP skipped: no attributes
...
cpu params.per_env PASS
cpu bounded.absorbing SKIP skipped: not applicable (the unit runs no ctx.bounded_loop with active=)
...
cpu hot_path.topology PASS
cpu hot_path.buffers PASS
In pytest, load the kit’s plugin from your top-level conftest.py and call the
track_gen_conformance fixture, which fails the test with the same list:
# conftest.py
pytest_plugins = ["track_gen.testing.pytest_plugin"]
# test_conformance.py
def test_wobble_conforms(track_gen_conformance):
track_gen_conformance("acme.wobble", params={"lobes": 5})
Run that job with TRACK_GEN_PLUGINS=strict, so a plugin that fails to load fails the
run instead of being quarantined with a warning (see Broken plugins: quarantine and TRACK_GEN_PLUGINS).
The checks¶
Each check has a stable name, so a test can assert which check a generator fails.
track_gen.testing.CHECKS["generator"] lists them in run order. Every check runs once
per device, except that the hot_path.* checks and host.require_caller_capturable
run on CPU devices only and the capture.* checks and bounded.lowering on CUDA
devices only. On CPU every generate() runs its Python
eagerly, so the guards see every call it makes; on CUDA the generator replays a graph,
and recording it into a real capture is the stronger test. A failed check’s message
names the env, the field or the call site (module.function:line), and the rule it
broke.
Check |
Passes when |
Usual causes of a failure |
|---|---|---|
|
The generator resolves and validates its params, allocates its buffers, runs
its warm-up pass and one more |
An exception in |
|
The raw centerline is |
Writing a flag other than |
|
The pipeline built with |
A kernel that writes outside |
|
The raw centerline is finite for every env the generator marked valid, and every
per-point |
A NaN from a zero-length normalisation or a division by zero; a scratch value read
before it is written; a point the kernel never writes (a launch one point short).
Write |
|
Every per-point |
The pipeline writes these fields from the raw output, so a failure points at a raw output the pipeline mishandles, or at a pipeline bug: report it with the seeds. |
|
|
As for |
|
Each valid env’s raw loop closes (the gap from its last point back to its first is at most 3 times its median step, or 1.5 times its longest step, so a straight that closes the loop passes when the loop has one as long elsewhere), and the pipeline accepts at least one env. |
Writing an open curve, or one that spans less than a full turn. A loop far larger
than |
|
Every slot of every per-point attribute the generator writes ( |
A launch that covers |
|
Every exported attribute of the result ( |
A NaN written into a raw attribute (a non-finite real point). It does not
catch |
|
Two generators built with the same seeds give bitwise identical batches: the raw
output and every |
A random number not drawn from the env’s seed (Python’s |
|
Each env changes when its seed changes, and no two envs of one batch are
identical. Only envs valid in both batches are compared: two rejected envs may be
alike. |
A kernel that ignores |
|
Env |
Reading a neighbour’s row ( |
|
|
State carried from one call to the next: a scratch buffer read before it is
written, an output accumulated into, an output slot left unwritten, a Python
counter, seeds copied at construction instead of read from |
|
For the stage under test (the generator, or the stage |
A kernel reading |
|
For a generator (or the stage under test) that runs a |
An iteration after the freeze that still writes an output (a kernel without the
frozen-env early return, a converged env that keeps moving); |
|
The same loop, in every lowering of Iterative stages (BOUNDED), gives the bytes
of an eager generation: the host poll (a), a Warp capture with conditional
windows (b, skipped with the reason when |
As for |
|
For a pipeline with |
Not a plugin bug: the check covers the planner’s refusal. Report a failure. |
|
Two |
A readback to branch on generated data or to size a launch. Branch on the device
instead (a per-env flag), with fixed bounds; stop a loop with
|
|
Two |
A temporary allocated per call, or an allocate-if- |
|
Every |
Launching with |
|
Two |
A generator that captures and replays its own graph: the facade does that already, and a graph launch cannot be recorded into a caller’s capture. |
|
The two calls emit the same launch sequence: the same kernels in the same order, with the same dimensions, buffers and scalar arguments. |
Launch dimensions or a scalar argument taken from a call counter or from host data; buffers swapped from one call to the next (ping-pong across calls). A recorded graph replays the first sequence forever. |
|
Every array the two calls launch on belongs to the generator (its result,
scratch, seed and parameter buffers, |
A buffer shared by every generator instance: a module global, a variable the spec’s functions close over, a class attribute (two generators then overwrite each other, and a graph keeps the pointer). |
|
|
Anything the |
|
|
An output the graph never writes (the poison survives the replay); a host value baked at capture. |
The hot_path.* checks run their two generations with every bounded_loop window
unrolled, the superset of what any lowering launches, so a sync or an allocation in a
late window is seen even when the kit’s seeds would stop the loop early; one more
generation with the default lowering reports the sanctioned polls.
HOST generators and skipped checks. A HOST generator (a Stage whose capture
class is HOST) runs eagerly on every call, so its launch sequence may depend on host
data: hot_path.topology and hot_path.buffers are reported as skipped (“waived:
HOST stage(s) […] run eagerly every call”). hot_path.host_sync and
hot_path.allocation still run: the ctx.host() operations pass (counted in
details["host_ops"]), while a direct .numpy(), a wp.synchronize*() or a
per-call allocation in the stage fails. A legacy GeneratorSpec registered with
capturable=False cannot use ctx.host(), so its hot_path.host_sync and
hot_path.allocation are skipped too (“waived: capturable=False”). hot_path.stream
and hot_path.capture always apply (see Opting out of capture: capturable=False). On CUDA,
capture.caller checks the refusal instead: recorded inside a caller’s capture after a
marker copy, generate() must raise NotCapturableError naming the HOST stages,
and the replayed caller graph must hold the marker and nothing of the generation (the
outputs keep their poison, the seeds and epochs are unchanged). capture.self_graph
checks the generator’s own graphs, one per segment: graph stays None, a graph
exists exactly for each capturable segment, and replays into poisoned outputs equal the
eager batch, before and after a reseed. When construct fails, every other check is reported as not run,
and the report fails.
Arguments¶
check_generator(name, params=None, *, devices=None, num_envs=8, seeds=(11, 1009),
config=None, capture_modes=("warp", "external", "torch"), skip=(),
modules=())
params: the generator’sgenerator_params, a Params instance or a mapping of its fields;Nonegives the defaults. A generator whose parameters select different code paths (a fallback, a mode) needs one call per path.devices: a device or a list of devices. The default is"cpu", plus"cuda:0"when Warp sees a CUDA device.num_envs: the batch sizeE, at least 2 (diversity and batch independence compare envs).seeds: two base seeds. Batch A seeds envewithseeds[0] + e, batch B (the “different seeds” batch, also used to reseed) withseeds[1] + e. The two ranges must not overlap.config: extraTrackGenConfigfields, for example{"num_points": 64}to keep a slow generator fast. The kit setsgenerator,deviceandnum_envsitself; pass the Params either asparams=or asconfig["generator_params"], not both.capture_modes: the modescapture.callerruns (see Recording into a caller’s graph: track_gen.testing.capture).skip: check names, or a prefix such as"hot_path", to report as skipped instead of running. Use it sparingly, and say why next to the call.modules: more dotted module prefixes forhot_path.streamto probe, besidestrack_genand the top-level package of the module that defines the generator (registry.describe(kind, name).origin; afunctools.partialcounts as the function it wraps). Name a helper package of yours with another top-level name here.
An unknown check name, overlapping seeds or an unknown capture mode raises
ValueError; an unknown generator name raises the registry’s error
(PluginLoadError for a quarantined plugin, carrying its load error).
Reading a report¶
A Report holds one CheckResult per check and device, in run order, and iterates
over them. report.ok is True when every check passed or was skipped,
report.failures lists the failed results and report.failed_names their sorted
names. report.result(name, device=None) returns one result (on the first device
when device is None). A CheckResult has name, device, passed,
skipped, message (why it failed or was skipped), details (measurements, such
as valid_fraction for output.closed) and status ("PASS", "FAIL" or
"SKIP").
Here is the report of a GeneratorSpec that reads its phase array back with
.numpy(), allocates a temporary with wp.empty and copies it into the output with
a bare wp.copy, all inside generate (messages shortened):
generator 'acme.broken': 14/17 checks passed, 5 skipped
cpu construct PASS
...
cpu statelessness PASS
cpu params.per_env SKIP skipped: no dynamic parameters
...
cpu hot_path.host_sync FAIL 2 host sync/readback(s) in two generate() calls: wp.array.numpy() from acme_broken.generate:35. A capturable hot path must not synchronize ...
cpu hot_path.allocation FAIL 2 allocation(s) in two generate() calls: wp.empty() from acme_broken.generate:37. A capturable hot path must not allocate ...
cpu hot_path.stream FAIL 2 call(s) off the call's stream in two generate() calls: wp.copy without stream= from acme_broken.generate:39. Emit device work through the LaunchContext ...
cpu hot_path.capture PASS
cpu hot_path.topology PASS
cpu hot_path.buffers PASS
The hot path runs twice, so each breach is counted twice. Each message names the line to fix and ends with the rule it breaks and how to comply.
Gate generators¶
check_gate_generator takes the same arguments (config holds GateGenConfig
fields) and runs the same checks on the GateSequence, listed in
CHECKS["gate_generator"]. It checks anchors Stages (The anchors stage contract) and
legacy GateGeneratorSpec objects alike (a legacy spec emits its
DeprecationWarning). The raw output is the generator’s anchors, checked before
count[e]: tap("anchors.xy"), [E*A] wp.vec2f with A the generator’s
capacity()["gates"] (a legacy spec: the [E*G] staging buffer, G =
config.max_gates). output.count asks for 0 <= count[e] <= max_gates for every
env and count[e] >= min_gates for valid ones, and for a Stage also
capacity()["min_gates"] <= anchors.count[e] <= capacity()["gates"] for every env;
output.accepted (the gate pipeline accepts at least one env) replaces
output.closed. Without a gate_ordering in config, the kit uses the
generator’s default order, as GateGenConfig(gate_ordering=None) does: its
preferred["order"] when it accepts it, else "ccw" when it accepts it, else the
first order it accepts (a legacy spec: "ccw" if supported, else the first of its
supported_orderings). With the gate generator guide’s acme.scatter
(Writing a Gate Generator):
from track_gen.testing import check_gate_generator
check_gate_generator("acme.scatter").raise_for_failures()
Pipeline stages: check_stage¶
check_stage checks a stage of an open pipeline kind (see
Writing a Stage) by running it where it belongs: in the default track pipeline
of presets.track, in the slot of its kind, built through TrackGenerator; or, for
the gate kinds, in the default gate pipeline of presets.gates, built through
GateGenerator:
an
orderstage runs asorder=of a stations pipeline (generator="bezier",source="stations"), so no generator’s ordering whitelist gets in the way;a
stationsstage runs asstations=, onsource="relaxed_stations"when it readscl.*channels, else onsource="stations";a
validitystage that reads agates.*oranchors.*channel is appended after"gate_default"in a gate pipeline; any other runs in the track pipeline;an
elevationstage runs in the track pipeline, or withpipeline="gates"aselevation=of the default gate pipeline (thebeziergate generator unlessupstreamnames another), its ports rewired togates.*: theoutput.*checks read theGateSequenceandstage.isolationre-runs the stage ongates.cum/gates.perim/gates.countintogates.z. Run it once per pipeline to check both.
With the stages of Writing a Stage and Writing an ordering registered:
from track_gen.testing import check_stage
check_stage("relax", "acme.taubin", {"passes": 4}).raise_for_failures()
check_stage("validity", "acme.max_curvature",
upstream={"elevation": "acme.sinusoid_hills"}).raise_for_failures()
check_stage("order", "acme.reverse").raise_for_failures()
check_stage("elevation", "acme.sinusoid_hills").raise_for_failures()
check_stage("elevation", "acme.sinusoid_hills",
pipeline="gates").raise_for_failures()
check_stage("width", "acme.ripple").raise_for_failures()
check_stage("width", "from_channel", ports={"src": "cl.acme.hw"},
upstream={"generator": "acme.ribbon"}).raise_for_failures()
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=())
kindis"relax","width","elevation","validity","order"or"stations"(ValueErrorotherwise);nameis looked up in the registry like a generator’s (an unknown name raises the registry’sValueError, a quarantined plugin itsPluginLoadError).paramsis the stage’s params: a Params instance, a mapping of its fields, orNonefor the defaults.pipelineisNone(the pipeline the kind belongs to, as above),"track"or"gates"."gates"takeselevation,validity,orderandstationsstages,"track"relax,width,elevationandvaliditystages; any other value or combination raisesValueError.upstreamholds morepresets.trackkeyword arguments, the rest of the pipeline the stage runs in:generator,generator_params,num_points,n_max,half_width, the other kinds’ stages,on_invalid.num_envs,deviceand the slot under test (relax,widthorelevation) are the kit’s and raiseValueError. For the gate kinds,upstreamholdspresets.gateskeyword arguments instead, with the same reserved keys (orderorstations). For a validity stage,upstream["checks"]lists the checks that run before it (default("default",), or("gate_default",)for a check that runs in the gate pipeline), and the stage is appended to them under its own name as id; a validity stage that writesvalid.failwithout reading it replaces them instead.portsis the stage’sStageRef.ports: it wires a required read port ({"src": "cl.acme.hw"}forwidth.from_channel) or rewires another; the channel must exist in the pipelineupstreamdescribes (here a generator that carriesraw.acme.hw).devices,num_envs(at least 2),seeds,capture_modes,skipandmodulesare those ofcheck_generator(see Arguments).
The report reads like a generator’s, under CHECKS[kind]. The checks are those of the
generator kit, on the whole pipeline’s Track or GateSequence (output.* checks it as
check_generator does after the pipeline), without diversity and
output.closed (the raw loop is the generator’s business; a stage need not add
diversity), plus:
Check |
Passes when |
Usual causes of a failure |
|---|---|---|
|
The pipeline with the stage accepts at least one env. When it accepts none, the
same pipeline with the slot’s neutral built-in ( |
A relax stage that leaves NaNs or self-crossings, a width under |
|
After a generation, the stage’s write buffers are filled with a poison byte and
only its |
Writing the buffer |
|
The stage owns at least one fail bit, the bits it does not own are exactly those
of the same pipeline without it, and |
OR-ing a built-in bit ( |
|
The stage’s own |
A negative or NaN width (a NaN source channel of |
determinism, batch_independence and statelessness compare whole Track
batches, fail_mask, half_width and every attribute included. A relax stage’s
bead-identity contract (point i out is point i in, moved; see
the relax contract) is documented, not checked: XPBD’s own
drift along the loop reaches the size of an off-by-one shift on some envs, so no
threshold separates the two. The hot_path.* checks guard two generations of the
whole Program and probe track_gen and the stage’s own top-level package; the
capture.* checks record the whole Program into a caller’s capture
(assert_caller_capturable) and replay its own graph. With a HOST stage the
Program runs that stage eagerly in its own segment, and the checks change as for a
HOST generator (see HOST generators above).
bounded.absorbing and bounded.lowering probe the loops the stage under test
runs, so a stage that calls bounded_loop(active=...) while declaring STATIC
fails construct instead (see Writing an iterative (BOUNDED) stage).
Runtime components: check_component¶
check_component checks a component of track_gen.tools (see
Writing a Component): your own, or one of the built-in adapters, which TrackGen’s
suite checks the same way.
check_component(factory, *, source="track", upstream=(), devices=None, num_envs=8,
steps=32, seeds=None, drive=None, modes=CAPTURE_MODES) -> Report
factory(cc)returns a fresh, unattached component;ccis theComponentContextit will be attached to (itslayoutandcapacity), so a component that owns device arrays from construction can build them for the kit’s device andnum_envs.upstreamholds the factories of the components it reads (after,source), attached first, in order.sourceis"track"(presets.track, the default),"gates"(presets.gates(gate_width=0.1): passable gates), or aPipelineSpec/TrackGenConfig/GateGenConfig(itsnum_envsanddevicereplaced). The kit builds the producer with explicit seeds (seeds=None:7 * e + 3) andon_invalid="keep", and attaches the chain throughComponentContext.for_producer.drive(step, inputs, source)writes the bound inputs before each step. Every input name the chain declares gets one[E * slots]array per run. By defaultpositionwalks along the course (the centerline, or between consecutive gates) with lateral offsets that leave the road (or the gate opening) now and then,box_positionis that point plus per-slot offsets (every fourth slot inactive, NaN, on some steps),orientationthe yaw of the motion andhalf_extents(0.03, 0.015); any other input needsdrive=(constructfails naming it).
Every check drives the whole chain in order, upstream first (on_generate, reset,
step of each component, inside the recorded graph too), and compares only the
component under test. The kit finds every wp.array reachable from its result
(dataclass fields, mapping values, named-tuple, tuple and list items, __slots__ and
attributes, recursively; nested names are dotted) and classifies each by pointer:
aliased when it belongs to the read target, an upstream result or the producer’s
result, owned otherwise. Only owned arrays are poisoned and compared; aliased ones must keep
their pointers.
Check |
What it does |
|---|---|
|
Attaches the chain; checks the declaration’s types ( |
|
A proxy of the read target exposing only the declared members ( |
|
A read target whose |
|
|
|
Two producers and chains with the same seeds and drive: identical owned bytes at every step. |
|
After a history of steps, |
|
|
|
The producer advanced with a mask, then the chain’s |
|
Stateful components that read a source whose read target carries |
|
Half the envs failed by the commit’s own kernel, once keeping their geometry
( |
|
Components that read a |
|
Attached to |
|
|
|
Two calls of each hot method emit the same launch trace; |
|
Every launched pointer belongs to the component, its upstream, its read target,
its inputs or the masks; and |
|
Per capture mode: on a component whose hot methods never ran eagerly, records
the chain a |
|
As |
A component that declares no read (DiscCollision on raw disc arrays) and gives the
same bytes on an empty read target reports requires.complete, topology and
source.view as skipped (“reads no source”), never passed; regenerate.masked and
invalid.sentinel run for every component (trivially passing for one that reads
nothing). One that reads another component’s result (source=) without declaring a
protocol of it skips requires.complete, topology and source.view (nothing to
check them against). CHECKS["component"] lists the names in run order.
The pytest plugin¶
track_gen.testing.pytest_plugin is opt-in: it has no pytest11 entry point, which
would load it into every pytest run of every environment that installs TrackGen. Load it
with pytest_plugins = ["track_gen.testing.pytest_plugin"] in your top-level
conftest.py (or -p track_gen.testing.pytest_plugin on the command line). It
provides:
the
track_gen_conformancefixture, a factory:track_gen_conformance(name, params=None, *, kind="generator", **kwargs)runscheck_generator(check_gate_generatorwithkind="gate_generator",check_stagewithkindone of"relax","width","elevation","validity","order","stations") withkwargs, fails the test listing every failed check, and returns theReport;the
track_gen_check_componentfixture, a factory:track_gen_check_component(factory, **kwargs)runscheck_componentthe same way;generator_names(kind="generator", *, builtin=None, exclude=()), the sorted names registered underkind, a generator kind or one of the six stage kinds (builtin=Falsekeeps the third-party ones,builtin=Truethe built-ins);parametrize_generators(kind="generator", *, argname="generator", builtin=None, exclude=()), the matchingpytest.mark.parametrize. It reads the registry at collection time, which loads the installed plugins then.
import pytest
from track_gen.testing.pytest_plugin import parametrize_generators
ACME = ["acme.wobble", "acme.wobbly", "acme.star"]
@pytest.mark.parametrize("name", ACME)
def test_acme_generator_conforms(name, track_gen_conformance):
track_gen_conformance(name)
def test_scatter_gates_conform(track_gen_conformance):
track_gen_conformance("acme.scatter", kind="gate_generator")
# Every third-party generator installed in this environment, yours included:
@parametrize_generators("generator", builtin=False)
def test_installed_plugins_conform(generator, track_gen_conformance):
track_gen_conformance(generator)
builtin=False selects every installed third-party plugin, not only yours: list your
names explicitly when the environment may hold others.
Building blocks¶
check_generator is assembled from three public pieces. Use them directly to guard a
code path the checks do not drive: a helper your package exports, a custom
enqueue, a generator wrapped in your own simulation step.
Contract guards: track_gen.testing.contract¶
hot_path_guard() watches a block for the four kinds of breach the hot_path
checks report: host syncs and readbacks, Warp allocations, work off the call’s stream,
and nested captures or graph launches. no_host_sync(), no_allocation() and
stream_explicit(modules) watch one kind each.
from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.testing import hot_path_guard
gen = TrackGenerator(presets.track(num_envs=8, device="cpu", generator="acme.wobble"))
# construction may allocate and sync; the generations after it may not
with hot_path_guard(modules=("track_gen", "acme_tracks")):
gen.generate() # raises HotPathViolation at the first breach
moduleslists the dotted module prefixes whose launches, copies, memsets and stream scopes the stream check covers ("acme_tracks"coversacme_tracks.kernelstoo). Host syncs, allocations and captures are checked from every caller.With
strict=True(the default) a breach raisesHotPathViolation(anAssertionError) before the call runs, and the block raises again on exit if a breach was recorded, so a broadexceptin the code under test cannot hide one. Withstrict=Falsethe guard only records:guard.violationslists everyViolation(category,api,site),guard.of(category)filters them,guard.report()formats them andguard.raise_for_violations()raises.allowlets known host syncs through, as(api, module, function)keys (("wp.array.numpy", "acme_tracks.gen", "generate")), and records them inguard.allowed.The host poll of a
ctx.bounded_loop(its lowering on an eagercpucall, see Iterative stages (BOUNDED)) and the operations of aHOSTstage’sctx.host()(HostOps) are sanctioned: they go through and are recorded inguard.sanctioned, in neitherallowednorviolations. Any other sync is reported, one made inside a loop body or directly in aHOSTstage included.
The guards patch Warp’s entry points while a guard is active, wherever a loaded module
binds them (from warp import launch at import time included), 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. Classes are not replaced (their __init__ is), so
isinstance and subclassing keep working everywhere. While the host-sync check is
active, TRACKGEN_DEBUG_SYNC is switched off: its sync is deliberate and opt-in. 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, guard
gen.enqueue(), since generate() replays the generator’s own graph, which the
guard reports as a wp.capture_launch.
Launch traces: track_gen.testing.trace¶
trace_launches() records every wp.launch, wp.launch_tiled,
wp.Launch.launch(), wp.copy and wp.array.assign in a block as a
LaunchRecord: the kernel, the dimensions, the stream handle passed, the pointer of
each array argument and the value of each other argument (a wp.struct argument is
flattened into its fields, keyed "<index>.<field>"). buffers_of(gen) collects
every wp.array reachable from a producer (its result, scratch, seed, parameter and
input buffers) into a BufferSet, against which pointers are checked. It does not
enter a registry spec (GeneratorSpec): what a spec’s functions close over is shared
by every generator built from it.
from track_gen.testing import (assert_pointers_within, assert_same_topology,
buffers_of, trace_launches)
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) # hot_path.topology
assert_pointers_within(first, owned) # hot_path.buffers
LaunchTrace.diff(other, buffers=...) returns the differences as text ([] when
there are none), and foreign_pointers(buffers) the arguments that point outside
buffers. With buffers, each pointer is compared as the buffer it falls in (its
label and byte offset), and a pointer outside them as "<foreign>".
On cpu a bounded_loop stops at the first window whose check finds every env
converged, so the length of an eager trace depends on the data: compare two calls on the
same seeds, as above. The kit’s hot_path.topology compares calls with every window
unrolled.
Recording into a caller’s graph: track_gen.testing.capture¶
assert_caller_capturable(producer) is the CUDA half of the hot-path contract. It
records producer.generate(stream=s) into a capture the caller owns, replays it and
compares every output with an eager twin, built fresh from the producer’s config. The
recording follows the discipline of TrackGen’s own capture tests:
the capture runs on a side stream that is not Warp’s current stream, so work that ignores the call’s stream escapes the capture and shows as a wrong replay;
the Warp mempool is disabled on the device, so an allocation inside the capture fails instead of being served from the pool;
wp.capture_launchraises, since a graph launch cannot be recorded into another capture (CUDA error 900), and the recording must not create or replace the producer’s own graph;the outputs are poisoned with finite, dtype-aware values before each replay (
-7for signed integers,-12345.0for every float component), so a slot the replay never writes is caught; a NaN poison would hide it, since outputs are NaN-padded;the producer is reseeded between replays, and the replay after the reseed must equal a twin built with the new seeds: the graph reads the live seeds.
from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.testing import assert_caller_capturable
gen = TrackGenerator(presets.track(num_envs=8, device="cuda:0", generator="acme.wobble"))
status = assert_caller_capturable(gen) # {"warp": "ok", "external": "ok", "torch": "ok"}
modes selects the captures: "warp" (a wp.ScopedCapture on the side stream),
"external" (a torch.cuda.graph capture registered with Warp through
track_gen.interop.torch.registered) and "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 and report
"skipped: torch with CUDA is not available" without it. The producer must be on a
CUDA device and never advanced (every epoch 0); otherwise, or for an unknown mode, it
raises ValueError. record= replaces the recorded call (default
producer.generate(stream=stream)), twin= the eager reference (default: a fresh
producer of the same type and config, run once with enqueue()), and reseed= the
seeds of the second comparison (default: the current seeds plus 7919). The helpers it is
built from (side_stream, mempool_disabled,
capture_launch_forbidden, recording, poison, poison_all and
result_arrays) are public in track_gen.testing.capture for custom capture
tests; CUDA Graph — capture, replay, and buffer reuse describes the execution model they test.
Testing your plugin’s own properties¶
The kit checks the contract; what your generator is meant to draw, or what your stage is
meant to compute, needs tests of its own. Write them as ordinary tests on CPU
(device="cpu": every rule of the contract holds there too), and wrap any test that
adds, replaces or removes specs in registry.scoped(), which snapshots every kind on
entry and restores it on exit, so nothing leaks into the next test. With the two quick
starts of Writing a Generator (acme.wobbly, a GeneratorSpec, and
acme.wobble, a Stage):
import dataclasses
import numpy as np
from track_gen import TrackGenerator
from track_gen.pipeline import presets
from track_gen.plugin import registry
def _gen(name, **params):
return TrackGenerator(presets.track(num_envs=8, device="cpu", generator=name,
generator_params=params or None))
def test_wobbly_runs_on_cpu():
assert _gen("acme.wobbly").generate().valid.numpy().all()
def test_a_variant_without_leaking_it():
with registry.scoped():
base = registry.get("generator", "acme.wobbly")
registry.add(dataclasses.replace(base, name="acme.wobbly_variant"))
assert "acme.wobbly_variant" in registry.available("generator")
assert "acme.wobbly_variant" not in registry.available("generator")
def test_amp_takes_effect_at_the_next_generation():
gen = _gen("acme.wobble")
gen.params["generator"].set("amp", 0.3)
np.testing.assert_array_equal(gen.generate().center.numpy(),
_gen("acme.wobble", amp=0.3).generate().center.numpy())
A variant of a Stage is a subclass with another name. Test a dynamic parameter
against a generator built with the same value, as above: a set followed by a
generation must give its output bit for bit. For a stage, compare a tap of its output
(Taps) with a NumPy reference on a small batch, and for a check, that it
sets its bit exactly where it should. That the installed entry point registers your
names is a test of the package (Packaging a Plugin).
A CI recipe¶
Two lanes. The CPU lane runs on any runner and catches every breach the kit can see
without a GPU; the CUDA lane runs the same tests on a GPU runner, where the defaults add
cuda:0 and the capture.* checks.
# CPU lane, every pull request
TRACK_GEN_PLUGINS=strict python -m pytest -q
# CUDA lane, on a GPU runner (install torch with CUDA for the torch capture modes)
TRACK_GEN_PLUGINS=strict python -m pytest -q
Keep the tests that check your plugin’s own properties (Testing your plugin’s own properties) next to the conformance test.
What the kit does not see¶
Calls outside Warp’s Python API. The guards patch Warp’s Python entry points. A sync inside a C extension, a torch operation, a call made from another thread, or a call through a Warp function bound to a local variable or a closure before the guard was entered is invisible to them; so is polling a stream or an event (
is_complete). Only the CUDA lane’s real capture catches their effect.Which stream, on CPU.
hot_path.streamchecks thatstream=is passed.stream=None, or a launch throughLaunchContext.current(device)instead of the call’sctx, passes on CPU and failscapture.calleron CUDA.Stream discipline outside the probed packages.
hot_path.streamcovers launches fromtrack_genand from the top-level package of the module that defines the generator (registry.describe(kind, name).origin). Add a helper package of yours with another top-level name withmodules=, or guard it withstream_explicit(modules=...).Native handles shared between instances.
hot_path.bufferscompares array pointers. Awp.HashGrid/wp.Meshid passed as auint64is a scalar to it.A launch sequence that changes after the first call.
hot_path.topologycompares two calls made long after construction’s warm-up pass: a one-time initialisation launch on the first call only is not seen.A loop that never stops early for the kit’s seeds.
bounded.absorbingcan only compare an early stop with the full bound when some window is skipped: with seeds that never converge beforen_itersit reports “not exercised”. Passseeds=(or params) under which the loop does stop.CPU against CUDA. Determinism is checked per device. CPU and CUDA results may differ (see Hard rules), and the kit never compares them.
One configuration per call. The checks exercise the parameters you pass. Run the kit once per setting that selects a different code path.
Masked regeneration and dynamic parameters. The kit does not drive
generate(advance=True, mask=...), andparams.per_envonly checks that each env reads its own value of each dynamic field, at one set of values, eagerly; test what a value change does (and a curriculum in a graph) as Testing your plugin’s own properties shows. (A masked regeneration is exact by construction: the commit restores the envs outside the mask; see The commit.)A stage in another position or pipeline.
check_stageruns the stage in its default slot ofpresets.trackwith theupstreamyou give. A stage you insert elsewhere, or rewire withports, needs a test of that pipeline.Track quality. One valid env is enough for
output.closed; yield, shape and speed are forbenchmarks/compare_generators.py.