Implement Agent Trajectory as a formal data structure.#969
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesThe PR adds typed trajectory storage with CSV and JSONL persistence, integrates it into environment and DSE execution, and updates reporting to load JSONL trajectories with CSV fallback. Tests cover trajectory behavior, caching, environment parameters, persistence, and reporting. Trajectory integration
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
alexmanle
requested review from
jeffnvidia,
podkidyshev and
srivatsankrishnan
as code owners
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 32-55: Make the identity-bearing data in
EnvParamsSample.env_params and TrialResult.action structurally immutable, not
only field-reassignment immutable, so values returned through
TrajectoryEntry.get() cannot alter cache matching. Deep-freeze nested mappings
and sequences when constructing or storing these components, or maintain a
private immutable canonical key used by find(); ensure persisted records and
future find() calls use that stable identity.
- Around line 104-118: Update TrajectoryWriter.append to inspect an existing
trajectory.csv header before accepting the first record schema: compare any
non-empty header with self._fields and raise ValueError on mismatch, while
allowing a pre-created empty file to be initialized with the header. Ensure
subsequent appends continue enforcing the established field order and do not
write duplicate headers.
In `@src/cloudai/systems/slurm/single_sbatch_runner.py`:
- Around line 208-225: Update unroll_dse in SingleSbatchRunner to support
declared tr.test.env_params by sampling an EnvParamsSample for each trial,
applying those same values to the execution test run and gym trajectory append.
Ensure the sampled environment parameters are passed wherever CloudAIGymEnv
requires them, and add an integration test covering DSE actions with environment
parameters.
In `@tests/test_cloudaigym.py`:
- Line 717: Suppress Ruff’s S311 finding for the deterministic Random usage in
the expected_sample setup, using the repository’s standard inline suppression
syntax and keeping the test’s deterministic behavior unchanged.
In `@tests/test_reporter.py`:
- Around line 50-59: Update test_load_trajectory_dataframe_prefers_json_lines to
also create a conflicting trajectory.csv in tmp_path with different data before
calling load_trajectory_dataframe. Keep the existing JSONL assertions and verify
the returned path and records still come from trajectory.jsonl, proving JSONL
takes precedence over CSV.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 469cd551-c61f-434b-afff-96cfbef6bf41
📒 Files selected for processing (13)
src/cloudai/configurator/__init__.pysrc/cloudai/configurator/cloudai_gym.pysrc/cloudai/configurator/env_params.pysrc/cloudai/configurator/trajectory.pysrc/cloudai/models/workload.pysrc/cloudai/report_generator/dse_report.pysrc/cloudai/reporter.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_cloudaigym.pytests/test_env_params.pytests/test_gymnasium_adapter_contract.pytests/test_reporter.pytests/test_trajectory.py
💤 Files with no reviewable changes (2)
- src/cloudai/configurator/env_params.py
- tests/test_env_params.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 258-264: The trajectory entry storage used by find() must snapshot
identity-contributing component state instead of comparing live mutable
instances. Update the append/store path and _identity_for handling so every
component with contributes_to_identity=True is frozen or represented by a
private immutable canonical key, while preserving existing identity matching for
persisted entries.
In `@src/cloudai/systems/slurm/single_sbatch_runner.py`:
- Around line 217-233: Apply the same constraint_check used by unroll_dse() in
the combination loop before creating the trial observation, computing reward, or
appending trajectory data; skip rejected combinations so only submitted runs are
recorded. Add a regression test covering a constraint-rejected DSE combination
and verifying it is absent from trajectory output and reports.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 3c0c2a29-883b-4460-86f3-6aba6ecf9382
📒 Files selected for processing (6)
src/cloudai/configurator/trajectory.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_cloudaigym.pytests/test_reporter.pytests/test_single_sbatch_runner.pytests/test_trajectory.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 352-356: Update _validate_components to reject any component
dataclass with contributes_to_identity=True unless its dataclass definition is
frozen. Preserve existing validation for non-dataclass components and duplicate
types, and raise a clear TypeError identifying the offending component.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: b339fc85-47e7-4b63-a8c0-19754eb9876c
📒 Files selected for processing (3)
src/cloudai/configurator/trajectory.pysrc/cloudai/systems/slurm/single_sbatch_runner.pytests/test_trajectory.py
| def _validate_components(components: Sequence[object]) -> None: | ||
| non_dataclasses = [type(component).__name__ for component in components if not dataclasses.is_dataclass(component)] | ||
| if non_dataclasses: | ||
| raise TypeError(f"trajectory components must be dataclass instances: {', '.join(non_dataclasses)}") | ||
| if len({type(component) for component in components}) != len(components): |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Python dataclasses __dataclass_params__ frozen attribute
💡 Result:
The dataclass_params attribute is an internal, undocumented class attribute added by the @dataclass decorator to a class [1][2]. It stores the configuration parameters passed to the decorator [3][4]. The frozen attribute is a member of the object stored in dataclass_params [1][5]. It reflects the value of the frozen parameter provided to the @dataclass decorator (which defaults to False) [6][7]. Key points regarding dataclass_params and the frozen attribute: - Undocumented Status: dataclass_params is considered an internal implementation detail of the Python dataclasses module [2]. Its structure and presence may change in future versions of Python [2][8]. - Introspection: Because it is attached to the class at runtime, it can be used for introspection to determine the configuration of a dataclass without attempting to modify instances or catching exceptions [1]. - Structure: The attribute holds an instance of the internal class _DataclassParams, which contains fields corresponding to the @dataclass parameters (e.g., init, repr, eq, order, unsafe_hash, frozen, etc.) [3][1]. - Purpose: When set to True via @dataclass(frozen=True), the dataclass decorator adds setattr and delattr methods to the class, which raise a FrozenInstanceError if an attempt is made to modify an instance [7][9]. The frozen attribute within dataclass_params tracks this configuration [1]. For production code, relying on dataclass_params is generally discouraged because it is an implementation detail rather than a public API [2]. If you need to check if a class is a dataclass, use the public function dataclasses.is_dataclass(obj) [10]. To check if an instance is frozen, the standard approach remains attempting the operation and handling the potential FrozenInstanceError, as there is no public "is_frozen" property [1][7].
Citations:
- 1: https://stackoverflow.com/questions/75958246/how-to-check-if-a-dataclass-is-frozen
- 2:
match_args,slotsand possiblykw_onlymissing from__dataclass_params__python/cpython#96142 - 3: https://github.com/python/cpython/blob/main/Lib/dataclasses.py
- 4: https://github.com/python/cpython/blob/5f9c0f5ddf441dedeb085b0d9f9c9488ca6bd44d/Lib/dataclasses.py
- 5: https://github.com/python/cpython/blob/3.11/Lib/dataclasses.py
- 6: https://peps.python.org/pep-0557/
- 7: https://docs.python.org/3.11/library/dataclasses.html
- 8: Update __dataclass_params__ to track actual changes made by decorator python/cpython#99713
- 9: https://docs.python.org/3.10/library/dataclasses.html
- 10: https://docs.python.org/3/library/dataclasses.html
Require frozen dataclasses for identity-contributing components. _freeze_component only snapshots field values; if a future contributes_to_identity=True dataclass isn’t frozen=True, callers can still mutate entry.get(Component).field and silently diverge the cache key from the persisted row. Add a _validate_components check for that case.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 355-355: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/cloudai/configurator/trajectory.py` around lines 352 - 356, Update
_validate_components to reject any component dataclass with
contributes_to_identity=True unless its dataclass definition is frozen. Preserve
existing validation for non-dataclass components and duplicate types, and raise
a clear TypeError identifying the offending component.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/cloudai/configurator/trajectory.py (1)
43-53: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winSnapshot the observation payload before storing it.
TrialResultfreezesactionbut leavesobservationcaller-owned. A mutable list can be changed after_store_entry()persists the row, causing cache hits to return observations that differ from the persisted trajectory. Freeze or deep-copy the observation as well.Proposed fix
def __post_init__(self) -> None: """Own an immutable snapshot of the action.""" object.__setattr__(self, "action", _freeze(self.action)) + object.__setattr__(self, "observation", _freeze(self.observation))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cloudai/configurator/trajectory.py` around lines 43 - 53, Update TrialResult.__post_init__ to snapshot observation before storing it, using the existing _freeze helper or an equivalent deep-copy mechanism. Preserve the observation’s sequence semantics while ensuring later caller mutations cannot alter the stored trajectory or cached results.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cloudai/configurator/base_agent.py`:
- Line 52: Update the Slurm DSE environment construction in
single_sbatch_runner.py, where CloudAIGymEnv is instantiated, to pass through
the configured trajectory_file_type alongside rewards. Preserve the selected
value, including "jsonl", instead of allowing the environment default of "csv"
to apply.
---
Outside diff comments:
In `@src/cloudai/configurator/trajectory.py`:
- Around line 43-53: Update TrialResult.__post_init__ to snapshot observation
before storing it, using the existing _freeze helper or an equivalent deep-copy
mechanism. Preserve the observation’s sequence semantics while ensuring later
caller mutations cannot alter the stored trajectory or cached results.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: 365ca573-91b2-4412-be87-68fada14bbbd
📒 Files selected for processing (3)
src/cloudai/cli/handlers.pysrc/cloudai/configurator/base_agent.pysrc/cloudai/configurator/trajectory.py
| default_factory=RewardOverrides, | ||
| description="Reward and observation overrides for the agent.", | ||
| ) | ||
| trajectory_file_type: Literal["csv", "jsonl"] = "csv" |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n## base_agent.py excerpt\n'
sed -n '1,140p' src/cloudai/configurator/base_agent.py
printf '\n## single_sbatch_runner.py excerpt\n'
sed -n '180,260p' src/cloudai/systems/slurm/single_sbatch_runner.py
printf '\n## search for trajectory_file_type and CloudAIGymEnv\n'
rg -n "trajectory_file_type|CloudAIGymEnv\(" srcRepository: NVIDIA/cloudai
Length of output: 9293
Forward trajectory_file_type in the Slurm DSE env setup. src/cloudai/systems/slurm/single_sbatch_runner.py:214 still constructs CloudAIGymEnv with only rewards, so jsonl configs silently fall back to CSV on this path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/cloudai/configurator/base_agent.py` at line 52, Update the Slurm DSE
environment construction in single_sbatch_runner.py, where CloudAIGymEnv is
instantiated, to pass through the configured trajectory_file_type alongside
rewards. Preserve the selected value, including "jsonl", instead of allowing the
environment default of "csv" to apply.
1 participant
Summary
Introduces a typed Trajectory data structure with extensible dataclass components, cache identity support, and CSV/JSONL persistence. Environment parameters are now stored as trajectory components, and CloudAI defaults to CSV output.
Features
The new Trajectory class manages in-memory trajectory entries and persists them to a file: trajectory.[csv/jsonl].
Initialize it with the iteration directory and any optional component types:
Append a completed trial using the fields required by the configured components:
Entries can be accessed as an ordered sequence or searched by trial identity:
Trajectory file type output can be selected using the new agent_config flag.
Test Plan
Tested locally with Python 3.14 using:
Also tested with CloudAI internal agents.
Notes