Formalize DSE Trajectory

src/cloudai/configurator/__init__.py

@@ -16,15 +16,16 @@
 
 from .base_agent import BaseAgent
 from .base_gym import BaseGym
-from .cloudai_gym import CloudAIGymEnv, TrajectoryEntry
+from .cloudai_gym import CloudAIGymEnv
 from .grid_search import GridSearchAgent
 from .gymnasium_adapter import GymnasiumAdapter
+from .trajectory import Trajectory
 
 __all__ = [
     "BaseAgent",
     "BaseGym",
     "CloudAIGymEnv",
     "GridSearchAgent",
     "GymnasiumAdapter",
-    "TrajectoryEntry",
+    "Trajectory",
 ]

src/cloudai/configurator/cloudai_gym.py

@@ -15,29 +15,20 @@
 # limitations under the License.
 
 import copy
-import csv
-import dataclasses
 import logging
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
 
 from cloudai.core import METRIC_ERROR, BaseRunner, Registry, TestRun
 from cloudai.util.lazy_imports import lazy
 
 from .base_agent import RewardOverrides
 from .base_gym import BaseGym
-from .env_params import EnvParams, ObsLeafDescriptor, write_env_params
+from .env_params import EnvParams, ObsLeafDescriptor
+from .trajectory import Trajectory
 
-
-@dataclasses.dataclass(frozen=True)
-class TrajectoryEntry:
-    """Represents a trajectory entry."""
-
-    step: int
-    action: dict[str, Any]
-    reward: float
-    observation: list
-    env_params: dict[str, Any] = dataclasses.field(default_factory=dict)
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class CloudAIGymEnv(BaseGym):
@@ -47,7 +38,12 @@ class CloudAIGymEnv(BaseGym):
     Uses the TestRun object and actual runner methods to execute jobs.
     """
 
-    def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrides):
+    def __init__(
+        self,
+        test_run: TestRun,
+        runner: BaseRunner,
+        rewards: RewardOverrides,
+    ):
         """
         Initialize the Gym environment using the TestRun object.
 
