Merge pull request #567 from NVIDIA/am/more-slurm-meta

Expand slurm meta to have per-step information

Author: amaslenn

src/cloudai/cli/handlers.py

@@ -168,7 +168,8 @@ def generate_reports(system: System, test_scenario: TestScenario, result_dir: Pa
             reporter = reporter_class(system, test_scenario, result_dir, cfg)
             reporter.generate()
         except Exception as e:
-            logging.warning(f"Error generating report: {e}")
+            logging.warning(f"Error generating report '{name}', see debug log for details")
+            logging.debug(e, stack_info=True)
 
 
 def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> None:

src/cloudai/systems/slurm/__init__.py

@@ -14,14 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .metadata import SlurmSystemMetadata
 from .single_sbatch_runner import SingleSbatchRunner
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
 from .slurm_installer import SlurmInstaller
 from .slurm_job import SlurmJob
+from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata, SlurmSystemMetadata
 from .slurm_node import SlurmNode, SlurmNodeState
 from .slurm_runner import SlurmRunner
-from .slurm_system import SlurmGroup, SlurmJobMetadata, SlurmPartition, SlurmSystem, parse_node_list
+from .slurm_system import SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list
 
 __all__ = [
     "SingleSbatchRunner",
@@ -34,6 +34,7 @@
     "SlurmNodeState",
     "SlurmPartition",
     "SlurmRunner",
+    "SlurmStepMetadata",
     "SlurmSystem",
     "SlurmSystemMetadata",
     "parse_node_list",

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -21,14 +21,13 @@
 from pathlib import Path
 from typing import Generator, Optional, cast
 
-import toml
-
 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_runner import SlurmJob, SlurmRunner
-from .slurm_system import SlurmJobMetadata, SlurmSystem
+from .slurm_system import SlurmSystem
 
 
 class SingleSbatchRunner(SlurmRunner):
@@ -184,7 +183,7 @@ 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)
 
-        await self.job_completion_callback(job)
+        self.on_job_completion(job)
 
     def _submit_test(self, tr: TestRun) -> SlurmJob:
         with open(self.scenario_root / "cloudai_sbatch_script.sh", "w") as f:
@@ -206,25 +205,20 @@ def _submit_test(self, tr: TestRun) -> SlurmJob:
         logging.info(f"Submitted slurm job: {job_id}")
         return SlurmJob(tr, id=job_id)
 
-    def store_job_metadata(self, job: SlurmJob):
-        logging.debug(f"Storing job metadata for job {job.id}")
-        res = None if self.mode == "dry-run" else self.system.get_job_status(job)
-        logging.debug(f"Job status ra: {res}")
-
-        job_name, job_state, time_sec = "unknown", "UNKNOWN", 0
-        if res:
-            job_name, job_state, time_sec = res[0], res[1], int(res[2])
-
-        job_meta = SlurmJobMetadata(
+    def _get_job_metadata(
+        self, job: SlurmJob, steps_metadata: list[SlurmStepMetadata]
+    ) -> tuple[Path, SlurmJobMetadata]:
+        return self.scenario_root / "slurm-job.toml", SlurmJobMetadata(
             job_id=int(job.id),
-            job_name=job_name,
-            job_state=job_state,
-            elapsed_time_sec=time_sec,
+            name=steps_metadata[0].name,
+            state=steps_metadata[0].state,
+            exit_code=steps_metadata[0].exit_code,
+            start_time=steps_metadata[0].start_time,
+            end_time=steps_metadata[0].end_time,
+            elapsed_time_sec=steps_metadata[0].elapsed_time_sec,
+            job_steps=steps_metadata[1:],
             srun_cmd="n/a for single sbatch run",
             test_cmd="n/a for single sbatch run",
+            is_single_sbatch=True,
+            job_root=self.scenario_root.absolute(),
         )
-
-        job_res = self.scenario_root / "slurm-job.toml"
-        with job_res.open("w") as job_file:
-            toml.dump(job_meta.model_dump(), job_file)
-        logging.debug(f"Saved job metadata: {job_res}")

src/cloudai/systems/slurm/metadata.py …/cloudai/systems/slurm/slurm_metadata.pysrc/cloudai/systems/slurm/metadata.py renamed to src/cloudai/systems/slurm/slurm_metadata.py src/cloudai/systems/slurm/metadata.py renamed to src/cloudai/systems/slurm/slurm_metadata.py

@@ -14,7 +14,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pydantic import BaseModel
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, field_serializer
+
+
+class _SlurmStepMetadataBase(BaseModel):
+    """Represents the metadata of a Slurm job step."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    job_id: int
+    name: str
+    state: str
+    start_time: str
+    end_time: str
+    elapsed_time_sec: int
+    exit_code: str
+
+
+class SlurmStepMetadata(_SlurmStepMetadataBase):
+    """Represents the metadata of a Slurm job step."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    step_id: str
+    submit_line: str
+
+    @classmethod
+    def from_sacct_single_line(cls, line: str, delimiter: str) -> SlurmStepMetadata:
+        data = line.split(delimiter)
+        if len(data) < 8:
+            raise ValueError(f"Invalid line: {line}")
+
+        job_id, step_id = data[0].split(".") if "." in data[0] else (data[0], "")
+
+        return cls(
+            job_id=int(job_id),
+            step_id=step_id,
+            name=data[1],
+            state=data[2],
+            exit_code=data[3],
+            start_time=data[4],
+            end_time=data[5],
+            elapsed_time_sec=int(data[6]),
+            submit_line=data[7],
+        )
+
+
+class SlurmJobMetadata(_SlurmStepMetadataBase):
+    """Represents the metadata of a Slurm job."""
+
+    srun_cmd: str
+    test_cmd: str
+    is_single_sbatch: bool = False
+    job_root: Path
+    job_steps: list[SlurmStepMetadata]
+
+    @field_serializer("job_root")
+    def _path_serializer(self, v: Path) -> str:
+        return str(v)
 
 
 class MetadataSystem(BaseModel):

src/cloudai/systems/slurm/slurm_runner.py

@@ -25,7 +25,8 @@
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
 from .slurm_job import SlurmJob
-from .slurm_system import SlurmJobMetadata, SlurmSystem
+from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata
+from .slurm_system import SlurmSystem
 
 
 class SlurmRunner(BaseRunner):
@@ -59,25 +60,46 @@ def _submit_test(self, tr: TestRun) -> SlurmJob:
         logging.info(f"Submitted slurm job: {job_id}")
         return SlurmJob(tr, id=job_id)
 
-    async def job_completion_callback(self, job: BaseJob) -> None:
-        self.store_job_metadata(job)
+    def on_job_completion(self, job: BaseJob) -> None:
+        logging.debug(f"Job completion callback for job {job.id}")
+        self.store_job_metadata(cast(SlurmJob, job))
 
-    def store_job_metadata(self, job):
-        jb = cast(SlurmJob, job)
-        system = cast(SlurmSystem, self.system)
-        cmd_gen = cast(SlurmCommandGenStrategy, jb.test_run.test.test_template.command_gen_strategy)
-        res = None if self.mode == "dry-run" else system.get_job_status(jb)
-        job_name, job_state, time_sec = "unknown", "UNKNOWN", 0
-        if res:
-            job_name, job_state, time_sec = res[0], res[1], int(res[2])
-        job_meta = SlurmJobMetadata(
-            job_id=int(jb.id),
-            job_name=job_name,
-            job_state=job_state,
-            elapsed_time_sec=time_sec,
-            srun_cmd=cmd_gen.gen_srun_command(jb.test_run),
-            test_cmd=" ".join(cmd_gen.generate_test_command({}, {}, jb.test_run)),
+    def _mock_job_metadata(self) -> SlurmStepMetadata:
+        return SlurmStepMetadata(
+            job_id=0,
+            step_id="",
+            name="unknown",
+            state="UNKNOWN",
+            exit_code="0",
+            start_time="",
+            end_time="",
+            elapsed_time_sec=0,
+            submit_line="dry-run test",
+        )
+
+    def _get_job_metadata(
+        self, job: SlurmJob, steps_metadata: list[SlurmStepMetadata]
+    ) -> tuple[Path, SlurmJobMetadata]:
+        cmd_gen = cast(SlurmCommandGenStrategy, job.test_run.test.test_template.command_gen_strategy)
+        return job.test_run.output_path / "slurm-job.toml", SlurmJobMetadata(
+            job_id=int(job.id),
+            name=steps_metadata[0].name,
+            state=steps_metadata[0].state,
+            exit_code=steps_metadata[0].exit_code,
+            start_time=steps_metadata[0].start_time,
+            end_time=steps_metadata[0].end_time,
+            elapsed_time_sec=steps_metadata[0].elapsed_time_sec,
+            job_steps=steps_metadata[1:],
+            srun_cmd=cmd_gen.gen_srun_command(job.test_run),
+            test_cmd=" ".join(cmd_gen.generate_test_command({}, {}, job.test_run)),
+            job_root=job.test_run.output_path.absolute(),
         )
 
-        with open(jb.test_run.output_path / "slurm-job.toml", "w") as job_file:
+    def store_job_metadata(self, job: SlurmJob):
+        system = cast(SlurmSystem, self.system)
+        steps_metadata = [self._mock_job_metadata()] if self.mode == "dry-run" else system.get_job_status(job)
+        slurm_job_file, job_meta = self._get_job_metadata(job, steps_metadata)
+
+        logging.debug(f"Storing job metadata for job {job.id} to {slurm_job_file}")
+        with slurm_job_file.open("w") as job_file:
             toml.dump(job_meta.model_dump(), job_file)

src/cloudai/systems/slurm/slurm_system.py

@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 import re
 from pathlib import Path
@@ -25,20 +27,10 @@
 from cloudai.models.scenario import ReportConfig, parse_reports_spec
 from cloudai.util import CommandShell
 
+from .slurm_metadata import SlurmStepMetadata
 from .slurm_node import SlurmNode, SlurmNodeState
 
 
-class SlurmJobMetadata(BaseModel):
-    """Represents the metadata of a Slurm job."""
-
-    job_id: int
-    job_name: str
-    job_state: str
-    elapsed_time_sec: int
-    srun_cmd: str
-    test_cmd: str
-
-
 class DataRepositoryConfig(BaseModel):
     """Configuration for a data repository."""
 
@@ -308,9 +300,12 @@ def is_job_completed(self, job: BaseJob, retry_threshold: int = 3) -> bool:
 
         return False
 
-    def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> Optional[tuple[str, str, str]]:
+    def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> list[SlurmStepMetadata]:
         retry_count = 0
-        command = f"sacct -j {job.id} --format=JobName,State,ElapsedRAW --delimiter=',' -p --noheader"
+        command = (
+            f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine "
+            "--delimiter='|' -p --noheader"
+        )
 
         while retry_count < retry_threshold:
             stdout, stderr = self.cmd_shell.execute(command).communicate()
@@ -326,12 +321,9 @@ def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> Optional[tup
                 logging.error(error_message)
                 raise RuntimeError(error_message)
 
-            # sacct produces a single line per job, first line is for overall job
-            job_states = stdout.strip().splitlines()[0]
-            data = job_states.split(",")
-            return data[0], data[1], data[2]
+            return [SlurmStepMetadata.from_sacct_single_line(line, "|") for line in stdout.splitlines()]
 
-        return None
+        return []
 
     def kill(self, job: BaseJob) -> None:
         """

tests/conftest.py

@@ -32,7 +32,7 @@ def create_autospec_dataclass(dataclass: type) -> Mock:
 def cleanup():
     yield
 
-    for f in {"env_vars.sh", "hostfile.txt"}:
+    for f in {"env_vars.sh", "hostfile.txt", "start_server_wrapper.sh"}:
         (Path.cwd() / f).unlink(missing_ok=True)
 
 

tests/test_single_sbatch_runner.py

@@ -43,7 +43,7 @@ def nccl_tr(slurm_system: SlurmSystem) -> TestRun:
         ),
         num_nodes=2,
         nodes=[],
-        output_path=slurm_system.output_path,
+        output_path=slurm_system.output_path / "nccl_test",
     )
     tr.test.test_template.command_gen_strategy = NcclTestSlurmCommandGenStrategy(slurm_system, {})
     return tr
@@ -61,7 +61,7 @@ def sleep_tr(slurm_system: SlurmSystem) -> TestRun:
         ),
         num_nodes=1,
         nodes=[],
-        output_path=slurm_system.output_path,
+        output_path=slurm_system.output_path / "sleep_test",
     )
     tr.test.test_template.command_gen_strategy = SleepSlurmCommandGenStrategy(slurm_system, {})
     tr.output_path.mkdir(parents=True, exist_ok=True)
@@ -480,8 +480,12 @@ def test_store_job_metadata(nccl_tr: TestRun, slurm_system: SlurmSystem) -> None
     assert out_file.exists()
     sjm = SlurmJobMetadata.model_validate(toml.load(out_file))
     assert sjm.job_id == 1
+    assert sjm.is_single_sbatch is True
     assert sjm.srun_cmd == "n/a for single sbatch run"
     assert sjm.test_cmd == "n/a for single sbatch run"
+    assert sjm.job_root == runner.scenario_root.absolute()
+
+    assert sjm == SlurmJobMetadata.model_validate(toml.loads(toml.dumps(sjm.model_dump())))
 
 
 def test_pre_test(nccl_tr: TestRun, sleep_tr: TestRun, slurm_system: SlurmSystem) -> None: