Skip to content

refactor(replace): run Substitute as one invocation - #336

Merged
binaryaaron merged 6 commits into
aagonzales/graph-sdk/feat/mainfrom
aagonzales/chore/graph-sdk-substitute-one-invocation
Oct 29, 2026
Merged

binaryaaron merged 6 commits into
aagonzales/graph-sdk/feat/mainfrom
aagonzales/chore/graph-sdk-substitute-one-invocation

Conversation

@binaryaaron

@binaryaaron binaryaaron commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

Substitute now runs as one call to the graph kernel instead of two. Its whole-frame replacement map becomes the first
step of the per-row workflow, so a row that DataDesigner drops is recorded as a failed target in the same record as
every other row. Outputs, errors and failed-record accounting for Anonymizer users are unchanged apart from the
trace-column dtypes listed below.

This PR also carries the kernel executor rewrite merged in from #337 (section below).

This branch stacks on #332 and is rebased onto its base branch when the stack moves.

What changes

  • One invocation (engine/replace/route.py). run_substitute_route lowers the input frame once and runs
    substitute-map, replace and verify-patches in one protect() call. substitute-map is an invocation-scoped
    operation: it calls LlmReplaceWorkflow.generate_map_only once over the whole frame, maps the returned rows back to
    their input positions with resolve_source_positions, and returns each row's replacement map. A row DataDesigner
    dropped gets None and fails at that step; its later steps are blocked.
  • replace reads the map from the graph. When a workflow has a map step, replace takes the map from its
    replacement_map input; Redact, Annotate and Hash still build their maps locally and behave as before.
  • Errors. The first error raised by any step other than the verifier is re-raised to the caller, as before: an
    error from generate_map_only or from applying replacements still reaches the caller once, as the same object.
  • Output frame. Rows are the targets whose output is available, in input order, taken from the input frame, with
    the replacement map, map source, replaced text and application columns read from the record. failed_records still
    comes from generate_map_only.
  • Tests. Two tests in tests/engine/test_replace_runner.py gave the runner an input frame with no entities while
    the mocked map result carried entities. The input frames now carry the same entities as the mocked result.

Net non-blank, non-comment code lines under src/: +24, all in route.py.

Known differences

  • In Substitute's trace_dataframe, _validated_entities and _augmented_entities keep the dtype they have in the
    input frame (a pyarrow struct type), where main has object. Their values are unchanged.
  • With an empty input frame, generate_map_only is no longer called; the result is still an empty frame with no
    failed records.

Checks

  • make check passes.
  • make test: 1,557 passed. test_wandb_snapshot_rejects_fifo_socket_and_device fails in the test environment with
    AF_UNIX path too long, as on the base.
  • uv run --group dev python -m tests.spec_programs.run_all: 70 pass, 8 retired, 0 fail.
  • Parity against main (f4f2b73d) on pandas 2 and 3, over all cases: no blocking difference beyond the differences
    the base already has; the dtype difference above is the only new non-blocking one.
  • Executed lines per case against main: median 1.20×, at most 1.23× on cases whose status matches main (base: 1.20×
    and 1.22×).
  • No new imports of private names across packages, and no new import cycles.

