Implement Agent Trajectory as a formal data structure.#969
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 |
![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
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
|
Despite the implementation is technically correct, I find it too complicated for the goal it achieves (apparently the PR summary misses the business goal!) As far as I could understand, the goal is to make the trajectory a better container:
My alternative suggestion for 2. would be to use the standard thing in data analysis - For the point 1. - our container should be just columnar flat data. So for any additional data that we'd like to populate into trajectory, we could use format what do you think? |
@podkidyshev I agree actually that this solution is to much for the problem. I think as long as it's easy to append new data and search elements (for caching) your proposed alternative is great with me! |
|
@alexmanle this look far better to me! I really like the new trajectory file look please make sure to thoroughly test these changes, including single-sbatch and reporter 🙏 potentially also let other team members be aware of it. |
2 participants
Summary
Introduces a pandas-native
Trajectoryfor storing, analyzing, caching, and persisting DSE trial data. Each trajectory is maintained as an in-memoryDataFrameand written incrementally totrajectory.csv.Trajectory data uses flat, namespaced columns such as
action.batch_size,observation.throughput, andenv_params.network_speed. This removes the previous entry/component, identity metadata, and writer abstractions.Features
Initialize a trajectory with an iteration directory:
An existing DataFrame can optionally be supplied for an in-memory warm start:
Append a completed trial using named domains:
Nested mappings are recursively flattened into pandas-friendly columns:
Search for cached trials using any subset of the stored domains:
Both
append()andfind()return copied pandasSeriesobjects. The complete trajectory is available as a copied DataFrame:CloudAIGym now keeps observations as metric-name mappings internally and passes them directly to the trajectory. They are converted to ordered lists only at the Gym interface boundary.
Reporting reads the new flat CSV format while remaining compatible with released trajectory CSVs containing serialized
actionandobservationcolumns.Test Plan
Also tested with CloudAI internal agents.
Notes
trajectory.csv.objectdtype to preserve Python type distinctions during cache lookup.env_params.*columns.