Merge pull request #609 from NVIDIA/andrei-bug

Fix to Buggy Implemention of PR 589

Author: srivatsankrishnan

src/cloudai/_core/base_runner.py

@@ -259,6 +259,7 @@ async def monitor_jobs(self) -> int:
                                 f"Job {job.id} for test {job.test_run.name} failed: {job_status_result.error_message}"
                             )
                             logging.error(error_message)
+                            await self.handle_job_completion(job)
                             await self.shutdown()
                             raise JobFailureError(job.test_run.name, error_message, job_status_result.error_message)
                     else:

src/cloudai/configurator/cloudai_gym.py

@@ -42,6 +42,7 @@ def __init__(self, test_run: TestRun, runner: Runner):
             runner (Runner): The runner object to execute jobs.
         """
         self.test_run = test_run
+        self.original_test_run = copy.deepcopy(test_run)  # Preserve clean state for DSE
         self.runner = runner
         self.max_steps = test_run.test.test_definition.agent_steps
         self.reward_function = Registry().get_reward_function(test_run.test.test_definition.agent_reward_function)
@@ -105,9 +106,24 @@ def step(self, action: Any) -> Tuple[list, float, bool, dict]:
 
         logging.info(f"Running step {self.test_run.step} with action {action}")
         new_tr = copy.deepcopy(self.test_run)
+        new_tr.output_path = self.runner.runner.get_job_output_path(new_tr)
         self.runner.runner.test_scenario.test_runs = [new_tr]
+
+        self.runner.runner.shutting_down = False
+        self.runner.runner.jobs.clear()
+        self.runner.runner.testrun_to_job_map.clear()
+
         asyncio.run(self.runner.run())
-        self.test_run = self.runner.runner.test_scenario.test_runs[0]
+
+        if (
+            self.runner.runner.test_scenario.test_runs
+            and self.runner.runner.test_scenario.test_runs[0].output_path.exists()
+        ):
+            self.test_run = self.runner.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(action)
         reward = self.compute_reward(observation)