Kernel executor rewrite (from #337)

The graph kernel's executor is rewritten to run its node sequence directly on mutable state, with structural pattern
matching for events and per-target results. Behaviour is unchanged: dispatch order, outcomes, ledgers, attempts,
recorded errors, rejection codes and cancellation all match the base. kernel.py goes from 333 to 272 code lines, and
five new spec-program checks pin interrupt propagation and cancellation from inside a callback.

What changes

  • New checks (tests/spec_programs/). Two lifecycle checks and three invocation-scoped checks pin behaviour no
    program covered (four are added before the rewrite, one after it):
    • a KeyboardInterrupt from a per-target callback reaches the caller of protect() as the same object, and later
      targets are not dispatched;
    • a callback that cancels the run stops the next per-target dispatch;
    • a KeyboardInterrupt from an invocation-scoped callback reaches the caller;
    • cancelling inside an invocation-scoped call completes that step for every target and dispatches nothing after it;
    • a cancel raised by the step just before an invocation-scoped step stops that step from being called.
  • Mutable state (engine/graph_core/kernel.py). The persistent vector, the linked-list event log and the copied
    state objects are gone. State holds a list of rows and a list of entries, and apply(state, event) changes it in
    place with one match over the event types. Its guards are unchanged: a dispatch to a node that is not pending, or
    after cancellation, raises; a completion or failure for a node that is not running is ignored; a second close is
    ignored.
  • A sequence, scheduled as one. execute runs the nodes in order, each over every ready unit in order. The stored
    scan cursor, next_ready and the per-node dependency sets are removed; a unit is ready at a node when it is not
    closed and every earlier node succeeded. Cancellation is checked before each per-target call and before each
    invocation-scoped call, as before.
  • Two dispatch helpers side by side. _run_target makes one call for one target; _run_invocation makes one call
    for every ready target. Both complete a target's output through the same helper, and an invocation-scoped callback's
    per-target entries are settled by a match: None fails the target with missing_output, an Exception
    fails it with raised, and any other value is checked as the target's output.
  • Types. Statuses are a StrEnum whose values are the previous strings, so ledgers and comparisons are unchanged.
    An input port's source is a Source(node, port) named tuple; seed values are kept apart from node outputs instead of
    sitting under node index -1. Events are named tuples.
  • Smaller changes. compile_sequence checks kinds without a side-effect loop; CancelToken.link replaces the
    direct write to a token's parent from run_in_frame; execution.py follows the renamed apply and the new row
    layout.

Net non-blank, non-comment code lines under src/: −58 (kernel.py and execution.py only).

Checks

  • make check passes.
  • make test: 1,557 passed. test_wandb_snapshot_rejects_fifo_socket_and_device fails in the test environment with
    AF_UNIX path too long, as on the base.
  • uv run --group dev python -m tests.spec_programs.run_all: 75 pass, 8 retired, 0 fail.
  • Each commit compared with its parent on probe scenarios through protect() (dispatch order, interrupt identity,
    cancellation before, during and after a call, every invocation failure form, ledgers, attempts and values):
    identical output at every step.
  • Parity against main (f4f2b73d) on pandas 2 and 3, over all cases: no blocking difference beyond the differences
    the base already has.
  • Executed lines per case against main: median 1.20×, at most 1.22× on cases whose status matches main (base: 1.20×
    and 1.23×).
  • No new imports of private names across packages, and no new import cycles.

Main merged (5155d6d)

This branch now includes main at 5155d6d through its base branch, including the Rewrite join change described on the remove-old-sdk PR.

Main merged (f294dc1)

This branch now also includes main at f294dc1 (#328, recovery from incomplete replacements in Rewrite) through its base branch.

🤖 Generated with Claude Code

binaryaaron and others added 6 commits October 26, 2026 19:50
Substitute lowers the input frame once and declares one workflow: an
invocation-scoped map step over all targets, then replace and verify-patches
per target. A row the map step drops fails at that step and is absent from
the output. The runner raises the first error recorded by a non-verifier
step, so a map-step or replace exception still reaches the caller.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Two Substitute runner tests fed an input frame with empty entities while the
mocked map frame carried the stub entities. The route now lowers the input
frame, so the two frames must agree.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
…odule

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
* test: pin interrupt propagation and in-call cancellation in the spec programs

Add E11.6 and E11.7 to the cancellation program and E14.9 and E14.10 to the
invocation-operation program. They pin that an interrupt raised by a callback
reaches the caller of protect(), and that a cancel() issued inside a callback
lets that call finish but stops the next dispatch.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

* refactor(graph): run the kernel's node sequence on mutable state

State holds rows and entries in dispatch order and changes in place through
apply(), a match over the event classes. The executor loops over nodes in
order and, within a node, over units in order; per-target and
invocation-scoped nodes settle results through one path. Plan.deps and the
persistent vector are gone.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

* refactor(graph): type the kernel's statuses and seed sources

Statuses are a StrEnum with the same string values; the ledger and attempts
still emit plain strings. A row keeps seed values and node outputs apart,
and a plan source names either a seed or a node output. CancelToken gains
link() so run_in_frame no longer writes a private attribute, and
compile_sequence validates kinds without a bare side-effect loop.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

* refactor(graph): simplify the kernel's events, readiness and frame tokens

Events become NamedTuples, readiness drops the redundant pending check,
frame tokens link through one call, and docstrings shrink. Behaviour is
unchanged.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

* refactor(graph): spell out the kernel's dispatch paths

execute now has one branch for invocation-scoped nodes and one for
per-target nodes, each calling its own helper. _run_invocation handles a
raised call, a wrong-shaped return and per-entry results as separate cases,
and the dense one-liners in compile_sequence, _inputs and apply are unfolded.
Behaviour is unchanged.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

* test: pin cancellation raised just before an invocation-scoped step

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

---------

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Co-authored-by: Claude <noreply@anthropic.com>
…/chore/graph-sdk-substitute-one-invocation

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
…/chore/graph-sdk-substitute-one-invocation

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Base automatically changed from aagonzales/chore/graph-sdk-invocation-operation to aagonzales/graph-sdk/feat/main October 29, 2026 21:15
@binaryaaron
binaryaaron merged commit 7881548 into aagonzales/graph-sdk/feat/main Oct 29, 2026
8 checks passed
@binaryaaron
binaryaaron deleted the aagonzales/chore/graph-sdk-substitute-one-invocation branch October 29, 2026 21:15
@binaryaaron

Copy link
Copy Markdown
Collaborator Author

Folded into #319: aagonzales/graph-sdk/feat/main was fast-forwarded to 44924e2e, which contains this branch unchanged, so this step (Substitute as one invocation) is now reviewed in #319. Nothing was rewritten; this description stays as the record of the step.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant