chore: import upstream snapshot with attribution

tests_load/suite/python_sdk/_helpers.py

@@ -0,0 +1,266 @@
+"""Helpers shared by the load-test suite.
+
+Functions here are intentionally simple so individual tests can be read
+top-to-bottom without jumping around.
+"""
+
+import base64
+import contextlib
+import json
+import logging
+import os
+import random
+import string
+import time
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, Iterable, Iterator, List, Set, Tuple
+
+import opik
+from opik import Opik
+from opik.rest_api.types.span_public import SpanPublic
+from opik.rest_api.types.trace_public import TracePublic
+
+
+PROJECT_PREFIX: str = "loadtest"
+REPORT_DIR: Path = Path(__file__).resolve().parents[2] / ".last_run"
+
+KB: int = 1_000
+MB: int = 1_000_000
+
+REQUIRED_SPAN_FIELDS: Tuple[str, ...] = ("name", "end_time", "input", "output")
+
+
+class Metrics:
+    """Tiny key/value recorder with a `timer(label)` context manager.
+
+    One instance per test; written to ``REPORT_DIR/<test_name>.json`` at
+    teardown by the ``metrics`` fixture.
+    """
+
+    test_name: str
+    _data: Dict[str, Any]
+
+    def __init__(self, test_name: str) -> None:
+        self.test_name = test_name
+        self._data = {"test_name": test_name}
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self._data[key] = value
+
+    @contextlib.contextmanager
+    def timer(self, label: str) -> Iterator[None]:
+        start: float = time.perf_counter()
+        yield
+        self._data[f"{label}_seconds"] = round(time.perf_counter() - start, 3)
+
+    def write(self) -> None:
+        REPORT_DIR.mkdir(parents=True, exist_ok=True)
+        out_path: Path = REPORT_DIR / f"{self.test_name}.json"
+        out_path.write_text(json.dumps(self._data, indent=2, default=str))
+        logging.getLogger(self.test_name).info(
+            "metrics: %s", json.dumps(self._data, default=str)
+        )
+
+
+def unique_project_name(scenario: str) -> str:
+    timestamp: str = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
+    return f"{PROJECT_PREFIX}-{scenario}-{timestamp}-{uuid.uuid4().hex[:6]}"
+
+
+def random_text(size_bytes: int) -> str:
+    return "".join(random.choices(string.ascii_letters + string.digits, k=size_bytes))
+
+
+def random_bytes(size_bytes: int) -> bytes:
+    return os.urandom(size_bytes)
+
+
+PNG_MAGIC: bytes = b"\x89PNG\r\n\x1a\n"
+
+
+def random_base64_png(size_bytes: int) -> str:
+    """Returns a standard-alphabet base64-encoded byte stream that the SDK's
+    attachment-extraction pipeline recognises as a PNG.
+
+    Two constraints from the SDK that the naïve approach (random bytes,
+    ``urlsafe_b64encode``) trips over:
+
+    1. The extraction regex is ``[A-Za-z0-9+/]``, so the encoded blob
+       must use the standard ``+/`` alphabet (``b64encode``), not the
+       ``-_`` url-safe one.
+    2. The decoder MIME-sniffs the raw bytes and skips anything that
+       comes back as ``application/octet-stream`` or ``text/plain``,
+       so the byte stream must start with a recognised magic header.
+
+    The bytes after the 8-byte PNG magic are pure noise — the test
+    doesn't render the image, it just exercises the extract/upload path.
+    """
+    raw_size: int = (size_bytes // 4) * 3
+    payload: bytes = PNG_MAGIC + os.urandom(max(0, raw_size - len(PNG_MAGIC)))
+    return base64.b64encode(payload).decode("ascii")
+
+
+def think_time() -> None:
+    """Sleeps a small randomized interval so submits aren't lockstep.
+
+    Real callers always have at least some gap between consecutive SDK calls
+    (LLM responses, request handling, etc.). The range is small on purpose so
+    it doesn't dominate runtime for high-count scenarios.
+    """
+    time.sleep(random.uniform(0.0005, 0.002))
+
+
+def opik_client() -> Opik:
+    return opik.Opik()
+
+
+def now_utc() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _assert_required_fields_set(
+    items: Iterable[Any], required_fields: Tuple[str, ...], kind: str
+) -> None:
+    """Fails the test if any item is missing one of `required_fields`.
+
+    Catches partial-write bugs where the entity lands but a later update
+    (e.g. ``end_time``) was dropped or never sent.
+    """
+    broken: List[str] = []
+    for item in items:
+        for field in required_fields:
+            if getattr(item, field) is None:
+                broken.append(f"{item.id}: {field} is None")
+                break
+    if broken:
+        raise AssertionError(
+            f"{len(broken)} {kind}(s) missing required fields "
+            f"({', '.join(required_fields)}); sample: {broken[:5]}"
+        )
+
+
+def verify_traces(
+    client: Opik,
+    project_name: str,
+    expected_count: int,
+    timeout_seconds: int = 900,
+) -> List[TracePublic]:
+    """Polls until at least `expected_count` traces are visible in the project.
+
+    Excludes ``input``/``output``/``metadata`` from the response. We only
+    need ``id``, ``name``, ``end_time`` for the load-test verification,
+    and excluding the bulky fields:
+
+    1. Streams much less data back for the high-count scenarios.
+    2. Side-steps OPIK-6651, where attachment-extracted traces fail to
+       stream because ``AttachmentService.list`` can't read
+       ``workspaceName`` from the reactor context during enrichment.
+       Excluding the fields skips the enrichment path entirely. Drop
+       the exclude list once OPIK-6651 is fixed.
+
+    ``end_time``/``name`` field checks still run; ``input``/``output``
+    are not included in the response so we can't assert on them here.
+    """
+    traces: List[TracePublic] = client.search_traces(
+        project_name=project_name,
+        max_results=expected_count,
+        wait_for_at_least=expected_count,
+        wait_for_timeout=timeout_seconds,
+        exclude=["input", "output", "metadata"],
+    )
+    _assert_required_fields_set(traces, ("name", "end_time"), "trace")
+    return traces
+
+
+def verify_spans_for_trace(
+    client: Opik,
+    project_name: str,
+    trace_id: str,
+    expected_count: int,
+    timeout_seconds: int = 600,
+) -> List[SpanPublic]:
+    """Polls until at least `expected_count` spans are visible for a trace."""
+    spans: List[SpanPublic] = client.search_spans(
+        project_name=project_name,
+        trace_id=trace_id,
+        max_results=expected_count,
+        wait_for_at_least=expected_count,
+        wait_for_timeout=timeout_seconds,
+    )
+    _assert_required_fields_set(spans, REQUIRED_SPAN_FIELDS, "span")
+    return spans
+
+
+def verify_exact_trace_ids(
+    client: Opik,
+    project_name: str,
+    expected_ids: Set[str],
+    timeout_seconds: int = 900,
+) -> Set[str]:
+    """Waits for all `expected_ids` to land and returns the delivered id set.
+
+    Raises ``AssertionError`` listing a sample of missing ids if some traces
+    are never delivered within the timeout. This catches dropped messages
+    (e.g. regression cases like OPIK-6444) rather than just under-counts.
+
+    Required-field validation runs inside ``verify_traces`` so any trace
+    landing without ``end_time``, ``name``, ``input``, or ``output`` also
+    fails the test.
+    """
+    traces: List[TracePublic] = verify_traces(
+        client,
+        project_name=project_name,
+        expected_count=len(expected_ids),
+        timeout_seconds=timeout_seconds,
+    )
+    delivered: Set[str] = {trace.id for trace in traces}
+    missing: Set[str] = expected_ids - delivered
+    if missing:
+        raise AssertionError(
+            f"{len(missing)}/{len(expected_ids)} traces missing after flush; "
+            f"sample missing ids: {sorted(missing)[:3]}"
+        )
+    return delivered
+
+
+def verify_attachments(
+    client: Opik,
+    project_name: str,
+    entity_type: str,
+    entity_id: str,
+    expected_count: int,
+    timeout_seconds: int = 600,
+) -> int:
+    """Polls the attachments REST endpoint until `expected_count` are listed.
+
+    The endpoint's ``path`` query parameter is a base64-encoded base URL
+    the backend uses to build download links — we pass the SDK's
+    configured base URL so the round-trip succeeds. Uses the standard
+    ``b64encode`` alphabet (not ``urlsafe_b64encode``) to match the
+    contract used by ``attachment/client.py`` and ``tests/e2e/verifiers.py``.
+    """
+    project = client.rest_client.projects.retrieve_project(name=project_name)
+    base_url: str = str(client.config.url_override).rstrip("/")
+    encoded_base_url: str = base64.b64encode(
+        base_url.encode("utf-8")
+    ).decode("ascii")
+    deadline: float = time.time() + timeout_seconds
+    last_seen: int = 0
+    while time.time() < deadline:
+        page = client.rest_client.attachments.attachment_list(
+            project_id=project.id,
+            entity_type=entity_type,
+            entity_id=entity_id,
+            path=encoded_base_url,
+            size=expected_count + 10,
+        )
+        last_seen = len(page.content or [])
+        if last_seen >= expected_count:
+            return last_seen
+        time.sleep(1)
+    raise TimeoutError(
+        f"Only {last_seen}/{expected_count} attachments visible after "
+        f"{timeout_seconds}s for {entity_type} {entity_id}"
+    )

tests_load/suite/python_sdk/conftest.py

@@ -0,0 +1,51 @@
+"""Pytest fixtures for the Python SDK load-test suite."""
+
+import logging
+import os
+from typing import Iterator
+
+import pytest
+from opik import context_storage
+
+from ._helpers import Metrics
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(levelname)s [%(asctime)s] %(name)s: %(message)s",
+)
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--load-scale",
+        type=float,
+        default=float(os.getenv("OPIK_LOAD_SCALE", "1.0")),
+        help="Multiplier applied to default trace/span counts in load tests.",
+    )
+
+
+@pytest.fixture(autouse=True)
+def _reset_opik_context_after_test() -> Iterator[None]:
+    """Clears SDK context between tests so leaks don't cross test boundaries.
+
+    The ``start_as_current_trace`` / ``start_as_current_span`` context
+    managers acquire ``context_storage`` project-name ownership on enter
+    but don't release it on exit; if the next test uses ``@opik.track``
+    with a new project, traces silently land in the leaked project. This
+    fixture neutralises that across our suite.
+    """
+    yield
+    context_storage.clear_all()
+
+
+@pytest.fixture
+def metrics(request: pytest.FixtureRequest) -> Iterator[Metrics]:
+    recorder = Metrics(test_name=request.node.name)
+    yield recorder
+    recorder.write()
+
+
+@pytest.fixture
+def load_scale(request: pytest.FixtureRequest) -> float:
+    return float(request.config.getoption("--load-scale"))

