Merge pull request #716 from NVIDIA/am/upds

Simplify internal hierarchy of classes

Author: amaslenn

src/cloudai/__init__.py

@@ -18,7 +18,6 @@
     Registry,
     Runner,
     System,
-    Test,
     TestDefinition,
     TestRun,
     TestScenario,
@@ -29,7 +28,6 @@
     "Registry",
     "Runner",
     "System",
-    "Test",
     "TestDefinition",
     "TestRun",
     "TestScenario",

src/cloudai/_core/base_reporter.py

@@ -67,7 +67,7 @@ def load_test_runs(self):
                         tr_file = tr.output_path / "test-run.toml"
                         if tr_file.exists():
                             tr_file = toml.load(tr_file)
-                            tr.test.test_definition = tr.test.test_definition.model_validate(tr_file["test_definition"])
+                            tr.test = tr.test.model_validate(tr_file["test_definition"])
                         self.trs.append(copy.deepcopy(tr))
                 else:
                     tr.current_iteration = int(iter.name)

src/cloudai/_core/base_runner.py

@@ -25,6 +25,7 @@
 from .command_gen_strategy import CommandGenStrategy
 from .exceptions import JobFailureError, JobSubmissionError
 from .job_status_result import JobStatusResult
+from .json_gen_strategy import JsonGenStrategy
 from .registry import Registry
 from .system import System
 from .test_scenario import TestRun, TestScenario
@@ -288,7 +289,7 @@ def get_job_status(self, job: BaseJob) -> JobStatusResult:
             JobStatusResult: The result containing the job status and an optional error message.
         """
         runner_job_status_result = self.get_runner_job_status(job)
-        workload_run_results = job.test_run.test.test_definition.was_run_successful(job.test_run)
+        workload_run_results = job.test_run.test.was_run_successful(job.test_run)
         if not runner_job_status_result.is_successful:
             return runner_job_status_result
         if not workload_run_results.is_successful:
@@ -375,5 +376,9 @@ async def delayed_kill_job(self, job: BaseJob, delay: int = 0):
         self.system.kill(job)
 
     def get_cmd_gen_strategy(self, system: System, test_run: TestRun) -> CommandGenStrategy:
-        strategy_cls = Registry().get_command_gen_strategy(type(system), type(test_run.test.test_definition))
+        strategy_cls = Registry().get_command_gen_strategy(type(system), type(test_run.test))
+        return strategy_cls(system, test_run)
+
+    def get_json_gen_strategy(self, system: System, test_run: TestRun) -> JsonGenStrategy:
+        strategy_cls = Registry().get_json_gen_strategy(type(system), type(test_run.test))
         return strategy_cls(system, test_run)

src/cloudai/_core/command_gen_strategy.py

@@ -18,16 +18,15 @@
 
 from .system import System
 from .test_scenario import TestRun
-from .test_template_strategy import TestTemplateStrategy
 
 
-class CommandGenStrategy(TestTemplateStrategy, ABC):
+class CommandGenStrategy(ABC):
     """Abstract base class defining the interface for command generation strategies across different systems."""
 
     TEST_RUN_DUMP_FILE_NAME: str = "test-run.toml"
 
     def __init__(self, system: System, test_run: TestRun) -> None:
-        super().__init__(system)
+        self.system = system
         self.test_run = test_run
         self._final_env_vars: dict[str, str | list[str]] = {}
 

src/cloudai/_core/grader.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,6 +19,8 @@
 from pathlib import Path
 from typing import Dict, List
 
+from .registry import Registry
+from .system import System
 from .test_scenario import TestRun, TestScenario
 
 
@@ -31,8 +33,9 @@ class Grader:
         logger (logging.Logger): Logger for the class, used to log messages related to the grading process.
     """
 
-    def __init__(self, output_path: Path) -> None:
+    def __init__(self, output_path: Path, system: System) -> None:
         self.output_path = output_path
+        self.system = system
 
     def grade(self, test_scenario: TestScenario) -> str:
         """
@@ -80,9 +83,11 @@ def _get_perfs_from_subdirs(self, directory_path: Path, tr: TestRun) -> List[flo
             List[float]: A list of performance values.
         """
         perfs = []
+
         for subdir in directory_path.iterdir():
             if subdir.is_dir() and subdir.name.isdigit():
-                perf = tr.test.test_template.grade(subdir, tr.ideal_perf)
+                grading_strategy = Registry().get_grading_strategy(type(self.system), type(tr.test))()
+                perf = grading_strategy.grade(subdir, tr.ideal_perf)
                 perfs.append(perf)
         return perfs
 

src/cloudai/_core/json_gen_strategy.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,17 +18,21 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
+from .system import System
 from .test_scenario import TestRun
-from .test_template_strategy import TestTemplateStrategy
 
 
-class JsonGenStrategy(TestTemplateStrategy, ABC):
+class JsonGenStrategy(ABC):
     """
     Abstract base class for generating Kubernetes job specifications based on system and test parameters.
 
     It specifies how to generate JSON job specifications based on system and test parameters.
     """
 
+    def __init__(self, system: System, test_run: TestRun) -> None:
+        self.system = system
+        self.test_run = test_run
+
     def sanitize_k8s_job_name(self, job_name: str) -> str:
         """
         Sanitize the job name to ensure it follows Kubernetes naming rules.
@@ -51,7 +55,7 @@ def sanitize_k8s_job_name(self, job_name: str) -> str:
         return sanitized_name[:253]
 
     @abstractmethod
-    def gen_json(self, tr: TestRun) -> Dict[Any, Any]:
+    def gen_json(self) -> Dict[Any, Any]:
         """
         Generate the Kubernetes job specification based on the given parameters.
 

