feat(dse): add pin_nodes_to_first_step option to lock DSE steps to sa…

src/cloudai/configurator/cloudai_gym.py

@@ -60,6 +60,7 @@ def __init__(
         self.reward_function = Registry().get_reward_function(test_run.test.agent_reward_function)
         self.params: EnvParams | None = EnvParams.from_test(test_run.test)
         self.trajectory = Trajectory(iteration_dir=self.iteration_dir)
+        self._pinned_nodes: list[str] = []
         super().__init__()
 
     @property
@@ -158,6 +159,8 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
 
             new_tr = copy.deepcopy(self.test_run)
             new_tr.output_path = self.runner.get_job_output_path(new_tr)
+            if self.test_run.test.pin_nodelist and self._pinned_nodes:
+                new_tr.nodes = self._pinned_nodes
             self.runner.test_scenario.test_runs = [new_tr]
 
             self.runner.shutting_down = False
@@ -169,6 +172,19 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
             except Exception as e:
                 logging.error(f"Error running step {self.test_run.step}: {e}")
 
+            if self.test_run.test.pin_nodelist and not self._pinned_nodes:
+                get_job_id = getattr(self.runner, "get_job_id", None)
+                fetch_cmd = getattr(self.runner.system, "fetch_command_output", None)
+                for f in new_tr.output_path.rglob("*.stdout"):
+                    job_id = get_job_id(f.read_text(errors="ignore"), "") if get_job_id else None
+                    if job_id and fetch_cmd:
+                        out, _ = fetch_cmd(f"sacct -j {job_id} -p --noheader -X --format=NodeList")
+                        nodes = out.splitlines()[0].strip().replace("|", "") if out.splitlines() else ""
+                        if nodes and nodes != "Unknown":
+                            self._pinned_nodes = [nodes]
+                            logging.info(f"Pinned DSE nodes to: {nodes}")
+                        break
+
             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:

src/cloudai/models/workload.py

@@ -123,6 +123,10 @@ class TestDefinition(BaseModel, ABC):
     agent_metrics: list[str] = Field(default=["default"])
     agent_reward_function: str = "inverse"
     agent_config: dict[str, Any] | None = Field(default=None, description="Agent configuration.")
+    pin_nodelist: bool = Field(
+        default=False,
+        description="If True, all DSE steps after the first will be pinned to the same nodes as step 1.",
+    )
     env_params: dict[str, EnvParamSpec] = Field(
         default_factory=dict,
         description=(