Feat/training report metadata

doc/reporting.rst

@@ -8,6 +8,12 @@ This chapter describes the reporting system in CloudAI. In this chapter, we will
 - :ref:`Enabling, Disabling and Configuring Reports <enabling-disabling-and-configuring-reports>`
 - :ref:`Reporting Registration <reporting-registration>`
 - :ref:`Reporting Configuration Implementation <reporting-configuration-implementation>`
+- :doc:`Reports <reports>`
+
+.. toctree::
+   :hidden:
+
+   reports
 
 .. _overview:
 

doc/reports.rst

@@ -0,0 +1,9 @@
+Reports
+=======
+
+CloudAI report schemas describe the generated output for supported workload categories.
+
+.. toctree::
+   :maxdepth: 1
+
+   training-report-schema

doc/training-report-schema.rst

@@ -0,0 +1,67 @@
+Training Report Schema
+======================
+
+``training_report.json`` is the unified training output for NeMoRun, MegatronRun, and Megatron-Bridge.
+
+Versioning
+----------
+
+The top-level ``schema_version`` field uses ``MAJOR.MINOR`` versioning:
+
+- Increment **MAJOR** when removing, renaming, or incompatibly changing a field.
+- Increment **MINOR** when adding an optional field or otherwise making a backward-compatible schema change.
+- Do not change the schema version for implementation fixes that leave the JSON contract unchanged.
+
+Consumers should reject unsupported major versions and tolerate unknown fields within a supported major version.
+
+Version History
+---------------
+
+1.0 — 2026-07-14
+~~~~~~~~~~~~~~~~
+
+Added to ``root``:
+
+- ``schema_version``: ``str``
+
+Added to ``root.config``:
+
+- ``test_id``: ``str``
+- ``test_name``: ``str``
+- ``description``: ``str``
+- ``test_scenario_name``: ``str``
+- ``system_path``: ``str``
+- ``tests_dir_path``: ``str``
+- ``test_scenario_path``: ``str``
+- ``container_image``: ``str``
+- ``cloudai_execution_node``: ``str``
+- ``env_vars``: ``dict[str, Any]``
+- ``gpus_per_node``: ``Optional[int]``
+- ``nodes``: ``list[str]``
+- ``clique_size``: ``Optional[int]``
+- ``fp8``: ``Optional[str]``
+- ``fp8_recipe``: ``Optional[str]``
+- ``expert_tensor_parallel_size``: ``int``
+
+Training Report Models
+----------------------
+
+.. autoclass:: cloudai.report_generator.training.models.TrainingResults
+   :members:
+   :exclude-members: __init__
+
+.. autoclass:: cloudai.report_generator.training.models.TrainingConfig
+   :members:
+   :exclude-members: __init__
+
+.. autoclass:: cloudai.report_generator.training.models.TrainingStep
+   :members:
+   :exclude-members: __init__
+
+.. autoclass:: cloudai.report_generator.training.models.StepAggregation
+   :members:
+   :exclude-members: __init__, from_steps
+
+.. autoclass:: cloudai.report_generator.training.models.MetricStats
+   :members:
+   :exclude-members: __init__, from_values

src/cloudai/_core/test_scenario.py

