fix(python): resolve Ollama embedding serialization error in create_table (#3583)

This PR fixes a serialization error when using Ollama embeddings in
`create_table`.

The use of `@cached_property` for the Ollama client was causing issues
during serialization/pickling, which is required by certain LanceDB
operations (like when using multiprocessing or certain storage
backends). Switching to a standard `@property` ensures the client is
instantiated when needed without being stored in a way that breaks
serialization.

Verified with the following script:
```python
import lancedb
from lancedb.embeddings import get_registry
import pickle

registry = get_registry().get(\"ollama\")
model = registry(name=\"llama3\")

# This would fail before the fix
pickled = pickle.dumps(model)
unpickled = pickle.loads(pickled)
```

Fixes #2629 (or similar serialization issues reported).

---------

Co-authored-by: Unmilan Mukherjee <Missing-Identity@users.noreply.github.com>
This commit is contained in:
Unmilan Mukherjee
2026-07-09 23:28:41 +05:30
committed by GitHub
parent 291e9e37be
commit 75c5c83f12
2 changed files with 28 additions and 1 deletions

View File

@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from functools import cached_property
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
import numpy as np
@@ -56,6 +56,16 @@ class OllamaEmbeddings(TextEmbeddingFunction):
embeddings = self._compute_embedding(texts)
return list(embeddings)
def __getstate__(self) -> dict[str, Any]:
state = super().__getstate__()
state["__dict__"] = {
k: v for k, v in state["__dict__"].items() if k != "_ollama_client"
}
return state
def __setstate__(self, state: dict[str, Any]) -> None:
super().__setstate__(state)
@cached_property
def _ollama_client(self) -> "ollama.Client":
ollama = attempt_import_or_raise("ollama")

View File

@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import os
import pickle
from typing import List, Optional, Union
from unittest.mock import MagicMock, patch
@@ -591,6 +592,22 @@ def test_openai_no_retry_on_401(mock_sleep):
assert mock_sleep.call_count == 0
def test_ollama_embeddings_pickle():
"""OllamaEmbeddings must pickle even after the cached client is created."""
registry = get_registry()
model = registry.get("ollama").create(name="nomic-embed-text")
# Simulate accessing the cached client, which stores it on the instance.
model.__dict__["_ollama_client"] = MagicMock()
pickled = pickle.dumps(model)
restored = pickle.loads(pickled)
assert restored.name == "nomic-embed-text"
assert restored.host == "http://localhost:11434"
assert "_ollama_client" not in restored.__dict__
def test_url_retrieve_downloads_image():
"""
Embedding functions like open-clip, siglip, and jinaai use url_retrieve()