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

@@ -22,6 +22,7 @@
 
 from cloudai.core import METRIC_ERROR, Registry, Runner, TestRun
 from cloudai.util.lazy_imports import lazy
+from cloudai._core.exceptions import JobFailureError
 
 from .base_gym import BaseGym
 
@@ -42,6 +43,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 +107,21 @@ 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)