Merge pull request #598 from NVIDIA/am/cmd-gen-use-tr-member

Rely on member test run object instead of args

Author: amaslenn

src/cloudai/_core/command_gen_strategy.py

@@ -29,28 +29,35 @@ class CommandGenStrategy(TestTemplateStrategy, ABC):
     def __init__(self, system: System, test_run: TestRun) -> None:
         super().__init__(system)
         self.test_run = test_run
+        self._final_env_vars: dict[str, str | list[str]] = {}
 
     @abstractmethod
-    def gen_exec_command(self, tr: TestRun) -> str:
+    def gen_exec_command(self) -> str:
         """
         Generate the execution command for a test based on the given parameters.
 
-        Args:
-            tr (TestRun): Contains the test and its run-specific configurations.
-
         Returns:
             str: The generated execution command.
         """
         pass
 
     @abstractmethod
-    def store_test_run(self, tr: TestRun) -> None:
+    def store_test_run(self) -> None:
         """
         Store the test run information in output folder.
 
         Only at command generation time, CloudAI has all the information to store the test run.
-
-        Args:
-            tr (TestRun): The test run object to be stored.
         """
         pass
+
+    @property
+    def final_env_vars(self) -> dict[str, str | list[str]]:
+        if not self._final_env_vars:
+            final_env_vars = self.system.global_env_vars.copy()
+            final_env_vars.update(self.test_run.test.extra_env_vars)
+            self._final_env_vars = final_env_vars
+        return self._final_env_vars
+
+    @final_env_vars.setter
+    def final_env_vars(self, value: dict[str, str | list[str]]) -> None:
+        self._final_env_vars = value

src/cloudai/systems/lsf/lsf_command_gen_strategy.py

@@ -42,7 +42,7 @@ def __init__(self, system: System, test_run: TestRun) -> None:
         super().__init__(system, test_run)
         self.system = cast(LSFSystem, system)
 
-    def gen_exec_command(self, tr: TestRun) -> str:
+    def gen_exec_command(self) -> str:
         """
         Generate the execution command for the test run.
 
@@ -52,11 +52,13 @@ def gen_exec_command(self, tr: TestRun) -> str:
         Returns:
             str: The generated LSF command.
         """
-        env_vars = self._override_env_vars(self.system.global_env_vars, tr.test.extra_env_vars)
-        cmd_args = self._flatten_dict(tr.test.cmd_args)
-        lsf_args = self._parse_lsf_args(tr.test.test_template.__class__.__name__, env_vars, cmd_args, tr)
+        env_vars = self.final_env_vars
+        cmd_args = self._flatten_dict(self.test_run.test.cmd_args)
+        lsf_args = self._parse_lsf_args(
+            self.test_run.test.test_template.__class__.__name__, env_vars, cmd_args, self.test_run
+        )
 
-        bsub_command = self._gen_bsub_command(lsf_args, env_vars, cmd_args, tr)
+        bsub_command = self._gen_bsub_command(lsf_args, env_vars, cmd_args, self.test_run)
 
         return bsub_command.strip()
 

src/cloudai/systems/lsf/lsf_runner.py

@@ -54,7 +54,7 @@ def get_job_id(self, stdout: str, stderr: str) -> int | None:
     def _submit_test(self, tr: TestRun) -> LSFJob:
         logging.info(f"Running test: {tr.name}")
         tr.output_path = self.get_job_output_path(tr)
-        exec_cmd = self.get_cmd_gen_strategy(self.system, tr).gen_exec_command(tr)
+        exec_cmd = self.get_cmd_gen_strategy(self.system, tr).gen_exec_command()
         logging.debug(f"Executing command for test {tr.name}: {exec_cmd}")
         job_id = 0
         if self.mode == "run":

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -109,11 +109,11 @@ def aux_commands(self) -> list[str]:
         max_nodes, _ = self.extract_sbatch_nodes_spec()
         tr.num_nodes = max_nodes
         cmd_gen = cast(SlurmCommandGenStrategy, self.get_cmd_gen_strategy(self.system, tr))
-        return [cmd_gen._metadata_cmd(tr), cmd_gen._ranks_mapping_cmd(tr)]
+        return [cmd_gen._metadata_cmd(), cmd_gen._ranks_mapping_cmd()]
 
     def get_single_tr_block(self, tr: TestRun) -> str:
         cmd_gen = cast(SlurmCommandGenStrategy, self.get_cmd_gen_strategy(self.system, tr))
-        srun_cmd = cmd_gen.gen_srun_command(tr)
+        srun_cmd = cmd_gen.gen_srun_command()
         nnodes, node_list = self.system.get_nodes_by_spec(tr.nnodes, tr.nodes)
         node_arg = f"--nodelist={','.join(node_list)}" if node_list else f"-N{nnodes}"
         extra_args = (