Skip to content

Commit b2b6813

Browse files
gary-huanggary-huang
authored
chore(llmobs): use pagination and a higher timeout for pulling datasets (#14505)
pulling a dataset over the default page limit will result in pulling an incomplete dataset for example, creating this dataset ``` dataset = LLMObs.create_dataset("1-then-big-gh-09031128", "", [{"input_data": "first", "expected_output": "1"}]) print(dataset.as_dataframe()) print(dataset._version) print(dataset.url) for i in range(0, 30000): dataset.append({"input_data": "a"*5000, "expected_output": "b"*100}) dataset.push() ``` https://dd.datad0g.com/llm/datasets/627c39ff-da5c-437a-888a-c079b4c65343 then pulling it will result in this instead of 30000 records: ``` dataset = LLMObs.pull_dataset("1-then-big-gh-09021127") print(dataset.as_dataframe()) ``` ``` input_data expected_output 0 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 1 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 2 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 3 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... ... ... ... 4995 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4996 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4997 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4998 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4999 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... [5000 rows x 2 columns] ``` with this fix, the correct full dataset is returned ``` input_data expected_output 0 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 1 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 2 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 3 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 4 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... ... ... ... 29997 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 29998 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 29999 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 30000 aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa... bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb... 30001 first 1 [30002 rows x 2 columns] ``` ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
1 parent c06c53f commit b2b6813

8 files changed

+3156
-21
lines changed

‎ddtrace/llmobs/_experiment.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -266,9 +266,13 @@ def as_dataframe(self) -> None:
266266
flat_record[("expected_output", "")] = expected_output
267267
column_tuples.add(("expected_output", ""))
268268

269-
for metadata_col, metadata_val in record.get("metadata", {}).items():
270-
flat_record[("metadata", metadata_col)] = metadata_val
271-
column_tuples.add(("metadata", metadata_col))
269+
metadata = record.get("metadata", {})
270+
if isinstance(metadata, dict):
271+
for metadata_col, metadata_val in metadata.items():
272+
flat_record[("metadata", metadata_col)] = metadata_val
273+
column_tuples.add(("metadata", metadata_col))
274+
else:
275+
logger.warning("unexpected metadata format %s", type(metadata))
272276

273277
data_rows.append(flat_record)
274278

‎ddtrace/llmobs/_writer.py‎

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -303,9 +303,10 @@ class LLMObsExperimentsClient(BaseLLMObsWriter):
303303
ENDPOINT = ""
304304
TIMEOUT = 5.0
305305
BULK_UPLOAD_TIMEOUT = 60.0
306+
LIST_RECORDS_TIMEOUT = 20
306307
SUPPORTED_UPLOAD_EXTS = {"csv"}
307308

308-
def request(self, method: str, path: str, body: JSONType = None) -> Response:
309+
def request(self, method: str, path: str, body: JSONType = None, timeout=TIMEOUT) -> Response:
309310
headers = {
310311
"Content-Type": "application/json",
311312
"DD-API-KEY": self._api_key,
@@ -315,7 +316,7 @@ def request(self, method: str, path: str, body: JSONType = None) -> Response:
315316
headers[EVP_SUBDOMAIN_HEADER_NAME] = self.EVP_SUBDOMAIN_HEADER_VALUE
316317

317318
encoded_body = json.dumps(body).encode("utf-8") if body else b""
318-
conn = get_connection(url=self._intake, timeout=self.TIMEOUT)
319+
conn = get_connection(url=self._intake, timeout=timeout)
319320
try:
320321
url = self._intake + self._endpoint + path
321322
logger.debug("requesting %s", url)
@@ -450,23 +451,36 @@ def dataset_get_with_records(self, name: str) -> Dataset:
450451
dataset_description = data[0]["attributes"].get("description", "")
451452
dataset_id = data[0]["id"]
452453

453-
path = f"/api/unstable/llm-obs/v1/datasets/{dataset_id}/records"
454-
resp = self.request("GET", path)
455-
if resp.status != 200:
456-
raise ValueError(f"Failed to pull dataset {name}: {resp.status} {resp.get_json()}")
457-
records_data = resp.get_json()
458-
454+
list_base_path = f"/api/unstable/llm-obs/v1/datasets/{dataset_id}/records"
455+
has_next_page = True
459456
class_records: List[DatasetRecord] = []
460-
for record in records_data.get("data", []):
461-
attrs = record.get("attributes", {})
462-
class_records.append(
463-
{
464-
"record_id": record["id"],
465-
"input_data": attrs["input"],
466-
"expected_output": attrs.get("expected_output"),
467-
"metadata": attrs.get("metadata", {}),
468-
}
469-
)
457+
list_path = list_base_path
458+
page_num = 0
459+
while has_next_page:
460+
resp = self.request("GET", list_path, timeout=self.LIST_RECORDS_TIMEOUT)
461+
if resp.status != 200:
462+
raise ValueError(
463+
f"Failed to pull {page_num}th page of dataset records {name}: {resp.status} {resp.get_json()}"
464+
)
465+
records_data = resp.get_json()
466+
467+
for record in records_data.get("data", []):
468+
attrs = record.get("attributes", {})
469+
class_records.append(
470+
{
471+
"record_id": record["id"],
472+
"input_data": attrs["input"],
473+
"expected_output": attrs.get("expected_output"),
474+
"metadata": attrs.get("metadata", {}),
475+
}
476+
)
477+
next_cursor = records_data.get("meta", {}).get("after")
478+
has_next_page = False
479+
if next_cursor:
480+
has_next_page = True
481+
list_path = f"{list_base_path}?page[cursor]={next_cursor}"
482+
logger.debug("next list records request path %s", list_path)
483+
page_num += 1
470484
return Dataset(name, dataset_id, class_records, dataset_description, curr_version, _dne_client=self)
471485

472486
def dataset_bulk_upload(self, dataset_id: str, records: List[DatasetRecord]):

‎tests/llmobs/llmobs_cassettes/datadog/datadog_api_unstable_llm-obs_v1_datasets_delete_post_a66cc19d.yaml‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
interactions:
2+
- request:
3+
body: '{"data": {"type": "datasets", "attributes": {"type": "soft", "dataset_ids":
4+
["f572bfe7-98d4-42c2-90d9-4c50e554d62b"]}}}'
5+
headers:
6+
Accept:
7+
- '*/*'
8+
? !!python/object/apply:multidict._multidict.istr
9+
- Accept-Encoding
10+
: - identity
11+
Connection:
12+
- keep-alive
13+
Content-Length:
14+
- '119'
15+
? !!python/object/apply:multidict._multidict.istr
16+
- Content-Type
17+
: - application/json
18+
User-Agent:
19+
- python-requests/2.32.3
20+
method: POST
21+
uri: https://api.datadoghq.com/api/unstable/llm-obs/v1/datasets/delete
22+
response:
23+
body:
24+
string: '{"data":[{"id":"f572bfe7-98d4-42c2-90d9-4c50e554d62b","type":"datasets","attributes":{"author":{"id":"a7cd01e3-f412-11ed-a144-0aa89e224034"},"created_at":"2025-09-04T23:06:39.294177Z","current_version":1,"deleted_at":"2025-09-04T23:06:46.662435Z","description":"A
25+
test dataset with a large number of records","name":"test-dataset-large-num-records","updated_at":"2025-09-04T23:06:40.341225Z"}}]}'
26+
headers:
27+
content-length:
28+
- '395'
29+
content-security-policy:
30+
- frame-ancestors 'self'; report-uri https://logs.browser-intake-datadoghq.com/api/v2/logs?dd-api-key=pube4f163c23bbf91c16b8f57f56af9fc58&dd-evp-origin=content-security-policy&ddsource=csp-report&ddtags=site%3Adatadoghq.com
31+
content-type:
32+
- application/vnd.api+json
33+
date:
34+
- Thu, 04 Sep 2025 23:06:46 GMT
35+
strict-transport-security:
36+
- max-age=31536000; includeSubDomains; preload
37+
vary:
38+
- Accept-Encoding
39+
x-content-type-options:
40+
- nosniff
41+
x-frame-options:
42+
- SAMEORIGIN
43+
status:
44+
code: 200
45+
message: OK
46+
version: 1

0 commit comments

Comments
 (0)