tests_load/suite/python_sdk/requirements.txt

@@ -0,0 +1,7 @@
+# Requirements for the Python SDK load-test suite.
+# Install the Opik SDK separately (e.g. `pip install -e sdks/python` from
+# the repo root, or `pip install opik` for the published package) — this
+# file only pins the suite's runtime dependencies.
+pytest
+pytest-timeout
+pytest-xdist

tests_load/suite/python_sdk/test_attachments.py

@@ -0,0 +1,150 @@
+"""Attachment scenarios — explicit and implicit."""
+
+from typing import List, Set
+
+import opik
+from opik import Attachment
+
+from . import _helpers
+from ._helpers import KB, Metrics
+
+
+def test_traces_with_explicit_attachments(
+    metrics: Metrics, load_scale: float
+) -> None:
+    """Traces with explicit ``Attachment`` uploads, via ``@opik.track``.
+
+    Inside a ``@opik.track``-decorated handler, two 50 KB binary
+    attachments are added with ``opik.update_current_trace(attachments=...)``
+    — the public pattern for attaching arbitrary files to the active
+    trace from inside instrumented user code. Stresses the multipart
+    upload path and the ``flush_tracker()`` contract around in-flight
+    uploads.
+
+    Volume: 500 traces × 2 attachments × 50 KB ≈ 50 MB of attachment
+    payload total, plus 1k multipart uploads to coordinate.
+
+    Verifies every submitted trace id lands with required fields set, and
+    that the attachment-list endpoint reports both attachments on a
+    sampled trace.
+    """
+    trace_count: int = int(500 * load_scale)
+    attachments_per_trace: int = 2
+    attachment_bytes: int = 50 * KB
+    trace_input_bytes: int = 100
+    project_name: str = _helpers.unique_project_name("explicit-attachments")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["attachments_per_trace"] = attachments_per_trace
+    metrics["attachment_bytes"] = attachment_bytes
+    metrics["trace_input_bytes"] = trace_input_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track(project_name=project_name)
+    def handle_request(prompt: str) -> str:
+        opik.update_current_trace(
+            attachments=[
+                Attachment(
+                    data=_helpers.random_bytes(attachment_bytes),
+                    file_name=f"attachment-{j}.bin",
+                    content_type="application/octet-stream",
+                )
+                for j in range(attachments_per_trace)
+            ]
+        )
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return f"echo: {prompt}"
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            handle_request(prompt=_helpers.random_text(trace_input_bytes))
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    last_trace_id: str = submitted_trace_ids[-1]
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+        delivered_attachment_count: int = _helpers.verify_attachments(
+            client,
+            project_name=project_name,
+            entity_type="trace",
+            entity_id=last_trace_id,
+            expected_count=attachments_per_trace,
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+    metrics["delivered_attachments_on_sample_trace"] = delivered_attachment_count
+    assert delivered_attachment_count >= attachments_per_trace
+
+
+def test_traces_with_implicit_attachments(
+    metrics: Metrics, load_scale: float
+) -> None:
+    """Traces whose attachments are extracted from base64 input automatically.
+
+    The handler accepts an ``image`` argument whose value is a
+    ``data:image/png;base64,<~400 KB>`` URL. Because the embedded base64
+    blob exceeds ``min_base64_embedded_attachment_size`` (250 KB by
+    default), the SDK's attachment-extraction pipeline detects it and
+    uploads it as an attachment without any explicit user action. This
+    is the path most user code hits when logging multi-modal LLM I/O.
+
+    Volume: 500 traces × 400 KB of base64 ≈ 200 MB of payload that the
+    SDK has to scan, extract, and upload as 500 attachments.
+
+    Verifies every submitted trace id lands with required fields set, and
+    that at least one extracted attachment is reported on a sampled trace.
+    """
+    trace_count: int = int(500 * load_scale)
+    embedded_base64_bytes: int = 400 * KB
+    trace_prompt_bytes: int = 100
+    project_name: str = _helpers.unique_project_name("implicit-attachments")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["embedded_base64_bytes"] = embedded_base64_bytes
+    metrics["trace_prompt_bytes"] = trace_prompt_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track(project_name=project_name)
+    def handle_image_request(prompt: str, image: str) -> str:
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return f"caption for {prompt}: {image[:32]}..."
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            large_base64: str = _helpers.random_base64_png(embedded_base64_bytes)
+            handle_image_request(
+                prompt=_helpers.random_text(trace_prompt_bytes),
+                image=f"data:image/png;base64,{large_base64}",
+            )
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    last_trace_id: str = submitted_trace_ids[-1]
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+        delivered_attachment_count: int = _helpers.verify_attachments(
+            client,
+            project_name=project_name,
+            entity_type="trace",
+            entity_id=last_trace_id,
+            expected_count=1,
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+    metrics["delivered_attachments_on_sample_trace"] = delivered_attachment_count
+    assert delivered_attachment_count >= 1

tests_load/suite/python_sdk/test_bursts.py

@@ -0,0 +1,178 @@
+"""Burst, spread, and concurrent scenarios."""
+
+import threading
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import List, Set
+
+import opik
+
+from . import _helpers
+from ._helpers import Metrics
+
+
+def test_burst_single_loop(metrics: Metrics, load_scale: float) -> None:
+    """Fast-paced burst from a single thread.
+
+    Calls a ``@opik.track``-decorated handler 50k times in a tight loop
+    with only the minimal randomised think-time (0.5–2 ms) shared by the
+    rest of the suite — enough to keep the runner's docker-compose Opik
+    stack from being DoS-ed when multiple heavy scenarios run in
+    parallel under xdist, but still tight enough that the SDK's
+    in-process queue and batch flusher are kept busy throughout.
+
+    Volume: 50k traces, ~200 B input each.
+
+    Verifies every submitted trace id lands with required fields set.
+    """
+    trace_count: int = int(50_000 * load_scale)
+    trace_input_bytes: int = 200
+    project_name: str = _helpers.unique_project_name("burst")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["trace_input_bytes"] = trace_input_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track(project_name=project_name)
+    def handle_request(prompt: str) -> str:
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return f"echo: {prompt}"
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            handle_request(prompt=_helpers.random_text(trace_input_bytes))
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+
+
+def test_spread_over_time(metrics: Metrics, load_scale: float) -> None:
+    """Steady-rate workload paced over a long window.
+
+    Calls a ``@opik.track``-decorated handler 10k times evenly spaced
+    across a 10-minute window (~17 traces/sec sustained). Mirrors a real
+    moderate-rate production workload and exercises the periodic flush
+    path that fires on its interval rather than on batch-size triggers.
+
+    Volume: 10k traces over 600 s, ~200 B input each.
+
+    Verifies every submitted trace id lands with required fields set.
+    """
+    trace_count: int = int(10_000 * load_scale)
+    window_seconds: int = max(1, int(600 * load_scale))
+    trace_input_bytes: int = 200
+    project_name: str = _helpers.unique_project_name("spread")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["window_seconds"] = window_seconds
+    metrics["trace_input_bytes"] = trace_input_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track(project_name=project_name)
+    def handle_request(prompt: str) -> str:
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return f"echo: {prompt}"
+
+    interval: float = window_seconds / trace_count
+    next_log_time: float = time.perf_counter()
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            handle_request(prompt=_helpers.random_text(trace_input_bytes))
+            next_log_time += interval
+            sleep_for: float = next_log_time - time.perf_counter()
+            if sleep_for > 0:
+                time.sleep(sleep_for)
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+
+
+def test_concurrent_writers_share_one_client(
+    metrics: Metrics, load_scale: float
+) -> None:
+    """30 threads invoking the same ``@opik.track``-decorated handler.
+
+    Every thread calls into the same global Opik client (the one the
+    ``@opik.track`` decorator uses by default). Each invocation gets its
+    own trace via thread-local context — exactly how a real multi-thread
+    server uses the SDK. Realistic think-time prevents lockstep submits.
+    This is the configuration most likely to surface batcher races —
+    same shape as the OPIK-6444 unit regression, just one level up.
+
+    Volume: 30 threads × 1k traces = 30k traces, ~200 B input each.
+
+    Verifies that every submitted trace id lands with required fields
+    set. Any dropped message fails the test with a sample of missing ids.
+    """
+    thread_workers: int = 30
+    traces_per_worker: int = int(1_000 * load_scale)
+    total_traces: int = thread_workers * traces_per_worker
+    trace_input_bytes: int = 200
+    project_name: str = _helpers.unique_project_name("concurrent")
+
+    metrics["project_name"] = project_name
+    metrics["thread_workers"] = thread_workers
+    metrics["traces_per_worker"] = traces_per_worker
+    metrics["total_traces"] = total_traces
+    metrics["trace_input_bytes"] = trace_input_bytes
+
+    submitted_trace_ids: List[str] = []
+    submitted_lock: threading.Lock = threading.Lock()
+
+    @opik.track(project_name=project_name)
+    def handle_request(worker_id: int, prompt: str) -> str:
+        trace_id: str = opik.opik_context.get_current_trace_data().id
+        with submitted_lock:
+            submitted_trace_ids.append(trace_id)
+        return f"worker-{worker_id}: {prompt}"
+
+    def worker(worker_id: int) -> None:
+        for _ in range(traces_per_worker):
+            handle_request(
+                worker_id=worker_id,
+                prompt=_helpers.random_text(trace_input_bytes),
+            )
+            _helpers.think_time()
+
+    with metrics.timer("logging"):
+        with ThreadPoolExecutor(max_workers=thread_workers) as pool:
+            futures: List[Future[None]] = [
+                pool.submit(worker, w) for w in range(thread_workers)
+            ]
+            for future in futures:
+                future.result()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client,
+            project_name=project_name,
+            expected_ids=set(submitted_trace_ids),
+            timeout_seconds=1200,
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)

tests_load/suite/python_sdk/test_dataset_items.py

@@ -0,0 +1,90 @@
+"""Dataset-items upload scenarios.
+
+Each ``Dataset.insert()`` call creates a new dataset version on the
+backend; the BE snapshots the previous version's items into the new
+version via a ClickHouse ``INSERT … SELECT`` (``COPY_VERSION_ITEMS``).
+On multi-replica ClickHouse deployments that SELECT can non-
+deterministically return short, truncating the new version's row set;
+every subsequent version then cascades off the truncated baseline.
+Loss is purely server-side — single-thread sequential REST calls
+already trigger it.
+
+These tests can't *reproduce* the bug on a single-replica localhost
+install (Notion: "Dataset migration replay: silent data loss on the
+version chain"), but they:
+
+1. Provide a green baseline for environments where the bug can fire
+   (production, multi-replica staging) — running the suite there will
+   surface any short-COPY by way of the item-count assertion.
+2. Cover that ``Dataset.insert()`` + ``Dataset.get_items()`` round-trip
+   cleanly across many sequential versions on a single thread.
+3. Stay on the public, high-level API (``Dataset.insert`` /
+   ``Dataset.get_items``) rather than the lower-level REST client the
+   ``opik migrate dataset`` tool uses internally.
+"""
+
+from typing import Any, Dict, List
+
+from opik import Opik
+
+from . import _helpers
+from ._helpers import KB, Metrics
+
+
+def test_dataset_insert_many_versions(metrics: Metrics, load_scale: float) -> None:
+    """Sequential ``Dataset.insert()`` calls, single thread, many versions.
+
+    Mirrors the shape of the production repro from the Notion writeup
+    "Dataset migration replay: silent data loss on the version chain":
+    one dataset, many versions, modest payload per item, no client-side
+    concurrency. The test asserts that the dataset's latest version
+    streams back exactly the expected total — i.e. that no
+    ``COPY_VERSION_ITEMS`` truncation happened anywhere along the chain.
+
+    Volume at ``load_scale=1.0``:
+    - 50 versions × 50 items per version = 2500 items
+    - ~4 KB payload per item
+
+    Verifies via ``dataset.get_items()`` (which streams the latest
+    version's items, equivalent to ``stream_dataset_items`` with the
+    latest version hash) that the delivered count matches the expected
+    total. Catches both the metadata-vs-storage disagreement noted in
+    the repro (where ``items_total`` reports N but the stream returns
+    fewer) and the cascading truncation pattern.
+    """
+    versions: int = int(50 * load_scale)
+    items_per_version: int = 50
+    item_payload_bytes: int = 4 * KB
+    expected_total: int = versions * items_per_version
+    dataset_name: str = _helpers.unique_project_name("dataset-insert")
+
+    metrics["dataset_name"] = dataset_name
+    metrics["versions"] = versions
+    metrics["items_per_version"] = items_per_version
+    metrics["item_payload_bytes"] = item_payload_bytes
+    metrics["expected_total_items"] = expected_total
+
+    client: Opik = _helpers.opik_client()
+    dataset = client.create_dataset(name=dataset_name)
+
+    with metrics.timer("insert"):
+        for _ in range(versions):
+            items: List[Dict[str, Any]] = [
+                {
+                    "input": _helpers.random_text(item_payload_bytes),
+                    "expected_output": _helpers.random_text(100),
+                }
+                for _ in range(items_per_version)
+            ]
+            dataset.insert(items)
+
+    with metrics.timer("verify"):
+        delivered_items: List[Dict[str, Any]] = dataset.get_items()
+
+    metrics["delivered_item_count"] = len(delivered_items)
+    assert len(delivered_items) == expected_total, (
+        f"Dataset items lost: expected {expected_total}, got {len(delivered_items)}. "
+        "Likely a server-side COPY_VERSION_ITEMS truncation on multi-replica "
+        "ClickHouse — see Notion 'Dataset migration replay: silent data loss "
+        "on the version chain' for the failure mode."
+    )

tests_load/suite/python_sdk/test_heavy_payload.py

@@ -0,0 +1,130 @@
+"""Heavy payload scenarios."""
+
+from typing import List, Set
+
+import opik
+from opik.rest_api.types.span_public import SpanPublic
+
+from . import _helpers
+from ._helpers import KB, MB, Metrics
+
+
+def test_traces_with_one_megabyte_payload(
+    metrics: Metrics, load_scale: float
+) -> None:
+    """Low trace count, very large per-trace input/output, via ``@opik.track``.
+
+    Each call to ``handle_request`` produces a trace whose ``input``
+    contains ~1 MB of random text (the function argument) and whose
+    ``output`` is another ~1 MB (the return value). The decorator
+    captures both automatically — the same shape as wrapping an LLM call
+    that takes a long prompt and returns a long completion.
+
+    Volume: 500 traces × (1 MB in + 1 MB out) ≈ 1 GB of payload.
+
+    Verifies every submitted trace id lands with required fields set.
+    """
+    trace_count: int = int(500 * load_scale)
+    trace_input_bytes: int = 1 * MB
+    trace_output_bytes: int = 1 * MB
+    project_name: str = _helpers.unique_project_name("heavy-trace-payload")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["trace_input_bytes"] = trace_input_bytes
+    metrics["trace_output_bytes"] = trace_output_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track(project_name=project_name)
+    def handle_request(prompt: str) -> str:
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return _helpers.random_text(trace_output_bytes)
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            handle_request(prompt=_helpers.random_text(trace_input_bytes))
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+
+
+def test_spans_with_heavy_payload(metrics: Metrics, load_scale: float) -> None:
+    """Moderate fan-out of heavy spans under small trace shells, via context managers.
+
+    Uses ``start_as_current_trace`` / ``start_as_current_span`` so each
+    span explicitly carries ~500 KB in ``input`` and ``output``. The outer
+    trace itself stays small. Stresses span batching for large payloads
+    where the parent trace is light.
+
+    Volume: 200 traces × 5 spans × (500 KB in + 500 KB out) ≈ 1 GB of
+    span payload total. Trace shells are ~1 KB each.
+
+    Verifies every submitted trace id lands with required fields set, and
+    that the last trace's 5 spans are all visible and well-formed.
+    """
+    trace_count: int = int(200 * load_scale)
+    spans_per_trace: int = 5
+    trace_input_bytes: int = 1 * KB
+    trace_output_bytes: int = 1 * KB
+    span_input_bytes: int = 500 * KB
+    span_output_bytes: int = 500 * KB
+    project_name: str = _helpers.unique_project_name("heavy-span-payload")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["spans_per_trace"] = spans_per_trace
+    metrics["trace_input_bytes"] = trace_input_bytes
+    metrics["trace_output_bytes"] = trace_output_bytes
+    metrics["span_input_bytes"] = span_input_bytes
+    metrics["span_output_bytes"] = span_output_bytes
+
+    submitted_trace_ids: List[str] = []
+    last_trace_id: str = ""
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            with opik.start_as_current_trace(
+                name="handle_request",
+                project_name=project_name,
+                input={"prompt": _helpers.random_text(trace_input_bytes)},
+                output={"completion": _helpers.random_text(trace_output_bytes)},
+            ) as trace:
+                for j in range(spans_per_trace):
+                    with opik.start_as_current_span(
+                        name=f"heavy_tool_call_{j}",
+                        input={"prompt": _helpers.random_text(span_input_bytes)},
+                        output={"completion": _helpers.random_text(span_output_bytes)},
+                    ):
+                        pass
+                submitted_trace_ids.append(trace.id)
+                last_trace_id = trace.id
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+        sample_spans: List[SpanPublic] = _helpers.verify_spans_for_trace(
+            client,
+            project_name=project_name,
+            trace_id=last_trace_id,
+            expected_count=spans_per_trace,
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+    metrics["delivered_spans_on_sample_trace"] = len(sample_spans)
+    assert len(sample_spans) >= spans_per_trace

tests_load/suite/python_sdk/test_ingestion_rate.py

@@ -0,0 +1,133 @@
+"""High spans/traces ingestion rate scenarios."""
+
+from typing import List, Set
+
+import opik
+from opik.rest_api.types.span_public import SpanPublic
+
+from . import _helpers
+from ._helpers import Metrics
+
+
+def test_many_traces_one_span_each(metrics: Metrics, load_scale: float) -> None:
+    """High trace count, low spans-per-trace, via ``@opik.track``.
+
+    Mimics a user-facing handler (``handle_request``) that makes one
+    downstream call (``downstream_call``). Both are ``@opik.track``-
+    decorated so each invocation creates a trace with a nested span —
+    the same shape an instrumented LLM app would emit.
+
+    Volume: 100k traces × 1 span each ≈ 200k observations. Payloads are
+    intentionally small (100 B) so the test stresses message count, not
+    per-message size.
+
+    Verifies every submitted trace id lands with required fields set.
+    """
+    trace_count: int = int(100_000 * load_scale)
+    trace_input_bytes: int = 100
+    downstream_output_bytes: int = 100
+    project_name: str = _helpers.unique_project_name("many-traces")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["trace_input_bytes"] = trace_input_bytes
+    metrics["downstream_output_bytes"] = downstream_output_bytes
+
+    submitted_trace_ids: List[str] = []
+
+    @opik.track
+    def downstream_call(payload: str) -> str:
+        return _helpers.random_text(downstream_output_bytes)
+
+    @opik.track(project_name=project_name)
+    def handle_request(prompt: str) -> str:
+        submitted_trace_ids.append(opik.opik_context.get_current_trace_data().id)
+        return downstream_call(payload=prompt)
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            handle_request(prompt=_helpers.random_text(trace_input_bytes))
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+
+
+def test_many_spans_per_trace(metrics: Metrics, load_scale: float) -> None:
+    """Moderate trace count, heavy span fan-out per trace, via context managers.
+
+    Uses ``opik.start_as_current_trace`` and ``opik.start_as_current_span``
+    — the pattern a user reaches for when they want explicit control over
+    where a trace/span starts and ends rather than wrapping a function.
+
+    Volume: 5k traces × 50 spans = 250k spans. Payloads are small
+    (~100 B) so the test stresses span-batching and trace/span ordering
+    guarantees more than raw byte volume.
+
+    Verifies every submitted trace id lands with required fields set, and
+    that the last trace's 50 spans are all visible and well-formed.
+    """
+    trace_count: int = int(5_000 * load_scale)
+    spans_per_trace: int = 50
+    trace_input_bytes: int = 100
+    trace_output_bytes: int = 100
+    span_input_bytes: int = 100
+    span_output_bytes: int = 100
+    project_name: str = _helpers.unique_project_name("many-spans")
+
+    metrics["project_name"] = project_name
+    metrics["trace_count"] = trace_count
+    metrics["spans_per_trace"] = spans_per_trace
+    metrics["trace_input_bytes"] = trace_input_bytes
+    metrics["trace_output_bytes"] = trace_output_bytes
+    metrics["span_input_bytes"] = span_input_bytes
+    metrics["span_output_bytes"] = span_output_bytes
+
+    submitted_trace_ids: List[str] = []
+    last_trace_id: str = ""
+
+    with metrics.timer("logging"):
+        for _ in range(trace_count):
+            with opik.start_as_current_trace(
+                name="handle_request",
+                project_name=project_name,
+                input={"prompt": _helpers.random_text(trace_input_bytes)},
+                output={"completion": _helpers.random_text(trace_output_bytes)},
+            ) as trace:
+                for j in range(spans_per_trace):
+                    with opik.start_as_current_span(
+                        name=f"tool_call_{j}",
+                        input={"prompt": _helpers.random_text(span_input_bytes)},
+                        output={"completion": _helpers.random_text(span_output_bytes)},
+                    ):
+                        pass
+                submitted_trace_ids.append(trace.id)
+                last_trace_id = trace.id
+            _helpers.think_time()
+
+    with metrics.timer("flush"):
+        opik.flush_tracker()
+
+    client = _helpers.opik_client()
+    with metrics.timer("verify"):
+        delivered_trace_ids: Set[str] = _helpers.verify_exact_trace_ids(
+            client, project_name=project_name, expected_ids=set(submitted_trace_ids)
+        )
+        sample_spans: List[SpanPublic] = _helpers.verify_spans_for_trace(
+            client,
+            project_name=project_name,
+            trace_id=last_trace_id,
+            expected_count=spans_per_trace,
+        )
+
+    metrics["delivered_trace_count"] = len(delivered_trace_ids)
+    metrics["delivered_spans_on_sample_trace"] = len(sample_spans)
+    assert len(sample_spans) >= spans_per_trace