Stage 03 showed that a dask-backed Parcels run is ~327× slower than an in-memory one even though no disk I/O happens — the field is cache-resident and the process is CPU-bound. This directory profiles where that CPU goes, to confirm it is dask's client-side graph machinery (build → optimize → schedule), not array math.
profile_dask.py runs the exact Parcels gather —
field.isel(<scattered indices>).compute() — in a loop against a field that has
been .persist()ed into RAM, so every microsecond measured is pure scheduling
overhead.
uv pip install --python "$PARCELS_PYTHON" viztracer # one-time, into the parcels env
python profile_dask.py --ncalls 20 --tool viztracer
vizviewer profiling/results/dask_trace.json # opens a Perfetto UI in the browserVizTracer records every function call/return on a timeline. In the viewer, zoom
into a single compute() and you can watch it expand into the dask pipeline call
by call — graph construction (HighLevelGraph, blockwise/slicing), optimization
(optimize, cull, fuse), and the scheduler (get_async / execute_task).
That repeating block, once per gather, is the overhead. Keep --ncalls small
(~20) — the trace records everything, so it grows fast.
The JSON is a standard Chrome trace, so you can also open it at
https://ui.perfetto.dev or chrome://tracing instead of vizviewer.
python profile_dask.py --ncalls 200 --tool cprofile
snakeviz profiling/results/dask_sampling.prof # interactive icicle
# or: gprof2dot -f pstats profiling/results/dask_sampling.prof | dot -Tsvg -o callgraph.svgDeterministic, attributes wall time to named functions. The script prints the top
functions by self time (where the CPU actually is) and by cumulative time
(the call hierarchy), filtered to dask/xarray. See
results/cprofile_top.txt for a captured run.
cProfile and VizTracer only see Python frames; py-spy samples the whole process (including NumPy/Cython C frames) with negligible overhead and emits a speedscope trace:
pixi exec --spec py-spy py-spy record --format speedscope \
-o profiling/results/dask_sampling.speedscope.json \
-- "$PARCELS_PYTHON" profile_dask.py --ncalls 500 --tool none
# then drag the JSON onto https://www.speedscope.app(--tool none makes the script just run the loop; py-spy profiles it externally.)
scheduler_knobs.py (pixi run knobs) times the same
in-RAM gather under dask-internal configurations — scheduler="synchronous",
optimize_graph=False, fewer/larger chunks, and batching. Captured run in
results/scheduler_knobs.txt. Headline: staying inside dask,
synchronous + larger chunks cut the per-gather overhead ~6–12× (the thread-pool
lock/condition storm is the bulk of it). But every .compute() still rebuilds and
tokenizes a graph, so it bottoms out around ~50× NumPy — tuning narrows the gap,
the windowed approach (stage 05) closes it.
Parcels samples each velocity component separately — XLinear_Velocity calls
XLinear for U, then V, then W, and each ends in its own value.compute() — so a
3-D RK2 step issues ~6 .compute() calls. batched_compute.py
(pixi run batched) tests collapsing them into one dask.compute(u, v, w).
Captured run in results/batched_compute.txt.
Finding (worth flagging — it corrects the intuition): batching reduces the
call count (6 → 2 for RK2-3D) but not wall time (~1.05× threads, neutral
under synchronous). The cost is per-task, not per-.compute()-call: merging
three components into one graph just yields one graph with 3× the tasks, so the
scheduler does the same total work. The dask-only lever that does move the
needle is the scheduler (synchronous ≈ 2× here) plus fewer/larger chunks
(see scheduler_knobs.py) — not the number of calls. Batching is still worth
doing for cleaner code and to pair with dask.compute(..., scheduler=...), but
it isn't a speedup by itself. (You also can't batch across RK stages — they're
sequential.) Closing the gap still needs the in-RAM window (stage 05).
chunk_size_sweep.py (pixi run sweep) sweeps the chunk
edge for a scattered gather and reports both sides at once. Captured run in
results/chunk_size_sweep.txt:
chunk (t,y,x) #chunks #tasks ms/call read-amp
(1, 64, 64) 4332 10902 1676.20 1565x
(1, 128, 128) 1200 2385 476.42 1940x
(1, 300, 300) 192 385 99.71 1728x
(1, 600, 600) 48 97 39.11 1728x
(1, 1200, 1200) 12 25 22.66 1728x
Bigger chunks → far fewer tasks → ~74× faster in-RAM sampling (10902→25 tasks, 1676→23 ms). But read amplification doesn't improve (~1700×) — for scattered access each touched chunk still drags its whole self off disk, and the naive path re-reads it every step. So on the lazy-off-disk path, bigger chunks just trade scheduler overhead for read amplification; neither extreme wins.
Conclusion: large chunks only pay off if you read the chunk into memory
once, reuse it, and toss it — amortizing the big read over all particles and
sub-steps in that time interval, with low task count. And note dask.persist()
keeps it resident but still samples through the scheduler (~200× NumPy); the full
win is materializing to NumPy and indexing that. That combination — big
chunks + load-to-NumPy + reuse + evict — is exactly the stage-05 rolling window.
All three converge on the same story: with the data in RAM, time is spent
constructing and scheduling a fresh task graph on every .compute() — there
is no array work to speak of. That is why the fix is structural (load the needed
window into NumPy, or batch many gathers into one graph), not faster storage. See
../03_parcels/README.md for the options.