prepare: plan each program step once, and run its physical plan on every training step #202

🌀 Claude (xarray-sql ML demos)

Summary

A cached program builds nothing on a later training step. It still plans each of its steps from scratch: decoding the step's Substrait, optimizing it, and planning it physically. A prepared program would plan each step once and run that physical plan on every later step, as jax.jit traces once and runs the compiled function after that. xarray-sql will add Context.prepare(sql) for the same thing on its side, and it needs ddx's half for statements that call grad.

What it would save

Demo 2 of xarray-sql's ML benchmarks (a 1-layer GPT, xqlsystems/xarray-sql#269) runs one SGD step as 76 ddx steps. I timed DataFusionBackend.materialize at ddx 1a3074c on DataFusion 55.0.0, with target_partitions=1, on an Intel i7-8700. For the "reused" row, I patched materialize to keep each step's ExecutionPlan, keyed by its Substrait bytes, and to run it with ctx.execute(plan, 0). Nothing resets the plans between runs (see below), so the values are wrong, and the step count drifts to 67. The times are what running a planned step costs:

The 0.16 s that's left is outside materialize: plan_key and view_sql, and planning the rewritten statements. A prepared program would plan those once too. The profile behind this is in the comment on #186: #186 (comment).

What a reused plan needs

These are what I found trying it from Python on DataFusion 55:

1. The step's inputs, read when it runs. A register_record_batches table (a MemTable) copies its batches into the plan when the plan is made. Registering a new table under the same name doesn't reach a plan that's already made: a planned step keeps reading the previous run's weights and earlier steps' outputs. A plan does see new values written into the same Arrow buffers, because the scan holds the buffers themselves. Two ways to get there:
   • Keep each table's buffers for the program's life, and write each run's values into them. This needs row counts that don't change between runs: true of a model's weights and a fixed batch size. xarray-sql's in-memory tables share their NumPy arrays' memory, so xarray-sql can write a new weight into the arrays a plan reads. ddx would do the same with each step's output: copy the result into the table the next step's plan reads.
   • A table provider whose scan reads a slot when it runs, not when it's planned. Through datafusion-ffi this costs about 1.5 ms per scan (datafusion-ffi rebuilds the session and task context, with every function, on each scan and execute (~1.5 ms per scan) xarray-sql#273), more than the planning it saves. So it pays only for a provider in the same library as the plan, not through FFI.
2. A reset between runs. A hash join keeps its build side after the first run, so a reused plan reads its left input as it was then (a ten-line repro is in datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream xarray-sql#283). All 2,988 joins in the GPT's forward plan are HashJoinExec with mode=CollectLeft, and so are the steps' joins. DataFusion resets a plan with reset_plan_states, which is Rust only. datafusion-python doesn't bind it, FFI_ExecutionPlan doesn't forward reset_state, and with_new_children keeps the build side. datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream xarray-sql#283 is a draft request for ExecutionPlan.reset_state() in datafusion-python. Until then, a stopgap works in Python: decoding the plan again, ExecutionPlan.from_bytes(ctx, plan.to_bytes()), gives fresh nodes. It copies the in-memory tables' current values into the plan, so it works with item 1's in-place writes and not with re-registered tables. It fails on a plan nested too deep to decode ("recursion limit reached"), but every one of the GPT's steps decodes. With it, a step takes 1.5 ms, against 3.3 ms planned each time and 0.65 ms with a true reset.
3. The reset covers repartitioning too. A RepartitionExec keeps state as well. Running a plan with one again panics with "partition not used yet" (a PanicException), though running its partitions one after another works the first time. Separately, target_partitions=1 is faster for steps this small: a cached step drops from 0.66 s to 0.43 s with it (program: choose step boundaries by what they save (fewer, larger steps) #186).
4. Dynamic filters off. reset_plan_states doesn't reset dynamic filters (Support dynamic filters and work-table plans state reset apache/datafusion#19905). A prepared plan should be made with datafusion.optimizer.enable_dynamic_filter_pushdown=false until it does.
5. The same key, plus settings. A step's physical plan depends on what plan_key covers (the tables' schemas, the views' definitions), and on the session's settings. It also depends on the tables' row counts, through join selection. A plan made for other row counts still gives the right answer, but item 1 already fixes them.

A possible shape

• Backend gains an optional pair, for example prepare(name, plan) -> handle and run_prepared(handle). DataFusionBackend implements them with a kept ExecutionPlan, reset before each run; a backend without them is planned each time, as now. Programs keeps each step's handle next to its program.
• ddxdb.ad.prepare(ctx, statements) returns an object whose call runs every program's prepared steps, then the rewritten statements' kept plans, and returns their results. The contract: between calls, the tables the statements read keep their schemas and row counts, and their values may change in place. xarray-sql's Context.prepare would call it for statements that call grad, vjp or jvp.
• Runs of one prepared object can't overlap, as for a program today. This matters for run: run a program's independent steps concurrently #157 (running independent steps concurrently): concurrent steps are fine, concurrent runs aren't.

Related

• datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream xarray-sql#283: datafusion-python has no reset_plan_states (draft for upstream).
• DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth xarray-sql#267: DataFusion inlines a CTE at every reference, so a deeper GPT can't be planned at all; WITH … AS MATERIALIZED lands in DataFusion 56. Comment with the 55 numbers: DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth xarray-sql#267 (comment)
• Stateless execution plans for plan caching apache/datafusion#19351: stateless execution plans for plan caching, upstream's fix for item 2.
• MaterializedCteExec::reset_state clears its shared buffer in place apache/datafusion#26229 and fix: keep the output of a materialized CTE per run, not in the plan apache/datafusion#26270: a materialized CTE's output lives on its plan node, so DataFusion 56's MaterializedCteExec needs the same reset.
• rewrite: a cached program still optimizes each shared CTE's placeholder (15% of a cached GPT step) #184: a cached program still optimizes each shared CTE's placeholder. A prepared program wouldn't optimize at all after its first run.