diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py index cdafd5ce5..1d1e256b6 100644 --- a/src/cloudai/configurator/cloudai_gym.py +++ b/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: diff --git a/src/cloudai/models/workload.py b/src/cloudai/models/workload.py index 22c3c04ad..21d490276 100644 --- a/src/cloudai/models/workload.py +++ b/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=(