Implement Agent Trajectory as a formal data structure.

[alexmanle]

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:

trajectory = Trajectory(
    iteration_dir=iteration_dir,
    file_type="csv",
    components=(EnvParamsSample,),
)

Append a completed trial using the fields required by the configured components:

trajectory.append(
    step=1,
    action={"batch_size": 32},
    reward=0.95,
    observation=[120.0],
    env_params={"osl": 100},
)

Entries can be accessed as an ordered sequence or searched by trial identity:

entry = trajectory[0]
cached = trajectory.find(
    {"batch_size": 32},
    env_params={"network_speed": 100},
)

result = entry.get(TrialResult)
env_sample = entry.get(EnvParamsSample)

Trajectory file type output can be selected using the new agent_config flag.

[agent_config]
trajectory_file_type = "jsonl"

Test Plan

Tested locally with Python 3.14 using:

uv run --extra dev ruff check src tests
uv run --extra dev pyright
uv run --extra dev pytest -q
git diff --check

Also tested with CloudAI internal agents.

Notes

• Actions and identity-bearing extension values are immutable after insertion, while persistence is handled automatically by the trajectory writer.
• Env_Params extension is used for domain randomization.
• JSONL trajectory output is now possible.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

The 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

Layer / File(s)	Summary
Typed trajectory model and persistence src/cloudai/configurator/trajectory.py, tests/test_trajectory.py	Adds typed entries, component schemas, CSV/JSONL writers, ordered appends, exact cache lookup, and comprehensive tests.
Environment trajectory execution and caching src/cloudai/configurator/cloudai_gym.py, src/cloudai/configurator/env_params.py, src/cloudai/configurator/__init__.py, src/cloudai/models/workload.py, src/cloudai/systems/slurm/single_sbatch_runner.py, src/cloudai/cli/handlers.py, src/cloudai/configurator/base_agent.py, tests/test_cloudaigym.py, tests/test_env_params.py, tests/test_gymnasium_adapter_contract.py	Routes cached and executed results through Trajectory, records optional environment parameters in trajectory output, configures trajectory file types, updates exports, and removes the separate environment-parameter CSV writer.
Trajectory loading and DSE reporting src/cloudai/report_generator/dse_report.py, src/cloudai/reporter.py, tests/test_reporter.py, tests/test_single_sbatch_runner.py	Loads JSONL trajectories preferentially, falls back to CSV, and uses the shared loader in DSE reporting.

Estimated code review effort: 4 (Complex) | ~60 minutes

Suggested reviewers: srivatsankrishnan, podkidyshev, jeffnvidia

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: introducing a formal trajectory data structure.
Description check	✅ Passed	The description matches the changeset by covering typed trajectory components, persistence, cache identity, and configurable output.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between 0de6ca8 and d469b00.

📒 Files selected for processing (13)

• src/cloudai/configurator/__init__.py
• src/cloudai/configurator/cloudai_gym.py
• src/cloudai/configurator/env_params.py
• src/cloudai/configurator/trajectory.py
• src/cloudai/models/workload.py
• src/cloudai/report_generator/dse_report.py
• src/cloudai/reporter.py
• src/cloudai/systems/slurm/single_sbatch_runner.py
• tests/test_cloudaigym.py
• tests/test_env_params.py
• tests/test_gymnasium_adapter_contract.py
• tests/test_reporter.py
• tests/test_trajectory.py

💤 Files with no reviewable changes (2)

• src/cloudai/configurator/env_params.py
• tests/test_env_params.py

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between d469b00 and 7aefb83.

📒 Files selected for processing (6)

• src/cloudai/configurator/trajectory.py
• src/cloudai/systems/slurm/single_sbatch_runner.py
• tests/test_cloudaigym.py
• tests/test_reporter.py
• tests/test_single_sbatch_runner.py
• tests/test_trajectory.py

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between 7aefb83 and daebbb8.

📒 Files selected for processing (3)

• src/cloudai/configurator/trajectory.py
• src/cloudai/systems/slurm/single_sbatch_runner.py
• tests/test_trajectory.py

[coderabbitai]

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 win

Snapshot the observation payload before storing it.

TrialResult freezes action but leaves observation caller-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

📥 Commits

Reviewing files that changed from the base of the PR and between a0ad79c and ddab127.

📒 Files selected for processing (3)

• src/cloudai/cli/handlers.py
• src/cloudai/configurator/base_agent.py
• src/cloudai/configurator/trajectory.py