From 1f3093a51f9d23426798c97ecd2db31498a7fb7e Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:10:41 +0000 Subject: [PATCH] fix(python): reopen native tables in forked workers --- python/python/lancedb/permutation.py | 7 +++ python/python/lancedb/table.py | 76 ++++++++++++++++++++++++++-- python/python/tests/test_torch.py | 76 ++++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index b3d95c352..32ef87798 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -591,6 +591,13 @@ class Permutation: then the first split will be used. """ assert base_table is not None, "base_table is required" + # A PyTorch fork worker may construct its Permutation lazily from a + # table opened in the parent process. Reopen that table before the + # Rust reader clones its object-store clients and connection pools. + if hasattr(base_table, "_ensure_open"): + base_table._ensure_open() + if permutation_table is not None and hasattr(permutation_table, "_ensure_open"): + permutation_table._ensure_open() if split is not None: if permutation_table is None: raise ValueError( diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 31a70c298..188de239d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import inspect import deprecation +import os import warnings from abc import ABC, abstractmethod from dataclasses import dataclass @@ -2139,6 +2140,8 @@ class LanceTable(Table): namespace_path = [] self._conn = connection self._namespace_path = namespace_path + self._storage_options = storage_options + self._index_cache_size = index_cache_size self._location = location # Store location for use in _dataset_path self._namespace_client = namespace_client self._pushdown_operations = pushdown_operations or set() @@ -2164,10 +2167,65 @@ class LanceTable(Table): managed_versioning=managed_versioning, ) ) + self._initialize_reopen_state(name) + + def _initialize_reopen_state(self, name: str) -> None: + """Capture the state needed to replace inherited native handles.""" + self._name = name + self._pid = os.getpid() + self._branch = self._table.current_branch() + self._checkout_version: Optional[int] = None + + # A native table owns object-store clients and connection pools. Those + # handles must not be used after fork, so retain a process-independent + # connection description while it is still safe to inspect the parent + # connection. Some tables constructed from a bare Rust handle cannot + # be serialized; those retain the historical best-effort behavior. + try: + self._connection_state: Optional[str] = self._conn.serialize() + self._can_reopen_after_fork = not self._conn.uri.startswith("memory://") + except Exception: + self._connection_state = None + self._can_reopen_after_fork = False + + def _ensure_open(self) -> None: + """Reopen native table handles inherited from another process.""" + pid = os.getpid() + if self._pid == pid: + return + + if not self._can_reopen_after_fork or self._connection_state is None: + # In-memory and opaque Rust-only connections cannot be recreated + # from connection metadata. Their local handles retain the prior + # best-effort fork behavior. + self._pid = pid + return + + from lancedb import deserialize_conn + + connection = deserialize_conn(self._connection_state, for_worker=True) + reopened = connection.open_table( + self._name, + namespace_path=self._namespace_path or None, + storage_options=self._storage_options, + index_cache_size=self._index_cache_size, + branch=self._branch, + version=self._checkout_version, + ) + + # Keep this Python object stable because user datasets commonly retain + # it across fork. Replace every process-bound component with the fresh + # child's equivalent. + self._conn = reopened._conn + self._table = reopened._table + self._namespace_client = reopened._namespace_client + self._pushdown_operations = reopened._pushdown_operations + self._route_pushdown_to_rust = reopened._route_pushdown_to_rust + self._pid = pid @property def name(self) -> str: - return self._table.name + return self._name @property def namespace(self) -> List[str]: @@ -2379,18 +2437,20 @@ class LanceTable(Table): def _wrap_branch_handle( self, async_table: "AsyncTable", version: Optional[int] = None ) -> "LanceTable": - # version is unused locally: the pin already lives on async_table and a - # local handle is not reopened via a serialized connection. - return LanceTable( + table = LanceTable( self._conn, async_table.name, namespace_path=self._namespace_path, + storage_options=self._storage_options, + index_cache_size=self._index_cache_size, namespace_client=self._namespace_client, pushdown_operations=self._pushdown_operations, route_pushdown_to_rust=self._route_pushdown_to_rust, location=self._location, _async=async_table, ) + table._checkout_version = version + return table def checkout(self, version: Union[int, str]): """Checkout a version of the table. This is an in-place operation. @@ -2429,6 +2489,9 @@ class LanceTable(Table): 0 [1.1, 0.9] vector """ LOOP.run(self._table.checkout(version)) + # Resolve tags to their numeric version so a forked child can reopen + # the same pinned view through ``open_table(version=...)``. + self._checkout_version = self.version def checkout_latest(self): """Checkout the latest version of the table. This is an in-place operation. @@ -2437,6 +2500,7 @@ class LanceTable(Table): version of the table. """ LOOP.run(self._table.checkout_latest()) + self._checkout_version = None def restore(self, version: Optional[Union[int, str]] = None): """Restore a version of the table. This is an in-place operation. @@ -2485,6 +2549,7 @@ class LanceTable(Table): if version is not None: LOOP.run(self._table.checkout(version)) LOOP.run(self._table.restore()) + self._checkout_version = None def count_rows(self, filter: Optional[str] = None) -> int: return LOOP.run(self._table.count_rows(filter)) @@ -3595,6 +3660,7 @@ class LanceTable(Table): self = cls.__new__(cls) self._conn = db self._namespace_path = namespace_path + self._index_cache_size = None self._location = location self._namespace_client = namespace_client self._pushdown_operations = pushdown_operations or set() @@ -3623,6 +3689,7 @@ class LanceTable(Table): enable_v2_manifest_paths ) + self._storage_options = storage_options self._table = LOOP.run( self._conn._conn.create_table( name, @@ -3639,6 +3706,7 @@ class LanceTable(Table): namespace_client=namespace_client, ) ) + self._initialize_reopen_state(name) return self def delete(self, where: Union[str, Expr]) -> DeleteResult: diff --git a/python/python/tests/test_torch.py b/python/python/tests/test_torch.py index 38a367330..a54b139ba 100644 --- a/python/python/tests/test_torch.py +++ b/python/python/tests/test_torch.py @@ -342,6 +342,42 @@ def _multiworker_dataloader_target(db_uri: str, result_queue): result_queue.put(count) +class _LazyPermutationDataset(torch.utils.data.Dataset): + """Match applications that create their Permutation inside a fork worker.""" + + def __init__(self, table): + self._table = table + self._permutation = None + self._length = table.count_rows() + + def __len__(self): + return self._length + + def __getitems__(self, indices): + if self._permutation is None: + inherited_connection = self._table._conn + self._permutation = Permutation.identity(self._table) + if self._table._conn is inherited_connection: + raise RuntimeError("Permutation reused a connection inherited by fork") + return self._permutation.__getitems__(indices) + + +def _lazy_multiworker_dataloader_target(db_uri: str, result_queue): + table = lancedb.connect(db_uri).open_table("test_table") + dataset = _LazyPermutationDataset(table) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=10, + num_workers=2, + multiprocessing_context="fork", + ) + count = 0 + for batch in dataloader: + assert batch["a"].size(0) == 10 + count += 1 + result_queue.put(count) + + def _remote_multiworker_dataloader_target(port: int, result_queue): import lancedb from lancedb.permutation import Permutation @@ -410,6 +446,46 @@ def test_permutation_dataloader_fork_workers(tmp_path): assert queue.get() == 100 +@pytest.mark.skipif( + sys.platform != "linux", + reason=( + "fork() is unavailable on Windows and unsafe on macOS " + "(Apple frameworks/TLS are not fork-safe)" + ), +) +def test_lazy_permutation_reopens_inherited_table_in_fork_worker(tmp_path): + """A lazily built Permutation must not reuse an inherited table client. + + Object-store table handles contain HTTP connection pools that are unsafe + after fork. The local table makes the handle replacement deterministic + without requiring an S3 service in the unit-test environment. + """ + db_uri = str(tmp_path / "db") + db = lancedb.connect(db_uri) + db.create_table("test_table", pa.table({"a": list(range(1000))})) + + ctx = mp.get_context("spawn") + queue = ctx.Queue() + proc = ctx.Process( + target=_lazy_multiworker_dataloader_target, + args=(db_uri, queue), + ) + proc.start() + proc.join(timeout=30) + + if proc.is_alive(): + proc.terminate() + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + pytest.fail("Lazy Permutation hung in a fork-based DataLoader worker") + + assert proc.exitcode == 0, f"child exited with code {proc.exitcode}" + assert not queue.empty(), "child produced no batches" + assert queue.get() == 100 + + @pytest.mark.skipif( sys.platform != "linux", reason=(