Dataset Management
Relevant source files
The following files were used as context for generating this wiki page:
Dataset Management in the Osmosis SDK encompasses the utilities and services required to ingest, validate, and transform structured data for LLM training workflows. This system supports two primary modes: Test Mode, used for local agent loop evaluation, and CLI Evaluation, used for large-scale rubric-based scoring of conversation records.
The management layer ensures that datasets provided in various formats (JSON, JSONL, Parquet) conform to strict schemas before they are processed by the evaluation engines. This allows developers to maintain consistency between local testing and remote rollout environments.
Dataset Architecture and Formats
The SDK provides specialized readers to handle different file extensions and data structures. While both Test Mode and CLI services support JSON-based formats, they serve different internal consumers.
Supported File Formats
The system supports three primary data formats for input:
- JSON: An array of objects.
- JSONL: JSON Lines format, where each line is a valid JSON object.
- Parquet: Apache Parquet format (primarily supported in Test Mode).
Sources: osmosis_ai/rollout/test_mode/dataset.py:12-16, osmosis_ai/cli_services/dataset.py:100-110
Component Overview
The following diagram illustrates how raw files are processed into validated internal objects.
The logic separates raw parsing from semantic validation to ensure that missing required columns are caught early in the pipeline.
Sources: osmosis_ai/rollout/test_mode/dataset.py:48-60, osmosis_ai/cli_services/dataset.py:82-95
Test Mode Dataset Management
Test mode focuses on providing data for RolloutAgentLoop implementations. It utilizes the DatasetReader class to iterate through files and convert them into protocol-compliant requests.
Required Schema (Test Mode)
For local testing of agents, the dataset must contain specific columns to facilitate the interaction loop and subsequent reward calculation.
| Column | Type | Description |
|---|---|---|
ground_truth | String | The expected output used for reward calculation. |
user_prompt | String | The initial input from the user to the agent. |
system_prompt | String | The system instructions provided to the agent. |
Sources: osmosis_ai/rollout/test_mode/dataset.py:20-24, osmosis_ai/rollout/test_mode/dataset.py:41-45
Transformation Logic
Validated rows are converted into RolloutRequest objects via dataset_row_to_request. This function maps the dataset columns to the message history and embeds ground_truth and other extra columns into the metadata field.
# osmosis_ai/rollout/test_mode/dataset.py:309-322
metadata: Dict[str, Any] = {
"ground_truth": row["ground_truth"],
"test_mode": True,
"row_index": row_index,
}
# Preserve any extra columns from dataset
for key, value in row.items():
if key not in REQUIRED_COLUMNS_SET:
metadata[key] = value
Sources: osmosis_ai/rollout/test_mode/dataset.py:309-322
CLI Evaluation Dataset Management
The CLI services use a DatasetLoader to process JSONL records specifically for rubric evaluation. Unlike Test Mode, this pathway focuses on solution_str (the assistant's response) and rubric_id.
The DatasetRecord Object
Each entry in a CLI dataset is parsed into a DatasetRecord data class. This object handles the extraction of optional fields like original_input or extra_info.
| Field | Description |
|---|---|
rubric_id | Identifies which rubric configuration applies to this record. |
solution_str | The actual text generated by the assistant to be evaluated. |
conversation_id | Optional identifier for tracking specific conversations. |
metadata | Arbitrary key-value pairs preserved for reporting. |
Sources: osmosis_ai/cli_services/dataset.py:16-35
Filtering and Sessions
During an evaluation session, the EvaluationSession filters the loaded dataset based on the requested rubric_id. Records that do not match the ID are skipped.
Data Validation Rules
The SDK enforces strict validation across both management modules to prevent runtime failures during LLM calls.
Validation Constraints
- Non-Nullity: Required fields (like
solution_stroruser_prompt) cannot be null. - Type Enforcement: Required fields must be strings.
- Non-Empty: Strings cannot be empty or consist solely of whitespace.
- Case Insensitivity: Column names in input files (e.g.,
.parquetor.json) are matched case-insensitively but normalized to lowercase for internal use.
Sources: osmosis_ai/rollout/test_mode/dataset.py:257-285, osmosis_ai/cli_services/dataset.py:112-135
Error Handling
The system raises specific exceptions when validation fails:
DatasetParseError: Raised when the file format is incorrect or the JSON is malformed.DatasetValidationError: Raised when required columns are missing or contain invalid data types.CLIError: Used in CLI services to report missingsolution_stror empty datasets.
Sources: osmosis_ai/rollout/test_mode/dataset.py:90-100, osmosis_ai/cli_services/dataset.py:126-140
Implementation Details
Memory Efficiency
The DatasetReader in Test Mode provides an iter_rows method. This allows for memory-efficient iteration over large datasets (especially Parquet and JSONL) by yielding rows one by one rather than loading the entire file into memory.
Sources: osmosis_ai/rollout/test_mode/dataset.py:116-140
Extra Columns
Both systems are designed to be extensible. Any column present in the input file that is not a "required" column is automatically captured and placed into a metadata or extra_info dictionary. This ensures that custom user data travels with the record through the entire evaluation pipeline.
Sources: osmosis_ai/rollout/test_mode/dataset.py:287-293, osmosis_ai/cli_services/dataset.py:145-155
Conclusion
Dataset Management provides the foundational data layer for the Osmosis SDK. By combining flexible format support with rigid schema validation and normalization, it ensures that evaluation engines receive high-quality, predictable input. Whether running local agent tests or remote rubric evaluations, the management system preserves the integrity of the original data while preparing it for complex LLM workflows.