src/cloudai/_core/registry.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable, ClassVar, List, Set, Tuple, Type, Union
+from typing import TYPE_CHECKING, Callable, ClassVar, List, Set, Tuple, Type
 
 if TYPE_CHECKING:
     from ..configurator.base_agent import BaseAgent
@@ -49,16 +49,6 @@ class Registry(metaclass=Singleton):
     """Registry for implementations mappings."""
 
     runners_map: ClassVar[dict[str, Type[BaseRunner]]] = {}
-    strategies_map: ClassVar[
-        dict[
-            Tuple[
-                Type[Union[JsonGenStrategy, GradingStrategy]],
-                Type[System],
-                Type[TestDefinition],
-            ],
-            Type[Union[JsonGenStrategy, GradingStrategy]],
-        ]
-    ] = {}
     installers_map: ClassVar[dict[str, Type[BaseInstaller]]] = {}
     systems_map: ClassVar[dict[str, Type[System]]] = {}
     test_definitions_map: ClassVar[dict[str, Type[TestDefinition]]] = {}
@@ -68,6 +58,8 @@ class Registry(metaclass=Singleton):
     report_configs: ClassVar[dict[str, ReportConfig]] = {}
     reward_functions_map: ClassVar[dict[str, RewardFunction]] = {}
     command_gen_strategies_map: ClassVar[dict[tuple[Type[System], Type[TestDefinition]], Type[CommandGenStrategy]]] = {}
+    json_gen_strategies_map: ClassVar[dict[tuple[Type[System], Type[TestDefinition]], Type[JsonGenStrategy]]] = {}
+    grading_strategies_map: ClassVar[dict[Tuple[Type[System], Type[TestDefinition]], Type[GradingStrategy]]] = {}
 
     def add_runner(self, name: str, value: Type[BaseRunner]) -> None:
         """
@@ -94,30 +86,23 @@ def update_runner(self, name: str, value: Type[BaseRunner]) -> None:
         """
         self.runners_map[name] = value
 
-    def add_strategy(
-        self,
-        strategy_interface: Type[Union[JsonGenStrategy, GradingStrategy]],
-        system_types: List[Type[System]],
-        definition_types: List[Type[TestDefinition]],
-        strategy: Type[Union[JsonGenStrategy, GradingStrategy]],
+    def add_grading_strategy(
+        self, system_type: Type[System], tdef_type: Type[TestDefinition], strategy: Type[GradingStrategy]
     ) -> None:
-        for system_type in system_types:
-            for def_type in definition_types:
-                key = (strategy_interface, system_type, def_type)
-                if key in self.strategies_map:
-                    raise ValueError(f"Duplicating implementation for '{key}', use 'update()' for replacement.")
-                self.update_strategy(key, strategy)
-
-    def update_strategy(
-        self,
-        key: Tuple[
-            Type[Union[JsonGenStrategy, GradingStrategy]],
-            Type[System],
-            Type[TestDefinition],
-        ],
-        value: Type[Union[JsonGenStrategy, GradingStrategy]],
+        key = (system_type, tdef_type)
+        if key in self.grading_strategies_map:
+            raise ValueError(f"Duplicating implementation for '{key}', use 'update()' for replacement.")
+        self.update_grading_strategy(key, strategy)
+
+    def update_grading_strategy(
+        self, key: Tuple[Type[System], Type[TestDefinition]], value: Type[GradingStrategy]
     ) -> None:
-        self.strategies_map[key] = value
+        self.grading_strategies_map[key] = value
+
+    def get_grading_strategy(self, system_type: Type[System], tdef_type: Type[TestDefinition]) -> Type[GradingStrategy]:
+        if (system_type, tdef_type) not in self.grading_strategies_map:
+            raise KeyError(f"Grading gen strategy for '{system_type.__name__}, {tdef_type.__name__}' not found.")
+        return self.grading_strategies_map[(system_type, tdef_type)]
 
     def add_installer(self, name: str, value: Type[BaseInstaller]) -> None:
         """
@@ -274,3 +259,25 @@ def get_command_gen_strategy(
         if (system_type, tdef_type) not in self.command_gen_strategies_map:
             raise KeyError(f"Command gen strategy for '{system_type.__name__}, {tdef_type.__name__}' not found.")
         return self.command_gen_strategies_map[(system_type, tdef_type)]
+
+    def add_json_gen_strategy(
+        self, system_type: Type[System], tdef_type: Type[TestDefinition], value: Type[JsonGenStrategy]
+    ) -> None:
+        if (system_type, tdef_type) in self.json_gen_strategies_map:
+            raise ValueError(
+                f"Duplicating implementation for '{system_type.__name__}, {tdef_type.__name__}', use 'update()' "
+                "for replacement."
+            )
+        self.update_json_gen_strategy(system_type, tdef_type, value)
+
+    def update_json_gen_strategy(
+        self, system_type: Type[System], tdef_type: Type[TestDefinition], value: Type[JsonGenStrategy]
+    ) -> None:
+        self.json_gen_strategies_map[(system_type, tdef_type)] = value
+
+    def get_json_gen_strategy(
+        self, system_type: Type[System], tdef_type: Type[TestDefinition]
+    ) -> Type[JsonGenStrategy]:
+        if (system_type, tdef_type) not in self.json_gen_strategies_map:
+            raise KeyError(f"JSON gen strategy for '{system_type.__name__}, {tdef_type.__name__}' not found.")
+        return self.json_gen_strategies_map[(system_type, tdef_type)]