Merge pull request #661 from NVIDIA/am/scenario-reports-cfg

Configure reports via scenario config

Author: amaslenn

doc/reporting.md

@@ -10,6 +10,13 @@ Per-test reports are linked to a particular workload type (e.g. `NcclTest`). All
 To list all available reports, one can use `cloudai list-reports` command. Use verbose output to also print report configurations.
 
 
+## Notes and general flow
+1. All reports should be registered via `Registry()` (`.add_report()` or `.add_scenario_report()`).
+1. Scenario reports are configurable via system config (Slurm-only for now) and scenario config.
+1. Configuration in a scenario config has the highest priority. Next, system config is checked. Then it defaults to report config from the registry.
+1. Then report is generated (or not) according to this final config.
+
+
 ## Enable, disable and configure reports
 **NOTE** Only scenario-level reports can be configured today.
 

src/cloudai/_core/test_scenario.py

@@ -26,6 +26,7 @@
 from .test_template_strategy import TestTemplateStrategy
 
 if TYPE_CHECKING:
+    from ..models.scenario import ReportConfig
     from .report_generation_strategy import ReportGenerationStrategy
     from .test import Test
 
@@ -184,18 +185,26 @@ class TestScenario:
 
     __test__ = False
 
-    def __init__(self, name: str, test_runs: List[TestRun], job_status_check: bool = True) -> None:
+    def __init__(
+        self,
+        name: str,
+        test_runs: List[TestRun],
+        job_status_check: bool = True,
+        reports: dict[str, ReportConfig] | None = None,
+    ) -> None:
         """
         Initialize a TestScenario instance.
 
         Args:
             name (str): Name of the test scenario.
             test_runs (List[TestRun]): List of tests in the scenario with custom run options.
             job_status_check (bool): Flag indicating whether to check the job status or not.
+            reports (Optional[dict[str, ReportConfig]]): Reports to be generated for the scenario.
         """
         self.name = name
         self.test_runs = test_runs
         self.job_status_check = job_status_check
+        self.reports = reports or {}
 
     def __repr__(self) -> str:
         """

src/cloudai/cli/handlers.py

@@ -157,7 +157,9 @@ def generate_reports(system: System, test_scenario: TestScenario, result_dir: Pa
         logging.debug(f"Generating report '{name}' ({reporter_class.__name__})")
 
         cfg = registry.report_configs.get(name, ReportConfig(enable=False))
-        if isinstance(system, SlurmSystem) and system.reports and name in system.reports:
+        if scenario_cfg := test_scenario.reports.get(name):
+            cfg = scenario_cfg
+        elif isinstance(system, SlurmSystem) and system.reports and name in system.reports:
             cfg = system.reports[name]
         logging.debug(f"Report '{name}' config is: {cfg.model_dump_json(indent=None)}")
 

src/cloudai/models/scenario.py

@@ -155,6 +155,7 @@ class TestScenarioModel(BaseModel):
     tests: list[TestRunModel] = Field(alias="Tests", min_length=1)
     pre_test: Optional[str] = None
     post_test: Optional[str] = None
+    reports: dict[str, ReportConfig] = Field(default_factory=dict)
 
     @model_validator(mode="after")
     def check_no_self_dependency(self):
@@ -188,6 +189,11 @@ def check_all_dependencies_are_known(self):
 
         return self
 
+    @field_validator("reports", mode="before")
+    @classmethod
+    def parse_reports(cls, value: dict[str, Any] | None) -> dict[str, ReportConfig] | None:
+        return parse_reports_spec(value)
+
 
 class TestRunDetails(BaseModel):
     """

src/cloudai/test_scenario_parser.py

@@ -167,6 +167,7 @@ def _parse_data(self, data: Dict[str, Any]) -> TestScenario:
             name=ts_model.name,
             test_runs=list(test_runs_by_id.values()),
             job_status_check=ts_model.job_status_check,
