Test Mode Datasets
Relevant source files
The following files were used as context for generating this wiki page:
Test Mode Datasets serve as the foundational data layer for validating RolloutAgentLoop implementations locally. This system allows developers to run agent loops against structured datasets containing prompts and ground truth answers without requiring a live TrainGate server environment.
The module facilitates the ingestion, validation, and transformation of raw data into protocol-compliant requests. It supports multiple file formats, including JSON, JSONL, and Parquet, ensuring compatibility with common data science and LLM training workflows.
Sources: osmosis_ai/rollout/test_mode/init.py:1-15, osmosis_ai/rollout/test_mode/dataset.py:1-20
Dataset Architecture and Ingestion
The ingestion process is centered around the DatasetReader class, which handles file format detection and row-level validation. The system enforces strict schema requirements to ensure that every record contains the necessary context for an LLM rollout.
File Format Support
The system supports three primary data formats:
- JSON: A standard array of objects.
- JSONL: Line-delimited JSON objects, suitable for large datasets.
- Parquet: Columnar storage format, which uses metadata for efficient row counting.
Sources: osmosis_ai/rollout/test_mode/dataset.py:45-80, tests/test_test_mode_dataset.py:40-100
Ingestion Logic Flow
The following diagram illustrates how raw files are processed into validated DatasetRow objects.
The ingestion engine normalizes column names to lowercase to allow for flexible casing in source files (e.g., "USER_PROMPT" becomes "user_prompt").
Sources: osmosis_ai/rollout/test_mode/dataset.py:100-135, tests/test_test_mode_dataset.py:180-200
Data Schema and Validation
The system enforces a set of required columns designated in REQUIRED_COLUMNS.
Required Fields
| Field | Type | Description |
|---|---|---|
user_prompt | String | The primary input or question for the agent. |
system_prompt | String | The system instructions or persona for the LLM. |
ground_truth | String | The expected correct answer for evaluation. |
| Sources: osmosis_ai/rollout/test_mode/dataset.py:25-30 |
Validation Rules
The DatasetReader performs the following checks on every row:
- Existence: All required columns must be present.
- Type Enforcement: Values for required columns must be strings.
- Content: Values cannot be null or whitespace-only strings.
- Preservation: Extra columns (e.g., "difficulty", "category") are preserved and passed into the metadata of the resulting request.
Sources: osmosis_ai/rollout/test_mode/dataset.py:140-165, tests/test_test_mode_dataset.py:270-320
Data Transformation
Once validated, DatasetRow objects are converted into RolloutRequest objects via the dataset_row_to_request function. This transformation bridges the gap between raw data and the Osmosis rollout protocol.
Metadata Enrichment
During transformation, the system enriches the metadata field of the RolloutRequest:
test_mode: Set toTrueto signal local execution.row_index: The absolute index of the row in the dataset.ground_truth: The reference answer is stored here for later comparison by reward functions.- Custom Fields: Any additional columns found in the source file are automatically included in this dictionary.
Sources: osmosis_ai/rollout/test_mode/dataset.py:200-215, tests/test_test_mode_dataset.py:365-385
CLI Integration
The osmosis test command leverages these datasets to drive batch and interactive testing.
Dataset Loading Parameters
The CLI provides several arguments to control how datasets are read:
--limit: Maximum number of rows to process.--offset: Number of rows to skip at the beginning of the file.--row: In interactive mode, specifies the specific absolute index to start testing.
Sources: osmosis_ai/rollout/test_mode/cli.py:90-115, osmosis_ai/rollout/test_mode/cli.py:215-240
# Implementation of dataset reading in the CLI
reader = DatasetReader(args.dataset)
total_rows = len(reader)
rows = reader.read(limit=args.limit, offset=args.offset)
Sources: osmosis_ai/rollout/test_mode/cli.py:220-225
Error Handling
The dataset module defines specific exceptions to handle failures during the ingestion lifecycle:
DatasetParseError: Raised when the file format is invalid (e.g., malformed JSON).DatasetValidationError: Raised when the file is readable but violates schema constraints (e.g., missinguser_prompt).
Sources: osmosis_ai/rollout/test_mode/exceptions.py:1-20, tests/test_test_mode_dataset.py:230-260
Conclusion
Test Mode Datasets provide a robust and flexible mechanism for offline agent validation. By supporting standard formats and enforcing strict data quality rules, the system ensures that developers can reliably test agent logic against known scenarios before moving to remote training environments.