Engine-level cancellation safepoints; remove cancellation_requested sysop#3158
Engine-level cancellation safepoints; remove cancellation_requested sysop#3158antoniosarosi merged 8 commits intocanaryfrom
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR refactors cancellation handling in the BAML engine from distributed per-path checks to centralized engine-level safepoints. Distributed cancellation guards are removed from LLM orchestration paths and the Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Binary size checks passed✅ 7 passed
Generated by |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/bex_engine/src/lib.rs (1)
1075-1136: 🧹 Nitpick | 🔵 TrivialRedundant safepoints in
ScheduleFuturearm.The safepoint at line 1089 (after
execute_sys_op) already guards entry into bothReadyandAsyncbranches. The subsequent checks at lines 1093 and 1110 can never observe a newly-cancelled token because no.awaitpoint intervenes between them. Consider removing the per-branch checks to reduce noise.♻️ Proposed simplification
let sys_op_result = self.execute_sys_op(pending.operation, &args, call_id, cancel); Self::cancellation_safepoint(cancel, &abort_handles)?; match sys_op_result { SysOpResult::Ready(result) => { - Self::cancellation_safepoint(cancel, &abort_handles)?; - // Sync operation - set future to Ready without touching stack. ... } SysOpResult::Async(fut) => { - Self::cancellation_safepoint(cancel, &abort_handles)?; - // Async operation — wrap in Abortable and spawn. ... } }
antoniosarosi
dismissed stale reviews from coderabbitai[bot] and coderabbitai[bot]
Dismissing stale CodeRabbit CHANGES_REQUESTED after follow-up commits resolved all review threads and addressed actionable findings.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/bex_engine/src/lib.rs (1)
1033-1061:⚠️ Potential issue | 🟡 Minor
FunctionEndevent is never emitted under "cancel wins" semantics, orphaning theFunctionStart.Previously, the
Completearm was the only guarantee thatFunctionEndwould always be emitted (non-Completeexits on error paths already skipped it). This PR's "cancel wins" safepoint at Line 1037 introduces a new code path where a VM that has fully computed its result still does not emitFunctionEnd, leaving aFunctionStartwithout a pairedFunctionEndin the event store.If downstream consumers (e.g.,
bex_events::event_store::events_for_span, collectors) assume theCompletepath always produces aFunctionEnd, they will now silently receive an incomplete event stream for cancellations that race with VM completion.Consider emitting a
FunctionEndwith a cancellation-indicating result before returningErr(Cancelled), or document that consumers must handle unpaired spans:📌 One approach — emit FunctionEnd before returning Cancelled
VmExecState::Complete(value) => { // "Cancel wins" semantics: if cancellation races with a // completed VM step, report `Cancelled` rather than // returning a success value. - Self::cancellation_safepoint(cancel, &abort_handles)?; - - // Emit FunctionEnd for the root entry-point span if tracing - if let Some(state) = span_state.as_mut() { + let cancelled = cancel.is_cancelled(); + if cancelled { + Self::abort_inflight_tasks(&abort_handles); + } + + // Always emit FunctionEnd so consumers see paired span events. + if let Some(state) = span_state.as_mut() { if let Some(root_span) = state.stack.pop() { ... bex_events::event_store::emit(end_event); } } + + if cancelled { + return Err(EngineError::Cancelled); + } + return self.heap.with_gc_protection(...); }
Merging this PR will improve performance by 12.48%
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ⚡ | WallTime | bench_incremental_no_change |
121.5 µs | 109.9 µs | +10.53% |
| ⚡ | WallTime | bench_incremental_close_string |
1,022 µs | 914 µs | +11.82% |
| ⚡ | WallTime | bench_incremental_add_attribute |
1,018.3 µs | 914.9 µs | +11.3% |
| ⚡ | WallTime | bench_incremental_add_user_field |
1.1 ms | 1 ms | +10.14% |
| ⚡ | WallTime | bench_single_simple_file |
1,037.7 µs | 930 µs | +11.58% |
| ⚡ | WallTime | bench_incremental_add_string_char |
1,024.7 µs | 911 µs | +12.48% |
| ⚡ | WallTime | bench_empty_project |
957.9 µs | 853.7 µs | +12.2% |
Comparing codex/cancellation-safepoints (7adc0c4) with canary (a91edc3)
Footnotes
-
91 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
antoniosarosi
force-pushed
the
codex/cancellation-safepoints
branch
from
484f3e0 to
b94ad41
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/bex_engine/src/lib.rs (1)
1126-1152: 🧹 Nitpick | 🔵 TrivialSafepoints at lines 1134 and 1152 are redundant with the one at line 1129.
Between the post-
execute_sys_opsafepoint (line 1129) and the Ready/Async match arms (lines 1134, 1152), no async work, blocking, or significant computation occurs — just pattern matching. A concurrent cancellation in that window is caught at the next loop-top safepoint anyway.These extra checks cost only an atomic load each, so they're harmless. Just noting the redundancy in case you want to trim.
1 participant
centralize cancellation checks in bex_engine via a safepoint helper in the VM event loop
add safepoint checks at VM loop boundaries and ScheduleFuture boundaries so cancelled calls do not dispatch new sysops
remove baml.sys.cancellation_requested() from builtins, native sysop impl, and LLM orchestration BAML
update cancellation test coverage to drop the removed sysop cases
refresh baml_std snapshots for lexer/parser/HIR/MIR/codegen changes
Validation
cargo test -p bex_engine --test cancellation -- --nocapture
cargo test -p bex_engine --test orchestration -- --nocapture
cargo test -p baml_tests __baml_std -- --nocapture
Summary by CodeRabbit
Changes
cancellation_requested()API function; cancellation is now managed entirely by the engine.Tests