Remove TestTemplate concept

Previously it help all the strategies, but now it is obsolete.

Author: amaslenn

src/cloudai/_core/test.py

@@ -18,8 +18,6 @@
 
 from typing import TYPE_CHECKING, Dict, List, Union
 
-from .test_template import TestTemplate
-
 if TYPE_CHECKING:
     from cloudai.models.workload import TestDefinition
 
@@ -29,15 +27,13 @@ class Test:
 
     __test__ = False
 
-    def __init__(self, test_definition: TestDefinition, test_template: TestTemplate) -> None:
+    def __init__(self, test_definition: TestDefinition) -> None:
         """
         Initialize a Test instance.
 
         Args:
             test_definition (TestDefinition): The test definition object.
-            test_template (TestTemplate): The test template object
         """
-        self.test_template = test_template
         self.test_definition = test_definition
 
     @property

src/cloudai/_core/test_template.py

@@ -41,7 +41,6 @@
 from ._core.system import System
 from ._core.test import Test
 from ._core.test_scenario import METRIC_ERROR, TestDependency, TestRun, TestScenario
-from ._core.test_template import TestTemplate
 from .configurator.base_agent import BaseAgent
 from .configurator.cloudai_gym import CloudAIGymEnv
 from .configurator.grid_search import GridSearchAgent
@@ -95,7 +94,6 @@
     "TestScenario",
     "TestScenarioParser",
     "TestScenarioParsingError",
-    "TestTemplate",
     "case_name",
     "format_validation_error",
 ]

src/cloudai/core.py

@@ -55,9 +55,7 @@ def gen_exec_command(self) -> str:
         """
         env_vars = self.final_env_vars
         cmd_args = flatten_dict(self.test_run.test.cmd_args)
-        lsf_args = self._parse_lsf_args(
-            self.test_run.test.test_template.__class__.__name__, env_vars, cmd_args, self.test_run
-        )
+        lsf_args = self._parse_lsf_args(self.test_run.test.name, env_vars, cmd_args, self.test_run)
 
         bsub_command = self._gen_bsub_command(lsf_args, env_vars, cmd_args, self.test_run)
 

src/cloudai/systems/lsf/lsf_command_gen_strategy.py

@@ -118,7 +118,7 @@ def gen_srun_command(self) -> str:
         return self._gen_srun_command()
 
     def job_name_prefix(self) -> str:
-        return self.test_run.test.test_template.__class__.__name__
+        return self.test_run.test.name
 
     def job_name(self) -> str:
         job_name_prefix = self.job_name_prefix()

src/cloudai/systems/slurm/slurm_command_gen_strategy.py

@@ -21,14 +21,7 @@
 import toml
 from pydantic import ValidationError
 
-from .core import (
-    Registry,
-    System,
-    Test,
-    TestConfigParsingError,
-    TestTemplate,
-    format_validation_error,
-)
+from .core import Registry, System, Test, TestConfigParsingError, format_validation_error
 from .models.workload import TestDefinition
 
 
@@ -92,31 +85,15 @@ def load_test_definition(self, data: dict) -> TestDefinition:
 
         return test_def
 
-    def _get_test_template(self, name: str, tdef: TestDefinition) -> TestTemplate:
-        """
-        Dynamically retrieves the appropriate TestTemplate subclass based on the given name.
-
-        Args:
-            name (str): The name of the test template.
-            tdef (TestDefinition): The test definition.
-
-        Returns:
-            Type[TestTemplate]: A subclass of TestTemplate corresponding to the given name.
-        """
-        obj = TestTemplate(system=self.system)
-        return obj
-
     def _parse_data(self, data: Dict[str, Any]) -> Test:
         """
         Parse data for a Test object.
 
         Args:
             data (Dict[str, Any]): Data from a source (e.g., a TOML file).