@@ -240,6 +240,15 @@ def apply_params_set(self, action: dict[str, Any], env_params: dict[str, Any] |
         return new_tr
 
 
+@dataclass(frozen=True)
+class ConfigPaths:
+    """Source paths used to build a system and test scenario."""
+
+    system_path: Path
+    tests_dir_path: Optional[Path]
+    test_scenario_path: Path
+
+
 @dataclass
 class TestScenario:
     """
@@ -250,6 +259,7 @@ class TestScenario:
         tests (List[Test]): Tests in the scenario.
         job_status_check (bool): Flag indicating whether to check the job status or not.
         reports (dict[str, ReportConfig] | None): Report configurations for the scenario.
+        config_paths (Optional[ConfigPaths]): Source configuration paths, when parsed from files.
     """
 
     __test__ = False
@@ -258,6 +268,7 @@ class TestScenario:
     test_runs: list[TestRun]
     job_status_check: bool = True
     reports: dict[str, ReportConfig] = field(default_factory=dict)
+    config_paths: Optional[ConfigPaths] = None
 
     def __repr__(self) -> str:
         """

src/cloudai/core.py

@@ -48,7 +48,15 @@
 from ._core.report_generation_strategy import ReportGenerationStrategy
 from ._core.runner import Runner
 from ._core.system import System
-from ._core.test_scenario import METRIC_ERROR, MetricErrorSentinel, MetricValue, TestDependency, TestRun, TestScenario
+from ._core.test_scenario import (
+    METRIC_ERROR,
+    ConfigPaths,
+    MetricErrorSentinel,
+    MetricValue,
+    TestDependency,
+    TestRun,
+    TestScenario,
+)
 from .configurator.base_agent import BaseAgent, BaseAgentConfig, RewardOverrides
 from .configurator.cloudai_gym import CloudAIGymEnv
 from .configurator.env_params import (
@@ -77,6 +85,7 @@
     "CloudAIGymEnv",
     "CmdArgs",
     "CommandGenStrategy",
+    "ConfigPaths",
     "DockerImage",
     "Encoding",
     "File",

src/cloudai/parser.py

@@ -31,7 +31,7 @@
 )
 from ._core.registry import Registry
 from ._core.system import System
-from ._core.test_scenario import TestScenario
+from ._core.test_scenario import ConfigPaths, TestScenario
 from .test_parser import TestParser
 from .test_scenario_parser import TestScenarioParser
 from .toml_utils import format_toml_decode_error
@@ -125,6 +125,12 @@ def parse(
         except TestScenarioParsingError:
             exit(1)  # exit right away to keep error message readable for users
 
+        test_scenario.config_paths = ConfigPaths(
+            system_path=self.system_config_path.resolve(),
+            tests_dir_path=test_path.resolve() if test_path is not None else None,
+            test_scenario_path=test_scenario_path.resolve(),
+        )
+
         scenario_tests = {tr.test.name for tr in test_scenario.test_runs}
         hook_scenario_tests = {
             tr.test.name for hook_scenario in hook_test_scenario_mapping.values() for tr in hook_scenario.test_runs

src/cloudai/registration.py

@@ -38,7 +38,7 @@ def register_all():
     )
     from cloudai.core import Registry
     from cloudai.models.scenario import ReportConfig
-    from cloudai.report_generator.training import TrainingReportGenerationStrategy
+    from cloudai.report_generator.training import TrainingReporter
     from cloudai.reporter import DSEReporter, PerTestReporter, StatusReporter, TarballReporter
 
     # Import systems
@@ -305,14 +305,11 @@ def register_all():
     Registry().add_report(GrokTestDefinition, JaxToolboxReportGenerationStrategy)
     Registry().add_report(MegatronRunTestDefinition, CheckpointTimingReportGenerationStrategy)
     Registry().add_report(MegatronRunTestDefinition, MegatronRunReportGenerationStrategy)
-    Registry().add_report(MegatronRunTestDefinition, TrainingReportGenerationStrategy)
     Registry().add_report(MegatronBridgeTestDefinition, MegatronBridgeReportGenerationStrategy)
-    Registry().add_report(MegatronBridgeTestDefinition, TrainingReportGenerationStrategy)
     Registry().add_report(NCCLTestDefinition, NcclTestPerformanceReportGenerationStrategy)
     Registry().add_report(NeMoLauncherTestDefinition, NeMoLauncherReportGenerationStrategy)
     Registry().add_report(NeMoRunTestDefinition, NeMoRunReportGenerationStrategy)
     Registry().add_report(NeMoRunTestDefinition, NeMoRunDataStoreReportGenerationStrategy)
-    Registry().add_report(NeMoRunTestDefinition, TrainingReportGenerationStrategy)
     Registry().add_report(NemotronTestDefinition, JaxToolboxReportGenerationStrategy)
     Registry().add_report(UCCTestDefinition, UCCTestReportGenerationStrategy)
     Registry().add_report(TritonInferenceTestDefinition, TritonInferenceReportGenerationStrategy)
@@ -328,6 +325,7 @@ def register_all():
     Registry().add_report(VllmTestDefinition, VLLMBenchReportGenerationStrategy)
 
     Registry().add_scenario_report("per_test", PerTestReporter, ReportConfig(enable=True))
+    Registry().add_scenario_report("training", TrainingReporter, ReportConfig(enable=True))
     Registry().add_scenario_report(
         "moe_benchmark_throughput",
         MoEBenchmarkThroughputReporter,

src/cloudai/report_generator/training/__init__.py

@@ -14,6 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .report_generation_strategy import TrainingReportGenerationStrategy
+from .reporter import TrainingReporter
 
-__all__ = ["TrainingReportGenerationStrategy"]
+__all__ = ["TrainingReporter"]

src/cloudai/report_generator/training/mappings.py

@@ -55,15 +55,22 @@
 # Framework's resolved config artifact. (world_size, num_nodes, model_name) and computed data_parallel_size are not
 # mapped here.
 NEMO_MODEL_CONFIG: dict[str, str] = {
+    # Precision
+    "fp8": "model.fp8",
+    "fp8_recipe": "model.fp8_recipe",
+    # Batch
     "micro_batch_size": "data.micro_batch_size",
     "global_batch_size": "data.global_batch_size",
     "seq_length": "data.seq_length",
+    # Parallelism
     "tensor_parallel_size": "parallelism.tensor_model_parallel_size",
     "pipeline_parallel_size": "parallelism.pipeline_model_parallel_size",
     "context_parallel_size": "parallelism.context_parallel_size",
     "virtual_pipeline_parallel_size": "parallelism.virtual_pipeline_model_parallel_size",
     "sequence_parallel": "parallelism.sequence_parallel",
     "expert_parallel_size": "parallelism.expert_model_parallel_size",
+    "expert_tensor_parallel_size": "parallelism.expert_tensor_parallel_size",
+    # Model architecture
     "num_layers": "model.num_layers",
     "hidden_size": "model.hidden_size",
     "num_attention_heads": "model.num_attention_heads",
@@ -72,22 +79,30 @@
     "kv_channels": "model.kv_channels",
     "normalization": "model.normalization",
     "position_embedding_type": "model.position_embedding_type",
+    # MoE
     "num_experts": "model.num_moe_experts",
     "moe_router_topk": "model.moe_router_topk",
     "moe_ffn_hidden_size": "model.moe_ffn_hidden_size",
     "moe_grouped_gemm": "model.moe_grouped_gemm",
 }
 
 MEGATRON_MODEL_CONFIG: dict[str, str] = {
+    # Precision
+    "fp8": "fp8",
+    "fp8_recipe": "fp8_recipe",
+    # Batch
     "micro_batch_size": "micro_batch_size",
     "global_batch_size": "global_batch_size",
     "seq_length": "seq_length",
+    # Parallelism
     "tensor_parallel_size": "tensor_model_parallel_size",
     "pipeline_parallel_size": "pipeline_model_parallel_size",
     "context_parallel_size": "context_parallel_size",
     "virtual_pipeline_parallel_size": "virtual_pipeline_model_parallel_size",
     "sequence_parallel": "sequence_parallel",
     "expert_parallel_size": "expert_model_parallel_size",
+    "expert_tensor_parallel_size": "expert_tensor_parallel_size",
+    # Model architecture
     "num_layers": "num_layers",
     "hidden_size": "hidden_size",
     "num_attention_heads": "num_attention_heads",
@@ -96,22 +111,30 @@
     "kv_channels": "kv_channels",
     "normalization": "normalization",
     "position_embedding_type": "position_embedding_type",
+    # MoE
     "num_experts": "num_experts",
     "moe_router_topk": "moe_router_topk",
     "moe_ffn_hidden_size": "moe_ffn_hidden_size",
     "moe_grouped_gemm": "moe_grouped_gemm",
 }
 
 MEGATRON_BRIDGE_MODEL_CONFIG: dict[str, str] = {
+    # Precision
+    "fp8": "mixed_precision.fp8",
+    "fp8_recipe": "mixed_precision.fp8_recipe",
+    # Batch
     "micro_batch_size": "train.micro_batch_size",
     "global_batch_size": "train.global_batch_size",
     "seq_length": "model.seq_length",
+    # Parallelism
     "tensor_parallel_size": "model.tensor_model_parallel_size",
     "pipeline_parallel_size": "model.pipeline_model_parallel_size",
     "context_parallel_size": "model.context_parallel_size",
     "virtual_pipeline_parallel_size": "model.virtual_pipeline_model_parallel_size",
     "sequence_parallel": "model.sequence_parallel",
     "expert_parallel_size": "model.expert_model_parallel_size",
+    "expert_tensor_parallel_size": "model.expert_tensor_parallel_size",
+    # Model architecture
     "num_layers": "model.num_layers",
     "hidden_size": "model.hidden_size",
     "num_attention_heads": "model.num_attention_heads",
@@ -120,6 +143,7 @@
     "kv_channels": "model.kv_channels",
     "normalization": "model.normalization",
     "position_embedding_type": "model.position_embedding_type",
+    # MoE
     "num_experts": "model.num_moe_experts",
     "moe_router_topk": "model.moe_router_topk",
     "moe_ffn_hidden_size": "model.moe_ffn_hidden_size",
@@ -129,25 +153,37 @@
 
 # CloudAI TestDefinition (user TOML + defaults). TrainingConfig field -> dotted path in TestDefinition.model_dump().
 NEMO_TEST_CONFIG: dict[str, str] = {
+    # Environment
+    "container_image": "cmd_args.docker_image_url",
+    # Profiling
     "profiling_enabled": "nsys.enable",
     "profiling_start_step": "extra_cmd_args.*start_step",
     "profiling_stop_step": "extra_cmd_args.*end_step",
+    # Aggregation window
     "exclude_start_steps": "training_report.exclude_start_steps",
     "exclude_post_profiling_steps": "training_report.exclude_post_profiling_steps",
 }
 
 MEGATRON_TEST_CONFIG: dict[str, str] = {
+    # Environment
+    "container_image": "cmd_args.docker_image_url",
+    # Profiling
     "profiling_enabled": "nsys.enable",
     "profiling_start_step": "cmd_args.profile_step_start",
     "profiling_stop_step": "cmd_args.profile_step_end",
+    # Aggregation window
     "exclude_start_steps": "training_report.exclude_start_steps",
     "exclude_post_profiling_steps": "training_report.exclude_post_profiling_steps",
 }
 
 MEGATRON_BRIDGE_TEST_CONFIG: dict[str, str] = {
+    # Environment
+    "container_image": "cmd_args.container_image",
+    # Profiling
     "profiling_enabled": "cmd_args.enable_nsys",
     "profiling_start_step": "cmd_args.profiling_start_step",
     "profiling_stop_step": "cmd_args.profiling_stop_step",
+    # Aggregation window
     "exclude_start_steps": "training_report.exclude_start_steps",
     "exclude_post_profiling_steps": "training_report.exclude_post_profiling_steps",
 }