perf(core-py): release the GIL while blocking on the tokio runtime - #2485
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Walkthrough
ChangesConcurrent engine calls
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Python
participant PySessionContext
participant CallLock
participant Runtime
Python->>PySessionContext: invoke runtime method
PySessionContext->>Python: detach blocking work
PySessionContext->>CallLock: acquire per-context mutex
PySessionContext->>Runtime: execute runtime operation
Runtime-->>PySessionContext: return result
PySessionContext-->>Python: return result
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
![coderabbitai[bot]](https://cdn.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/wren-core-py/src/context.rs (1)
371-400: 🚀 Performance & Scalability | 🔵 TrivialLong-running queries serialize the whole context.
lock_calls()is held for the entiresql().await+collect().await, so any concurrenttransform_sql,list_tables, ordry_runon the sameSessionContextblocks until this query finishes. This is correct given the shared-catalog re-registration invariant, but for read-heavy or long-running query workloads consider routing heavyquerycalls through a dedicated context (or documenting this ordering guarantee) so a single slow query doesn't stall lightweight metadata calls on the same context.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/wren-core-py/src/context.rs` around lines 371 - 400, Document the ordering guarantee in query and related SessionContext operations, or route heavy query execution through a dedicated context so sql().await and collect().await do not block lightweight transform_sql, list_tables, and dry_run calls. Preserve the shared-catalog re-registration invariant when choosing the implementation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@core/wren-core-py/src/context.rs`:
- Around line 371-400: Document the ordering guarantee in query and related
SessionContext operations, or route heavy query execution through a dedicated
context so sql().await and collect().await do not block lightweight
transform_sql, list_tables, and dry_run calls. Preserve the shared-catalog
re-registration invariant when choosing the implementation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: ec67f681-9631-4181-ae10-0ac406a834d0
📒 Files selected for processing (2)
core/wren-core-py/src/context.rscore/wren-core-py/tests/test_modeling_core.py
|
The shared CI failure here ( |
abcf935 to
8d605a0
Compare
Two modes: separate contexts per thread, and one shared context with mixed methods. A barrier forces the threads to overlap. The shared-context mode catches the same-context catalog race if the per-context call lock is ever removed while the GIL is released. Correctness-only by design - performance evidence lives in the PR benchmark, not in CI.
Every blocking call in PySessionContext (transform_sql, query, register_*, dry_run, load_mdl, and the constructor) held the GIL for the whole planning/execution call, so concurrent calls from Python threads serialized: 4-thread throughput scaled at 1.04x on a 16-core machine, and any embedding asyncio application froze entirely while a call was in flight. Wrap each runtime.block_on in Python::detach. Python-borrowed arguments are copied to owned values before entering the detached closure, so no GIL-dependent data crosses the boundary. The futures never call back into Python (session properties are converted eagerly at construction; UDFs are ByPass* stubs), so releasing the GIL cannot deadlock. Releasing the GIL also removes the accidental serialization that hid a real race: transform_sql re-applies the MDL catalog onto the context's shared catalog list on every call, so two concurrent calls on the same context could observe each other's half-registered catalogs (planning failed with 'table not found'). Preserve the previous per-context ordering with an internal call lock, taken only while detached; calls on different contexts now run in parallel. list_tables has no block_on but reads the same shared catalog list, so it takes the lock too. The unused Clone derive is removed so a Rust clone can't silently share the lock and catalog state while diverging on ctx/exec_ctx/mdl. Behavior note: methods that take &mut self (load_mdl) can now observe PyO3's borrow error if called concurrently with an in-flight call on the same object, instead of silently waiting for the GIL - a loud failure rather than an implicit queue.
8d605a0 to
b7d0ac1
Compare
|
Related core fix: #2495 While validating this PR's same-context concurrency path, I traced the This PR keeps the per-context |
|
Thanks for taking the time to review and merge this. Really appreciate it! |
2 participants
Summary
Release the Python GIL while
PySessionContextwaits on TokioRuntime::block_on, allowing calls on differentSessionContextinstances to run concurrently.
This covers all 13 Python-reachable
block_onsites, includingconstruction, SQL transformation, query execution, table registration,
dry runs, MDL loading, and function discovery.
Problem
Every blocking
SessionContextcall previously held the GIL for theentire planning or execution operation. Python threads therefore
serialized even when they used independent contexts and independent
Tokio runtimes.
This also prevented an asyncio event loop from progressing while a
synchronous Wren call was offloaded to a worker thread: the worker still
held the GIL until the call completed.
Implementation
Each blocking operation now runs inside
Python::detach.Python-borrowed arguments such as SQL, function names, and paths are
converted to Rust-owned values before entering the detached closure.
Session properties are also converted eagerly during construction, and
the futures do not call back into Python.
Releasing the GIL exposed a pre-existing same-context catalog race.
transform_sql_with_ctxre-applies the MDL catalog to shared contextstate, so concurrent calls could observe a partially registered catalog.
A per-context
call_lockpreserves the previous ordering for calls onthe same context. The lock is acquired only after releasing the GIL to
avoid a lock/GIL deadlock. Calls using different contexts can run in
parallel.
list_tablesdoes not useblock_on, but it reads the same catalogstate and therefore participates in the same locking protocol.
The unused
Cloneimplementation is removed because cloned wrapperscould share catalog and lock state while later diverging in their
ctx,exec_ctx, andmdlfields.Behavioral note
Methods taking
&mut self, currentlyload_mdl, may now return a PyO3borrow error when called concurrently with another in-flight operation
on the same object. Previously such calls implicitly waited for the GIL.
Same-context engine calls otherwise retain their previous serialized
ordering.
Benchmark
Release build on an Apple M5 Pro with 18 cores. Each thread used its own
SessionContextand performed 200transform_sqlcalls. Measurementswere collected with Python 3.14 using the abi3-py311 wheel.
Single-thread throughput remained effectively unchanged (
956versus953calls/s). At 16 threads, throughput improved from991to5,142calls/s, a 5.19x improvement at the same concurrency level.
Scaling remains sublinear because every context currently owns a
multi-thread Tokio runtime, which oversubscribes the machine at higher
thread counts. Sharing one process-wide runtime is planned as a
follow-up. I plan to follow up with a PR for this.
Tests
assertions.
Validation:
cargo test: 37 passedcargo clippy --all-targets --all-features -- -D warningspytest: 34 passedSummary by CodeRabbit
Summary by CodeRabbit
Bug Fixes
Tests
transform_sqlacross threads (both separate and shared session contexts).