[NA] [SDK] perf: stream Dataset.insert so a large insert holds batches, not items - #8122
[NA] [SDK] perf: stream Dataset.insert so a large insert holds batches, not items#8122alexkuzmik wants to merge 4 commits into
Dataset.insert so a large insert holds batches, not items#8122Conversation
…hes, not items `insert` read everything before sending anything. A call built the dataclasses, the survivors of deduplication, the REST models and the list of batches - four structures as long as the input - and then handed every batch to the thread pool at once, so the whole dataset stayed resident until the last upload finished, however small the individual batches were. It now consumes its argument lazily and uploads as it reads. `insert` takes any iterable, including a generator over a file or a database cursor, so a dataset larger than memory no longer has to be materialised as a list first. Two changes make that hold end to end: - `sequence_splitter.stream_into_batches` yields batches as items arrive. `split_into_batches` is now `list()` around it, so the eager callers keep their list and the size accounting has one implementation. `max_length=None` means "no count limit" rather than `len(items)`, which no iterable can answer and which meant the same thing. - `_send_batches` waits before pulling instead of submitting everything, so at most `num_threads` batches exist at any moment and the producer only runs as far ahead as the uploads drain. Deduplication still works, lazily, but is not free: the content hashes it keeps grow with the number of items inserted and outlive the call, so `deduplication=False` is what actually bounds memory. Said plainly in the docstring and the docs rather than implied. One behaviour changes deliberately. Reading and converting used to finish before the first upload, so a source that failed halfway meant nothing was sent; the failure now reaches the caller with the batches before it already persisted, and the batch still accumulating is lost with it. That matches the contract `num_threads` already documents for a batch that fails to upload, and it is pinned by a test rather than left to be discovered. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
documentation
⏱️ pre-commit per-hook timing
⏭️ 39 skipped (no matching files changed)
|
|
🌿 Preview your docs: https://opik-preview-01a06799-4a73-729f-8674-277ae8caeef5.docs.buildwithfern.com/docs/opik No broken links found Unverified links (timeout / rate-limited / server error — not failing the check)• https://aistudio.google.com/apikey (401) 📌 Results for commit f048ef4 |
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| def read_items(path): | ||
| with open(path) as file: | ||
| for line in file: | ||
| yield {"input": json.loads(line)} |
There was a problem hiding this comment.
Example fails without json import
The copyable Python example calls json.loads(line) without importing json, so it raises NameError on the first input line — should we add import json to the snippet?
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
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
apps/opik-documentation/documentation/fern/docs-v2/evaluation/advanced/manage_datasets.mdx
around lines 306-312, update the `read_items` Python example to import the `json` module
before calling `json.loads(line)`. Keep the snippet self-contained and directly runnable
as shown.
There was a problem hiding this comment.
Commit c18d743 addressed this comment by adding import json to the Python example before json.loads(line) is called.
There was a problem hiding this comment.
Fixed in c18d743 — added import json to the snippet.
🤖 Reply posted via /address-github-pr-comments
| with futures.ThreadPoolExecutor(max_workers=num_threads) as pool: | ||
| submitted = [ | ||
| pool.submit( | ||
| self._insert_batch_with_retry, | ||
| batch, | ||
| batch_group_id=batch_group_id, | ||
| in_flight: Set[futures.Future] = set() |
There was a problem hiding this comment.
Concurrent inserts can exhaust threads
Each parallel Dataset.insert creates its own ThreadPoolExecutor, so overlapping calls can allocate up to num_threads workers each and grow the process thread count without a shared bound — should we use a dataset/client-owned bounded executor with explicit shutdown and tracked futures?
Want Baz to fix this for you? Activate Fixer
Other fix methods
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 705-706, refactor
`_send_batches` so parallel inserts do not create a new `ThreadPoolExecutor` for every
call. Use a shared, bounded executor owned by the `Dataset` or its client, retain and
manage submitted futures so exceptions still propagate correctly, and add an explicit
lifecycle/shutdown path that safely waits for outstanding uploads.
There was a problem hiding this comment.
Skipping — pre-existing: the per-insert ThreadPoolExecutor came with num_threads and this PR does not change how it is created, only how much work is queued into it (strictly less — at most num_threads batches now, rather than every batch up front). Concurrent inserts each allocating a pool is worth addressing, but it is the same before and after this change, so it belongs in its own ticket rather than being conflated with the memory bound.
🤖 Reply posted via /address-github-pr-comments
There was a problem hiding this comment.
Thanks for the clarification—I’ll keep the distinction between pre-existing executor lifecycle concerns and PR-introduced queueing changes in mind, and save it to memory once the PR is merged.
| self._hashes.add(item_hash) | ||
| self._id_to_hash[item.id] = item_hash | ||
|
|
||
| return deduplicated_items | ||
| yield item |
There was a problem hiding this comment.
_deduplicate writes item hashes to self._hashes/self._id_to_hash before _insert_batch_with_retry succeeds, so failed batches remain cached and a later insert with deduplication=True skips them — should we stage these entries and commit them only after persistence succeeds?
Want Baz to fix this for you? Activate Fixer
Other fix methods
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 762-764 (and
`__internal_api__insert_items_as_dataclasses__` around lines 807-810), fix
`_deduplicate` so hashes are not treated as persisted before the corresponding item has
actually been uploaded. Refactor the insert pipeline to stage deduplication hashes per
batch and commit them to `self._hashes`/`self._id_to_hash` (and set `_hashes_synced`
accordingly) only after `_insert_batch_with_retry` succeeds; any exception while
streaming, batching, converting, or sending should invalidate/roll back the staged
entries for that batch rather than leaving them cached, so a later deduplicated insert
can retry stranded items. Preserve duplicate filtering within the same in-progress
insert without leaving unsaved items in the long-lived cache, and add a regression test
covering a source/upload failure followed by a retry.
There was a problem hiding this comment.
Skipping — pre-existing and unchanged by this PR: the hashes were recorded before any upload here too, since deduplication ran to completion before the first batch was sent. Streaming interleaves the two but does not alter which side of the upload the hash lands on.
It is a real product issue, though — a failed batch leaves its items looking inserted to the next deduplicated insert. Out of scope for a memory change; worth its own ticket, and _deduplicate's docstring now says the cache is not transactional.
🤖 Reply posted via /address-github-pr-comments
There was a problem hiding this comment.
Thanks for clarifying—this is pre-existing behavior and out of scope for this PR. I’ll save that context to memory once the PR is merged.
| # Generator, not a list: conversion happens one item ahead of the batch | ||
| # being filled, so the REST models for the whole insert are never all | ||
| # alive at once. | ||
| rest_items = (self._convert_to_rest_item(item) for item in items_to_send) |
There was a problem hiding this comment.
Partial failures leave count cache stale
_send_batches() can upload earlier batches before the lazy producer raises, so count invalidation is skipped and dataset_items_count/TestSuite.items_count returns the stale pre-insert value — should we clear the cached count in a finally or equivalent failure path whenever a batch may have been uploaded?
Want Baz to fix this for you? Activate Fixer
Other fix methods
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 807-810, update
__internal_api__insert_items_as_dataclasses__ so the lazy streaming upload invalidates
_dataset_items_count even when producing a later batch raises after earlier batches were
persisted. Wrap _send_batches and the count invalidation in a finally-based cleanup (or
equivalent failure path), preserving the original exception while ensuring subsequent
dataset_items_count accesses refresh from the backend.
There was a problem hiding this comment.
Commit c18d743 addressed this comment by invalidating _dataset_items_count in a finally block around _send_batches, including partial-upload failures.
There was a problem hiding this comment.
Fixed in c18d743 — the invalidation now sits in a finally, so a source that fails part-way or a batch that fails after earlier ones landed no longer leaves items_count reporting the number from before an insert that did add items. Test: test_source_raises_mid_stream__cached_item_count_is_invalidated.
🤖 Reply posted via /address-github-pr-comments
- `next(batch_iterator, None)` did not typecheck: the sentinel is not a batch. `StopIteration` says the same thing without one, and CI caught it because my own verification never ran the hooks (the file list reached pre-commit as a single argument, so every hook reported nothing to check). - Invalidate the cached item count in a `finally`. A partial insert still changed the dataset, so a failure part-way used to skip the invalidation and leave `items_count` reporting the number from before it. - Import `json` in the docs example, which called `json.loads` without it. - Cover the generator path against a real backend, sequentially and fanned out, since the existing e2e insert tests all pass a materialised list. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Referenced in a PR reply, so it needs to be true in the code: an item's hash is recorded before its batch is known to have landed, which predates streaming but was written down nowhere. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Every other test here shrinks `DATASET_ITEMS_MAX_BATCH_SIZE` to 2 so batching is observable, which means none of them would notice the streaming splitter mishandling either cap at its actual value. Two cases, both fed by a one-shot generator: - 2,500 items on the default worker count, filling two whole batches and leaving a partial one. Asserts the sizes are exactly `[1000, 1000, 500]`, that every item arrives once with a unique id - nothing lost at a boundary, nothing sent twice by the pool - and that one batch_group_id covers the whole stream. - 120 items of 100KB, so a batch fills on bytes long before it fills on count. Asserts no batch exceeds the 5MB the uploads are built around and that the count cap is not what split them; without the size accounting all 120 would arrive as a single batch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Python SDK Unit Tests Results (Python 3.14)5 163 tests +19 5 161 ✅ +19 1m 57s ⏱️ -25s Results for commit b6ee2a1. ± Comparison against base commit c9eb439. This pull request removes 1 and adds 20 tests. Note that renamed tests count towards both.♻️ This comment has been updated with latest results. |
| def counting_source(): | ||
| nonlocal pulled | ||
| for item in _make_items(item_count): | ||
| pulled += 1 | ||
| yield item | ||
|
|
||
| def tracked_upload(*args, **kwargs): | ||
| nonlocal started, peak_gap | ||
| with lock: | ||
| started += 1 | ||
| peak_gap = max(peak_gap, pulled - started * batch_size) |
There was a problem hiding this comment.
Unsynchronized lookahead bound measurement
counting_source increments pulled outside lock while tracked_upload reads it under the lock, so peak_gap measures an unsynchronized producer/worker interleaving rather than a deterministic bound — should we protect both accesses with lock or use a queue/event-based measurement?
Want Baz to fix this for you? Activate Fixer
Other fix methods
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
859-869, fix the `counting_source` and `tracked_upload` concurrency measurement so
`pulled` is not incremented outside the lock while being read under it. Protect both the
increment and read with the same lock, or replace the counter-based approach with a
queue/event-based measurement that deterministically tracks producer progress relative
to uploads. Keep the `peak_gap` assertion focused on the intended bounded-prefetch
behavior.
| assert list(sequence_splitter.stream_into_batches(items, **limits)) == ( | ||
| sequence_splitter.split_into_batches(items, **limits) | ||
| ) |
There was a problem hiding this comment.
💡 suggestion | Testing
This compatibility test cannot fail. split_into_batches is now defined as list(stream_into_batches(...)), so both sides of the assertion call the same generator and the test reduces to list(x) == list(x) — it holds for any implementation, including a broken one.
That matters more than the usual tautology nit, because this is the test positioned to catch a boundary change. I differential-tested the refactor against main over 300 randomized inputs and the batch boundaries are byte-identical, so the code is fine — but that confidence came from outside the suite, not from here.
Pinning the expected batches literally makes it a real regression test and keeps the documented boundary rules (oversized item emitted alone and ahead of the accumulating batch; inclusive size bound) readable in one place:
| assert list(sequence_splitter.stream_into_batches(items, **limits)) == ( | |
| sequence_splitter.split_into_batches(items, **limits) | |
| ) | |
| @pytest.mark.parametrize( | |
| "items, limits, expected", | |
| [ | |
| ( | |
| [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], | |
| {"max_length": 4}, | |
| [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]], | |
| ), | |
| ([1, 2, 3], {"max_length": 10}, [[1, 2, 3]]), | |
| ([], {"max_length": 4}, []), | |
| ], | |
| ) | |
| def test_stream__yields_the_documented_batches(self, items, limits, expected): | |
| assert list(sequence_splitter.stream_into_batches(items, **limits)) == expected |
The two megabyte-object cases are worth keeping as their own case with expected batch lengths asserted (comparing the objects themselves is noisy), so the size-limit and oversized-item paths stay covered.
🤖 Review posted via /review-github-pr
|
|
||
| try: | ||
| batch = next(batch_iterator) | ||
| except StopIteration: | ||
| break | ||
|
|
There was a problem hiding this comment.
❓ question | Performance
Worth a follow-up ticket rather than anything for this PR — the memory fix here is real and I verified the in-flight bound holds.
Waiting before pulling is what caps in-flight batches at num_threads, but it also means next(batch_iterator) runs on the submitting thread: the source read, DatasetItem(**item), content_hash(), _convert_to_rest_item, and the splitter's JSON size walk all execute between submissions. So the loop alternates wait → produce → submit, and workers that free up during the produce step stay idle until it finishes.
Simulating this loop's semantics (100 batches, 1.0s upload):
| threads | producer cost/batch | worker utilization |
|---|---|---|
| 8 | 0s | 96% |
| 8 | 0.1s | 28% |
| 4 | 0.1s | 63% |
The diagnostic detail is that more threads make it worse — a serialized producer can't feed them — and makespan has a floor at n × produce_cost regardless of pool size, so raising num_threads can't recover it. Under the GIL the producer's CPU work also competes with the upload threads' serialization.
The standard fix decouples the two concerns that are currently fused, and gets bounded memory and work-conserving uploads instead of trading one for the other:
q = queue.Queue(maxsize=num_threads * 2) # capacity = backpressure
# producer thread: for batch in batches: q.put(batch)
# N workers: while (b := q.get()) is not SENTINEL: upload(b)put() blocks when full, so peak memory stays bounded by maxsize while a freed worker picks up work immediately. In the same simulation, raising capacity from 8 to 16 cut makespan from 14.6 to 8.2 — capacity and parallelism want to be separate knobs.
🤖 Review posted via /review-github-pr
There was a problem hiding this comment.
Correction to the comment above — I compared against the wrong baseline, and this PR is a throughput improvement, not a regression.
The ~28% figure measures this submit loop in isolation, against an implicit 100% baseline. That baseline is wrong. On main, all the producer work — building the four full-length lists — ran to completion before _send_batches was called, so during that entire phase no uploads were in flight and every worker sat idle. Interleaving production with uploads beats front-loading it.
Same parameters as above (100 batches, 1.0s upload, 0.1s produce per batch, 8 threads), measured across the whole insert() call rather than just the loop:
| producer phase | upload phase | total | utilization | |
|---|---|---|---|---|
main |
10.0s (workers idle) | 12.5s | 22.5s | 56% |
| this PR | interleaved | 14.6s | 86% |
Across the range:
| threads | produce/batch | main |
this PR | delta |
|---|---|---|---|---|
| 4 | 0.05s | 30.0s | 26.4s | −12% |
| 4 | 0.1s | 35.0s | 27.8s | −21% |
| 4 | 0.3s | 55.0s | 33.4s | −39% |
| 8 | 0.1s | 22.5s | 14.6s | −35% |
| 8 | 0.3s | 42.5s | 17.8s | −58% |
| 8 | 0s (free producer) | 12.5s | 13.0s | +4% |
| 1 | any | — | — | 0% |
The only case this PR loses is a zero-cost producer, which doesn't occur — conversion and the splitter's size walk both cost something.
So the corrected picture for this change: CPU complexity unchanged (O(n), same constant — verified by interleaved in-process benchmark, flat ns/item across a 16× size range), peak memory improved O(n) → O(num_threads), and wall-clock improved 12–58% in realistic cases.
The bounded-queue suggestion still stands on its own merits — 86% isn't 100%, and decoupling queue depth from worker count closes the remainder — but it's an incremental improvement on something this PR already made substantially better, not a fix for anything introduced here. Apologies for the misframing.
🤖 Review posted via /review-github-pr
|
👋 Review summary What looks good
Overall One structural note, filed inline as a follow-up rather than a change request: this fixes the memory half of the problem but leaves the throughput half, and both come from the same design decision — producer and consumers sharing a thread. Because Separately, and pre-existing rather than introduced here: count-only callers ( Inline comments: 1 suggestion, 1 question — nothing blocking. This review does not constitute an approval; a human reviewer should still approve. 🤖 Review posted via /review-github-pr |
Follow-up: 14 of the 16
|
| Site | Why |
|---|---|
api_objects/experiment/experiment.py:245 |
len(batches) twice — a log line, and worker_count = min(num_threads, len(batches), MAX_THREADS). Sizing the pool needs the count up front. |
cli/imports/experiment.py:665 |
enumerate(item_batches, start=1) plus len(item_batches) inside the loop for a batch N/M progress message. |
Single forward pass — the list is pure overhead (14):
message_processing/batching/batchers.py:28,:65—for batch in ...: batches.append(Message(batch=batch))api_objects/opik_client.py:478,:849,:900,:961— inlinefor batch in split_into_batches(...):api_objects/annotation_queue/annotation_queue.py:230,:259,:348,:376—for batch in batches:api_objects/threads/threads_client.py:154— inline for-loopapi_objects/experiment/experiment.py:146—for batch in batches:cli/migrate/datasets/version_replay.py:392—for batch in item_batches:api_objects/dataset/dataset.py:908(delete) —for batch in batches:
None of these index, re-iterate, or take a length. Each batch is consumed and discarded before the next is pulled — exactly the shape stream_into_batches was written for.
Several run on inputs that get large: delete_traces, the four annotation-queue ID paths, and dataset.delete. Pointing those at stream_into_batches gives them the same O(n) → O(num_threads) peak-memory win insert just got, for free, with no behaviour change since a single forward pass is identical either way.
Two options for the ticket, and the first seems lower-risk:
- Repoint the 14 at
stream_into_batches. Mechanical, no behaviour change, leavessplit_into_batchesfor the 2 that need a list. - Make
split_into_batcheslazy and fix up the 2. Cleaner end state, but it changes the return type of a widely-used helper, andexperiment.py:245wants the count for pool sizing anyway — so you'd reintroduce alist()there regardless.
Also worth folding into the same ticket: count-only callers still pay a full jsonable_encoder.encode() plus recursive size walk per item for a limit that's never consulted (~7× slower than necessary on flat ID lists). Pre-existing — on main the 0.0 if max_payload_size_MB is None guard was already dead code, since the parameter gets reassigned to float("inf") before the loop — but now that the dead branch is gone, restoring the short-circuit behind a real flag is straightforward. delete_traces and the annotation-queue paths are the ones paying for it.
🤖 Review posted via /review-github-pr
2 participants
Details
Dataset.insertread everything before sending anything: it built the dataclasses, the survivors of deduplication, the REST models and the list of batches — four structures as long as the input — then handed every batch to the thread pool at once, so the whole dataset stayed resident until the last upload finished, however small the batches were.It now consumes its argument lazily and uploads as it reads, so
insertaccepts any iterable — a generator over a file or a database cursor included — and a dataset larger than memory no longer has to be materialised as a list first. Two changes make that hold end to end:sequence_splitter.stream_into_batchesyields batches as items arrive (withsplit_into_batchesnowlist()around it, so the eager callers keep their list and the size accounting has one implementation), and_send_batcheswaits before pulling instead of submitting everything, so at mostnum_threadsbatches exist at once and the producer only runs as far ahead as the uploads drain.Deduplication still works, lazily, but is not free: its content hashes grow with the number of items inserted and outlive the call, so
deduplication=Falseis what actually bounds memory. That is now stated in the docstring and the docs instead of implied.One behaviour changes deliberately: a source that fails halfway used to mean nothing was sent, and now the failure reaches the caller with the earlier batches already persisted, the batch still accumulating lost with it. That matches the contract
num_threadsalready documents for a failed batch, and a test pins it.Change checklist
Issues
AI-WATERMARK
AI-WATERMARK: yes
Testing
pytest tests/unit— 5158 passed, 3 skippedmake precommit— ruff, ruff-format, mypy on the changed filesNew tests cover the splitter and the insert path. For the splitter: streaming yields exactly what the eager version returns across the existing scenarios (count limit, size limit, both, empty, oversized item), a generator input needs no
len(), only the first batch plus the item that closed it is read to produce one batch, and an oversized item still comes out ahead of the batch accumulating around it.For
insert: a generator uploads the same items a list does and is consumed exactly once; dataclass items work from a generator; deduplication still drops duplicates lazily; the source is read no further ahead than the uploads (measured as the gap between items read and items whose upload has started, parametrized over 1 and 4 threads — barrier-free, so it cannot flake); no more batches are in flight thannum_threads; and a source that raises midway leaves the earlier batches persisted.Documentation
manage_datasets.mdx— a section on inserting more items than fit in memory: the generator form, that the iterable is consumed once, what happens if reading fails partway, and thatdeduplication=Falseis what bounds memory.