@@ -62,15 +58,10 @@ def __init__(self, test_run: TestRun, runner: BaseRunner, rewards: RewardOverrid
         self.rewards = rewards
         self.max_steps = test_run.test.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
-        self.trajectory: dict[int, list[TrajectoryEntry]] = {}
         self.params: EnvParams | None = EnvParams.from_test(test_run.test)
+        self.trajectory = Trajectory(iteration_dir=self.iteration_dir)
         super().__init__()
 
-    @property
-    def env_params_record_path(self) -> Path:
-        """``env.csv`` lives alongside ``trajectory.csv`` so a plain ``merge`` joins them."""
-        return self.iteration_dir / "env.csv"
-
     @property
     def upcoming_trial(self) -> int:
         """
@@ -104,7 +95,7 @@ def define_observation_space(self) -> list:
     def reset(
         self,
         seed: Optional[int] = None,
-        options: Optional[dict[str, Any]] = None,  # noqa: Vulture
+        options: Optional[dict[str, Any]] = None,
     ) -> Tuple[list, dict[str, Any]]:
         """
         Reset the environment and reinitialize the TestRun.
@@ -118,6 +109,7 @@ def reset(
                 - observation (list): Initial observation.
                 - info (dict): Additional info for debugging.
         """
+        del options
         if seed is not None:
             lazy.np.random.seed(seed)
         self.test_run.current_iteration = 0
@@ -152,59 +144,50 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
         if cached_result is not None:
             logging.info(
                 "Retrieved cached result from trajectory with reward %s (from step %s). Skipping execution.",
-                cached_result.reward,
-                cached_result.step,
+                cached_result["reward"],
+                cached_result["step"],
             )
-            self.write_trajectory(
-                TrajectoryEntry(
-                    step=self.test_run.step,
-                    action=action,
-                    reward=cached_result.reward,
-                    observation=cached_result.observation,
-                    env_params=sampled_env_params,
-                )
-            )
-
-            return cached_result.observation, cached_result.reward, False, info
-
-        if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
-            logging.info("Constraint check failed. Skipping step.")
-            return [-1.0], self.rewards.constraint_failure, True, info
-
-        new_tr = copy.deepcopy(self.test_run)
-        new_tr.output_path = self.runner.get_job_output_path(new_tr)
-        self.runner.test_scenario.test_runs = [new_tr]
-
-        self.runner.shutting_down = False
-        self.runner.jobs.clear()
-        self.runner.testrun_to_job_map.clear()
-
-        try:
-            self.runner.run()
-        except Exception as e:
-            logging.error(f"Error running step {self.test_run.step}: {e}")
-
-        if self.runner.test_scenario.test_runs and self.runner.test_scenario.test_runs[0].output_path.exists():
-            self.test_run = self.runner.test_scenario.test_runs[0]
+            observation = {
+                metric: cached_result[f"observation.{metric}"] for metric in self.test_run.test.agent_metrics
+            }
+            reward = cast(float, cached_result["reward"])
         else:
-            self.test_run = copy.deepcopy(self.original_test_run)
-            self.test_run.step = new_tr.step
-            self.test_run.output_path = new_tr.output_path
-
-        metrics = self.get_observation(action)
-        reward = self.compute_reward(metrics)
-
-        self.write_trajectory(
-            TrajectoryEntry(
-                step=self.test_run.step,
-                action=action,
-                reward=reward,
-                observation=metrics,
-                env_params=sampled_env_params,
-            )
+            if not self.test_run.test.constraint_check(self.test_run, self.runner.system):
+                logging.info("Constraint check failed. Skipping step.")
+                return [-1.0], self.rewards.constraint_failure, True, info
+
+            new_tr = copy.deepcopy(self.test_run)
+            new_tr.output_path = self.runner.get_job_output_path(new_tr)
+            self.runner.test_scenario.test_runs = [new_tr]
+
+            self.runner.shutting_down = False
+            self.runner.jobs.clear()
+            self.runner.testrun_to_job_map.clear()
+
+            try:
+                self.runner.run()
+            except Exception as e:
+                logging.error(f"Error running step {self.test_run.step}: {e}")
+
+            if self.runner.test_scenario.test_runs and self.runner.test_scenario.test_runs[0].output_path.exists():
+                self.test_run = self.runner.test_scenario.test_runs[0]
+            else:
+                self.test_run = copy.deepcopy(self.original_test_run)
+                self.test_run.step = new_tr.step
+                self.test_run.output_path = new_tr.output_path
+
+            observation = self.get_observation()
+            reward = self.compute_reward(observation)
+
+        self.trajectory.append(
+            step=self.test_run.step,
+            action=action,
+            reward=reward,
+            observation=observation,
+            env_params=sampled_env_params,
         )
 
-        return metrics, reward, False, info
+        return list(observation.values()), reward, False, info
 
     def render(self, mode: str = "human"):
         """
@@ -225,38 +208,35 @@ def seed(self, seed: Optional[int] = None):
         if seed is not None:
             lazy.np.random.seed(seed)
 
-    def compute_reward(self, observation: list) -> float:
+    def compute_reward(self, observation: dict[str, Any]) -> float:
         """
         Compute a reward based on the TestRun result.
 
         Args:
-            observation (list): The observation list containing the average value.
+            observation: Metric values keyed by configured metric name.
 
         Returns:
             float: Reward value.
         """
-        return self.reward_function(observation)
+        return self.reward_function(list(observation.values()))
 
-    def get_observation(self, action: Any) -> list:
+    def get_observation(self) -> dict[str, Any]:
         """
         Get the observation from the TestRun object.
 
-        Args:
-            action (Any): Action taken by the agent.
-
         Returns:
-            list: The observation.
+            Metric values keyed by configured metric name.
         """
         all_metrics = self.test_run.test.agent_metrics
         if not all_metrics:
             raise ValueError("No agent metrics defined for the test run")
 
-        observation = []
+        observation: dict[str, Any] = {}
         for metric in all_metrics:
             v = self.test_run.get_metric_value(self.runner.system, metric)
             if v is METRIC_ERROR:
                 v = self.rewards.metric_failure
-            observation.append(v)
+            observation[metric] = v
         return observation
 
     def structured_observation_descriptors(self) -> Optional[Dict[str, ObsLeafDescriptor]]:
@@ -277,83 +257,15 @@ def encode_env_params(self, env_params: dict[str, Any]) -> Dict[str, Any]:
         """
         return self.params.encode(env_params) if self.params is not None else {}
 
-    def write_trajectory(self, entry: TrajectoryEntry):
-        """
-        Append the entry to the in-memory cache and trajectory.csv (plus env.csv when declared).
-
-        ``trajectory.csv`` and the ``env.csv`` projection are sunk from the same
-        ``TrajectoryEntry`` here, so a trial that never produces an entry (e.g. a
-        constraint failure returns before this call) lands in neither file and the
-        two stay 1:1 step-aligned.
-        """
-        self.current_trajectory.append(entry)
-
-        file_exists = self.trajectory_file_path.exists()
-        logging.debug(f"Writing trajectory into {self.trajectory_file_path} (exists: {file_exists})")
-        self.trajectory_file_path.parent.mkdir(parents=True, exist_ok=True)
-
-        with open(self.trajectory_file_path, mode="a", newline="") as file:
-            writer = csv.writer(file)
-            if not file_exists:
-                writer.writerow(["step", "action", "reward", "observation"])
-            writer.writerow([entry.step, entry.action, entry.reward, entry.observation])
-
-        write_env_params(self.env_params_record_path, entry.step, entry.env_params)
-
     @property
     def iteration_dir(self) -> Path:
-        """Per-iteration output dir; trajectory.csv and env.csv both live here, step-aligned."""
+        """Per-iteration output directory containing the trajectory output."""
         return self.runner.scenario_root / self.test_run.name / f"{self.test_run.current_iteration}"
 
     @property
     def trajectory_file_path(self) -> Path:
-        return self.iteration_dir / "trajectory.csv"
-
-    @property
-    def current_trajectory(self) -> list[TrajectoryEntry]:
-        return self.trajectory.setdefault(self.test_run.current_iteration, [])
-
-    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> TrajectoryEntry | None:
-        """
-        Return a cached entry only when the full trial identity matches.
-
-        Trial identity is ``(action, env_params)``: env-randomized parameters
-        change the workload's behaviour, so a trial repeating the same action
-        under a different ``env_params`` sample must miss and re-run. Empty
-        env_params on both sides is the back-compat path for workloads that
-        do not declare any ``[env_params.*]`` block. The sample is passed in (a
-        per-trial local owned by ``step``), exactly like ``action``.
-        """
-        for entry in self.current_trajectory:
-            action_match = self._values_match_exact(entry.action, action)
-            env_params_match = self._values_match_exact(entry.env_params, env_params)
-            if action_match and env_params_match:
-                return entry
+        return self.trajectory.output_path
 
-        return None
-
-    @classmethod
-    def _values_match_exact(cls, left: Any, right: Any) -> bool:
-        if type(left) is not type(right):
-            return False
-
-        elif isinstance(left, dict):
-            left_keys = set(left.keys())
-            right_keys = set(right.keys())
-            if left_keys != right_keys:
-                return False
-
-            return all(cls._values_match_exact(left[key], right[key]) for key in left_keys)
-
-        elif isinstance(left, (list, tuple)):
-            if len(left) != len(right):
-                return False
-
-            for left_item, right_item in zip(left, right, strict=True):
-                if not cls._values_match_exact(left_item, right_item):
-                    return False
-
-            return True
-
-        else:
-            return left == right
+    def get_cached_trajectory_result(self, action: Any, env_params: dict[str, Any]) -> "pd.Series | None":
+        """Return the first trajectory row matching the action and environment parameters."""
+        return self.trajectory.find(action=action, env_params=env_params)

src/cloudai/configurator/env_params.py

@@ -27,11 +27,9 @@
 
 from __future__ import annotations
 
-import csv
 import dataclasses
 import math
 import random
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
@@ -239,30 +237,6 @@ def observation_descriptors(self) -> Dict[str, ObsLeafDescriptor]:
         return {name: param.observation_descriptor() for name, param in self.params.items()}
 
 
-def write_env_params(path: Path, step: int, sample: Dict[str, Any]) -> None:
-    """
-    Append one trial's env_params sample to a step-aligned CSV.
-
-    The CSV mirrors how ``trajectory.csv`` serialises its ``action`` column
-    (one row per env.step(), sample dict stringified in a single cell) so the
-    two files align 1:1 on ``step`` and a plain ``merge`` joins them.
-
-    Empty samples are skipped, so a run without env_params writes nothing and
-    callers can sink every trial unconditionally.
-    """
-    if step < 1:
-        raise ValueError(f"step must be a positive trial index (cloudai DSE is 1-based); got {step}")
-    if not sample:
-        return
-    new_file = not path.exists()
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("a", newline="") as f:
-        writer = csv.writer(f)
-        if new_file:
-            writer.writerow(("step", "env"))
-        writer.writerow([step, sample])
-
-
 def validate_domain_randomization_active(test_scenario: "TestScenario") -> None:
     """
     Reject prepped configs that declare domain randomization no agent will run.