diff --git a/doc/reporting.rst b/doc/reporting.rst index cb5916739..4eed11ed4 100644 --- a/doc/reporting.rst +++ b/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 ` - :ref:`Reporting Registration ` - :ref:`Reporting Configuration Implementation ` +- :doc:`Reports ` + +.. toctree:: + :hidden: + + reports .. _overview: diff --git a/doc/reports.rst b/doc/reports.rst new file mode 100644 index 000000000..dca865a1c --- /dev/null +++ b/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 diff --git a/doc/training-report-schema.rst b/doc/training-report-schema.rst new file mode 100644 index 000000000..f589890a2 --- /dev/null +++ b/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 diff --git a/src/cloudai/_core/test_scenario.py b/src/cloudai/_core/test_scenario.py index ac75959e4..32fa25c62 100644 --- a/src/cloudai/_core/test_scenario.py +++ b/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: """ diff --git a/src/cloudai/core.py b/src/cloudai/core.py index 5bb739f99..06601f7f4 100644 --- a/src/cloudai/core.py +++ b/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", diff --git a/src/cloudai/parser.py b/src/cloudai/parser.py index 890709d8a..b76dc6787 100644 --- a/src/cloudai/parser.py +++ b/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 diff --git a/src/cloudai/registration.py b/src/cloudai/registration.py index 4ff418cfa..9ead24bbf 100644 --- a/src/cloudai/registration.py +++ b/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, diff --git a/src/cloudai/report_generator/training/__init__.py b/src/cloudai/report_generator/training/__init__.py index 1c7f33425..926472ba0 100644 --- a/src/cloudai/report_generator/training/__init__.py +++ b/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"] diff --git a/src/cloudai/report_generator/training/mappings.py b/src/cloudai/report_generator/training/mappings.py index 2588d5c29..62ffe4517 100644 --- a/src/cloudai/report_generator/training/mappings.py +++ b/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,6 +79,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", @@ -79,15 +87,22 @@ } 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,6 +111,7 @@ "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", @@ -103,15 +119,22 @@ } 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", } diff --git a/src/cloudai/report_generator/training/models.py b/src/cloudai/report_generator/training/models.py index 43b9d64da..c5681d35f 100644 --- a/src/cloudai/report_generator/training/models.py +++ b/src/cloudai/report_generator/training/models.py @@ -21,6 +21,8 @@ from dataclasses import MISSING, dataclass, fields from typing import Any, List, Optional +SCHEMA_VERSION = "1.0" # training report schema; bump on breaking changes to TrainingConfig/TrainingResults + @dataclass class MetricStats: @@ -111,10 +113,41 @@ class TrainingConfig: """ Resolved training configuration from the framework artifact + CloudAI. - The CloudAI-computed fields (test_template_name, data_parallel_size, model_name, world_size, num_nodes) default - here and are filled in by the parser after construction. + CloudAI-computed fields are supplied by the parser during construction. """ + # Test identity + test_id: str # scenario section id + test_name: str # test definition name + description: str + test_scenario_name: str + test_template_name: str + + # Configuration sources + system_path: str + tests_dir_path: str + test_scenario_path: str + + # Environment + container_image: str = "" + cloudai_execution_node: str + env_vars: dict[str, Any] # system global env + test extra env + + # Hardware + # Depends on: num_nodes, gpus_per_node + world_size: Optional[int] = None + # Populated from TestRun.nnodes. + num_nodes: int + gpus_per_node: Optional[int] = None + nodes: list[str] # compressed nodelist from the scenario + # Depends on: env_vars["CLIQUE_SIZE"] + clique_size: Optional[int] = None + + # Precision + fp8: Optional[str] = None + # Depends on: fp8 + fp8_recipe: Optional[str] = None + # Batch micro_batch_size: int global_batch_size: int @@ -127,7 +160,9 @@ class TrainingConfig: virtual_pipeline_parallel_size: Optional[int] sequence_parallel: bool expert_parallel_size: int - data_parallel_size: Optional[int] = None # CloudAI-computed (None when gpus_per_node is unavailable) + expert_tensor_parallel_size: int + # Depends on: world_size / (tensor_parallel_size * pipeline_parallel_size * context_parallel_size) + data_parallel_size: Optional[int] = None # Model architecture num_layers: int @@ -146,27 +181,24 @@ class TrainingConfig: moe_ffn_hidden_size: Optional[int] moe_grouped_gemm: Optional[bool] - # Hardware - world_size: Optional[int] = None # CloudAI-computed (None when gpus_per_node is unavailable) - num_nodes: int = 0 # CloudAI-computed - - # Profiling (CloudAI-computed from the run's nsys/profiler settings) + # Profiling profiling_enabled: bool = False + # Depends on: profiling_enabled profiling_start_step: Optional[int] = None + # Depends on: profiling_enabled profiling_stop_step: Optional[int] = None # Aggregation window (steps dropped before computing the top-level aggregation) exclude_start_steps: int = 5 + # Depends on: profiling_enabled, profiling_stop_step exclude_post_profiling_steps: int = 2 - # Identity - test_template_name: str = "" # CloudAI-computed - -@dataclass +@dataclass(kw_only=True) class TrainingResults: """Container for parsed training output.""" + schema_version: str = SCHEMA_VERSION config: TrainingConfig steps: List[TrainingStep] aggregation: Optional[StepAggregation] = None # None when no steps remain after exclusions diff --git a/src/cloudai/report_generator/training/parser.py b/src/cloudai/report_generator/training/parser.py index ccff5ab55..0447b3527 100644 --- a/src/cloudai/report_generator/training/parser.py +++ b/src/cloudai/report_generator/training/parser.py @@ -24,13 +24,14 @@ import fnmatch import json import logging +import socket from abc import ABC, abstractmethod from pathlib import Path from typing import Any, ClassVar, Optional import yaml -from cloudai.core import System, TestRun +from cloudai.core import System, TestRun, TestScenario from .mappings import ( MEGATRON_BRIDGE_MODEL_CONFIG, @@ -80,18 +81,24 @@ def _has_tb_event_files(tb_dir: Path) -> bool: def can_parse(self, tr: TestRun) -> bool: """Return True when the run produced the TB events and config artifact this parser needs.""" name = type(self).__name__ + + # Verify there's an existing TensorBoard directory containing at least one event file. tb_dir = self.get_tb_dir(tr) if not (tb_dir.is_dir() and self._has_tb_event_files(tb_dir)): logging.warning(f"{name}: no TensorBoard events at '{tb_dir}'; skipping training report") return False + + # If the workload expects a configuration file, verify that it exists. config_path = self.get_config_path(tr) if config_path is None or not config_path.is_file(): logging.warning(f"{name}: config artifact not found under '{tr.output_path}'; skipping training report") return False + + # If the configuration file is structured, verify that it parses into a non-empty mapping. if config_path.suffix.lower() in {".json", ".yaml", ".yml"}: try: config = self.get_model_config(tr) - except (json.JSONDecodeError, yaml.YAMLError) as exc: + except (OSError, UnicodeDecodeError, json.JSONDecodeError, yaml.YAMLError) as exc: logging.warning(f"{name}: invalid config artifact at '{config_path}' ({exc}); skipping training report") return False if not isinstance(config, dict) or not config: @@ -101,10 +108,10 @@ def can_parse(self, tr: TestRun) -> bool: return False return True - def parse(self, tr: TestRun, system: System) -> TrainingResults: + def parse(self, tr: TestRun, system: System, test_scenario: TestScenario) -> TrainingResults: """Read TB scalars + the config artifact and assemble TrainingResults.""" steps: list[TrainingStep] = self._build_steps(self._read_scalars(tr)) - config: TrainingConfig = self._build_config(tr, system) + config: TrainingConfig = self._build_config(tr, system, test_scenario) aggregation = self._aggregate(steps, config) return TrainingResults(config=config, steps=steps, aggregation=aggregation) @@ -149,16 +156,30 @@ def _build_step(self, step: int, step_scalars: list[Scalar]) -> Optional[Trainin } return TrainingStep(iteration=step, **field_values) - def _build_config(self, tr: TestRun, system: System) -> TrainingConfig: + def _build_config(self, tr: TestRun, system: System, test_scenario: TestScenario) -> TrainingConfig: """Map the framework + test config into TrainingConfig, then fill the CloudAI-computed fields.""" - field_values: dict[str, Any] = {} - field_values.update(self._resolve_model_config(tr)) - field_values.update(self._resolve_test_config(tr)) - config = TrainingConfig(**field_values) - config.test_template_name = tr.test.test_template_name - config.num_nodes = tr.nnodes - config.model_name = self.get_model_name(tr) + env_vars = {**getattr(system, "global_env_vars", {}), **tr.test.extra_env_vars} + config_paths = test_scenario.config_paths + config = TrainingConfig( + test_id=tr.name, + test_name=tr.test.name, + description=tr.test.description, + test_scenario_name=test_scenario.name, + test_template_name=tr.test.test_template_name, + system_path=str(config_paths.system_path) if config_paths is not None else "", + tests_dir_path=str(config_paths.tests_dir_path) if config_paths and config_paths.tests_dir_path else "", + test_scenario_path=str(config_paths.test_scenario_path) if config_paths is not None else "", + cloudai_execution_node=socket.gethostname(), + env_vars=env_vars, + num_nodes=tr.nnodes, + nodes=list(tr.nodes), + **self._resolve_model_config(tr), + **self._resolve_test_config(tr), + ) + + # Hardware gpus_per_node = getattr(system, "gpus_per_node", None) or getattr(system, "ntasks_per_node", None) + config.gpus_per_node = gpus_per_node if gpus_per_node: world_size = config.num_nodes * gpus_per_node config.world_size = world_size @@ -168,8 +189,21 @@ def _build_config(self, tr: TestRun, system: System) -> TrainingConfig: f"{type(self).__name__}: system has no gpus_per_node/ntasks_per_node; " "world_size and data_parallel_size left unset" ) + + config.clique_size = self._get_clique_size(config.env_vars) + + # Model architecture + config.model_name = self.get_model_name(tr) return config + @staticmethod + def _get_clique_size(env_vars: dict[str, Any]) -> Optional[int]: + clique_size = env_vars.get("CLIQUE_SIZE") + try: + return int(clique_size) if clique_size is not None else None + except (TypeError, ValueError): + return None + def _read_scalars(self, tr: TestRun) -> list[Scalar]: """Read the run's scalar events (subclasses may drop workload-specific noise).""" return read_scalars(self.get_tb_dir(tr)) diff --git a/src/cloudai/report_generator/training/report_generation_strategy.py b/src/cloudai/report_generator/training/reporter.py similarity index 58% rename from src/cloudai/report_generator/training/report_generation_strategy.py rename to src/cloudai/report_generator/training/reporter.py index 73a4bc118..39c0b51f5 100644 --- a/src/cloudai/report_generator/training/report_generation_strategy.py +++ b/src/cloudai/report_generator/training/reporter.py @@ -19,13 +19,13 @@ from dataclasses import asdict from typing import ClassVar -from cloudai.core import ReportGenerationStrategy +from cloudai.core import Reporter from .parser import MegatronBridgeParser, MegatronParser, NeMoRunParser, TrainingParser -class TrainingReportGenerationStrategy(ReportGenerationStrategy): - """Writes training_report.json for training workloads (NeMoRun, MegatronRun, MegatronBridge).""" +class TrainingReporter(Reporter): + """Generates a training report for each supported test run in a scenario.""" REPORT_FILE_NAME = "training_report.json" @@ -35,18 +35,25 @@ class TrainingReportGenerationStrategy(ReportGenerationStrategy): "MegatronBridge": MegatronBridgeParser, } - def can_handle_directory(self) -> bool: - parser_cls = self.PARSERS.get(self.test_run.test.test_template_name) - return parser_cls is not None and parser_cls().can_parse(self.test_run) + def generate(self) -> None: + self.load_test_runs() - def generate_report(self) -> None: - parser_cls = self.PARSERS[self.test_run.test.test_template_name] - training_results = parser_cls().parse(self.test_run, self.system) + for tr in self.trs: + parser_cls = self.PARSERS.get(tr.test.test_template_name) + if parser_cls is None: + continue + parser = parser_cls() + if not parser.can_parse(tr): + continue + try: + training_results = parser.parse(tr, self.system, self.test_scenario) - report_path = self.test_run.output_path / self.REPORT_FILE_NAME - report_path.write_text(json.dumps(asdict(training_results), indent=2, default=self._json_default)) + report_path = tr.output_path / self.REPORT_FILE_NAME + report_path.write_text(json.dumps(asdict(training_results), indent=2, default=self._json_default)) - logging.info(f"Generated training report for '{self.test_run.name}' at {report_path}") + logging.info(f"Generated training report for '{tr.name}' at {report_path}") + except Exception as exc: + logging.warning(f"Error generating training report for '{tr.output_path}': {exc}") @staticmethod def _json_default(value: object) -> object: diff --git a/tests/report_generator/training/test_training_parser.py b/tests/report_generator/training/test_training_parser.py index f8d8ddc57..3fb17b53a 100644 --- a/tests/report_generator/training/test_training_parser.py +++ b/tests/report_generator/training/test_training_parser.py @@ -16,24 +16,36 @@ import logging import types +from dataclasses import asdict from pathlib import Path from typing import Any import pytest +from cloudai.core import ConfigPaths +from cloudai.models.scenario import ReportConfig from cloudai.report_generator.training import parser as parser_mod +from cloudai.report_generator.training import reporter as report_mod from cloudai.report_generator.training import tb_reader -from cloudai.report_generator.training.models import Scalar, TrainingStep +from cloudai.report_generator.training.models import SCHEMA_VERSION, Scalar, TrainingResults, TrainingStep from cloudai.report_generator.training.parser import MegatronBridgeParser, MegatronParser, NeMoRunParser -from cloudai.report_generator.training.report_generation_strategy import TrainingReportGenerationStrategy +from cloudai.report_generator.training.reporter import TrainingReporter def _scalars(rows: list[tuple]) -> list[Scalar]: return [Scalar(tag=tag, step=step, value=value, wall_time=wall_time) for tag, step, value, wall_time in rows] -def _system(gpus_per_node: int | None = 4, ntasks_per_node: int | None = None) -> Any: - return types.SimpleNamespace(gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node) +def _system( + gpus_per_node: int | None = 4, ntasks_per_node: int | None = None, global_env_vars: dict[str, Any] | None = None +) -> Any: + return types.SimpleNamespace( + gpus_per_node=gpus_per_node, ntasks_per_node=ntasks_per_node, global_env_vars=global_env_vars or {} + ) + + +def _scenario(name: str = "scenario", config_paths: ConfigPaths | None = None) -> Any: + return types.SimpleNamespace(name=name, config_paths=config_paths) class _Test(types.SimpleNamespace): @@ -54,6 +66,9 @@ def _tr( nsys: Any = None, extra_cmd_args: dict[str, Any] | None = None, training_report: dict[str, Any] | None = None, + description: str = "", + extra_env_vars: dict[str, Any] | None = None, + nodes: list[str] | None = None, **cmd_args: Any, ) -> Any: test = _Test( @@ -63,8 +78,10 @@ def _tr( nsys=nsys, extra_cmd_args=extra_cmd_args or {}, training_report=training_report, + description=description, + extra_env_vars=extra_env_vars or {}, ) - return types.SimpleNamespace(output_path=Path(output_path), nnodes=nnodes, test=test) + return types.SimpleNamespace(output_path=Path(output_path), nnodes=nnodes, name=name, nodes=nodes or [], test=test) def _nsys(enable: bool) -> Any: @@ -233,18 +250,27 @@ def text(self): def test_build_config_resolves_paths_and_computes_fields(): raw = { "data": {"micro_batch_size": 1, "global_batch_size": 8}, - "parallelism": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1, "context_parallel_size": 1}, - "model": {"num_layers": 30}, + "parallelism": { + "tensor_model_parallel_size": 4, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "expert_tensor_parallel_size": 4, + }, + "model": {"num_layers": 30, "fp8": "hybrid", "fp8_recipe": "tensorwise"}, } parser = NeMoRunParser() parser.get_model_config = lambda tr: raw - config = parser._build_config(_tr(nnodes=8, template="NeMoRun", recipe_name="gpt3"), _system(gpus_per_node=4)) + config = parser._build_config( + _tr(nnodes=8, template="NeMoRun", recipe_name="gpt3"), _system(gpus_per_node=4), _scenario() + ) assert config.test_template_name == "NeMoRun" # CloudAI-computed assert config.micro_batch_size == 1 # nested dotted-path resolve assert config.num_layers == 30 assert config.tensor_parallel_size == 4 + assert config.expert_tensor_parallel_size == 4 assert config.model_name == "gpt3" # CloudAI-computed + assert (config.fp8, config.fp8_recipe) == ("hybrid", "tensorwise") assert (config.num_nodes, config.world_size) == (8, 32) # 8 nodes x 4 gpus assert config.data_parallel_size == 8 # 32 / (tp4 * pp1 * cp1) @@ -254,7 +280,7 @@ def test_build_config_leaves_world_size_none_without_gpus_per_node(): raw = {"parallelism": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1}} parser = NeMoRunParser() parser.get_model_config = lambda tr: raw - config = parser._build_config(_tr(nnodes=8, recipe_name="gpt3"), _system(gpus_per_node=None)) + config = parser._build_config(_tr(nnodes=8, recipe_name="gpt3"), _system(gpus_per_node=None), _scenario()) assert config.world_size is None assert config.data_parallel_size is None @@ -270,15 +296,20 @@ def test_megatron_config_parses_string_literals(monkeypatch): "pipeline_model_parallel_size": "2", "context_parallel_size": "1", "sequence_parallel": "True", + "expert_tensor_parallel_size": "1", + "fp8": "e4m3", + "fp8_recipe": "mxfp8", } monkeypatch.setattr(parser_mod, "read_text", lambda _tb_dir: text) parser = MegatronParser() - config = parser._build_config(_tr(nnodes=16, name="dsv3"), _system(gpus_per_node=4)) + config = parser._build_config(_tr(nnodes=16, name="dsv3"), _system(gpus_per_node=4), _scenario()) assert config.micro_batch_size == 1 # "1" -> int assert config.sequence_parallel is True # "True" -> bool + assert config.expert_tensor_parallel_size == 1 assert config.model_name == "dsv3" + assert (config.fp8, config.fp8_recipe) == ("e4m3", "mxfp8") assert config.data_parallel_size == 16 # 64 / (tp2 * pp2 * cp1) @@ -287,10 +318,10 @@ def test_compute_data_parallel_size_rejects_invalid_topology(): parallel = {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1, "context_parallel_size": 1} parser.get_model_config = lambda tr: {"parallelism": parallel} with pytest.raises(ValueError, match="world_size"): # world_size 34 is not a multiple of tp*pp*cp=4 - parser._build_config(_tr(nnodes=17, recipe_name="x"), _system(gpus_per_node=2)) + parser._build_config(_tr(nnodes=17, recipe_name="x"), _system(gpus_per_node=2), _scenario()) parser.get_model_config = lambda tr: {"parallelism": {}} with pytest.raises(ValueError, match="tensor_parallel_size"): # tp missing from the parsed config - parser._build_config(_tr(recipe_name="x"), _system()) + parser._build_config(_tr(recipe_name="x"), _system(), _scenario()) # --- profiling ----------------------------------------------------------------------------------- @@ -348,28 +379,103 @@ def test_megatron_bridge_profiling_reads_typed_fields(): def test_build_config_sets_profiling_fields(): # M-Bridge exposes enable + step bounds as typed cmd_args, so _build_config folds all three into the config. - raw = {"model": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1}} + raw = { + "model": { + "tensor_model_parallel_size": 4, + "pipeline_model_parallel_size": 1, + "expert_tensor_parallel_size": 4, + }, + "mixed_precision": {"fp8": "hybrid", "fp8_recipe": "tensorwise"}, + } parser = MegatronBridgeParser() parser.get_model_config = lambda tr: raw tr = _tr(model_recipe_name="gpt3", enable_nsys=True, profiling_start_step=3, profiling_stop_step=7) - config = parser._build_config(tr, _system(gpus_per_node=4)) + config = parser._build_config(tr, _system(gpus_per_node=4), _scenario()) assert config.profiling_enabled is True assert (config.profiling_start_step, config.profiling_stop_step) == (3, 7) + assert config.expert_tensor_parallel_size == 4 + assert (config.fp8, config.fp8_recipe) == ("hybrid", "tensorwise") def test_build_config_reads_aggregation_flags_from_toml(): parser = NeMoRunParser() parser.get_model_config = lambda tr: {} tr = _tr(recipe_name="gpt3", training_report={"exclude_start_steps": 10, "exclude_post_profiling_steps": 3}) - config = parser._build_config(tr, _system(gpus_per_node=None)) + config = parser._build_config(tr, _system(gpus_per_node=None), _scenario()) assert (config.exclude_start_steps, config.exclude_post_profiling_steps) == (10, 3) +def test_build_config_sets_identity_hardware_and_env(monkeypatch): + monkeypatch.setattr(parser_mod.socket, "gethostname", lambda: "cloudai-host") + parser = NeMoRunParser() + parser.get_model_config = lambda tr: { + "parallelism": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1} + } + tr = _tr( + name="dsv3_run", + template="NeMoRun", + recipe_name="gpt3", + description="a proxy run", + nodes=["node-[01-08]"], + extra_env_vars={"NCCL_MNNVL_ENABLE": "1", "CLIQUE_SIZE": "8"}, + docker_image_url="nvcr.io/nvidia/nemo:24.12", + ) + config_paths = ConfigPaths( + system_path=Path("/configs/system.toml"), + tests_dir_path=Path("/configs/tests"), + test_scenario_path=Path("/configs/scenario.toml"), + ) + config = parser._build_config( + tr, + _system(gpus_per_node=4, global_env_vars={"CUDA_HOME": "/usr/local/cuda", "CLIQUE_SIZE": "4"}), + _scenario("nightly", config_paths), + ) + + results = TrainingResults(config=config, steps=[]) + assert results.schema_version == SCHEMA_VERSION + assert next(iter(asdict(results))) == "schema_version" + assert config.test_id == "dsv3_run" + assert config.test_name == "dsv3_run" + assert config.description == "a proxy run" + assert config.test_scenario_name == "nightly" + assert config.system_path == "/configs/system.toml" + assert config.tests_dir_path == "/configs/tests" + assert config.test_scenario_path == "/configs/scenario.toml" + assert config.container_image == "nvcr.io/nvidia/nemo:24.12" + assert config.cloudai_execution_node == "cloudai-host" + assert config.nodes == ["node-[01-08]"] + assert config.gpus_per_node == 4 + assert config.clique_size == 8 + assert config.env_vars == {"CUDA_HOME": "/usr/local/cuda", "CLIQUE_SIZE": "8", "NCCL_MNNVL_ENABLE": "1"} + + +def test_build_config_uses_empty_paths_without_config_provenance(): + parser = NeMoRunParser() + parser.get_model_config = lambda tr: {} + + config = parser._build_config(_tr(recipe_name="gpt3"), _system(gpus_per_node=None), _scenario()) + + assert (config.system_path, config.tests_dir_path, config.test_scenario_path) == ("", "", "") + + +def test_build_config_invalid_clique_size_stays_none(caplog): + parser = NeMoRunParser() + parser.get_model_config = lambda tr: {} + config = parser._build_config( + _tr(recipe_name="gpt3", extra_env_vars={"CLIQUE_SIZE": "invalid"}), + _system(gpus_per_node=None), + _scenario(), + ) + + assert config.clique_size is None + assert "CLIQUE_SIZE" not in caplog.text + + def test_build_config_aggregation_flags_default_when_absent(): parser = NeMoRunParser() parser.get_model_config = lambda tr: {} - config = parser._build_config(_tr(recipe_name="gpt3"), _system(gpus_per_node=None)) + config = parser._build_config(_tr(recipe_name="gpt3"), _system(gpus_per_node=None), _scenario()) assert (config.exclude_start_steps, config.exclude_post_profiling_steps) == (5, 2) @@ -543,5 +649,31 @@ class StringFallback: def __str__(self): return "fallback" - assert TrainingReportGenerationStrategy._json_default(ScalarWithItem()) == 7 - assert TrainingReportGenerationStrategy._json_default(StringFallback()) == "fallback" + assert TrainingReporter._json_default(ScalarWithItem()) == 7 + assert TrainingReporter._json_default(StringFallback()) == "fallback" + + +def test_training_reporter_passes_scenario(monkeypatch, tmp_path): + calls = [] + + class FakeParser: + def can_parse(self, tr): + calls.append(("can_parse", tr)) + return True + + def parse(self, tr, system, test_scenario): + calls.append(("parse", tr, system, test_scenario)) + return object() + + monkeypatch.setattr(TrainingReporter, "PARSERS", {"NeMoRun": FakeParser}) + monkeypatch.setattr(report_mod, "asdict", lambda _: {"generated": True}) + system = _system() + tr = _tr(output_path=tmp_path) + scenario = _scenario("nightly") + reporter = TrainingReporter(system, scenario, tmp_path, ReportConfig(enable=True)) + monkeypatch.setattr(reporter, "load_test_runs", lambda: reporter.trs.append(tr)) + + reporter.generate() + + assert calls == [("can_parse", tr), ("parse", tr, system, scenario)] + assert (tmp_path / "training_report.json").read_text() == '{\n "generated": true\n}' diff --git a/tests/test_init.py b/tests/test_init.py index b2d361cc2..466fbae66 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -16,6 +16,7 @@ from cloudai.core import Registry +from cloudai.report_generator.training import TrainingReporter from cloudai.reporter import DSEReporter, PerTestReporter, StatusReporter, TarballReporter from cloudai.systems.kubernetes import KubernetesInstaller, KubernetesSystem from cloudai.systems.lsf import LSFInstaller, LSFSystem @@ -280,6 +281,7 @@ def test_scenario_reports(): scenario_reports = Registry().scenario_reports assert list(scenario_reports.keys()) == [ "per_test", + "training", "moe_benchmark_throughput", "status", "dse", @@ -294,6 +296,7 @@ def test_scenario_reports(): ] assert list(scenario_reports.values()) == [ PerTestReporter, + TrainingReporter, MoEBenchmarkThroughputReporter, StatusReporter, DSEReporter, @@ -312,6 +315,7 @@ def test_report_configs(): configs = Registry().report_configs assert list(configs.keys()) == [ "per_test", + "training", "moe_benchmark_throughput", "status", "dse", diff --git a/tests/test_parser.py b/tests/test_parser.py index 17e1c11c9..8e8a9a17e 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -22,7 +22,15 @@ import toml from pydantic_core import ErrorDetails -from cloudai.core import Parser, Registry, Reporter, TestConfigParsingError, TestParser, format_validation_error +from cloudai.core import ( + ConfigPaths, + Parser, + Registry, + Reporter, + TestConfigParsingError, + TestParser, + format_validation_error, +) from cloudai.models.scenario import ReportConfig, parse_reports_spec from cloudai.systems.slurm.slurm_system import SlurmSystem @@ -66,6 +74,37 @@ def test_custom_hook_root_is_used_by_parse(self, parser: Parser, tmp_path: Path) assert "custom_hook_test" in {test.name for test in tests} + @patch("cloudai.parser.Parser.parse_test_scenario") + def test_parse_links_config_paths(self, parse_test_scenario: Mock, parser: Parser, tmp_path: Path): + tests_dir = tmp_path / "tests" + tests_dir.mkdir() + test_scenario_path = tmp_path / "test_scenario.toml" + parse_test_scenario.return_value = Mock(test_runs=[]) + parser = Parser(parser.system_config_path, tmp_path / "hooks") + + _, _, test_scenario = parser.parse(tests_dir, test_scenario_path) + + assert test_scenario is not None + assert test_scenario.config_paths == ConfigPaths( + system_path=parser.system_config_path.resolve(), + tests_dir_path=tests_dir.resolve(), + test_scenario_path=test_scenario_path.resolve(), + ) + + @patch("cloudai.parser.Parser.parse_test_scenario") + def test_parse_links_config_paths_without_tests_dir( + self, parse_test_scenario: Mock, parser: Parser, tmp_path: Path + ): + test_scenario_path = tmp_path / "test_scenario.toml" + parse_test_scenario.return_value = Mock(test_runs=[]) + parser = Parser(parser.system_config_path, tmp_path / "hooks") + + _, _, test_scenario = parser.parse(None, test_scenario_path) + + assert test_scenario is not None + assert test_scenario.config_paths is not None + assert test_scenario.config_paths.tests_dir_path is None + def test_custom_hook_root_is_used_for_hook_scenario_resolution(self, parser: Parser, tmp_path: Path): """A hook *scenario* toml (referenced via pre_test) must also be resolved from the custom hook_root, not just hook test tomls under `/test`.""" diff --git a/tests/test_test_scenario.py b/tests/test_test_scenario.py index a02a48b24..64eb2c844 100644 --- a/tests/test_test_scenario.py +++ b/tests/test_test_scenario.py @@ -35,7 +35,6 @@ TestScenarioParser, ) from cloudai.models.scenario import TestRunModel, TestScenarioModel -from cloudai.report_generator.training import TrainingReportGenerationStrategy from cloudai.systems.slurm.slurm_system import SlurmSystem from cloudai.test_scenario_parser import calculate_total_time_limit, get_reporters from cloudai.workloads.ai_dynamo import AIDynamoReportGenerationStrategy, AIDynamoTestDefinition @@ -720,12 +719,11 @@ def test_default_reporters_size(self): { CheckpointTimingReportGenerationStrategy, MegatronRunReportGenerationStrategy, - TrainingReportGenerationStrategy, }, ), ( MegatronBridgeTestDefinition, - {MegatronBridgeReportGenerationStrategy, TrainingReportGenerationStrategy}, + {MegatronBridgeReportGenerationStrategy}, ), (NCCLTestDefinition, {NcclTestPerformanceReportGenerationStrategy}), (NeMoLauncherTestDefinition, {NeMoLauncherReportGenerationStrategy}), @@ -734,7 +732,6 @@ def test_default_reporters_size(self): { NeMoRunReportGenerationStrategy, NeMoRunDataStoreReportGenerationStrategy, - TrainingReportGenerationStrategy, }, ), (NemotronTestDefinition, {JaxToolboxReportGenerationStrategy}),