Merge pull request #671 from NVIDIA/am/single-sbatch-trajectory

Write trajectory file for DSE jobs in single-sbatch mode

Author: amaslenn

src/cloudai/cli/handlers.py

@@ -121,7 +121,7 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace):
 
     for tr in runner.runner.test_scenario.test_runs:
         test_run = copy.deepcopy(tr)
-        env = CloudAIGymEnv(test_run=test_run, runner=runner)
+        env = CloudAIGymEnv(test_run=test_run, runner=runner.runner)
         agent_type = test_run.test.test_definition.agent
 
         agent_class = registry.agents_map.get(agent_type)

src/cloudai/configurator/cloudai_gym.py

@@ -20,7 +20,7 @@
 import logging
 from typing import Any, Dict, Optional, Tuple
 
-from cloudai.core import METRIC_ERROR, Registry, Runner, TestRun
+from cloudai.core import METRIC_ERROR, BaseRunner, Registry, TestRun
 from cloudai.util.lazy_imports import lazy
 
 from .base_gym import BaseGym
@@ -33,13 +33,13 @@ class CloudAIGymEnv(BaseGym):
     Uses the TestRun object and actual runner methods to execute jobs.
     """
 
-    def __init__(self, test_run: TestRun, runner: Runner):
+    def __init__(self, test_run: TestRun, runner: BaseRunner):
         """
         Initialize the Gym environment using the TestRun object.
 
         Args:
             test_run (TestRun): A test run object that encapsulates cmd_args, extra_cmd_args, etc.
-            runner (Runner): The runner object to execute jobs.
+            runner (BaseRunner): The runner object to execute jobs.
         """
         self.test_run = test_run
         self.original_test_run = copy.deepcopy(test_run)  # Preserve clean state for DSE
@@ -106,20 +106,20 @@ 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]
+        new_tr.output_path = self.runner.get_job_output_path(new_tr)
+        self.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()
+        self.runner.shutting_down = False
+        self.runner.jobs.clear()
+        self.runner.testrun_to_job_map.clear()
 
-        asyncio.run(self.runner.run())
+        try:
+            asyncio.run(self.runner.run())
+        except Exception as e:
+            logging.error(f"Error running step {self.test_run.step}: {e}")
 
-        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]
+        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:
             self.test_run = copy.deepcopy(self.original_test_run)
             self.test_run.step = new_tr.step
@@ -179,7 +179,7 @@ def get_observation(self, action: Any) -> list:
 
         observation = []
         for metric in all_metrics:
-            v = self.test_run.get_metric_value(self.runner.runner.system, metric)
+            v = self.test_run.get_metric_value(self.runner.system, metric)
             if v == METRIC_ERROR:
                 v = -1.0
             observation.append(v)
@@ -196,10 +196,7 @@ def write_trajectory(self, step: int, action: Any, reward: float, observation: l
             observation (list): The observation after taking the action.
         """
         trajectory_file_path = (
-            self.runner.runner.scenario_root
-            / self.test_run.name
-            / f"{self.test_run.current_iteration}"
-            / "trajectory.csv"
+            self.runner.scenario_root / self.test_run.name / f"{self.test_run.current_iteration}" / "trajectory.csv"
         )
 
         file_exists = trajectory_file_path.exists()

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -21,11 +21,12 @@
 from pathlib import Path
 from typing import Generator, Optional, cast
 
+from cloudai.configurator.cloudai_gym import CloudAIGymEnv
 from cloudai.core import JobIdRetrievalError, System, TestRun, TestScenario
-from cloudai.systems.slurm.slurm_metadata import SlurmJobMetadata, SlurmStepMetadata
 from cloudai.util import CommandShell, format_time_limit, parse_time_limit
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
+from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata
 from .slurm_runner import SlurmJob, SlurmRunner
 from .slurm_system import SlurmSystem
 
@@ -194,8 +195,25 @@ async def run(self):
             is_completed = True if self.mode == "dry-run" else self.system.is_job_completed(job)
             await asyncio.sleep(self.system.monitor_interval)
 
+        self.handle_dse()
+
         self.on_job_completion(job)
 