-            strict (bool): Whether to enforce strict validation for test definition.
 
         Returns:
             Test: Parsed Test object.
         """
         test_def = self.load_test_definition(data)
-        test_template = self._get_test_template(test_def.test_template_name, test_def)
-        return Test(test_definition=test_def, test_template=test_template)
+        return Test(test_definition=test_def)

src/cloudai/test_parser.py

@@ -18,7 +18,7 @@
 import logging
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Set, Tuple, Type
+from typing import Any, Dict, List, Optional, Set, Type
 
 import toml
 from pydantic import ValidationError
@@ -186,9 +186,7 @@ def _create_test_run(
         Raises:
             ValueError: If the test or nodes are not found within the system.
         """
-        original_test, tdef = self._prepare_tdef(test_info)
-
-        test = Test(test_definition=tdef, test_template=original_test.test_template)
+        test = Test(test_definition=self._prepare_tdef(test_info))
 
         hooks = [hook for hook in [pre_test, post_test] if hook is not None]
         total_time_limit = calculate_total_time_limit(test_hooks=hooks, time_limit=test_info.time_limit)
@@ -210,7 +208,7 @@ def _create_test_run(
 
         return tr
 
-    def _prepare_tdef(self, test_info: TestRunModel) -> Tuple[Test, TestDefinition]:
+    def _prepare_tdef(self, test_info: TestRunModel) -> TestDefinition:
         tp = TestParser([self.file_path], self.system)
         tp.current_file = self.file_path
 
@@ -231,4 +229,4 @@ def _prepare_tdef(self, test_info: TestRunModel) -> Tuple[Test, TestDefinition]:
                 f"Cannot configure test case '{test_info.id}' with both 'test_name' and 'test_template_name'."
             )
 
-        return test, test.test_definition
+        return test.test_definition

src/cloudai/test_scenario_parser.py

@@ -23,7 +23,7 @@
 import toml
 import yaml
 
-from cloudai.core import CommandGenStrategy, Test, TestRun, TestTemplate
+from cloudai.core import CommandGenStrategy, Test, TestRun
 from cloudai.models.scenario import TestRunDetails
 from cloudai.models.workload import CmdArgs, TestDefinition
 from cloudai.systems.kubernetes import KubernetesSystem
@@ -125,8 +125,7 @@ def base_tr(slurm_system: SlurmSystem) -> TestRun:
     return TestRun(
         name="tr-name",
         test=Test(
-            test_definition=TestDefinition(name="n", description="d", test_template_name="tt", cmd_args=CmdArgs()),
-            test_template=TestTemplate(slurm_system),
+            test_definition=TestDefinition(name="n", description="d", test_template_name="tt", cmd_args=CmdArgs())
         ),
         num_nodes=1,
         nodes=[],
@@ -161,8 +160,7 @@ def benchmark_tr(slurm_system: SlurmSystem) -> TestRun:
         test_template_name="NcclTest",
         cmd_args=NCCLCmdArgs(docker_image_url="fake://url/nccl"),
     )
-    test_template = TestTemplate(system=slurm_system)
-    test = Test(test_definition=test_definition, test_template=test_template)
+    test = Test(test_definition=test_definition)
     tr = TestRun(name="benchmark", test=test, num_nodes=1, nodes=["node1"], iterations=3)
     create_test_directories(slurm_system, tr)
     return tr
@@ -178,8 +176,7 @@ def dse_tr(slurm_system: SlurmSystem) -> TestRun:
         extra_env_vars={"VAR1": ["value1", "value2"]},
         agent_steps=12,
     )
-    test_template = TestTemplate(system=slurm_system)
-    test = Test(test_definition=test_definition, test_template=test_template)
+    test = Test(test_definition=test_definition)
 
     tr = TestRun(name="dse", test=test, num_nodes=1, nodes=["node1"], iterations=12)
     create_test_directories(slurm_system, tr)

tests/conftest.py

@@ -17,7 +17,7 @@
 
 import pytest
 
-from cloudai.core import Test, TestRun, TestTemplate
+from cloudai.core import Test, TestRun
 from cloudai.systems.kubernetes import KubernetesSystem
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition, NcclTestKubernetesJsonGenStrategy
 