+            reports=ts_model.reports,
         )
 
     def _create_test_run(

src/cloudai/workloads/nccl_test/nccl_comparisson_report.py

@@ -47,6 +47,7 @@
 class NcclComparissonReportConfig(ReportConfig):
     """Configuration for NCCL comparisson report."""
 
+    enable: bool = True
     group_by: list[str] = Field(default_factory=lambda: ["subtest_name"])
 
 

tests/test_reporter.py

@@ -208,3 +208,49 @@ def test_disabled_on_system_level(self, slurm_system: SlurmSystem) -> None:
         slurm_system.reports = {"sr1": ReportConfig(enable=False)}
         generate_reports(slurm_system, TestScenario(name="ts", test_runs=[]), slurm_system.output_path)
         assert MY_REPORT_CALLED == 0
+
+
+class TestGenerateReportPriority:
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        reg = Registry()
+        orig_reports = copy.deepcopy(reg.scenario_reports)
+        reg.scenario_reports.clear()
+
+        global MY_REPORT_CALLED
+        MY_REPORT_CALLED = 0
+
+        yield
+
+        reg.scenario_reports.clear()
+        reg.scenario_reports.update(orig_reports)
+
+    def test_non_registered_report_is_ignored(self, slurm_system: SlurmSystem) -> None:
+        generate_reports(slurm_system, TestScenario(name="ts", test_runs=[]), slurm_system.output_path)
+        assert MY_REPORT_CALLED == 0
+
+    def test_report_is_enabled_on_system_level(self, slurm_system: SlurmSystem) -> None:
+        Registry().add_scenario_report("sr1", MyReporter, ReportConfig(enable=True))
+        slurm_system.reports = {"sr1": ReportConfig(enable=True)}
+        generate_reports(slurm_system, TestScenario(name="ts", test_runs=[]), slurm_system.output_path)
+        assert MY_REPORT_CALLED == 1
+
+    def test_report_is_enabled_on_scenario_level(self, slurm_system: SlurmSystem) -> None:
+        Registry().add_scenario_report("sr1", MyReporter, ReportConfig(enable=True))
+        slurm_system.reports = {}
+        generate_reports(
+            slurm_system,
+            TestScenario(name="ts", test_runs=[], reports={"sr1": ReportConfig(enable=True)}),
+            slurm_system.output_path,
+        )
+        assert MY_REPORT_CALLED == 1
+
+    def test_report_scenario_has_highest_priority(self, slurm_system: SlurmSystem) -> None:
+        Registry().add_scenario_report("sr1", MyReporter, ReportConfig(enable=True))
+        slurm_system.reports = {"sr1": ReportConfig(enable=False)}
+        generate_reports(
+            slurm_system,
+            TestScenario(name="ts", test_runs=[], reports={"sr1": ReportConfig(enable=True)}),
+            slurm_system.output_path,
+        )
+        assert MY_REPORT_CALLED == 1

tests/test_test_scenario_parser.py

@@ -15,13 +15,16 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, cast
 from unittest.mock import create_autospec
 
 import pytest
+import toml
 
 from cloudai.core import Test, TestDefinition, TestRun, TestScenario
+from cloudai.models.scenario import TestScenarioModel
 from cloudai.test_scenario_parser import calculate_total_time_limit
+from cloudai.workloads.nccl_test.nccl_comparisson_report import NcclComparissonReportConfig
 
 
 class DummyTestRun(TestRun):
@@ -71,3 +74,27 @@ def test_calculate_total_time_limit(
     test_hooks: List[TestScenario], time_limit: Optional[str], expected: Optional[str]
 ) -> None:
     assert calculate_total_time_limit(test_hooks, time_limit) == expected
+
+
+def test_report_spec_is_parsed() -> None:
+    model = TestScenarioModel.model_validate(
+        toml.loads("""
+    name = "scenario"
+
+    [reports]
+    nccl_comparisson = { enable = false, group_by = ["my_field"] }
+
+    [[Tests]]
+    id = "1"
+    num_nodes = 2
+
+    name = "name"
+    description = "desc"
+    test_template_name = "NcclTest"
+    """)
+    )
+
+    assert len(model.reports) == 1
+    cfg = cast(NcclComparissonReportConfig, model.reports["nccl_comparisson"])
+    assert cfg.enable is False
+    assert cfg.group_by == ["my_field"]