mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 22:26:53 +00:00
f5f6fbe69ad4b0a77ab0dec85050fa756bb2bea3
5
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
f5f6fbe69a |
fix(python): stop using Arc::strong_count to decide when the runtime is idle
Review on this PR found three more real bugs in the previous commit's exit-lifecycle design, all traceable to the same root cause: using Arc::strong_count as a proxy for "is anything still using the runtime" is the wrong signal. It reflects who currently holds a reference, not who's logically still in flight, and it let a spawned task become the final, worker-thread-side owner of the Runtime -- which reproduced unprompted in this branch's own test suite (a panic: "cannot drop a runtime in a context where blocking is not allowed"). Specifically: - A task holding an Arc for its own whole lifetime could end up being the one to drop the Runtime when it completes, which panics if that drop happens on one of the runtime's own worker threads. - shutdown() bailed out as soon as it saw an empty slot, so an install already in progress at that exact moment could finish afterward, leaving a live, un-waited-for runtime behind as the interpreter moved on into teardown. - Once the bounded wait elapsed, shutdown() simply returned without touching the runtime at all, recreating the original race for any call slower than the grace period. Replaces Arc-based tracking with an explicit OUTSTANDING counter that every top-level spawn/spawn_blocking/block_on call increments before it starts and decrements only once it has truly finished, via an OutstandingGuard token that carries no reference to the runtime itself. shutdown() waits for this to reach zero (not for Arc::strong_count to look low) before ever touching the runtime, and forces the retirement attempt anyway once its bound elapses instead of leaving the runtime untouched. Also narrows what the fork-child handler is allowed to do. ArcSwapOption turned out not to be safe to call from there either: swap/compare_and_swap reconcile reader "debts" internally (thread-local state, potentially an allocation), none of which is safe in a forked child that may have inherited another thread's lock mid-acquisition. The handler now only bumps a bare GENERATION counter; get_runtime() compares its installed runtime's generation against the live counter from ordinary (non-signal) context and rebuilds on a mismatch, exactly the check the handler used to perform directly. Verified against the same concurrent reproducer as the previous commit: 50/50 clean runs (vs. hangs and, separately, an unprompted panic under the prior design), plus a 30-second-join variant and a targeted test forcing the install-vs-shutdown race directly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
4206a727d2 |
fix(python): close two use-after-free/hang windows in runtime shutdown
The previous commit's shutdown() freed the shared Tokio runtime as soon as it looked unused (an AtomicPtr swap + Box::from_raw), but callers only ever held a bare reference into it, with no way to tell shutdown they were still using it. A concurrent caller could dereference memory shutdown had already freed. Switching to an ArcSwapOption<Runtime> so every caller holds an owned Arc clone for as long as it's using the runtime fixes the use-after-free, but isn't sufficient on its own, for two reasons: - future_into_py spawns a task that, once running, spawns a second one to do the real work and awaits its JoinHandle. If both nested calls independently re-resolved "the current runtime" via get_runtime(), a reclaim landing between them could bind the outer and inner task to two different runtime instances, permanently stalling the outer task if its own runtime was the one torn down while it awaited the inner one. - Runtime::shutdown_timeout gives spawned (non-blocking) tasks no grace period at all: a task keeps running until it next yields, then is dropped. An outermost caller that only held its Arc for the instant it submitted a task let the runtime's reference count fall back to baseline immediately, long before the task itself finished, so shutdown could reclaim the runtime while a top-level task was still genuinely in flight, abandoning it silently. spawn/spawn_blocking now use Handle::try_current() to pin any nested call to whatever runtime is already executing it, and the outermost call moves its Arc clone into the spawned task itself so the reference count stays elevated for the task's whole lifetime, not just its submission. Verified with two new targeted unit tests plus the existing concurrent stress test, and empirically against a concurrent reproducer built from this PR's own review feedback: 100+ clean runs with no hangs or crashes, versus a consistent ~35-55% failure rate beforehand. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
2dec579641 |
fix(python): shut down the shared Tokio runtime on interpreter exit
The Python bindings route every async call through a single, lazily created Tokio runtime, stored behind a raw pointer so it can survive being reset (not dropped) across fork() -- dropping it there would hang joining now-dead worker threads. That leak-on-fork path is deliberate and documented, but normal process exit had no equivalent handling at all: nothing ever tells the runtime to shut down, so its worker threads keep running, completely uncoordinated with the interpreter, until the process ends. If a worker thread is mid-task exactly as Py_Finalize begins tearing down interpreter state, it can panic on state that's already gone. That panic happens on a background thread with no PyO3-wrapped call frame to catch it -- and if it lands while another unwind is already in progress, Rust aborts the whole process outright rather than merely failing that one call. This reproduces as an intermittent SIGABRT on otherwise-successful short-lived processes that open a connection, do a handful of operations, and exit normally. Give the runtime a coordinated, bounded shutdown at the one point it was missing one: register a Python atexit callback that calls runtime::shutdown(), which takes the runtime through the same atomic slot the fork handler uses and calls shutdown_timeout() rather than a bare drop -- bounded, so a wedged background task can never hang process exit, but real in-flight work gets a genuine chance to finish first. Runs with the GIL released (Python::detach), since the wait would otherwise be able to deadlock against any in-flight task that itself needs the GIL to complete. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> |
||
|
|
8a5cd74e48 |
fix: ensure read freshness provider is built into namespace client (#3571)
By default the read freshness provider was not included in the namespace client, preventing the read freshness headers from being included in the request. This prevents checkout_latest() from working as expected when using the namespace client. This fix ensures the provided is built into the client when the namespace impl and properties are provided. |
||
|
|
a17c241e86 |
feat(python): make Permutation fork-safe for PyTorch DataLoader workers (#3339)
## Summary PyTorch's `DataLoader` uses fork-based multiprocessing by default on Linux, but threads do not survive `fork()`. LanceDB's Python bindings drive async work through two threaded layers, both of which become inert in a forked child: - `BackgroundEventLoop` runs an asyncio loop on a Python `threading.Thread`. - `pyo3-async-runtimes::tokio` holds a global multi-threaded tokio runtime whose worker threads also die on fork — and its runtime lives in a `OnceLock` that cannot be replaced after first use. As a result, any `Permutation` (or other async API) used inside a fork-based `DataLoader` worker hangs indefinitely. This PR makes both layers fork-safe so `Permutation` works as a `torch.utils.data.Dataset` with `num_workers > 0`. ## Approach ### Rust — new `python/src/runtime.rs` Mirrors the pattern used in [Lance's Python bindings](https://github.com/lance-format/lance/blob/456198cd6f42be07f99617a6d7e39d6209cdf3cc/python/src/lib.rs#L139), adapted for the async-bridge use case. - `LanceRuntime` implements `pyo3_async_runtimes::generic::Runtime + ContextExt`, backed by an `AtomicPtr<tokio::runtime::Runtime>` we own (sidestepping `pyo3-async-runtimes`'s frozen `OnceLock` global). - A `pthread_atfork(after_in_child)` handler nulls the pointer; the next `spawn` rebuilds the runtime in the child. The previous runtime is intentionally **leaked** — calling `Drop` would try to join now-dead worker threads and hang. - `runtime::future_into_py` is a drop-in for `pyo3_async_runtimes::tokio::future_into_py`. All ~80 call sites in `arrow.rs` / `connection.rs` / `permutation.rs` / `query.rs` / `table.rs` are updated to route through it. - `python/Cargo.toml` adds `libc = "0.2"` and the tokio `rt-multi-thread` feature. ### Python — `lancedb/background_loop.py` - Refactors `BackgroundEventLoop.__init__` to a reusable `_start()` method. - An `os.register_at_fork(after_in_child=…)` hook calls `LOOP._start()` to give the singleton a fresh asyncio loop and thread **in place**. This matters because the rest of the codebase imports `LOOP` via `from .background_loop import LOOP` — rebinding the module attribute would leave those references holding the dead loop. ### Python — `lancedb/__init__.py` Removes the `__warn_on_fork` pre-fork warning (and the now-unused `import warnings`). Fork is supported. ## Test plan - [x] New `test_permutation_dataloader_fork_workers` in `python/tests/test_torch.py`: runs a `Permutation` through `torch.utils.data.DataLoader(num_workers=2, multiprocessing_context="fork")` inside a spawn-isolated child with a 30s hang detector. **Pre-fix**: timed out at 36s. **Post-fix**: passes in ~3.6s. - [x] New `test_remote_connection_after_fork` in `python/tests/test_remote_db.py`: forks a child that creates a fresh `lancedb.connect(...)` against a mock HTTP server and calls `table_names()`; passes in <1s, validates the runtime reset is sufficient for fresh remote clients. - [x] All 62 tests in `test_torch.py` + `test_permutation.py` pass. - [x] All 35 tests in `test_remote_db.py` pass. - [x] `test_table.py` (87) + `test_db.py` + `test_query.py` (157, minus one unrelated `sentence_transformers` import skip) — 244 passing. - [x] `cargo clippy -p lancedb-python --tests` clean. - [x] `cargo fmt`, `ruff check`, `ruff format` all clean. ## Known limitation (follow-up) This PR makes a **freshly-built** `lancedb.connect(...)` work in a forked child. An **inherited** `Connection` from the parent still carries an inherited `reqwest::Client` whose hyper connection pool references socket FDs and TCP/TLS state shared with the parent — using it from the child after fork is unsafe (especially with HTTP/1.1 keep-alive). The recommended pattern for fork-based `DataLoader` workers that hit a remote DB is to construct a new connection inside the worker. Auto-clearing inherited HTTP client pools on fork would require tracking live `Connection` instances in `lancedb` core and is left for a follow-up PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> |