Generate test-run.toml in Dynamo k8s runs

Author: amaslenn

src/cloudai/_core/json_gen_strategy.py

@@ -18,6 +18,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
+import toml
+
 from .system import System
 from .test_scenario import TestRun
 
@@ -29,6 +31,8 @@ class JsonGenStrategy(ABC):
     It specifies how to generate JSON job specifications based on system and test parameters.
     """
 
+    TEST_RUN_DUMP_FILE_NAME: str = "test-run.toml"
+
     def __init__(self, system: System, test_run: TestRun) -> None:
         self.system = system
         self.test_run = test_run
@@ -54,6 +58,14 @@ def sanitize_k8s_job_name(self, job_name: str) -> str:
         sanitized_name = re.sub(r"[^a-z0-9]+$", "", sanitized_name)
         return sanitized_name[:253]
 
+    def store_test_run(self) -> None:
+        from cloudai.models.scenario import TestRunDetails
+
+        test_cmd, srun_cmd = ("", "n/a")
+        with (self.test_run.output_path / self.TEST_RUN_DUMP_FILE_NAME).open("w") as f:
+            trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=srun_cmd)
+            toml.dump(trd.model_dump(), f)
+
     @abstractmethod
     def gen_json(self) -> Dict[Any, Any]:
         """

src/cloudai/systems/kubernetes/kubernetes_runner.py

@@ -42,6 +42,10 @@ def _submit_test(self, tr: TestRun) -> KubernetesJob:
 
         return job
 
+    def on_job_submit(self, tr: TestRun) -> None:
+        json_gen = self.get_json_gen_strategy(self.system, tr)
+        json_gen.store_test_run()
+
     def on_job_completion(self, job: BaseJob) -> None:
         k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)
         k_job = cast(KubernetesJob, job)