Commit Graph
3 Commits
Author SHA1 Message Date
Bruno RamirezandClaude Sonnet 5 3fca33fcb5 fix: shut down the shared Tokio runtime on interpreter exit (#4175)
Short-lived Python processes using this client can occasionally crash
with SIGABRT during interpreter shutdown, even after every operation
they ran completed successfully. The cause is the shared Tokio runtime
backing every async call: it's never told to shut down at normal process
exit, only reset (and deliberately leaked) on `fork()`. Its worker
threads keep running, uncoordinated with the interpreter, until the
process actually ends, and if one is mid-task exactly as `Py_Finalize`
starts 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, so Rust aborts the whole process
instead of just failing that one call. This PR gives the runtime a
coordinated, bounded shutdown by registering a Python `atexit` callback
that runs while the interpreter is still fully valid.

Getting the exit lifecycle right took a few rounds of review. Earlier
versions freed the runtime as soon as `Arc::strong_count` looked low,
but that's the wrong signal — it reflects who currently holds a
reference, not who's logically still in flight. That mistake showed up
three ways: a caller could dereference memory already freed out from
under it; an install already in progress could finish invisibly after
`shutdown()` had already decided there was nothing to do; and a spawned
task could end up as the final owner of the `Runtime`, so completing it
dropped the runtime from inside one of its own worker threads, which
Tokio itself forbids and panics on (this reproduced unprompted in this
branch's own test suite). Fixing all three meant replacing
reference-count-based tracking with an explicit counter of in-flight
top-level calls that `shutdown()` waits on directly.

This was accomplished with the following changes:
- The runtime lives in an `ArcSwapOption<Tagged>`, where `Tagged` pairs
the `Runtime` with the fork generation it was built in.
- An `OUTSTANDING` counter, incremented before a top-level
`spawn`/`spawn_blocking`/`block_on` call does anything else and
decremented only once it has truly finished (via an `OutstandingGuard`
token that carries no reference to the runtime), is what `shutdown()`
waits on — not `Arc::strong_count` or whether the slot looks empty. This
closes the install-race and makes it impossible for a task's own
completion to be the runtime's final drop.
- Once `shutdown()`'s bound elapses, it stops waiting and attempts
retirement anyway, rather than returning with the runtime and its
workers left fully alive.
- `spawn`/`spawn_blocking` use `Handle::try_current()` to pin any nested
spawn (`future_into_py` spawns a task that itself spawns a second one
for the real work) to whichever runtime is already executing it, so a
reclaim landing between the two calls can't split one logical operation
across two different runtime instances.
- The fork-child handler now only bumps a bare `GENERATION` counter — no
`ArcSwapOption` call of any kind from that context, since
`swap`/`compare_and_swap` do real reader-reconciliation work
(thread-local state, potentially an allocation) that isn't safe in a
forked child. `get_runtime()` compares its installed runtime's
generation against the live counter from ordinary context and rebuilds
on a mismatch.
- Registered `shutdown_runtime` as a Python `atexit` callback in the
`_lancedb` module init, running with the GIL released (`Python::detach`)
since the bounded wait could otherwise deadlock against any in-flight
task that itself needs the GIL.

### Testing
- Unit tests in `runtime.rs` cover: shutdown with no runtime created,
shutdown after use and lazy rebuild afterward, calling shutdown twice in
a row, a concurrent stress test racing many threads against shutdown, a
nested-spawn test reproducing `future_into_py`'s own
spawn-within-a-spawn shape under concurrent shutdown, a test confirming
a top-level task in flight survives a concurrent shutdown reclaim, and a
test forcing the install-vs-shutdown race directly.
- Built the wheel and ran a concurrent reproducer (many threads
hammering the client while `atexit` fires) over 100 times with no hangs
or crashes, plus a 30-second-join variant and repeated runs of a
short-lived process confirming clean exits with no added latency.

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-15 12:48:31 -07:00
Ryan Green 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.
2026-06-25 21:47:55 -07:00
Weston PaceandClaude Opus 4.7 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>
2026-05-05 13:44:10 -07:00