diff --git a/.github/workflows/benchmark-report.yml b/.github/workflows/benchmark-report.yml index 1bdbc8f1..20b2d7b8 100644 --- a/.github/workflows/benchmark-report.yml +++ b/.github/workflows/benchmark-report.yml @@ -29,7 +29,7 @@ jobs: run: uv sync --no-dev --group testing - name: Run benchmarks - run: uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json output.json + run: uv run pytest tests/benchmarks/ --benchmark-json output.json - name: Store benchmark results uses: benchmark-action/github-action-benchmark@v1 diff --git a/.github/workflows/benchmark_pr.py b/.github/workflows/benchmark_pr.py index 2e8ed445..42a811b9 100644 --- a/.github/workflows/benchmark_pr.py +++ b/.github/workflows/benchmark_pr.py @@ -31,15 +31,18 @@ def make_report(old_path, new_path, out_file, header=None): old = load_stats(old_path) new = load_stats(new_path) - # Merge on benchmark name - df = old[-1].merge(new[-1], on="Benchmark", suffixes=("_old", "_new")) + # Merge on benchmark name. Outer join so benchmarks that exist on only one + # side (added or removed by the PR) still show up in the report. + df = old[-1].merge(new[-1], on="Benchmark", suffixes=("_old", "_new"), how="outer") pct_change = 100 * (df["mean_new"] - df["mean_old"]) / df["mean_old"] - df["Percent Change"] = pct_change.map("{:+.2f}".format) + df["Percent Change"] = pct_change.map("{:+.2f}".format).where( + pct_change.notna(), "n/a" + ) # Format runtimes - df["mean_old"] = df["mean_old"].map("{:.5f}".format) - df["mean_new"] = df["mean_new"].map("{:.5f}".format) + for col in ("mean_old", "mean_new"): + df[col] = df[col].map("{:.5f}".format).where(df[col].notna(), "-") # Change column names to commit ids df = df.rename( diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 86ef968f..7f1d9fbe 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -85,7 +85,9 @@ jobs: if: steps.cache_baseline.outputs.cache-hit != 'true' run: | git checkout ${{ github.event.pull_request.base.sha }} - uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json baseline.json + # -o python_files: the base commit's pyproject.toml may predate the + # bench_*.py collection config, so set it on the command line instead. + uv run pytest tests/benchmarks/ -o "python_files=test_*.py bench_*.py" --benchmark-json baseline.json - name: Cache baseline results uses: actions/cache/save@v5 @@ -97,7 +99,7 @@ jobs: - name: Run benchmark on PR head run: | git checkout ${{ github.event.pull_request.head.sha }} - uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json pr.json + uv run pytest tests/benchmarks/ -o "python_files=test_*.py bench_*.py" --benchmark-json pr.json - name: Generate report id: report diff --git a/pyproject.toml b/pyproject.toml index 804ddce3..68070227 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,6 +74,10 @@ dev = [ [tool.setuptools_scm] +[tool.pytest.ini_options] +python_files = ["test_*.py", "bench_*.py"] +norecursedirs = ["tests/benchmarks"] + [tool.ruff] line-length = 90 target-version = "py310" diff --git a/tests/benchmarks/__init__.py b/tests/benchmarks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/benchmarks/_graph_builders.py b/tests/benchmarks/_graph_builders.py new file mode 100644 index 00000000..a0e7ec06 --- /dev/null +++ b/tests/benchmarks/_graph_builders.py @@ -0,0 +1,127 @@ +"""Builders for the synthetic graphs the benchmarks run against. + +A plain module rather than conftest.py, because benchmark modules import these builders +directly and a conftest is not importable. The conftest next door holds only fixtures. +""" + +import numpy as np +import polars as pl +import tracksdata as td +from skimage.draw import disk +from tracksdata.nodes import Mask + +from funtracks.data_model import Tracks +from funtracks.utils.tracksdata_utils import create_empty_graph + + +def make_tracks( + n_frames: int, + cells_per_frame: int, + frame_shape: tuple[int, int], + seed: int = 42, +) -> Tracks: + """Build a synthetic Tracks with segmentation masks. + + Produces a *solution* graph, not a candidate graph: cells are linked 1-to-1 between + adjacent frames so every node has out-degree <= 1. This matters because the user + actions assume a tracking solution -- UserDeleteEdge rejects a removal that would + leave out-degree > 1, and UserAddEdge rejects merges -- so they raise + InvalidActionError on the multi-successor candidate graphs that + compute_graph_from_seg produces. + + Args: + n_frames: Number of time points. + cells_per_frame: Cells generated per time point. + frame_shape: Spatial (y, x) shape of each frame. + seed: Seed for the random generator, for reproducible geometry. + + Returns: + A Tracks with pos, area, mask and bbox node attrs, iou edge attrs, and a + writable unmanaged "score" node attr (populated on every node) for + attribute-update benchmarks. + """ + rng = np.random.default_rng(seed) + graph = create_empty_graph( + node_attributes=[ + "pos", + "area", + td.DEFAULT_ATTR_KEYS.MASK, + td.DEFAULT_ATTR_KEYS.BBOX, + ], + node_default_values=[0.0, 0.0, 0.0, 0.0], + edge_attributes=["iou"], + position_attrs=["pos"], + ndim=3, + ) + + nodes: list[dict] = [] + node_ids: list[int] = [] + edges: list[dict] = [] + node_id = 1 + prev_frame_ids: list[int] = [] + + for t in range(n_frames): + frame_ids = [] + for _ in range(cells_per_frame): + cy = int(rng.integers(25, frame_shape[0] - 25)) + cx = int(rng.integers(25, frame_shape[1] - 25)) + radius = int(rng.integers(8, 15)) + + # Mask is stored in its own local frame with an offset bbox, so the + # disk is drawn centered in a (2r-1, 2r-1) box. + mask_shape = (2 * radius - 1, 2 * radius - 1) + rr, cc = disk( + center=(radius - 1, radius - 1), radius=radius, shape=mask_shape + ) + mask_arr = np.zeros(mask_shape, dtype=bool) + mask_arr[rr, cc] = True + bbox = np.array([cy - radius + 1, cx - radius + 1, cy + radius, cx + radius]) + + nodes.append( + { + "t": t, + "pos": [float(cy), float(cx)], + "area": float(mask_arr.sum()), + "solution": True, + td.DEFAULT_ATTR_KEYS.MASK: Mask(mask_arr, bbox=bbox), + td.DEFAULT_ATTR_KEYS.BBOX: bbox, + } + ) + node_ids.append(node_id) + frame_ids.append(node_id) + node_id += 1 + + # Link 1-to-1 with the previous frame to keep out-degree <= 1. + for prev_id, cur_id in zip(prev_frame_ids, frame_ids, strict=False): + edges.append( + { + "source_id": prev_id, + "target_id": cur_id, + "solution": True, + "iou": float(rng.uniform(0.1, 0.9)), + } + ) + prev_frame_ids = frame_ids + + # Register "score" before adding the nodes so every node gets a stored value. Adding + # the key afterwards leaves existing nodes null, and tracksdata applies the default + # lazily on read (_maybe_fill_null in graph/_rustworkx_graph.py, guarded by + # s.has_nulls()). That backfill costs two extra polars collects per read until the + # node is first written, which made the first update of any node ~2x slower than + # subsequent ones and turned the attribute benchmarks into a measure of null + # backfill rather than of the update path. + graph.add_node_attr_key("score", dtype=pl.Float64, default_value=0.0) + for node in nodes: + node["score"] = 0.0 + + graph.bulk_add_nodes(nodes=nodes, indices=node_ids) + graph.bulk_add_edges(edges) + graph._update_metadata(segmentation_shape=(n_frames, *frame_shape)) + + return Tracks( + graph, + time_attr="t", + pos_attr="pos", + tracklet_attr="tracklet_id", + lineage_attr="lineage_id", + ) diff --git a/tests/benchmarks/bench_actions.py b/tests/benchmarks/bench_actions.py new file mode 100644 index 00000000..e643283e --- /dev/null +++ b/tests/benchmarks/bench_actions.py @@ -0,0 +1,180 @@ +import pytest + +from funtracks.user_actions import ( + UserAddEdge, + UserDeleteEdge, + UserDeleteNode, + UserUpdateNodeAttrs, + UserUpdateNodesAttrs, + UserUpdateSegmentation, +) +from funtracks.utils.tracksdata_utils import td_mask_to_pixels + +from ._graph_builders import make_tracks + +NUM_FRAMES = 20 +FRAME_SHAPE = (500, 500) +CELLS_PER_FRAME = 50 + +# Operations per timed call. Even, so the self-inverting benchmarks return the graph to +# its starting state. Raise this (rather than ROUNDS) if a benchmark reads as noisy. +N_OPS = 50 +ROUNDS = 3 + +# Fraction of a node's mask painted in the segmentation benchmark. Must stay < 1: a +# patch that fully covers the mask makes UserUpdateSegmentation delete the node instead +# of updating its mask, which is a different (and much more expensive) code path. +PATCH_FRACTION = 3 + + +@pytest.fixture(scope="module") +def tracks(_warm_jit): + """One Tracks object shared by every benchmark in this module.""" + return make_tracks( + n_frames=NUM_FRAMES, + cells_per_frame=CELLS_PER_FRAME, + frame_shape=FRAME_SHAPE, + ) + + +@pytest.fixture(scope="module") +def solution_edges(tracks): + """Every existing (source, target) edge, in node id order. + + Node ids are assigned frame by frame, so this is ordered by time: consecutive entries + are cells in the same frame, and the benchmarks' per-round slices therefore walk + forward through the movie rather than repeatedly hitting one frame. + """ + edges = [] + for node in sorted(int(n) for n in tracks.graph_solution.node_ids()): + successors = tracks.successors(node) + if successors: + edges.append((node, int(successors[0]))) + return edges + + +def _node_batches(solution_edges, n_rounds=ROUNDS, per_round=N_OPS): + """Split source nodes into one distinct batch per round. + + Each round gets its own nodes, and callers offset their slices so no two benchmarks + share nodes. That keeps test_delete_nodes from removing nodes the attribute + benchmarks reported on, and keeps every round of a destructive benchmark doing the + same amount of work. + + Note this is no longer needed to avoid a measurement artifact: "score" is populated + on every node at build time (see _graph_builders.make_tracks), so a node + that has never been written costs the same to update as one that has. + """ + nodes = [source for source, _ in solution_edges] + batches = [nodes[r * per_round : (r + 1) * per_round] for r in range(n_rounds)] + assert all(len(batch) == per_round for batch in batches), ( + f"need {n_rounds * per_round} source nodes to give each round a distinct batch, " + f"only have {len(nodes)}" + ) + return batches + + +def test_update_node_attrs_single(benchmark, tracks, solution_edges): + """N single-node attribute updates, each its own history entry.""" + batches = iter(_node_batches(solution_edges)) + + def run(): + for i, node in enumerate(next(batches)): + UserUpdateNodeAttrs(tracks, node, {"score": float(i)}) + + benchmark.pedantic(run, rounds=ROUNDS, iterations=1) + + +def test_update_node_attrs_bulk(benchmark, tracks, solution_edges): + """The same N updates as one batched action. + + UserUpdateNodesAttrs exists to collapse history and refresh into a single entry for + the whole batch; this pairs with test_update_node_attrs_single to show what that + batching is worth. + """ + # Offset past the batches test_update_node_attrs_single consumed, so these updates + # also land on nodes not yet touched in this session. + batches = iter(_node_batches(solution_edges[ROUNDS * N_OPS :])) + + def run(): + nodes = next(batches) + UserUpdateNodesAttrs( + tracks, nodes, {"score": [float(i) for i in range(len(nodes))]} + ) + + benchmark.pedantic(run, rounds=ROUNDS, iterations=1) + + +def test_add_delete_edges(benchmark, tracks, solution_edges): + """Alternating edge delete and re-add on one edge. + + Both directions trigger UpdateTrackIDs: deleting orphans the downstream segment onto + a fresh track/lineage id, and re-adding merges it back. Self-inverting, so the graph + is unchanged afterwards. + """ + edge = solution_edges[len(solution_edges) // 2] + + def run(): + for _ in range(N_OPS // 2): + UserDeleteEdge(tracks, edge) + UserAddEdge(tracks, edge) + + benchmark.pedantic(run, rounds=ROUNDS, iterations=1) + + +def test_update_segmentation(benchmark, tracks, solution_edges): + """Alternating paint-out and paint-back of one sub-mask patch. + + Simulates the paint stroke path: UpdateNodeSeg plus the regionprops recompute + (position, area) and the IoU refresh on adjacent edges. Removing then restoring the + same patch keeps every operation on the UpdateNodeSeg branch and leaves the mask + exactly as it started -- the patch is a strict subset of the mask so the node is + never fully erased, and the paint-back targets an existing node id so it never falls + through to UserAddNode. + """ + node = solution_edges[len(solution_edges) // 2][0] + mask_pixels = td_mask_to_pixels( + tracks.get_mask(node), tracks.get_time(node), ndim=tracks.ndim + ) + n_patch = len(mask_pixels[0]) // PATCH_FRACTION + patch = tuple(dim_pixels[:n_patch] for dim_pixels in mask_pixels) + + def run(): + for i in range(N_OPS): + if i % 2 == 0: + # Erase the patch: old value is the node, new value is background. + UserUpdateSegmentation( + tracks, + new_value=0, + updated_pixels=[(patch, node)], + current_track_id=1, + ) + else: + # Paint it back: old value is background, new value is the node. + UserUpdateSegmentation( + tracks, + new_value=node, + updated_pixels=[(patch, 0)], + current_track_id=1, + ) + + benchmark.pedantic(run, rounds=ROUNDS, iterations=1) + + +def test_delete_nodes(benchmark, tracks, solution_edges): + """N node deletions, each cascading to its adjacent edges and track ids. + + Destructive and not self-inverting, so this runs last in the module and takes a + fresh slice of nodes per round. Deleting a mid-track node makes UserDeleteNode + delete both adjacent edges and reconnect the neighbours with a skip edge, which is + the expensive case. + """ + # Start past the nodes the attribute benchmarks consumed so deletions do not remove + # nodes those benchmarks already reported on. + batches = iter(_node_batches(solution_edges[2 * ROUNDS * N_OPS :])) + + def run(): + for node in next(batches): + UserDeleteNode(tracks, node) + + benchmark.pedantic(run, rounds=ROUNDS, iterations=1) diff --git a/tests/benchmarks/bench_candidate_graph.py b/tests/benchmarks/bench_candidate_graph.py index 33078e87..574e3a6f 100644 --- a/tests/benchmarks/bench_candidate_graph.py +++ b/tests/benchmarks/bench_candidate_graph.py @@ -43,28 +43,6 @@ def seg_data(): return _generate_segmentation() -@pytest.fixture(scope="module") -def _warm_spatial_compile(): - """Warm the spatial_graph rtree JIT compile before the timed benchmark. - - Building a Tracks with segmentation creates a GraphArrayView, which builds a - spatial_graph rtree whose Cython module is JIT-compiled by witty on first use and - cached process-wide. That compile is a one-time cost (seconds-to-tens-of-seconds on - a cold Windows runner) that would otherwise land inside the timed region of - test_graph_to_solution. Trigger it here (untimed) on a tiny graph with the same - dimensionality so the benchmark measures tracking work, not compilation. - """ - tiny = _generate_segmentation(num_frames=2, frame_shape=(64, 64), cells_per_frame=2) - graph = compute_graph_from_seg(tiny, MAX_EDGE_DISTANCE, iou=True) - Tracks( - graph, - time_attr="t", - pos_attr="pos", - tracklet_attr="tracklet_id", - lineage_attr="lineage_id", - ) - - def test_compute_graph_from_seg(benchmark, seg_data): benchmark.pedantic( compute_graph_from_seg, @@ -75,15 +53,16 @@ def test_compute_graph_from_seg(benchmark, seg_data): ) -def test_graph_to_solution(benchmark, seg_data, _warm_spatial_compile): +def test_graph_to_solution(benchmark, seg_data, _warm_jit): """Benchmark candidate graph -> Tracks with track ids (tracklet/lineage assignment). Candidate-graph construction is benchmarked separately above and is built here in (untimed) setup, so only the Tracks construction -- dominated by TrackAnnotator._assign_tracklet_ids and _assign_lineage_ids -- is measured. Setting tracklet_attr/lineage_attr registers a TrackAnnotator and triggers the - track-id computation on construction. The _warm_spatial_compile fixture ensures the - one-time spatial_graph rtree JIT compile is not measured here. + track-id computation on construction. The session-scoped _warm_jit fixture (see + conftest.py) ensures the one-time spatial_graph rtree JIT compile is not measured + here. """ def setup(): diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py new file mode 100644 index 00000000..851b9c9b --- /dev/null +++ b/tests/benchmarks/conftest.py @@ -0,0 +1,50 @@ +"""Shared fixtures for the benchmark suite.""" + +import pytest + +from funtracks.user_actions import UserUpdateNodeAttrs, UserUpdateSegmentation +from funtracks.utils.tracksdata_utils import td_mask_to_pixels + +from ._graph_builders import make_tracks + + +@pytest.fixture(scope="session") +def _warm_jit(): + """Pay the one-time JIT compile costs before any timed benchmark. + + Several lazy compiles would otherwise land inside whichever timed region hits them + first, and they are large enough to skew or dwarf the measurement: + + 1. Building a Tracks with segmentation creates a GraphArrayView, which builds + a spatial_graph rtree whose Cython module is JIT-compiled by witty on first use + (seconds-to-tens-of-seconds on a cold Windows runner). + 2. UserUpdateSegmentation calls Mask.intersection to decide whether a paint stroke + fully covers a node, which lands in tracksdata's numba kernels + (``@njit fast_intersection_with_bbox`` / ``fast_iou_with_bbox`` in + tracksdata/functional/_iou.py). Nothing in funtracks itself is numba-compiled. + First call ~450ms against ~0.5ms steady state -- an 800x first-call penalty. + + Both caches are keyed on types, not sizes: numba specialises per argument type + signature and witty per generated source. That is what makes a tiny warm-up graph + sufficient -- it installs the ``(int64[:], int64[:], bool[:,:], bool[:,:])`` signature + that the full-size benchmarks then reuse. It also means the warm-up only works while + it shares the real graph's dtypes, so make_tracks is the single source of + geometry for both; do not hand-roll a graph for this fixture. + """ + tracks = make_tracks(n_frames=2, cells_per_frame=2, frame_shape=(64, 64)) + + # Exercise the update-segmentation path too. Building the tracks is not enough: the + # numba compile happens on the first Mask.intersection call, not during construction. + node = next(iter(sorted(int(n) for n in tracks.graph_solution.node_ids()))) + mask_pixels = td_mask_to_pixels( + tracks.get_mask(node), tracks.get_time(node), ndim=tracks.ndim + ) + n_patch = max(1, len(mask_pixels[0]) // 3) + patch = tuple(dim_pixels[:n_patch] for dim_pixels in mask_pixels) + UserUpdateSegmentation( + tracks, new_value=0, updated_pixels=[(patch, node)], current_track_id=1 + ) + + # Same story, smaller magnitude (~2x rather than ~800x): the first attribute update + # is slower than steady state, enough to skew a 3-round benchmark. + UserUpdateNodeAttrs(tracks, node, {"score": 1.0})