mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 07:58:31 +00:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 1825f6df9e | |||
| 6a1dacb1fe | |||
| 62dea8acd8 | |||
| ae9e8e8f8d | |||
| 36e44ab7a9 | |||
| feccabd739 | |||
| 1f3093a51f |
@@ -804,7 +804,7 @@ class LanceDBConnection(DBConnection):
|
||||
"manifest_enabled": self._manifest_enabled,
|
||||
"namespace_client_properties": self._namespace_client_properties,
|
||||
"read_consistency_interval_seconds": (
|
||||
rci.total_seconds() if rci else None
|
||||
rci.total_seconds() if rci is not None else None
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -482,6 +482,16 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
def serialize(self) -> str:
|
||||
import json
|
||||
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot serialize a namespace connection constructed from an "
|
||||
"opaque namespace client. Pass namespace_client_impl and "
|
||||
"namespace_client_properties when constructing the connection."
|
||||
)
|
||||
|
||||
return json.dumps(
|
||||
{
|
||||
"connection_type": "namespace",
|
||||
@@ -493,7 +503,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"storage_options": self.storage_options or None,
|
||||
"read_consistency_interval_seconds": (
|
||||
self.read_consistency_interval.total_seconds()
|
||||
if self.read_consistency_interval
|
||||
if self.read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
}
|
||||
@@ -569,6 +579,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
self,
|
||||
name,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
@@ -607,6 +618,8 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
self,
|
||||
name,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
@@ -899,10 +912,13 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
self,
|
||||
name,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
location=table_uri,
|
||||
namespace_client=namespace_client,
|
||||
managed_versioning=managed_versioning,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
_async=async_table,
|
||||
)
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -6,6 +6,8 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import inspect
|
||||
import deprecation
|
||||
import os
|
||||
import threading
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
@@ -163,7 +165,7 @@ def _maybe_add_fts_error_note(
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db import LanceDBConnection
|
||||
from .db import DBConnection, LanceDBConnection
|
||||
from ._lancedb import (
|
||||
Table as LanceDBTable,
|
||||
OptimizeStats,
|
||||
@@ -2105,6 +2107,23 @@ class Table(ABC):
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LanceTableReopenState:
|
||||
"""Process-independent coordinates for reopening a native table."""
|
||||
|
||||
connection_state: Optional[str]
|
||||
can_reopen_after_fork: bool
|
||||
fork_reopen_error: Optional[str]
|
||||
name: str
|
||||
namespace_path: List[str]
|
||||
storage_options: Optional[Dict[str, str]]
|
||||
index_cache_size: Optional[int]
|
||||
location: Optional[str]
|
||||
managed_versioning: Optional[bool]
|
||||
branch: Optional[str]
|
||||
checkout_version: Optional[int]
|
||||
|
||||
|
||||
class LanceTable(Table):
|
||||
"""
|
||||
A table in a LanceDB database.
|
||||
@@ -2139,7 +2158,10 @@ 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._managed_versioning = managed_versioning
|
||||
self._namespace_client = namespace_client
|
||||
self._pushdown_operations = pushdown_operations or set()
|
||||
# When the connection built the namespace client natively (e.g. an
|
||||
@@ -2164,9 +2186,206 @@ 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._native_state_guard = (self._pid, threading.RLock())
|
||||
|
||||
# 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. Connections without reconstructible metadata are not
|
||||
# safe to reuse in a forked child, so retain a clear diagnostic rather
|
||||
# than advertising them as reopenable based on JSON encoding alone.
|
||||
try:
|
||||
connection_uri: Optional[str] = self._conn.uri
|
||||
except Exception:
|
||||
connection_uri = None
|
||||
|
||||
fork_reopen_error: Optional[str] = None
|
||||
try:
|
||||
connection_state: Optional[str] = self._conn.serialize()
|
||||
can_reopen_after_fork = connection_uri is not None and not (
|
||||
connection_uri.startswith("memory://")
|
||||
)
|
||||
except Exception as error:
|
||||
connection_state = None
|
||||
can_reopen_after_fork = False
|
||||
if connection_uri is not None and not connection_uri.startswith(
|
||||
"memory://"
|
||||
):
|
||||
fork_reopen_error = (
|
||||
f"Cannot reopen table {name!r} in a forked process: {error}"
|
||||
)
|
||||
|
||||
self._reopen_state = _LanceTableReopenState(
|
||||
connection_state=connection_state,
|
||||
can_reopen_after_fork=can_reopen_after_fork,
|
||||
fork_reopen_error=fork_reopen_error,
|
||||
name=name,
|
||||
namespace_path=list(self._namespace_path),
|
||||
storage_options=(
|
||||
dict(self._storage_options)
|
||||
if self._storage_options is not None
|
||||
else None
|
||||
),
|
||||
index_cache_size=self._index_cache_size,
|
||||
location=self._location,
|
||||
managed_versioning=self._managed_versioning,
|
||||
branch=self._table.current_branch(),
|
||||
checkout_version=None,
|
||||
)
|
||||
|
||||
@property
|
||||
def _connection_state(self) -> Optional[str]:
|
||||
"""Serialized connection retained for worker reconstruction."""
|
||||
return self._reopen_state.connection_state
|
||||
|
||||
@property
|
||||
def _can_reopen_after_fork(self) -> bool:
|
||||
return self._reopen_state.can_reopen_after_fork
|
||||
|
||||
@property
|
||||
def _branch(self) -> Optional[str]:
|
||||
state = getattr(self, "_reopen_state", None)
|
||||
if state is not None:
|
||||
return state.branch
|
||||
return getattr(self, "_legacy_branch", None)
|
||||
|
||||
@_branch.setter
|
||||
def _branch(self, value: Optional[str]) -> None:
|
||||
state = getattr(self, "_reopen_state", None)
|
||||
if state is not None:
|
||||
state.branch = value
|
||||
else:
|
||||
self._legacy_branch = value
|
||||
|
||||
@property
|
||||
def _checkout_version(self) -> Optional[int]:
|
||||
state = getattr(self, "_reopen_state", None)
|
||||
if state is not None:
|
||||
return state.checkout_version
|
||||
return getattr(self, "_legacy_checkout_version", None)
|
||||
|
||||
@_checkout_version.setter
|
||||
def _checkout_version(self, value: Optional[int]) -> None:
|
||||
state = getattr(self, "_reopen_state", None)
|
||||
if state is not None:
|
||||
state.checkout_version = value
|
||||
else:
|
||||
self._legacy_checkout_version = value
|
||||
|
||||
def _native_state_lock(self):
|
||||
"""Return the per-process lock coordinating native mode and reopen state."""
|
||||
pid = os.getpid()
|
||||
guard = getattr(self, "_native_state_guard", None)
|
||||
if guard is None:
|
||||
candidate = (pid, threading.RLock())
|
||||
guard = self.__dict__.setdefault("_native_state_guard", candidate)
|
||||
elif guard[0] != pid:
|
||||
# A lock inherited while another parent thread held it cannot be
|
||||
# safely acquired in the child. Child state starts single-threaded,
|
||||
# so replace it before coordinating the first reopen.
|
||||
guard = (pid, threading.RLock())
|
||||
self._native_state_guard = guard
|
||||
return guard[1]
|
||||
|
||||
@classmethod
|
||||
def _open_from_reopen_state(
|
||||
cls,
|
||||
connection: "DBConnection",
|
||||
state: "_LanceTableReopenState",
|
||||
) -> "LanceTable":
|
||||
"""Open a table from its complete process-independent descriptor."""
|
||||
async_connection = getattr(connection, "_conn", None)
|
||||
if async_connection is None:
|
||||
async_connection = connection._inner
|
||||
|
||||
namespace_client = getattr(connection, "_namespace_client", None)
|
||||
async_table = LOOP.run(
|
||||
async_connection.open_table(
|
||||
state.name,
|
||||
namespace_path=state.namespace_path,
|
||||
storage_options=state.storage_options,
|
||||
index_cache_size=state.index_cache_size,
|
||||
location=state.location,
|
||||
namespace_client=namespace_client,
|
||||
managed_versioning=state.managed_versioning,
|
||||
)
|
||||
)
|
||||
table = cls(
|
||||
connection,
|
||||
state.name,
|
||||
namespace_path=state.namespace_path,
|
||||
storage_options=state.storage_options,
|
||||
index_cache_size=state.index_cache_size,
|
||||
location=state.location,
|
||||
namespace_client=namespace_client,
|
||||
managed_versioning=state.managed_versioning,
|
||||
pushdown_operations=getattr(
|
||||
connection, "_namespace_client_pushdown_operations", None
|
||||
),
|
||||
route_pushdown_to_rust=getattr(
|
||||
connection, "_route_pushdown_to_rust", False
|
||||
),
|
||||
_async=async_table,
|
||||
)
|
||||
if state.branch is not None:
|
||||
table = table.branches.checkout(state.branch, state.checkout_version)
|
||||
elif state.checkout_version is not None:
|
||||
table.checkout(state.checkout_version)
|
||||
return table
|
||||
|
||||
def _ensure_open(self) -> None:
|
||||
"""Reopen native table handles inherited from another process."""
|
||||
with self._native_state_lock():
|
||||
pid = os.getpid()
|
||||
if getattr(self, "_pid", pid) == pid:
|
||||
return
|
||||
|
||||
state = getattr(self, "_reopen_state", None)
|
||||
fork_reopen_error = getattr(state, "fork_reopen_error", None)
|
||||
if fork_reopen_error is not None:
|
||||
raise RuntimeError(fork_reopen_error)
|
||||
if (
|
||||
state is None
|
||||
or not state.can_reopen_after_fork
|
||||
or state.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(state.connection_state, for_worker=True)
|
||||
reopened = self._open_from_reopen_state(
|
||||
connection,
|
||||
state,
|
||||
)
|
||||
|
||||
# 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._reopen_state = reopened._reopen_state
|
||||
self._pid = pid
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
if hasattr(self, "_name"):
|
||||
return self._name
|
||||
# Preserve compatibility with lightweight / legacy instances that
|
||||
# were constructed without running ``LanceTable.__init__``.
|
||||
return self._table.name
|
||||
|
||||
@property
|
||||
@@ -2383,18 +2602,75 @@ 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,
|
||||
managed_versioning=self._managed_versioning,
|
||||
_async=async_table,
|
||||
)
|
||||
table._checkout_version = version
|
||||
return table
|
||||
|
||||
def _resolve_checkout_version(self, version: Union[int, str]) -> int:
|
||||
if isinstance(version, int):
|
||||
return version
|
||||
try:
|
||||
return self.tags.get_version(version)
|
||||
except RuntimeError as err:
|
||||
# Native checkout historically exposes an unknown tag as ValueError.
|
||||
# Preserve that contract while resolving tags before mutating the table.
|
||||
if "Ref not found" in str(err) and "does not exist" in str(err):
|
||||
raise ValueError(str(err)) from err
|
||||
raise
|
||||
|
||||
async def _commit_native_state(
|
||||
self,
|
||||
transition,
|
||||
version: Optional[int],
|
||||
started: threading.Event,
|
||||
finished: threading.Event,
|
||||
):
|
||||
"""Commit a native transition and its fork coordinate as one task."""
|
||||
started.set()
|
||||
try:
|
||||
task = asyncio.ensure_future(transition)
|
||||
try:
|
||||
result = await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
# BackgroundEventLoop cancels its submitted task when the
|
||||
# waiting caller is interrupted. Let an already-started native
|
||||
# transition reach its authoritative terminal state before the
|
||||
# per-table boundary is released.
|
||||
result = await task
|
||||
self._checkout_version = version
|
||||
raise
|
||||
self._checkout_version = version
|
||||
return result
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
def _run_native_state_transition(self, transition, version: Optional[int]):
|
||||
started = threading.Event()
|
||||
finished = threading.Event()
|
||||
try:
|
||||
return LOOP.run(
|
||||
self._commit_native_state(transition, version, started, finished)
|
||||
)
|
||||
except BaseException:
|
||||
if started.is_set():
|
||||
while not finished.is_set():
|
||||
try:
|
||||
finished.wait()
|
||||
except BaseException: # noqa: PERF203
|
||||
continue
|
||||
raise
|
||||
|
||||
def checkout(self, version: Union[int, str]):
|
||||
"""Checkout a version of the table. This is an in-place operation.
|
||||
@@ -2432,7 +2708,14 @@ class LanceTable(Table):
|
||||
vector type
|
||||
0 [1.1, 0.9] vector
|
||||
"""
|
||||
LOOP.run(self._table.checkout(version))
|
||||
# Resolve tags before mutating the native handle. This leaves the live
|
||||
# handle and reopen descriptor aligned if tag lookup fails, and avoids a
|
||||
# second fallible version lookup after checkout succeeds.
|
||||
with self._native_state_lock():
|
||||
resolved_version = self._resolve_checkout_version(version)
|
||||
self._run_native_state_transition(
|
||||
self._table.checkout(resolved_version), resolved_version
|
||||
)
|
||||
|
||||
def checkout_latest(self):
|
||||
"""Checkout the latest version of the table. This is an in-place operation.
|
||||
@@ -2440,7 +2723,8 @@ class LanceTable(Table):
|
||||
The table will be set back into standard mode, and will track the latest
|
||||
version of the table.
|
||||
"""
|
||||
LOOP.run(self._table.checkout_latest())
|
||||
with self._native_state_lock():
|
||||
self._run_native_state_transition(self._table.checkout_latest(), None)
|
||||
|
||||
def restore(self, version: Optional[Union[int, str]] = None):
|
||||
"""Restore a version of the table. This is an in-place operation.
|
||||
@@ -2486,9 +2770,13 @@ class LanceTable(Table):
|
||||
>>> len(table.list_versions())
|
||||
4
|
||||
"""
|
||||
if version is not None:
|
||||
LOOP.run(self._table.checkout(version))
|
||||
LOOP.run(self._table.restore())
|
||||
with self._native_state_lock():
|
||||
if version is not None:
|
||||
resolved_version = self._resolve_checkout_version(version)
|
||||
self._run_native_state_transition(
|
||||
self._table.checkout(resolved_version), resolved_version
|
||||
)
|
||||
self._run_native_state_transition(self._table.restore(), None)
|
||||
|
||||
def count_rows(self, filter: Optional[str] = None) -> int:
|
||||
return LOOP.run(self._table.count_rows(filter))
|
||||
@@ -3599,7 +3887,9 @@ class LanceTable(Table):
|
||||
self = cls.__new__(cls)
|
||||
self._conn = db
|
||||
self._namespace_path = namespace_path
|
||||
self._index_cache_size = None
|
||||
self._location = location
|
||||
self._managed_versioning = None
|
||||
self._namespace_client = namespace_client
|
||||
self._pushdown_operations = pushdown_operations or set()
|
||||
self._route_pushdown_to_rust = route_pushdown_to_rust
|
||||
@@ -3627,6 +3917,7 @@ class LanceTable(Table):
|
||||
enable_v2_manifest_paths
|
||||
)
|
||||
|
||||
self._storage_options = storage_options
|
||||
self._table = LOOP.run(
|
||||
self._conn._conn.create_table(
|
||||
name,
|
||||
@@ -3643,6 +3934,7 @@ class LanceTable(Table):
|
||||
namespace_client=namespace_client,
|
||||
)
|
||||
)
|
||||
self._initialize_reopen_state(name)
|
||||
return self
|
||||
|
||||
def delete(self, where: Union[str, Expr]) -> DeleteResult:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import json
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
@@ -101,6 +102,17 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey
|
||||
assert db_from_inner.read_consistency_interval == consistency_interval
|
||||
|
||||
|
||||
def test_serialize_preserves_zero_read_consistency_interval(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(0))
|
||||
table = db.create_table("items", pa.table({"x": [1]}))
|
||||
|
||||
encoded = json.loads(table._connection_state)
|
||||
assert encoded["read_consistency_interval_seconds"] == 0.0
|
||||
|
||||
restored = lancedb.deserialize_conn(table._connection_state)
|
||||
assert restored.read_consistency_interval == timedelta(0)
|
||||
|
||||
|
||||
def test_ingest_pd(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
|
||||
@@ -6,9 +6,13 @@
|
||||
import tempfile
|
||||
import shutil
|
||||
import importlib
|
||||
import multiprocessing as mp
|
||||
import sys
|
||||
from datetime import timedelta
|
||||
import pytest
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
from lance_namespace import connect as namespace_connect
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb.namespace import _MAX_QUERY_K
|
||||
from lancedb.table import AsyncTable, LanceTable
|
||||
@@ -72,6 +76,16 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable:
|
||||
return table
|
||||
|
||||
|
||||
def _direct_namespace_fork_child(table, result_queue):
|
||||
from lancedb.permutation import Permutation
|
||||
|
||||
try:
|
||||
permutation = Permutation.identity(table)
|
||||
result_queue.put(("ok", permutation.num_rows))
|
||||
except Exception as error:
|
||||
result_queue.put((type(error).__name__, str(error)))
|
||||
|
||||
|
||||
class TestNamespaceConnection:
|
||||
"""Test namespace-based LanceDB connection using DirectoryNamespace."""
|
||||
|
||||
@@ -419,7 +433,95 @@ class TestNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
db.create_table("test_table", schema=schema, storage_options=table_opts)
|
||||
created = db.create_table(
|
||||
"test_table", schema=schema, storage_options=table_opts
|
||||
)
|
||||
assert created._storage_options == table_opts
|
||||
|
||||
opened = db.open_table(
|
||||
"test_table",
|
||||
storage_options={"allow_http": "true"},
|
||||
index_cache_size=17,
|
||||
)
|
||||
assert opened._storage_options == {"allow_http": "true"}
|
||||
assert opened._index_cache_size == 17
|
||||
opened._pid = -1
|
||||
opened._ensure_open()
|
||||
assert opened.count_rows() == 0
|
||||
|
||||
def test_serialize_preserves_zero_read_consistency_interval(self):
|
||||
db = lancedb.connect_namespace(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
read_consistency_interval=timedelta(0),
|
||||
)
|
||||
|
||||
restored = lancedb.deserialize_conn(db.serialize())
|
||||
assert restored.read_consistency_interval == timedelta(0)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform != "linux",
|
||||
reason="fork() is only supported safely for this test on Linux",
|
||||
)
|
||||
def test_direct_namespace_with_descriptor_reopens_after_fork(self):
|
||||
properties = {"root": self.temp_dir}
|
||||
namespace = namespace_connect("dir", properties)
|
||||
db = lancedb.LanceNamespaceDBConnection(
|
||||
namespace,
|
||||
namespace_client_impl="dir",
|
||||
namespace_client_properties=properties,
|
||||
)
|
||||
table = db.create_table("items", pa.table({"id": [1]}))
|
||||
|
||||
ctx = mp.get_context("fork")
|
||||
result_queue = ctx.Queue()
|
||||
process = ctx.Process(
|
||||
target=_direct_namespace_fork_child,
|
||||
args=(table, result_queue),
|
||||
)
|
||||
process.start()
|
||||
process.join(10)
|
||||
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(5)
|
||||
pytest.fail("Direct namespace table hung while reopening after fork")
|
||||
|
||||
assert process.exitcode == 0
|
||||
assert result_queue.get(timeout=2) == ("ok", 1)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform != "linux",
|
||||
reason="fork() is only supported safely for this test on Linux",
|
||||
)
|
||||
def test_opaque_direct_namespace_reports_unsupported_fork(self):
|
||||
namespace = namespace_connect("dir", {"root": self.temp_dir})
|
||||
db = lancedb.LanceNamespaceDBConnection(namespace)
|
||||
table = db.create_table("items", pa.table({"id": [1]}))
|
||||
|
||||
with pytest.raises(ValueError, match="opaque namespace client"):
|
||||
db.serialize()
|
||||
assert not table._can_reopen_after_fork
|
||||
|
||||
ctx = mp.get_context("fork")
|
||||
result_queue = ctx.Queue()
|
||||
process = ctx.Process(
|
||||
target=_direct_namespace_fork_child,
|
||||
args=(table, result_queue),
|
||||
)
|
||||
process.start()
|
||||
process.join(10)
|
||||
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(5)
|
||||
pytest.fail("Opaque namespace table hung after fork")
|
||||
|
||||
assert process.exitcode == 0
|
||||
error_type, message = result_queue.get(timeout=2)
|
||||
assert error_type == "RuntimeError"
|
||||
assert "Cannot reopen table 'items' in a forked process" in message
|
||||
assert "namespace_client_impl and namespace_client_properties" in message
|
||||
|
||||
def test_namespace_operations(self):
|
||||
"""Test namespace management operations."""
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import asyncio
|
||||
import ctypes
|
||||
import gc
|
||||
import os
|
||||
@@ -9,7 +10,7 @@ import sys
|
||||
import threading
|
||||
import warnings
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures import CancelledError, ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from time import sleep
|
||||
from typing import List
|
||||
@@ -2145,6 +2146,260 @@ def test_restore(mem_db: DBConnection):
|
||||
table.restore(0)
|
||||
|
||||
|
||||
def test_restore_tracks_checkout_when_restore_fails():
|
||||
class FailingRestore:
|
||||
def __init__(self):
|
||||
self.live_version = None
|
||||
|
||||
async def checkout(self, version):
|
||||
self.live_version = version
|
||||
|
||||
async def restore(self):
|
||||
raise RuntimeError("injected restore failure")
|
||||
|
||||
inner = FailingRestore()
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = None
|
||||
|
||||
with pytest.raises(RuntimeError, match="injected restore failure"):
|
||||
table.restore(7)
|
||||
|
||||
assert table._checkout_version == inner.live_version
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("operation", "expected_descriptor", "expected_restore_calls"),
|
||||
[("checkout", 11, 0), ("restore", None, 1)],
|
||||
)
|
||||
def test_string_tag_resolves_before_checkout(
|
||||
operation, expected_descriptor, expected_restore_calls
|
||||
):
|
||||
class Tags:
|
||||
async def get_version(self, tag):
|
||||
assert tag == "tag-v1"
|
||||
return 11
|
||||
|
||||
class NoPostCheckoutVersionLookup:
|
||||
def __init__(self):
|
||||
self.tags = Tags()
|
||||
self.checkout_versions = []
|
||||
self.restore_calls = 0
|
||||
|
||||
async def checkout(self, version):
|
||||
self.checkout_versions.append(version)
|
||||
|
||||
async def version(self):
|
||||
raise RuntimeError("post-checkout version lookup must not run")
|
||||
|
||||
async def restore(self):
|
||||
self.restore_calls += 1
|
||||
|
||||
inner = NoPostCheckoutVersionLookup()
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = 3
|
||||
|
||||
getattr(table, operation)("tag-v1")
|
||||
|
||||
assert inner.checkout_versions == [11]
|
||||
assert table._checkout_version == expected_descriptor
|
||||
assert inner.restore_calls == expected_restore_calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["checkout", "restore"])
|
||||
def test_string_tag_resolution_failure_does_not_mutate_handle(operation):
|
||||
class FailingTags:
|
||||
async def get_version(self, tag):
|
||||
assert tag == "missing-tag"
|
||||
raise RuntimeError("injected tag lookup failure")
|
||||
|
||||
class UnchangedTable:
|
||||
def __init__(self):
|
||||
self.tags = FailingTags()
|
||||
self.checkout_calls = 0
|
||||
self.restore_calls = 0
|
||||
|
||||
async def checkout(self, version):
|
||||
self.checkout_calls += 1
|
||||
|
||||
async def restore(self):
|
||||
self.restore_calls += 1
|
||||
|
||||
inner = UnchangedTable()
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = 3
|
||||
|
||||
with pytest.raises(RuntimeError, match="injected tag lookup failure"):
|
||||
getattr(table, operation)("missing-tag")
|
||||
|
||||
assert table._checkout_version == 3
|
||||
assert inner.checkout_calls == 0
|
||||
assert inner.restore_calls == 0
|
||||
|
||||
|
||||
def test_native_state_transitions_are_serialized(monkeypatch):
|
||||
from lancedb.background_loop import LOOP
|
||||
|
||||
class Inner:
|
||||
def __init__(self):
|
||||
self.live_version = None
|
||||
|
||||
async def checkout(self, version):
|
||||
self.live_version = version
|
||||
|
||||
inner = Inner()
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = None
|
||||
|
||||
first_native_done = threading.Event()
|
||||
release_first_call = threading.Event()
|
||||
second_call_started = threading.Event()
|
||||
second_call_done = threading.Event()
|
||||
errors = []
|
||||
original_run = LOOP.run
|
||||
|
||||
def delayed_delivery(awaitable):
|
||||
result = original_run(awaitable)
|
||||
if threading.current_thread().name == "checkout-1":
|
||||
first_native_done.set()
|
||||
assert release_first_call.wait(5)
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", delayed_delivery)
|
||||
|
||||
def checkout(version):
|
||||
if version == 2:
|
||||
second_call_started.set()
|
||||
try:
|
||||
table.checkout(version)
|
||||
except BaseException as err:
|
||||
errors.append(err)
|
||||
finally:
|
||||
if version == 2:
|
||||
second_call_done.set()
|
||||
|
||||
first = threading.Thread(target=checkout, args=(1,), name="checkout-1")
|
||||
first.start()
|
||||
assert first_native_done.wait(5)
|
||||
|
||||
second = threading.Thread(target=checkout, args=(2,), name="checkout-2")
|
||||
second.start()
|
||||
assert second_call_started.wait(5)
|
||||
assert not second_call_done.wait(0.1)
|
||||
|
||||
release_first_call.set()
|
||||
first.join(5)
|
||||
second.join(5)
|
||||
|
||||
assert not first.is_alive()
|
||||
assert not second.is_alive()
|
||||
assert errors == []
|
||||
assert inner.live_version == 2
|
||||
assert table._checkout_version == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("operation", "args", "initial_version", "expected_version"),
|
||||
[
|
||||
("checkout", (11,), 3, 11),
|
||||
("checkout_latest", (), 3, None),
|
||||
("restore", (), 11, None),
|
||||
],
|
||||
)
|
||||
def test_native_state_commits_before_success_delivery(
|
||||
monkeypatch, operation, args, initial_version, expected_version
|
||||
):
|
||||
from lancedb.background_loop import LOOP
|
||||
|
||||
class Inner:
|
||||
def __init__(self, live_version):
|
||||
self.live_version = live_version
|
||||
|
||||
async def checkout(self, version):
|
||||
self.live_version = version
|
||||
|
||||
async def checkout_latest(self):
|
||||
self.live_version = None
|
||||
|
||||
async def restore(self):
|
||||
self.live_version = None
|
||||
|
||||
inner = Inner(initial_version)
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = initial_version
|
||||
|
||||
original_run = LOOP.run
|
||||
|
||||
def success_then_interrupt(awaitable):
|
||||
original_run(awaitable)
|
||||
raise KeyboardInterrupt("injected after native success")
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", success_then_interrupt)
|
||||
|
||||
with pytest.raises(KeyboardInterrupt, match="injected after native success"):
|
||||
getattr(table, operation)(*args)
|
||||
|
||||
assert inner.live_version == expected_version
|
||||
assert table._checkout_version == expected_version
|
||||
|
||||
|
||||
def test_native_state_waits_for_cancelled_delivery(monkeypatch):
|
||||
from lancedb.background_loop import LOOP
|
||||
|
||||
class Inner:
|
||||
def __init__(self):
|
||||
self.live_version = 3
|
||||
|
||||
async def checkout(self, version):
|
||||
await asyncio.sleep(0.01)
|
||||
self.live_version = version
|
||||
|
||||
inner = Inner()
|
||||
table = LanceTable.__new__(LanceTable)
|
||||
table._table = inner
|
||||
table._checkout_version = 3
|
||||
|
||||
original_run = LOOP.run
|
||||
|
||||
def cancel_while_running(awaitable):
|
||||
async def cancel_after_start():
|
||||
task = asyncio.create_task(awaitable)
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
return await task
|
||||
|
||||
return original_run(cancel_after_start())
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", cancel_while_running)
|
||||
|
||||
with pytest.raises(CancelledError):
|
||||
table.checkout(11)
|
||||
|
||||
assert inner.live_version == 11
|
||||
assert table._checkout_version == 11
|
||||
|
||||
|
||||
def test_reopen_preserves_explicit_table_location(tmp_path):
|
||||
db = lancedb.connect(tmp_path / "db")
|
||||
location = str(tmp_path / "physical-table")
|
||||
table = LanceTable.create(
|
||||
db,
|
||||
"items",
|
||||
pa.table({"x": [1]}),
|
||||
location=location,
|
||||
)
|
||||
|
||||
table._pid = -1
|
||||
table._ensure_open()
|
||||
|
||||
assert table.count_rows() == 1
|
||||
assert table._location == location
|
||||
|
||||
|
||||
def test_restore_with_tags(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
|
||||
@@ -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=(
|
||||
|
||||
Reference in New Issue
Block a user