+    def handle_dse(self):
+        for tr in self.test_scenario.test_runs:
+            if not tr.is_dse_job:
+                continue
+
+            for idx, combination in enumerate(tr.all_combinations):
+                next_tr = tr.apply_params_set(combination)
+                next_tr.step = idx + 1
+                next_tr.output_path = self.get_job_output_path(next_tr)
+
+                gym = CloudAIGymEnv(next_tr, self)
+                observation = gym.get_observation({})
+                reward = gym.compute_reward(observation)
+                gym.write_trajectory(idx, combination, reward, observation)
+
     def _submit_test(self, tr: TestRun) -> SlurmJob:
         with open(self.scenario_root / "cloudai_sbatch_script.sh", "w") as f:
             f.write(self.gen_sbatch_content())

tests/test_cloudaigym.py

@@ -21,7 +21,7 @@
 import pytest
 
 from cloudai.configurator import CloudAIGymEnv, GridSearchAgent
-from cloudai.core import Runner, Test, TestRun, TestScenario, TestTemplateStrategy
+from cloudai.core import BaseRunner, Runner, Test, TestRun, TestScenario, TestTemplateStrategy
 from cloudai.systems.slurm import SlurmSystem
 from cloudai.workloads.nemo_run import (
     Data,
@@ -45,7 +45,7 @@ def nemorun() -> NeMoRunTestDefinition:
 
 
 @pytest.fixture
-def setup_env(slurm_system: SlurmSystem, nemorun: NeMoRunTestDefinition) -> tuple[TestRun, Runner]:
+def setup_env(slurm_system: SlurmSystem, nemorun: NeMoRunTestDefinition) -> tuple[TestRun, BaseRunner]:
     tdef = nemorun.model_copy(deep=True)
     tdef.cmd_args.trainer = Trainer(
         max_steps=[1000, 2000],
@@ -81,10 +81,10 @@ def setup_env(slurm_system: SlurmSystem, nemorun: NeMoRunTestDefinition) -> tupl
 
     runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario)
 
-    return test_run, runner
+    return test_run, runner.runner
 
 
-def test_observation_space(setup_env):
+def test_observation_space(setup_env: tuple[TestRun, BaseRunner]):
     test_run, runner = setup_env
     env = CloudAIGymEnv(test_run=test_run, runner=runner)
     observation_space = env.define_observation_space()
@@ -147,7 +147,7 @@ def test_compute_reward_invalid():
     assert "Available functions: ['inverse', 'negative', 'identity']" in str(exc_info.value)
 
 
-def test_tr_output_path(setup_env: tuple[TestRun, Runner]):
+def test_tr_output_path(setup_env: tuple[TestRun, BaseRunner]):
     test_run, runner = setup_env
     test_run.test.test_definition.cmd_args.data.global_batch_size = 8  # avoid constraint check failure
     env = CloudAIGymEnv(test_run=test_run, runner=runner)
@@ -160,7 +160,7 @@ def test_tr_output_path(setup_env: tuple[TestRun, Runner]):
     assert env.test_run.output_path.name == "42"
 
 
-def test_action_space(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, Runner]):
+def test_action_space(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, BaseRunner]):
     tr, _ = setup_env
     nemorun.cmd_args.trainer = Trainer(
         max_steps=[1000, 2000], strategy=TrainerStrategy(tensor_model_parallel_size=[1, 2])
@@ -185,7 +185,7 @@ def test_action_space(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun,
 
 
 @pytest.mark.parametrize("num_nodes", (1, [1, 2], [3]))
-def test_all_combinations(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, Runner], num_nodes: int):
+def test_all_combinations(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, BaseRunner], num_nodes: int):
     tr, _ = setup_env
     nemorun.cmd_args.trainer = Trainer(max_steps=[1000], strategy=TrainerStrategy(tensor_model_parallel_size=[1, 2]))
     nemorun.extra_env_vars["DSE_VAR"] = ["1", "2", "3"]
@@ -224,7 +224,7 @@ def test_all_combinations(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestR
         assert expected in real_combinations, f"Expected {expected} in all_combinations"
 
 
-def test_all_combinations_non_dse(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, Runner]):
+def test_all_combinations_non_dse(nemorun: NeMoRunTestDefinition, setup_env: tuple[TestRun, BaseRunner]):
     tr, _ = setup_env
     tr.test.test_definition = nemorun
     assert len(tr.all_combinations) == 0