@@ -27,8 +27,7 @@ class TestNcclTestKubernetesJsonGenStrategy:
     def basic_test_run(self, kubernetes_system: KubernetesSystem) -> TestRun:
         cmd_args = NCCLCmdArgs.model_validate({"subtest_name": "all_reduce_perf", "docker_image_url": "fake_image_url"})
         nccl = NCCLTestDefinition(name="name", description="desc", test_template_name="NcclTest", cmd_args=cmd_args)
-        test_template = TestTemplate(kubernetes_system)
-        t = Test(test_definition=nccl, test_template=test_template)
+        t = Test(test_definition=nccl)
         return TestRun(name="t1", test=t, nodes=["node1", "node2"], num_nodes=2)
 
     @pytest.fixture
@@ -41,8 +40,7 @@ def test_run_with_env_vars(self, kubernetes_system: KubernetesSystem) -> TestRun
             cmd_args=cmd_args,
             extra_env_vars={"TEST_VAR": "test_value", "LIST_VAR": ["item1", "item2"]},
         )
-        test_template = TestTemplate(kubernetes_system)
-        t = Test(test_definition=nccl, test_template=test_template)
+        t = Test(test_definition=nccl)
         return TestRun(name="t1", test=t, nodes=["node1"], num_nodes=1)
 
     @pytest.fixture
@@ -64,8 +62,7 @@ def test_run_with_extra_args(self, kubernetes_system: KubernetesSystem) -> TestR
             cmd_args=cmd_args,
             extra_cmd_args={"extra-flag": "value"},
         )
-        test_template = TestTemplate(kubernetes_system)
-        t = Test(test_definition=nccl, test_template=test_template)
+        t = Test(test_definition=nccl)
         return TestRun(name="t1", test=t, nodes=["node1"], num_nodes=1)
 
     def json_gen_strategy(

tests/json_gen_strategy/test_nccl_kubernetes_json_gen_strategy.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import Mock
 
 from cloudai.core import Test, TestRun
 from cloudai.systems.runai.runai_system import RunAISystem
@@ -28,7 +27,7 @@ def json_gen_strategy(self, runai_system: RunAISystem, tr: TestRun) -> NcclTestR
     def test_gen_json(self, runai_system: RunAISystem) -> None:
         cmd_args = NCCLCmdArgs.model_validate({"subtest_name": "all_reduce_perf", "docker_image_url": "fake_image_url"})
         nccl = NCCLTestDefinition(name="name", description="desc", test_template_name="tt", cmd_args=cmd_args)
-        t = Test(test_definition=nccl, test_template=Mock())
+        t = Test(test_definition=nccl)
         tr = TestRun(name="t1", test=t, nodes=["node1", "node2"], num_nodes=2)
         json_payload = self.json_gen_strategy(runai_system, tr).gen_json()
 
@@ -63,7 +62,7 @@ def test_gen_json_with_cmd_args(self, runai_system: RunAISystem) -> None:
             }
         )
         nccl = NCCLTestDefinition(name="name", description="desc", test_template_name="tt", cmd_args=cmd_args)
-        t = Test(test_definition=nccl, test_template=Mock())
+        t = Test(test_definition=nccl)
         tr = TestRun(name="t1", test=t, nodes=["node1", "node2"], num_nodes=2)
 
         json_payload = self.json_gen_strategy(runai_system, tr).gen_json()
@@ -107,7 +106,7 @@ def test_gen_json_with_extra_cmd_args(self, runai_system: RunAISystem) -> None:
             cmd_args=cmd_args,
             extra_cmd_args={"--extra-arg": "value"},
         )
-        t = Test(test_definition=nccl, test_template=Mock())
+        t = Test(test_definition=nccl)
         tr = TestRun(name="t1", test=t, nodes=["node1", "node2"], num_nodes=2)
 
         json_payload = self.json_gen_strategy(runai_system, tr).gen_json()