Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 67 additions & 22 deletions sdks/python/src/opik/api_objects/dataset/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,11 @@ def __init__(
# one-shot sync on the first `insert()` instead of paying an N+1
# sync at list time.
self._hashes_synced: bool = True
# None until the backend version has actually been determined. Only a
# conclusive answer is stored, so a probe that failed to reach the
# backend is retried instead of pinning this dataset to sequential
# uploads for the rest of the session.
self._parallel_insert_supported_cache: Optional[bool] = None

@classmethod
def from_public(
Expand Down Expand Up @@ -611,44 +616,60 @@ def _insert_batch_with_retry(
)
LOGGER.debug("Successfully sent dataset items batch of size %d", len(batch))

@functools.cached_property
@property
def _parallel_insert_supported(self) -> bool:
"""Whether the backend tolerates concurrent batches sharing a batch_group_id.

Older backends race on them, so parallelism is only safe from
``constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT`` onwards. When the
version cannot be determined at all — unreachable endpoint, non-semver
build string — we report unsupported rather than risk the race.

Cached because the backend cannot change version mid-session, and
parallel upload is the default: probing per ``insert`` would add a
round trip to every call in a loop.
version cannot be determined we report unsupported rather than risk the
race.

The answer is cached because the backend cannot change version
mid-session and parallel upload is the default: probing per ``insert``
would add a round trip to every call in a loop. Only a conclusive
answer is cached — an unreachable backend is re-probed on the next
insert so parallel upload resumes once it recovers.
"""
if self._parallel_insert_supported_cache is not None:
return self._parallel_insert_supported_cache

try:
backend_version = self._rest_client.version()["version"]
except Exception:
LOGGER.warning(
"Could not reach the Opik backend to determine its version, "
"falling back to a sequential dataset upload for this insert.",
exc_info=True,
)
return False

try:
supported = (
semantic_version.SemanticVersion.parse(backend_version)
>= constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT
)
except Exception:
LOGGER.warning(
"Could not determine the Opik backend version, falling back to a "
"Could not parse the Opik backend version %s, falling back to a "
"sequential dataset upload. Parallel upload requires backend %s "
"or newer.",
constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT,
exc_info=True,
)
return False

if not supported:
LOGGER.warning(
"Opik backend %s does not support parallel dataset upload, falling "
"back to a sequential upload. Upgrade to backend %s or newer to use "
"num_threads.",
backend_version,
constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT,
exc_info=True,
)
supported = False
else:
if not supported:
LOGGER.warning(
"Opik backend %s does not support parallel dataset upload, "
"falling back to a sequential upload. Upgrade to backend %s or "
"newer to use num_threads.",
backend_version,
constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT,
)

self._parallel_insert_supported_cache = supported
return supported

def _send_batches(
Expand Down Expand Up @@ -718,6 +739,25 @@ def __internal_api__insert_items_as_dataclasses__(
num_threads: int = 1,
deduplication: bool = True,
) -> None:
# Validated here rather than in each public entry point: every insert
# path funnels through this method. A truthy string or None would
# otherwise silently pick the wrong duplicate-checking behaviour, and
# a non-integer worker count would fail on the comparison below with a
# TypeError instead of naming the offending argument.
if not isinstance(deduplication, bool):
raise ValueError("deduplication must be a bool")
Comment on lines +747 to +748

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Internal inserts bypass backend safety gate

The internal insertion funnel validates deduplication but forwards num_threads to _send_batches without applying the _parallel_insert_supported gate, so direct calls with num_threads > 1 use ThreadPoolExecutor against old backends instead of the fallback in insert() — should we apply the same worker-count compatibility check here?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/api_objects/dataset/dataset.py` around lines 745-746, update
`__internal_api__insert_items_as_dataclasses__` so it applies the same
`_parallel_insert_supported` compatibility gate as `insert()`. Before forwarding
`num_threads` to `_send_batches`, reduce values greater than 1 to 1 when the backend
does not support parallel insertion, ensuring direct internal callers cannot bypass the
sequential-upload fallback.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit 00c79c8 addressed this comment by applying the _parallel_insert_supported gate inside the internal insertion funnel before forwarding num_threads to batch sending. Unsupported backends now force worker count to 1, including for direct internal callers.

if isinstance(num_threads, bool) or not isinstance(num_threads, int):
raise ValueError("num_threads must be a positive integer")
if num_threads < 1:
raise ValueError("num_threads must be a positive integer")
Comment thread
alexkuzmik marked this conversation as resolved.

# Gated here rather than in `insert` so every caller of this funnel is
# covered: older backends race on concurrent batches that share a
# batch_group_id, and a direct caller asking for workers must not be
# able to skip that check.
if num_threads > 1 and not self._parallel_insert_supported:
num_threads = 1
Comment on lines +758 to +759

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Direct callers get incidental type errors

The internal funnel compares num_threads > 1 before validating its type, so direct callers of __internal_api__insert_items_as_dataclasses__ with non-orderable values get TypeError instead of the documented ValueError — should we move the same positive-integer validation to this boundary or normalize before comparing?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/api_objects/dataset/dataset.py` around lines 752-753, update
`__internal_api__insert_items_as_dataclasses__` so `num_threads` is validated before the
`> 1` comparison; currently `None` or non-orderable values raise an incidental
`TypeError` instead of the documented `ValueError`. Apply the same positive-integer and
boolean-exclusion checks used by `insert`, preferably through a shared private
validator, while preserving the sequential fallback for unsupported parallel inserts.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit 9c86938 addressed this comment by validating num_threads before the > 1 comparison, rejecting booleans, non-integers, and values below 1 with ValueError.


if deduplication:
items_to_send = self._deduplicate(items)
else:
Expand Down Expand Up @@ -756,22 +796,27 @@ def insert(
deduplication: Whether to skip items whose content already exists
in the dataset. Pass ``False`` to insert every item as-is
without any duplicate checking, which is significantly faster
on large datasets.
on large datasets. The next insert that does deduplicate has to
re-read the dataset's items to account for what was skipped.
num_threads: Number of worker threads used to upload the item
batches. Must be a positive integer, defaults to ``4``; pass
``1`` to upload sequentially. All batches land in a single
dataset version. If a batch fails the call raises, and the
batches that already succeeded stay persisted. Older Opik
backends do not support parallel upload and fall back to a
sequential one.

Raises:
ValueError: If ``num_threads`` is not a positive integer, or
``deduplication`` is not a bool.
"""
if isinstance(num_threads, bool) or not isinstance(num_threads, int):
raise ValueError("num_threads must be a positive integer")
if num_threads < 1:
raise ValueError("num_threads must be a positive integer")

if num_threads > 1 and not self._parallel_insert_supported:
num_threads = 1
# Checked here too so bad input raises before any item is converted.
if not isinstance(deduplication, bool):
raise ValueError("deduplication must be a bool")

dataset_items: List[dataset_item.DatasetItem] = [ # type: ignore
(dataset_item.DatasetItem(**item) if isinstance(item, dict) else item)
Expand Down
4 changes: 4 additions & 0 deletions sdks/python/src/opik/api_objects/dataset/rest_operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,10 @@ def get_test_suites(
dataset_items_count=dataset_fern.dataset_items_count,
client=client,
)
# This suite already holds items on the backend that we have not
# hashed locally, so the first insert must sync before it can tell
# a duplicate from a new item.
suite_dataset.__internal_api__hashes_synced__ = False
Comment on lines +287 to +290

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Listed suites lose owning project scope

In get_test_suites, the lazy-sync suite_dataset keeps the caller's project_name even when dataset_fern resolves via project_id, so workspace-wide or fallback listings can stream or upsert same-named suites into the wrong dataset — should we resolve and retain the returned suite's owning project name before enabling __internal_api__hashes_synced__, as Dataset.from_public already does?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/api_objects/dataset/rest_operations.py` around lines 287-290 in
`get_test_suites`, fix the lazy-sync setup so each listed suite retains its actual
owning project rather than the caller’s `project_name`. Resolve the project name from
`dataset_fern.project_id` (as `Dataset.from_public` does) and pass that resolved name
when constructing `suite_dataset` before enabling `__internal_api__hashes_synced__`,
ensuring streaming and upserts target the correct dataset.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not taking this one here, because it is pre-existing rather than introduced by this change.

get_datasets uses the identical pattern — it passes the caller's project_name straight into the Dataset while deriving project_id from that same name, and it sets __internal_api__hashes_synced__ = False too. So the scope behaviour of a listed suite is now exactly the scope behaviour of a listed dataset; before this PR the suite was inconsistent with it in the other direction (it never synced at all, which is the resubmission bug this PR fixes).

Resolving the owning project per listed item, the way from_public does, means a get_project_by_id call for every dataset in the listing. That is precisely the N+1 the lazy-sync design was built to avoid, so it is a deliberate trade-off to change, not a local fix — and it would need to change get_datasets at the same time to keep the two consistent. Worth doing as its own change if we want listings to be project-accurate.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the context—this is an intentional lazy-sync trade-off shared with get_datasets, rather than a regression in this PR. I’ll save this to memory once the PR is merged.

Comment on lines +287 to +290

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing listed-suite deduplication coverage

The regression test only checks the private __internal_api__hashes_synced__ flag, so it doesn't verify that a duplicate item is actually omitted from the REST batch after retrieving a populated suite via get_test_suites() — should we extend it, or add an integration/E2E test, to assert the batch excludes the duplicate or the backend item count stays unchanged?

Severity

Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/api_objects/dataset/rest_operations.py` around lines 287-290,
extend the tests for `get_test_suites()` to cover the cross-file synchronization
behavior introduced by `__internal_api__hashes_synced__ = False`. Populate a suite,
retrieve it through `get_test_suites()`, insert an identical item with deduplication
enabled, and assert that the REST batch omits the duplicate or that the backend item
count remains unchanged; do not limit the test to checking the private flag.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit 00c79c8 addressed this comment by adding an integration-style regression test that retrieves a populated suite, inserts a duplicate and a new item, and asserts only the new item is submitted.


suites.append(
test_suite_module.TestSuite(
Expand Down
142 changes: 142 additions & 0 deletions sdks/python/tests/unit/api_objects/dataset/test_dataset_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import pytest

from opik.api_objects import constants
from opik.api_objects.dataset import dataset_item
from opik.api_objects.dataset.dataset import Dataset


Expand Down Expand Up @@ -551,6 +552,147 @@ def test_insert__version_endpoint_unreachable__uploads_sequentially(monkeypatch)
), "A failing version probe must not break insert; it falls back to sequential"


def test_insert__version_probe_recovers__parallel_upload_resumes(monkeypatch):
"""A transient probe failure must not pin the dataset to sequential uploads."""
_small_batches(monkeypatch, size=_GATE_BATCH_SIZE)
mock_rest_client = Mock()
mock_rest_client.version.side_effect = [
ConnectionError("backend unreachable"),
{"version": constants.MIN_BACKEND_VERSION_FOR_PARALLEL_INSERT},
]
dataset = Dataset(
name="test_dataset",
description="Test description",
project_name="Test project",
rest_client=mock_rest_client,
)

# Spying on the upload layer keeps the assertion on the worker count that
# actually reached it, so a gate that stops honouring the probe still fails
# this test.
workers_per_insert = []
original_send = Dataset._send_batches

def spy_send_batches(self, batches, batch_group_id, num_threads):
workers_per_insert.append(num_threads)
return original_send(self, batches, batch_group_id, num_threads)

monkeypatch.setattr(Dataset, "_send_batches", spy_send_batches)

dataset.insert(_make_items(4), deduplication=False)
dataset.insert(_make_items(4), deduplication=False)

assert mock_rest_client.version.call_count == 2, (
"The failed probe must be retried rather than cached as unsupported"
)
assert workers_per_insert[0] == 1, (
"While the backend is unreachable the upload must stay sequential"
)
assert workers_per_insert[1] > 1, (
"Once the backend answers, the upload must go back to using workers"
)
Comment on lines +582 to +593

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parallel upload regression goes undetected

test_insert__version_probe_recovers__parallel_upload_resumes only checks the probe count and _parallel_insert_supported, so a sequential second insert would still pass — should we add a barrier or active-worker counter to the batch upload mock to assert overlap while retaining the probe-retry assertions?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/tests/unit/api_objects/dataset/test_dataset_client.py` around lines
569-577, strengthen `test_insert__version_probe_recovers__parallel_upload_resumes` so it
verifies that the second insert actually uses parallel/concurrent batch uploading, not
merely that the `_parallel_insert_supported` cache flag is set and the probe call count
is correct. Instrument the batch upload mock with a barrier, event, or deterministic
active-concurrent-call counter, and assert that overlap occurs among batch uploads
during the second insert, while retaining the existing version-probe retry assertions.
Alternatively, if adding this determinism proves too complex, consider renaming the test
to reflect that it only covers probe-cache recovery.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 00c79c8. test_insert__version_probe_recovers__parallel_upload_resumes now installs a spy on _send_batches that calls through to the real implementation and records the worker count each insert actually used, then asserts [0] == 1 (sequential while the backend is unreachable) and [1] > 1 (workers again once it answers). A second insert that stayed sequential now fails the test, which was the gap you flagged.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed. The spy now verifies the effective worker count reaching _send_batches: the first insert uses 1, while the recovered second insert uses multiple workers, and _send_batches fans those workers out through ThreadPoolExecutor.



def test_insert__unparseable_version__probed_once(monkeypatch):
"""An unparseable version is a conclusive answer, so it must still cache."""
_small_batches(monkeypatch, size=_GATE_BATCH_SIZE)
mock_rest_client = _mock_rest_client("dev-local")
dataset = Dataset(
name="test_dataset",
description="Test description",
project_name="Test project",
rest_client=mock_rest_client,
)

for _ in range(3):
dataset.insert(_make_items(4), deduplication=False)

assert mock_rest_client.version.call_count == 1, (
"A version the SDK cannot parse will not change, so it must not be re-probed"
)


def test_internal_insert__old_backend__worker_count_still_gated(monkeypatch):
"""The gate lives in the funnel, so a direct caller cannot skip it."""
_small_batches(monkeypatch, size=_GATE_BATCH_SIZE)
mock_rest_client = _mock_rest_client("2.2.7")
dataset = Dataset(
name="test_dataset",
description="Test description",
project_name="Test project",
rest_client=mock_rest_client,
)

# The spy calls through, so the upload still happens and the assertions
# below cover the items actually reaching the backend, not just the
# argument the gate computed.
used_workers = []
original_send = Dataset._send_batches

def spy_send_batches(self, batches, batch_group_id, num_threads):
used_workers.append(num_threads)
return original_send(self, batches, batch_group_id, num_threads)

monkeypatch.setattr(Dataset, "_send_batches", spy_send_batches)

dataset.__internal_api__insert_items_as_dataclasses__(
[dataset_item.DatasetItem(**item) for item in _make_items(4)],
num_threads=4,
)

assert used_workers == [1], (
"A backend that predates parallel insert must force a sequential upload "
"even when the internal API is called directly"
)

create_or_update = mock_rest_client.datasets.create_or_update_dataset_items
submitted = sorted(
item.data["input"]["i"]
for call in create_or_update.call_args_list
for item in call.kwargs["items"]
)
assert submitted == [0, 1, 2, 3], (
"Forcing a sequential upload must still deliver every item exactly once"
)


@pytest.mark.parametrize("bad_value", [0, -1, 1.5, "2", True, None])
def test_internal_insert__invalid_num_threads__raises_value_error(bad_value):
"""Direct callers get the named ValueError, not a TypeError from the gate."""
mock_rest_client = Mock()
dataset = Dataset(
name="test_dataset",
description="Test description",
project_name="Test project",
rest_client=mock_rest_client,
)

with pytest.raises(ValueError, match="num_threads must be a positive integer"):
dataset.__internal_api__insert_items_as_dataclasses__(
[dataset_item.DatasetItem(**item) for item in _make_items(2)],
num_threads=bad_value,
)

mock_rest_client.datasets.create_or_update_dataset_items.assert_not_called()


@pytest.mark.parametrize("bad_value", ["false", 0, 1, None, "", []])
def test_insert__non_bool_deduplication__raises_before_any_request(bad_value):
mock_rest_client = Mock()
dataset = Dataset(
name="test_dataset",
description="Test description",
project_name="Test project",
rest_client=mock_rest_client,
)

with pytest.raises(ValueError, match="deduplication must be a bool"):
dataset.insert(_make_items(3), deduplication=bad_value)

mock_rest_client.datasets.create_or_update_dataset_items.assert_not_called()
mock_rest_client.version.assert_not_called()


def test_insert__sequential__uploads_sequentially_without_probing_version(monkeypatch):
mock_rest_client = _mock_rest_client()

Expand Down
Loading
Loading