From e02cab6dd6112241d5fe98e85a19582b6611d292 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 18:31:26 +0000 Subject: [PATCH 01/17] test: characterize measurement module surfaces Signed-off-by: Aaron Gonzales --- .importlinter | 4 +- .../test_measurement_module_compatibility.py | 231 ++++++++++++++++++ 2 files changed, 234 insertions(+), 1 deletion(-) create mode 100644 tests/tools/test_measurement_module_compatibility.py diff --git a/.importlinter b/.importlinter index 6e8d170a..6bf5262d 100644 --- a/.importlinter +++ b/.importlinter @@ -1,5 +1,7 @@ [importlinter] -root_package = anonymizer +root_packages = + anonymizer + measurement_tools # The NDD adapter layer is a leaf boundary: it must never import the engine # sub-workflows that call it. diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py new file mode 100644 index 00000000..7859878a --- /dev/null +++ b/tests/tools/test_measurement_module_compatibility.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +import importlib.util +import os +import pickle +import subprocess +import sys +from collections.abc import Iterator +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest +from pydantic import BaseModel + +REPO_ROOT = Path(__file__).resolve().parents[2] +MEASUREMENT_ROOT = REPO_ROOT / "tools/measurement" + + +@pytest.fixture(autouse=True) +def _measurement_import_path(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.syspath_prepend(str(MEASUREMENT_ROOT)) + + +def _load_module(path: Path, name: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except BaseException: + sys.modules.pop(name, None) + raise + return module + + +@pytest.mark.parametrize( + ("relative_path", "defined_class", "canonical_import"), + [ + ("measurement_tools/wandb_models.py", "WandbMode", "StrictFrozenModel"), + ("measurement_tools/wandb_setup.py", "WandbPublisher", "ResolvedWandbConfig"), + ("create_wandb_report.py", "WandbReportResult", "WandbProjectPath"), + ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer"), + ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath"), + ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult"), + ], +) +def test_exact_tool_paths_support_repeated_dynamic_loading( + relative_path: str, + defined_class: str, + canonical_import: str, +) -> None: + path = MEASUREMENT_ROOT / relative_path + module_names = (f"compat_first_{path.stem}", f"compat_second_{path.stem}") + first = _load_module(path, module_names[0]) + second = _load_module(path, module_names[1]) + try: + first_class = getattr(first, defined_class) + second_class = getattr(second, defined_class) + assert first_class is not second_class + assert first_class.__module__ == module_names[0] + assert second_class.__module__ == module_names[1] + assert getattr(first, canonical_import) is getattr(second, canonical_import) + finally: + for module_name in module_names: + sys.modules.pop(module_name, None) + + +def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: + models = importlib.import_module("measurement_tools.wandb_models") + setup = importlib.import_module("measurement_tools.wandb_setup") + report_models = importlib.import_module("measurement_tools.wandb_report_models") + report = _load_module(MEASUREMENT_ROOT / "create_wandb_report.py", "compat_report_facade") + runner = _load_module(MEASUREMENT_ROOT / "run_benchmarks.py", "compat_runner_facade") + sweep = _load_module(MEASUREMENT_ROOT / "sweep_benchmarks.py", "compat_sweep_facade") + try: + assert setup.WandbMode is models.WandbMode + assert setup.ResolvedWandbConfig is models.ResolvedWandbConfig + assert report.WandbProjectPath is report_models.WandbProjectPath + assert report.WandbRunPath is report_models.WandbRunPath + assert runner.WandbMode is models.WandbMode + assert runner.ResolvedWandbConfig is models.ResolvedWandbConfig + assert sweep.WandbProjectPath is report_models.WandbProjectPath + assert sweep.run_benchmarks.is_remote_input_source is runner.is_remote_input_source + + assert models.WandbMode.__module__ == "measurement_tools.wandb_models" + assert models.ResolvedWandbConfig.__module__ == "measurement_tools.wandb_models" + settings = models.ResolvedWandbConfig() + reconstructed = models.ResolvedWandbConfig.model_validate(settings.model_dump()) + assert reconstructed == settings + assert pickle.loads(pickle.dumps(models.WandbMode.offline)) is models.WandbMode.offline + finally: + for module_name in ("compat_report_facade", "compat_runner_facade", "compat_sweep_facade"): + sys.modules.pop(module_name, None) + + +@pytest.mark.parametrize( + ("script", "help_text"), + [ + ("create_wandb_report.py", "Create a W&B benchmark report"), + ("run_benchmarks.py", "Run Anonymizer benchmark suites"), + ("sweep_benchmarks.py", "Run a benchmark suite sweep"), + ("analyze_benchmark_output.py", "Analyze joined benchmark measurements"), + ], +) +def test_executable_tool_paths_preserve_help_and_main_behavior(script: str, help_text: str) -> None: + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + filter(None, (str(MEASUREMENT_ROOT), env.get("PYTHONPATH"))) + ) + result = subprocess.run( + [sys.executable, str(MEASUREMENT_ROOT / script), "--help"], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0 + assert help_text in result.stdout + assert "Usage:" in result.stdout + + +def test_observed_monkeypatch_namespaces_remain_available() -> None: + setup = importlib.import_module("measurement_tools.wandb_setup") + report = _load_module(MEASUREMENT_ROOT / "create_wandb_report.py", "compat_report_patches") + runner = _load_module(MEASUREMENT_ROOT / "run_benchmarks.py", "compat_runner_patches") + sweep = _load_module(MEASUREMENT_ROOT / "sweep_benchmarks.py", "compat_sweep_patches") + surfaces: list[tuple[Any, tuple[str, ...]]] = [ + ( + setup, + ( + "_PUBLISHER_WANDB_MODULE", + "prepare_wandb_staging_dir", + "_effective_wandb_tags", + "publish_benchmark_wandb_best_effort", + ), + ), + ( + setup.os, + ("open", "close", "fstat"), + ), + ( + report, + ( + "require_wandb_report_sdk", + "require_wandb_workspace_sdk", + "build_benchmark_report", + "build_benchmark_group_report", + "build_benchmark_workspace", + "_save_report", + "_save_workspace", + "create_benchmark_report", + "create_benchmark_group_report", + "create_benchmark_workspace", + ), + ), + ( + runner, + ( + "Anonymizer", + "_run_case", + "_execute_case", + "_run_case_error", + "configured_measurement_session", + "export_measurement_tables", + "combine_detection_artifact_analysis", + "_git_metadata", + "_package_versions", + "publish_benchmark_wandb_best_effort", + "run_suite", + "run_or_plan", + ), + ), + ( + sweep, + ( + "create_benchmark_workspace", + "create_benchmark_group_report", + "run_benchmarks", + "run_sweep", + ), + ), + ] + try: + for namespace, names in surfaces: + for name in names: + assert hasattr(namespace, name), f"{namespace.__name__}.{name}" + finally: + for module_name in ("compat_report_patches", "compat_runner_patches", "compat_sweep_patches"): + sys.modules.pop(module_name, None) + + +def _inherited_policy_names( + model: type[BaseModel], + registry: dict[type[BaseModel], dict[str, Any]], +) -> Iterator[str]: + for parent in reversed(model.mro()): + if issubclass(parent, BaseModel): + yield from registry.get(parent, {}) + + +def test_outbound_field_policy_validation_is_complete_at_import() -> None: + models = importlib.import_module("measurement_tools.wandb_models") + + models.validate_outbound_field_policies() + for model, _policies in models.OUTBOUND_FIELD_POLICIES.items(): + assert set(model.model_fields) == set(_inherited_policy_names(model, models.OUTBOUND_FIELD_POLICIES)) + + +@pytest.mark.parametrize( + ("relative_path", "anchor"), + [ + ("tools/measurement/import_wandb_run.py", "from measurement_tools.wandb_models import"), + ("tools/measurement/measurement_tools/wandb_logging.py", "from measurement_tools.wandb_models import"), + ("tools/measurement/measurement_tools/wandb_report_models.py", "from measurement_tools.wandb_models import"), + ("tools/measurement/sweep_benchmarks.py", "import run_benchmarks"), + ("tests/tools/test_measurement_strict_import_publisher.py", "publish_benchmark_wandb_best_effort.__module__"), + (".github/workflows/benchmark-ci.yml", "uv run python tools/measurement/run_benchmarks.py"), + ("tools/measurement/examples/run-repo-data-smoke-with-dd-traces.sh", "tools/measurement/run_benchmarks.py"), + ("docs/development/observability.md", "tools/measurement/create_wandb_report.py"), + ], +) +def test_direct_consumer_anchors_remain_at_observed_paths(relative_path: str, anchor: str) -> None: + assert anchor in (REPO_ROOT / relative_path).read_text(encoding="utf-8") From 8202e5cfb27df62e12b8e3e4f818fe5d1db2875e Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 18:54:26 +0000 Subject: [PATCH 02/17] refactor: split wandb settings and metadata Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 54 +- .../measurement_tools/wandb_metadata.py | 402 ++++++++++++ .../measurement_tools/wandb_models.py | 588 ++---------------- .../measurement_tools/wandb_policy.py | 31 + .../measurement_tools/wandb_settings.py | 187 ++++++ 5 files changed, 707 insertions(+), 555 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_metadata.py create mode 100644 tools/measurement/measurement_tools/wandb_policy.py create mode 100644 tools/measurement/measurement_tools/wandb_settings.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 7859878a..d363c381 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -41,20 +41,26 @@ def _load_module(path: Path, name: str) -> ModuleType: @pytest.mark.parametrize( - ("relative_path", "defined_class", "canonical_import"), + ("relative_path", "defined_class", "canonical_import", "canonical_module"), [ - ("measurement_tools/wandb_models.py", "WandbMode", "StrictFrozenModel"), - ("measurement_tools/wandb_setup.py", "WandbPublisher", "ResolvedWandbConfig"), - ("create_wandb_report.py", "WandbReportResult", "WandbProjectPath"), - ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer"), - ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath"), - ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult"), + ( + "measurement_tools/wandb_models.py", + "WandbMode", + "StrictFrozenModel", + "measurement_tools.wandb_settings", + ), + ("measurement_tools/wandb_setup.py", "WandbPublisher", "ResolvedWandbConfig", None), + ("create_wandb_report.py", "WandbReportResult", "WandbProjectPath", None), + ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer", None), + ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath", None), + ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult", None), ], ) def test_exact_tool_paths_support_repeated_dynamic_loading( relative_path: str, defined_class: str, canonical_import: str, + canonical_module: str | None, ) -> None: path = MEASUREMENT_ROOT / relative_path module_names = (f"compat_first_{path.stem}", f"compat_second_{path.stem}") @@ -63,9 +69,13 @@ def test_exact_tool_paths_support_repeated_dynamic_loading( try: first_class = getattr(first, defined_class) second_class = getattr(second, defined_class) - assert first_class is not second_class - assert first_class.__module__ == module_names[0] - assert second_class.__module__ == module_names[1] + if canonical_module is None: + assert first_class is not second_class + assert first_class.__module__ == module_names[0] + assert second_class.__module__ == module_names[1] + else: + assert first_class is second_class + assert first_class.__module__ == canonical_module assert getattr(first, canonical_import) is getattr(second, canonical_import) finally: for module_name in module_names: @@ -89,8 +99,8 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: assert sweep.WandbProjectPath is report_models.WandbProjectPath assert sweep.run_benchmarks.is_remote_input_source is runner.is_remote_input_source - assert models.WandbMode.__module__ == "measurement_tools.wandb_models" - assert models.ResolvedWandbConfig.__module__ == "measurement_tools.wandb_models" + assert models.WandbMode.__module__ == "measurement_tools.wandb_settings" + assert models.ResolvedWandbConfig.__module__ == "measurement_tools.wandb_settings" settings = models.ResolvedWandbConfig() reconstructed = models.ResolvedWandbConfig.model_validate(settings.model_dump()) assert reconstructed == settings @@ -100,6 +110,22 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: sys.modules.pop(module_name, None) +@pytest.mark.parametrize( + "canonical_module", + [ + "measurement_tools.wandb_policy", + "measurement_tools.wandb_settings", + "measurement_tools.wandb_metadata", + ], +) +def test_wandb_model_facade_reexports_every_canonical_symbol(canonical_module: str) -> None: + models = importlib.import_module("measurement_tools.wandb_models") + canonical = importlib.import_module(canonical_module) + + for name in canonical.__all__: + assert getattr(models, name) is getattr(canonical, name) + + @pytest.mark.parametrize( ("script", "help_text"), [ @@ -111,9 +137,7 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: ) def test_executable_tool_paths_preserve_help_and_main_behavior(script: str, help_text: str) -> None: env = os.environ.copy() - env["PYTHONPATH"] = os.pathsep.join( - filter(None, (str(MEASUREMENT_ROOT), env.get("PYTHONPATH"))) - ) + env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(MEASUREMENT_ROOT), env.get("PYTHONPATH")))) result = subprocess.run( [sys.executable, str(MEASUREMENT_ROOT / script), "--help"], cwd=REPO_ROOT, diff --git a/tools/measurement/measurement_tools/wandb_metadata.py b/tools/measurement/measurement_tools/wandb_metadata.py new file mode 100644 index 00000000..a627bd3a --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_metadata.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sanitized benchmark metadata and W&B config projection.""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import ( + Field, + FiniteFloat, + StrictBool, + StrictInt, + StrictStr, + field_validator, + model_validator, +) + +from measurement_tools.validation import ( + NonNegativeFloat, + NonNegativeInt, + Probability, + StrictFrozenModel, + VisibleIdentifier, + VisibleSlurmIdentifier, +) +from measurement_tools.wandb_settings import WANDB_TAG_MAX_LENGTH, ResolvedWandbConfig, WandbMode + +__all__ = [ + "BenchmarkMetadata", + "ConfigMetadata", + "DetectMetadata", + "ExecutionMetadata", + "GitMetadata", + "ImportedRunMetadata", + "MatrixMetadata", + "ModelSourcesMetadata", + "ReplaceMetadata", + "RewriteMetadata", + "RuntimeMetadata", + "SafeScalar", + "SlurmJobMetadata", + "SlurmMetadata", + "SweepMetadata", + "WandbConfigPayload", + "WandbRunMetadata", + "WandbTag", + "WorkloadMetadata", + "WorkloadSourceMetadata", +] + +SafeScalar = StrictBool | StrictInt | FiniteFloat | StrictStr +WandbTag = Annotated[StrictStr, Field(min_length=1, max_length=WANDB_TAG_MAX_LENGTH)] +_NUMERIC_SWEEP_PARAMETER_TAILS = frozenset( + { + "detect_gliner_threshold", + "detect_validation_excerpt_window_chars", + "detect_validation_max_entities_per_call", + "replace_digest_length", + "rewrite_max_repair_iterations", + } +) +_BOOLEAN_SWEEP_PARAMETER_TAILS = frozenset( + {"emit_telemetry", "evaluate", "replace_normalize_label", "rewrite_strict_entity_protection"} +) +_ENUM_SWEEP_PARAMETER_VALUES = { + "replace_algorithm": frozenset({"md5", "sha1", "sha256"}), + "rewrite_risk_tolerance": frozenset({"high", "low", "minimal", "moderate"}), +} + + +class BenchmarkMetadata(StrictFrozenModel): + metadata_schema_version: NonNegativeInt | None = None + suite_schema_version: NonNegativeInt | None = None + wandb_sanitizer_version: NonNegativeInt | None = None + measurement_schema_version: NonNegativeInt | None = None + suite_id: VisibleIdentifier | None = None + workload_count: NonNegativeInt | None = None + config_count: NonNegativeInt | None = None + matrix_entry_count: NonNegativeInt | None = None + case_count: NonNegativeInt | None = None + case_retries: NonNegativeInt | None = None + case_retry_backoff_sec: NonNegativeFloat | None = None + suite_file_hash: StrictStr | None = None + + +class SlurmJobMetadata(StrictFrozenModel): + role: VisibleSlurmIdentifier + job_id: VisibleSlurmIdentifier + + +class SlurmMetadata(StrictFrozenModel): + job_id: VisibleIdentifier | None = None + array_job_id: VisibleIdentifier | None = None + array_task_id: VisibleIdentifier | None = None + array_task_count: NonNegativeInt | None = None + restart_count: NonNegativeInt | None = None + ntasks: NonNegativeInt | None = None + job_num_nodes: NonNegativeInt | None = None + jobs: Annotated[list[SlurmJobMetadata], Field(max_length=64, exclude_if=lambda value: not value)] = Field( + default_factory=list + ) + + +class ExecutionMetadata(StrictFrozenModel): + backend: Literal["local", "slurm"] | None = None + export: StrictBool | None = None + fail_fast: StrictBool | None = None + dd_trace: Literal["none", "last_message", "all_messages"] | None = None + dd_task_trace: StrictBool | None = None + slurm: SlurmMetadata | None = None + + +class RuntimeMetadata(StrictFrozenModel): + anonymizer_version: StrictStr | None = None + datadesigner_version: StrictStr | None = None + wandb_version: StrictStr | None = None + python_version: StrictStr | None = None + platform_machine: StrictStr | None = None + platform_system: StrictStr | None = None + + +class GitMetadata(StrictFrozenModel): + commit: StrictStr | None = None + branch: StrictStr | None = None + dirty: StrictBool | None = None + + +class ModelSourcesMetadata(StrictFrozenModel): + has_model_configs: StrictBool | None = None + has_model_providers: StrictBool | None = None + has_artifact_path: StrictBool | None = None + + +class WorkloadSourceMetadata(StrictFrozenModel): + kind: Literal["local_file", "remote_file"] | None = None + suffix: StrictStr | None = None + + +class WorkloadMetadata(StrictFrozenModel): + id: VisibleIdentifier | None = None + text_column: VisibleIdentifier | None = None + has_id_column: StrictBool | None = None + has_data_summary: StrictBool | None = None + row_limit: NonNegativeInt | None = None + row_offset: NonNegativeInt | None = None + source: WorkloadSourceMetadata | None = None + + +class DetectMetadata(StrictFrozenModel): + entity_label_source: Literal["custom", "default"] | None = None + entity_label_count: NonNegativeInt | None = None + entity_label_set_hash: StrictStr | None = None + gliner_threshold: Probability | None = None + validation_max_entities_per_call: NonNegativeInt | None = None + validation_excerpt_window_chars: NonNegativeInt | None = None + + +class ReplaceMetadata(StrictFrozenModel): + strategy: Literal["annotate", "hash", "redact", "substitute"] | None = None + normalize_label: StrictBool | None = None + algorithm: Literal["md5", "sha1", "sha256"] | None = None + digest_length: NonNegativeInt | None = None + has_format_template: StrictBool | None = None + has_instructions: StrictBool | None = None + + +class RewriteMetadata(StrictFrozenModel): + risk_tolerance: Literal["minimal", "low", "moderate", "high"] | None = None + max_repair_iterations: NonNegativeInt | None = None + strict_entity_protection: StrictBool | None = None + has_privacy_goal: StrictBool | None = None + has_protect: StrictBool | None = None + has_preserve: StrictBool | None = None + has_instructions: StrictBool | None = None + + +class ConfigMetadata(StrictFrozenModel): + id: VisibleIdentifier | None = None + mode: Literal["replace", "rewrite"] | None = None + evaluate: StrictBool | None = None + emit_telemetry: StrictBool | None = None + detect: DetectMetadata | None = None + replace: ReplaceMetadata | None = None + rewrite: RewriteMetadata | None = None + + +class MatrixMetadata(StrictFrozenModel): + workload: VisibleIdentifier | None = None + config: VisibleIdentifier | None = None + repetitions: NonNegativeInt | None = None + + +class SweepMetadata(StrictFrozenModel): + id: VisibleIdentifier + arm_id: VisibleIdentifier + params: dict[StrictStr, SafeScalar] + + @field_validator("params") + @classmethod + def validate_params(cls, params: dict[str, SafeScalar]) -> dict[str, SafeScalar]: + for key, value in params.items(): + if not _sweep_parameter_is_safe(key, value): + raise ValueError(f"unsupported W&B sweep parameter: {key}") + return params + + @classmethod + def from_arm(cls, *, sweep_id: str, arm_id: str, params: dict[str, Any]) -> SweepMetadata: + safe_params = { + key: value + for key, value in params.items() + if isinstance(value, bool | int | float | str) and _sweep_parameter_is_safe(key, value) + } + return cls(id=sweep_id, arm_id=arm_id, params=safe_params) + + +class ImportedRunMetadata(StrictFrozenModel): + completion_seal_schema_version: NonNegativeInt + completion_seal_sha256: StrictStr = Field(pattern=r"^[0-9a-f]{64}$") + producer_repository: VisibleIdentifier + producer_commit: StrictStr = Field(pattern=r"^[0-9a-f]{40,64}$") + phase: VisibleIdentifier + case_id: VisibleIdentifier + + +class WandbRunMetadata(StrictFrozenModel): + """Validated benchmark metadata that is safe to project into W&B config.""" + + run_kind: Literal["native_suite", "sweep_arm", "imported_case"] = "native_suite" + benchmark: BenchmarkMetadata | None = None + execution: ExecutionMetadata | None = None + runtime: RuntimeMetadata | None = None + git: GitMetadata | None = None + model_sources: ModelSourcesMetadata | None = None + workloads: tuple[WorkloadMetadata, ...] = () + configs: tuple[ConfigMetadata, ...] = () + matrix: tuple[MatrixMetadata, ...] = () + sweep: SweepMetadata | None = None + imported: ImportedRunMetadata | None = None + + @field_validator("workloads", "configs", "matrix", mode="before") + @classmethod + def freeze_sequences(cls, value: Any) -> tuple[Any, ...]: + if not isinstance(value, list | tuple): + raise ValueError("W&B metadata sequences must be lists or tuples") + return tuple(value) + + @model_validator(mode="after") + def validate_run_kind(self) -> WandbRunMetadata: + if (self.sweep is not None) != (self.run_kind == "sweep_arm"): + raise ValueError("sweep metadata and sweep_arm run kind must be set together") + if (self.imported is not None) != (self.run_kind == "imported_case"): + raise ValueError("import metadata and imported_case run kind must be set together") + if self.benchmark is None or self.benchmark.suite_id is None: + raise ValueError("W&B run metadata requires a benchmark suite identity") + if self.run_kind == "imported_case": + config_ids = [config.id for config in self.configs if config.id is not None] + if len(config_ids) != 1: + raise ValueError("imported W&B metadata requires exactly one config identity") + return self + + def comparer_values(self) -> dict[str, SafeScalar]: + values: dict[str, SafeScalar] = {} + benchmark = self.benchmark + if benchmark is not None: + _set_scalar(values, "benchmark_suite_id", benchmark.suite_id) + _set_scalar(values, "benchmark_case_count", benchmark.case_count) + _set_scalar(values, "benchmark_workload_ids", _compact_values(item.id for item in self.workloads)) + _set_scalar( + values, + "benchmark_workload_row_limits", + _compact_values(item.row_limit for item in self.workloads), + ) + _set_scalar( + values, + "benchmark_workload_source_kinds", + _compact_values(item.source.kind if item.source else None for item in self.workloads), + ) + _set_scalar( + values, + "benchmark_workload_source_suffixes", + _compact_values(item.source.suffix if item.source else None for item in self.workloads), + ) + _set_scalar(values, "benchmark_config_ids", _compact_values(item.id for item in self.configs)) + _set_scalar(values, "benchmark_modes", _compact_values(item.mode for item in self.configs)) + _set_scalar(values, "benchmark_strategies", _compact_values(_config_strategy(item) for item in self.configs)) + _set_scalar( + values, + "benchmark_gliner_thresholds", + _compact_values(item.detect.gliner_threshold if item.detect else None for item in self.configs), + ) + _set_scalar( + values, + "benchmark_entity_label_counts", + _compact_values(item.detect.entity_label_count if item.detect else None for item in self.configs), + ) + _set_scalar( + values, + "benchmark_risk_tolerances", + _compact_values(item.rewrite.risk_tolerance if item.rewrite else None for item in self.configs), + ) + return values + + +class WandbConfigPayload(WandbRunMetadata): + suite_id: VisibleIdentifier + wandb_mode: WandbMode + wandb_log_tables: StrictBool + benchmark_suite_id: VisibleIdentifier | None = None + benchmark_case_count: NonNegativeInt | None = None + benchmark_workload_ids: SafeScalar | None = None + benchmark_workload_row_limits: SafeScalar | None = None + benchmark_workload_source_kinds: SafeScalar | None = None + benchmark_workload_source_suffixes: SafeScalar | None = None + benchmark_config_ids: SafeScalar | None = None + benchmark_modes: SafeScalar | None = None + benchmark_strategies: SafeScalar | None = None + benchmark_gliner_thresholds: SafeScalar | None = None + benchmark_entity_label_counts: SafeScalar | None = None + benchmark_risk_tolerances: SafeScalar | None = None + native_suite_id: VisibleIdentifier | None = None + sweep_id: VisibleIdentifier | None = None + sweep_arm_id: VisibleIdentifier | None = None + imported_config_id: VisibleIdentifier | None = None + sweep_params: dict[StrictStr, SafeScalar] = Field(default_factory=dict) + + def sdk_values(self) -> dict[str, Any]: + values = self.model_dump(mode="json", exclude_none=True) + values.update({f"sweep_param_{key}": value for key, value in self.sweep_params.items()}) + return values + + @classmethod + def from_run_metadata( + cls, + settings: ResolvedWandbConfig, + *, + suite_id: str, + metadata: WandbRunMetadata, + ) -> WandbConfigPayload: + values: dict[str, Any] = { + "suite_id": suite_id, + "wandb_mode": settings.wandb_mode, + "wandb_log_tables": settings.wandb_log_tables, + } + if metadata.benchmark is None or metadata.benchmark.suite_id != suite_id: + raise ValueError("W&B metadata suite identity does not match publication suite") + values.update({name: getattr(metadata, name) for name in WandbRunMetadata.model_fields}) + values.update(metadata.comparer_values()) + if metadata.run_kind == "native_suite": + values["native_suite_id"] = metadata.benchmark.suite_id + if metadata.sweep is not None: + values.update( + { + "sweep_id": metadata.sweep.id, + "sweep_arm_id": metadata.sweep.arm_id, + "sweep_params": metadata.sweep.params, + } + ) + if metadata.run_kind == "imported_case": + values["imported_config_id"] = metadata.configs[0].id + return cls.model_validate(values, strict=True) + + +def _set_scalar(mapping: dict[str, SafeScalar], key: str, value: Any) -> None: + if isinstance(value, bool | int | float | str): + mapping[key] = value + + +def _compact_values(values: Any) -> SafeScalar | None: + compacted: list[SafeScalar] = [] + seen: set[tuple[str, SafeScalar]] = set() + for value in values: + if not isinstance(value, bool | int | float | str): + continue + marker = (type(value).__name__, value) + if marker in seen: + continue + compacted.append(value) + seen.add(marker) + if len(compacted) == 1: + return compacted[0] + if len(compacted) > 1: + return ",".join(str(value) for value in compacted) + return None + + +def _config_strategy(config: ConfigMetadata) -> str | None: + if config.replace is not None: + return config.replace.strategy + return "rewrite" if config.rewrite is not None else None + + +def _sweep_parameter_is_safe(key: str, value: SafeScalar) -> bool: + if not key.startswith("configs_") or not key.isascii() or any(not (char.isalnum() or char == "_") for char in key): + return False + if any(key.endswith(f"_{tail}") for tail in _NUMERIC_SWEEP_PARAMETER_TAILS): + if isinstance(value, bool) or not isinstance(value, int | float) or value < 0: + return False + return not key.endswith("_detect_gliner_threshold") or value <= 1 + if any(key.endswith(f"_{tail}") for tail in _BOOLEAN_SWEEP_PARAMETER_TAILS): + return isinstance(value, bool) + return any(key.endswith(f"_{tail}") and value in allowed for tail, allowed in _ENUM_SWEEP_PARAMETER_VALUES.items()) diff --git a/tools/measurement/measurement_tools/wandb_models.py b/tools/measurement/measurement_tools/wandb_models.py index bc5a07ae..c72ea37e 100644 --- a/tools/measurement/measurement_tools/wandb_models.py +++ b/tools/measurement/measurement_tools/wandb_models.py @@ -4,54 +4,81 @@ from __future__ import annotations -import hashlib -import re from collections.abc import Mapping from enum import StrEnum -from ipaddress import ip_address from pathlib import Path from types import MappingProxyType -from typing import Annotated, Any, Literal, cast -from urllib.parse import urlsplit +from typing import Annotated, Literal, cast from pydantic import ( BaseModel, - ConfigDict, Field, - FiniteFloat, StrictBool, StrictInt, StrictStr, field_validator, model_validator, ) -from pydantic_settings import BaseSettings, SettingsConfigDict from measurement_tools.validation import ( MeasurementStrategy, NonNegativeFloat, NonNegativeInt, Probability, - RedactedStrictFrozenModel, StrictFrozenModel, VisibleIdentifier, - VisibleSlurmIdentifier, +) +from measurement_tools.wandb_metadata import ( + BenchmarkMetadata, + ConfigMetadata, + DetectMetadata, + ExecutionMetadata, + GitMetadata, + ImportedRunMetadata, + MatrixMetadata, + ModelSourcesMetadata, + ReplaceMetadata, + RewriteMetadata, + RuntimeMetadata, + SafeScalar, + SlurmJobMetadata, + SlurmMetadata, + SweepMetadata, + WandbConfigPayload, + WandbRunMetadata, + WandbTag, + WorkloadMetadata, + WorkloadSourceMetadata, ) from measurement_tools.wandb_metric_schema import AGGREGATED_MEASUREMENT_FIELDS, BENCHMARK_METRIC_NAMES +from measurement_tools.wandb_policy import DataClass, Exposure, FieldPolicy +from measurement_tools.wandb_settings import ( + DEFAULT_WANDB_PROJECT as DEFAULT_WANDB_PROJECT, +) +from measurement_tools.wandb_settings import ( + WANDB_TAG_MAX_LENGTH as WANDB_TAG_MAX_LENGTH, +) +from measurement_tools.wandb_settings import ( + ResolvedWandbConfig as ResolvedWandbConfig, +) +from measurement_tools.wandb_settings import ( + WandbInputs as WandbInputs, +) +from measurement_tools.wandb_settings import ( + WandbMode, +) +from measurement_tools.wandb_settings import ( + generated_wandb_tag as generated_wandb_tag, +) +from measurement_tools.wandb_settings import ( + is_safe_wandb_tag as is_safe_wandb_tag, +) +from measurement_tools.wandb_settings import ( + wandb_tag_value_is_sensitive as wandb_tag_value_is_sensitive, +) -DEFAULT_WANDB_PROJECT = "nemo-anonymizer-benchmarks" PUBLICATION_COMPLETE_KEY = "publication/complete" PUBLICATION_SEAL_DIGEST_KEY = "publication/completion_seal_sha256" -WANDB_TAG_MAX_LENGTH = 64 -_WANDB_TAG_DIGEST_LENGTH = 12 -_SENSITIVE_WANDB_TAG_PARTS = frozenset({"api_key", "credential", "credentials", "password", "secret", "token"}) -_SENSITIVE_WANDB_TAG_VALUE_PREFIXES = ("sk-", "ghp_", "github_pat_", "glpat-", "xoxb-", "xoxp-", "xoxa-") - - -class WandbMode(StrEnum): - online = "online" - offline = "offline" - disabled = "disabled" class WandbPublicationState(StrEnum): @@ -60,192 +87,6 @@ class WandbPublicationState(StrEnum): already_complete = "already_complete" -class WandbInputs(BaseSettings): - """Optional operator inputs; dedicated measurement variables only.""" - - wandb_mode: WandbMode | None = None - wandb_project: str | None = None - wandb_base_url: str | None = None - wandb_entity: str | None = None - wandb_group: str | None = None - wandb_job_type: str | None = None - wandb_run_name: str | None = None - wandb_tags: str | None = None - wandb_log_tables: bool | None = None - - model_config = SettingsConfigDict( - env_prefix="ANONYMIZER_MEASUREMENT_", - extra="forbid", - env_ignore_empty=True, - populate_by_name=True, - hide_input_in_errors=True, - ) - - -class ResolvedWandbConfig(RedactedStrictFrozenModel): - """Immutable, fully resolved publisher configuration.""" - - wandb_mode: WandbMode = WandbMode.disabled - wandb_project: str = DEFAULT_WANDB_PROJECT - wandb_base_url: str | None = None - wandb_entity: str | None = None - wandb_group: str | None = None - wandb_job_type: str | None = None - wandb_run_name: str | None = None - wandb_tags: str = "" - wandb_log_tables: bool = False - - @field_validator( - "wandb_project", - "wandb_base_url", - "wandb_entity", - "wandb_group", - "wandb_job_type", - "wandb_run_name", - ) - @classmethod - def nonempty_optional_strings(cls, value: str | None) -> str | None: - if value is None: - return None - stripped = value.strip() - if not stripped: - raise ValueError("W&B string settings cannot be empty") - return stripped - - @field_validator("wandb_base_url") - @classmethod - def validate_base_url(cls, value: str | None) -> str | None: - if value is None: - return None - parsed = urlsplit(value) - if ( - parsed.scheme not in {"http", "https"} - or parsed.hostname is None - or parsed.username is not None - or parsed.password is not None - or parsed.query - or parsed.fragment - ): - raise ValueError("W&B base URL must be a credential-free HTTP(S) URL") - if parsed.scheme == "http" and not _is_loopback_host(parsed.hostname): - raise ValueError("W&B base URL must use HTTPS unless it targets loopback") - return value.rstrip("/") - - @field_validator("wandb_tags") - @classmethod - def validate_tags(cls, value: str) -> str: - tags = [tag for tag in (part.strip() for part in value.split(",")) if tag] - if any(not is_safe_wandb_tag(tag) for tag in tags): - raise ValueError("W&B tags must be safe identifiers between 1 and 64 characters") - return value - - @property - def enabled(self) -> bool: - return self.wandb_mode != WandbMode.disabled - - @property - def effective_wandb_project(self) -> str: - return self.wandb_project - - @property - def effective_wandb_tags(self) -> list[str]: - return [tag for tag in (part.strip() for part in self.wandb_tags.split(",")) if tag] - - @classmethod - def from_env_and_overrides( - cls, - *, - defaults: Mapping[str, Any] | None = None, - **overrides: Any, - ) -> ResolvedWandbConfig: - source = WandbInputs() - values = dict(defaults or {}) - values.update(source.model_dump(exclude_none=True)) - values.update({key: value for key, value in overrides.items() if value is not None}) - return cls.model_validate(values) - - def validated_update(self, **updates: Any) -> ResolvedWandbConfig: - values = self.model_dump() - values.update(updates) - return type(self).model_validate(values) - - -def _is_loopback_host(hostname: str) -> bool: - if hostname.lower() == "localhost": - return True - try: - return ip_address(hostname).is_loopback - except ValueError: - return False - - -def is_safe_wandb_tag(value: str) -> bool: - if not value or len(value) > WANDB_TAG_MAX_LENGTH or "://" in value or value.startswith("/"): - return False - return not wandb_tag_value_is_sensitive(value) - - -def generated_wandb_tag(namespace: str, value: str) -> str | None: - candidate = f"{namespace}:{value}" - if wandb_tag_value_is_sensitive(candidate) or "://" in candidate or candidate.startswith("/"): - return None - if len(candidate) <= WANDB_TAG_MAX_LENGTH: - return candidate - digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:_WANDB_TAG_DIGEST_LENGTH] - prefix_length = WANDB_TAG_MAX_LENGTH - len(namespace) - len(digest) - 2 - return f"{namespace}:{value[:prefix_length]}-{digest}" - - -def wandb_tag_value_is_sensitive(value: str) -> bool: - normalized = value.lower() - if normalized.startswith(_SENSITIVE_WANDB_TAG_VALUE_PREFIXES): - return True - parts = {part for part in re.split(r"[^a-z0-9]+", normalized) if part} - if "api" in parts and "key" in parts: - return True - return bool(parts & _SENSITIVE_WANDB_TAG_PARTS) - - -class DataClass(StrEnum): - operational = "operational" - pseudonymous = "pseudonymous" - sensitive = "sensitive" - - -class Exposure(StrEnum): - local_only = "local_only" - aggregate = "aggregate" - table_opt_in = "table_opt_in" - never = "never" - - -class FieldPolicy(BaseModel): - data_class: DataClass - exposure: Exposure - - model_config = ConfigDict(extra="forbid", frozen=True) - - -SafeScalar = StrictBool | StrictInt | FiniteFloat | StrictStr -WandbTag = Annotated[StrictStr, Field(min_length=1, max_length=WANDB_TAG_MAX_LENGTH)] -_NUMERIC_SWEEP_PARAMETER_TAILS = frozenset( - { - "detect_gliner_threshold", - "detect_validation_excerpt_window_chars", - "detect_validation_max_entities_per_call", - "replace_digest_length", - "rewrite_max_repair_iterations", - } -) -_BOOLEAN_SWEEP_PARAMETER_TAILS = frozenset( - {"emit_telemetry", "evaluate", "replace_normalize_label", "rewrite_strict_entity_protection"} -) -_ENUM_SWEEP_PARAMETER_VALUES = { - "replace_algorithm": frozenset({"md5", "sha1", "sha256"}), - "rewrite_risk_tolerance": frozenset({"high", "low", "minimal", "moderate"}), -} - - class WandbInitPayload(StrictFrozenModel): run_id: VisibleIdentifier resume: Literal["allow", "never"] = "never" @@ -259,298 +100,6 @@ class WandbInitPayload(StrictFrozenModel): tags: tuple[WandbTag, ...] = () -class BenchmarkMetadata(StrictFrozenModel): - metadata_schema_version: NonNegativeInt | None = None - suite_schema_version: NonNegativeInt | None = None - wandb_sanitizer_version: NonNegativeInt | None = None - measurement_schema_version: NonNegativeInt | None = None - suite_id: VisibleIdentifier | None = None - workload_count: NonNegativeInt | None = None - config_count: NonNegativeInt | None = None - matrix_entry_count: NonNegativeInt | None = None - case_count: NonNegativeInt | None = None - case_retries: NonNegativeInt | None = None - case_retry_backoff_sec: NonNegativeFloat | None = None - suite_file_hash: StrictStr | None = None - - -class SlurmJobMetadata(StrictFrozenModel): - role: VisibleSlurmIdentifier - job_id: VisibleSlurmIdentifier - - -class SlurmMetadata(StrictFrozenModel): - job_id: VisibleIdentifier | None = None - array_job_id: VisibleIdentifier | None = None - array_task_id: VisibleIdentifier | None = None - array_task_count: NonNegativeInt | None = None - restart_count: NonNegativeInt | None = None - ntasks: NonNegativeInt | None = None - job_num_nodes: NonNegativeInt | None = None - jobs: Annotated[list[SlurmJobMetadata], Field(max_length=64, exclude_if=lambda value: not value)] = Field( - default_factory=list - ) - - -class ExecutionMetadata(StrictFrozenModel): - backend: Literal["local", "slurm"] | None = None - export: StrictBool | None = None - fail_fast: StrictBool | None = None - dd_trace: Literal["none", "last_message", "all_messages"] | None = None - dd_task_trace: StrictBool | None = None - slurm: SlurmMetadata | None = None - - -class RuntimeMetadata(StrictFrozenModel): - anonymizer_version: StrictStr | None = None - datadesigner_version: StrictStr | None = None - wandb_version: StrictStr | None = None - python_version: StrictStr | None = None - platform_machine: StrictStr | None = None - platform_system: StrictStr | None = None - - -class GitMetadata(StrictFrozenModel): - commit: StrictStr | None = None - branch: StrictStr | None = None - dirty: StrictBool | None = None - - -class ModelSourcesMetadata(StrictFrozenModel): - has_model_configs: StrictBool | None = None - has_model_providers: StrictBool | None = None - has_artifact_path: StrictBool | None = None - - -class WorkloadSourceMetadata(StrictFrozenModel): - kind: Literal["local_file", "remote_file"] | None = None - suffix: StrictStr | None = None - - -class WorkloadMetadata(StrictFrozenModel): - id: VisibleIdentifier | None = None - text_column: VisibleIdentifier | None = None - has_id_column: StrictBool | None = None - has_data_summary: StrictBool | None = None - row_limit: NonNegativeInt | None = None - row_offset: NonNegativeInt | None = None - source: WorkloadSourceMetadata | None = None - - -class DetectMetadata(StrictFrozenModel): - entity_label_source: Literal["custom", "default"] | None = None - entity_label_count: NonNegativeInt | None = None - entity_label_set_hash: StrictStr | None = None - gliner_threshold: Probability | None = None - validation_max_entities_per_call: NonNegativeInt | None = None - validation_excerpt_window_chars: NonNegativeInt | None = None - - -class ReplaceMetadata(StrictFrozenModel): - strategy: Literal["annotate", "hash", "redact", "substitute"] | None = None - normalize_label: StrictBool | None = None - algorithm: Literal["md5", "sha1", "sha256"] | None = None - digest_length: NonNegativeInt | None = None - has_format_template: StrictBool | None = None - has_instructions: StrictBool | None = None - - -class RewriteMetadata(StrictFrozenModel): - risk_tolerance: Literal["minimal", "low", "moderate", "high"] | None = None - max_repair_iterations: NonNegativeInt | None = None - strict_entity_protection: StrictBool | None = None - has_privacy_goal: StrictBool | None = None - has_protect: StrictBool | None = None - has_preserve: StrictBool | None = None - has_instructions: StrictBool | None = None - - -class ConfigMetadata(StrictFrozenModel): - id: VisibleIdentifier | None = None - mode: Literal["replace", "rewrite"] | None = None - evaluate: StrictBool | None = None - emit_telemetry: StrictBool | None = None - detect: DetectMetadata | None = None - replace: ReplaceMetadata | None = None - rewrite: RewriteMetadata | None = None - - -class MatrixMetadata(StrictFrozenModel): - workload: VisibleIdentifier | None = None - config: VisibleIdentifier | None = None - repetitions: NonNegativeInt | None = None - - -class SweepMetadata(StrictFrozenModel): - id: VisibleIdentifier - arm_id: VisibleIdentifier - params: dict[StrictStr, SafeScalar] - - @field_validator("params") - @classmethod - def validate_params(cls, params: dict[str, SafeScalar]) -> dict[str, SafeScalar]: - for key, value in params.items(): - if not _sweep_parameter_is_safe(key, value): - raise ValueError(f"unsupported W&B sweep parameter: {key}") - return params - - @classmethod - def from_arm(cls, *, sweep_id: str, arm_id: str, params: dict[str, Any]) -> SweepMetadata: - safe_params = { - key: value - for key, value in params.items() - if isinstance(value, bool | int | float | str) and _sweep_parameter_is_safe(key, value) - } - return cls(id=sweep_id, arm_id=arm_id, params=safe_params) - - -class ImportedRunMetadata(StrictFrozenModel): - completion_seal_schema_version: NonNegativeInt - completion_seal_sha256: StrictStr = Field(pattern=r"^[0-9a-f]{64}$") - producer_repository: VisibleIdentifier - producer_commit: StrictStr = Field(pattern=r"^[0-9a-f]{40,64}$") - phase: VisibleIdentifier - case_id: VisibleIdentifier - - -class WandbRunMetadata(StrictFrozenModel): - """Validated benchmark metadata that is safe to project into W&B config.""" - - run_kind: Literal["native_suite", "sweep_arm", "imported_case"] = "native_suite" - benchmark: BenchmarkMetadata | None = None - execution: ExecutionMetadata | None = None - runtime: RuntimeMetadata | None = None - git: GitMetadata | None = None - model_sources: ModelSourcesMetadata | None = None - workloads: tuple[WorkloadMetadata, ...] = () - configs: tuple[ConfigMetadata, ...] = () - matrix: tuple[MatrixMetadata, ...] = () - sweep: SweepMetadata | None = None - imported: ImportedRunMetadata | None = None - - @field_validator("workloads", "configs", "matrix", mode="before") - @classmethod - def freeze_sequences(cls, value: Any) -> tuple[Any, ...]: - if not isinstance(value, list | tuple): - raise ValueError("W&B metadata sequences must be lists or tuples") - return tuple(value) - - @model_validator(mode="after") - def validate_run_kind(self) -> WandbRunMetadata: - if (self.sweep is not None) != (self.run_kind == "sweep_arm"): - raise ValueError("sweep metadata and sweep_arm run kind must be set together") - if (self.imported is not None) != (self.run_kind == "imported_case"): - raise ValueError("import metadata and imported_case run kind must be set together") - if self.benchmark is None or self.benchmark.suite_id is None: - raise ValueError("W&B run metadata requires a benchmark suite identity") - if self.run_kind == "imported_case": - config_ids = [config.id for config in self.configs if config.id is not None] - if len(config_ids) != 1: - raise ValueError("imported W&B metadata requires exactly one config identity") - return self - - def comparer_values(self) -> dict[str, SafeScalar]: - values: dict[str, SafeScalar] = {} - benchmark = self.benchmark - if benchmark is not None: - _set_scalar(values, "benchmark_suite_id", benchmark.suite_id) - _set_scalar(values, "benchmark_case_count", benchmark.case_count) - _set_scalar(values, "benchmark_workload_ids", _compact_values(item.id for item in self.workloads)) - _set_scalar( - values, - "benchmark_workload_row_limits", - _compact_values(item.row_limit for item in self.workloads), - ) - _set_scalar( - values, - "benchmark_workload_source_kinds", - _compact_values(item.source.kind if item.source else None for item in self.workloads), - ) - _set_scalar( - values, - "benchmark_workload_source_suffixes", - _compact_values(item.source.suffix if item.source else None for item in self.workloads), - ) - _set_scalar(values, "benchmark_config_ids", _compact_values(item.id for item in self.configs)) - _set_scalar(values, "benchmark_modes", _compact_values(item.mode for item in self.configs)) - _set_scalar(values, "benchmark_strategies", _compact_values(_config_strategy(item) for item in self.configs)) - _set_scalar( - values, - "benchmark_gliner_thresholds", - _compact_values(item.detect.gliner_threshold if item.detect else None for item in self.configs), - ) - _set_scalar( - values, - "benchmark_entity_label_counts", - _compact_values(item.detect.entity_label_count if item.detect else None for item in self.configs), - ) - _set_scalar( - values, - "benchmark_risk_tolerances", - _compact_values(item.rewrite.risk_tolerance if item.rewrite else None for item in self.configs), - ) - return values - - -class WandbConfigPayload(WandbRunMetadata): - suite_id: VisibleIdentifier - wandb_mode: WandbMode - wandb_log_tables: StrictBool - benchmark_suite_id: VisibleIdentifier | None = None - benchmark_case_count: NonNegativeInt | None = None - benchmark_workload_ids: SafeScalar | None = None - benchmark_workload_row_limits: SafeScalar | None = None - benchmark_workload_source_kinds: SafeScalar | None = None - benchmark_workload_source_suffixes: SafeScalar | None = None - benchmark_config_ids: SafeScalar | None = None - benchmark_modes: SafeScalar | None = None - benchmark_strategies: SafeScalar | None = None - benchmark_gliner_thresholds: SafeScalar | None = None - benchmark_entity_label_counts: SafeScalar | None = None - benchmark_risk_tolerances: SafeScalar | None = None - native_suite_id: VisibleIdentifier | None = None - sweep_id: VisibleIdentifier | None = None - sweep_arm_id: VisibleIdentifier | None = None - imported_config_id: VisibleIdentifier | None = None - sweep_params: dict[StrictStr, SafeScalar] = Field(default_factory=dict) - - def sdk_values(self) -> dict[str, Any]: - values = self.model_dump(mode="json", exclude_none=True) - values.update({f"sweep_param_{key}": value for key, value in self.sweep_params.items()}) - return values - - @classmethod - def from_run_metadata( - cls, - settings: ResolvedWandbConfig, - *, - suite_id: str, - metadata: WandbRunMetadata, - ) -> WandbConfigPayload: - values: dict[str, Any] = { - "suite_id": suite_id, - "wandb_mode": settings.wandb_mode, - "wandb_log_tables": settings.wandb_log_tables, - } - if metadata.benchmark is None or metadata.benchmark.suite_id != suite_id: - raise ValueError("W&B metadata suite identity does not match publication suite") - values.update({name: getattr(metadata, name) for name in WandbRunMetadata.model_fields}) - values.update(metadata.comparer_values()) - if metadata.run_kind == "native_suite": - values["native_suite_id"] = metadata.benchmark.suite_id - if metadata.sweep is not None: - values.update( - { - "sweep_id": metadata.sweep.id, - "sweep_arm_id": metadata.sweep.arm_id, - "sweep_params": metadata.sweep.params, - } - ) - if metadata.run_kind == "imported_case": - values["imported_config_id"] = metadata.configs[0].id - return cls.model_validate(values, strict=True) - - class WandbHistoryPayload(StrictFrozenModel): metrics: dict[StrictStr, SafeScalar] @@ -819,47 +368,6 @@ class WandbPublishResult(StrictFrozenModel): record_count: StrictInt = 0 -def _set_scalar(mapping: dict[str, SafeScalar], key: str, value: Any) -> None: - if isinstance(value, bool | int | float | str): - mapping[key] = value - - -def _compact_values(values: Any) -> SafeScalar | None: - compacted: list[SafeScalar] = [] - seen: set[tuple[str, SafeScalar]] = set() - for value in values: - if not isinstance(value, bool | int | float | str): - continue - marker = (type(value).__name__, value) - if marker in seen: - continue - compacted.append(value) - seen.add(marker) - if len(compacted) == 1: - return compacted[0] - if len(compacted) > 1: - return ",".join(str(value) for value in compacted) - return None - - -def _config_strategy(config: ConfigMetadata) -> str | None: - if config.replace is not None: - return config.replace.strategy - return "rewrite" if config.rewrite is not None else None - - -def _sweep_parameter_is_safe(key: str, value: SafeScalar) -> bool: - if not key.startswith("configs_") or not key.isascii() or any(not (char.isalnum() or char == "_") for char in key): - return False - if any(key.endswith(f"_{tail}") for tail in _NUMERIC_SWEEP_PARAMETER_TAILS): - if isinstance(value, bool) or not isinstance(value, int | float) or value < 0: - return False - return not key.endswith("_detect_gliner_threshold") or value <= 1 - if any(key.endswith(f"_{tail}") for tail in _BOOLEAN_SWEEP_PARAMETER_TAILS): - return isinstance(value, bool) - return any(key.endswith(f"_{tail}") and value in allowed for tail, allowed in _ENUM_SWEEP_PARAMETER_VALUES.items()) - - def _aggregate_policies( *names: str, data_class: DataClass = DataClass.operational, diff --git a/tools/measurement/measurement_tools/wandb_policy.py b/tools/measurement/measurement_tools/wandb_policy.py new file mode 100644 index 00000000..77309f43 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_policy.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dependency-light exposure vocabulary for W&B fields.""" + +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel, ConfigDict + +__all__ = ["DataClass", "Exposure", "FieldPolicy"] + + +class DataClass(StrEnum): + operational = "operational" + pseudonymous = "pseudonymous" + sensitive = "sensitive" + + +class Exposure(StrEnum): + local_only = "local_only" + aggregate = "aggregate" + table_opt_in = "table_opt_in" + never = "never" + + +class FieldPolicy(BaseModel): + data_class: DataClass + exposure: Exposure + + model_config = ConfigDict(extra="forbid", frozen=True) diff --git a/tools/measurement/measurement_tools/wandb_settings.py b/tools/measurement/measurement_tools/wandb_settings.py new file mode 100644 index 00000000..acad5409 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_settings.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Resolved W&B settings and tag safety.""" + +from __future__ import annotations + +import hashlib +import re +from collections.abc import Mapping +from enum import StrEnum +from ipaddress import ip_address +from typing import Any +from urllib.parse import urlsplit + +from pydantic import field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + +from measurement_tools.validation import RedactedStrictFrozenModel + +__all__ = [ + "DEFAULT_WANDB_PROJECT", + "WANDB_TAG_MAX_LENGTH", + "ResolvedWandbConfig", + "WandbInputs", + "WandbMode", + "generated_wandb_tag", + "is_safe_wandb_tag", + "wandb_tag_value_is_sensitive", +] + +DEFAULT_WANDB_PROJECT = "nemo-anonymizer-benchmarks" +WANDB_TAG_MAX_LENGTH = 64 +_WANDB_TAG_DIGEST_LENGTH = 12 +_SENSITIVE_WANDB_TAG_PARTS = frozenset({"api_key", "credential", "credentials", "password", "secret", "token"}) +_SENSITIVE_WANDB_TAG_VALUE_PREFIXES = ("sk-", "ghp_", "github_pat_", "glpat-", "xoxb-", "xoxp-", "xoxa-") + + +class WandbMode(StrEnum): + online = "online" + offline = "offline" + disabled = "disabled" + + +class WandbInputs(BaseSettings): + """Optional operator inputs; dedicated measurement variables only.""" + + wandb_mode: WandbMode | None = None + wandb_project: str | None = None + wandb_base_url: str | None = None + wandb_entity: str | None = None + wandb_group: str | None = None + wandb_job_type: str | None = None + wandb_run_name: str | None = None + wandb_tags: str | None = None + wandb_log_tables: bool | None = None + + model_config = SettingsConfigDict( + env_prefix="ANONYMIZER_MEASUREMENT_", + extra="forbid", + env_ignore_empty=True, + populate_by_name=True, + hide_input_in_errors=True, + ) + + +class ResolvedWandbConfig(RedactedStrictFrozenModel): + """Immutable, fully resolved publisher configuration.""" + + wandb_mode: WandbMode = WandbMode.disabled + wandb_project: str = DEFAULT_WANDB_PROJECT + wandb_base_url: str | None = None + wandb_entity: str | None = None + wandb_group: str | None = None + wandb_job_type: str | None = None + wandb_run_name: str | None = None + wandb_tags: str = "" + wandb_log_tables: bool = False + + @field_validator( + "wandb_project", + "wandb_base_url", + "wandb_entity", + "wandb_group", + "wandb_job_type", + "wandb_run_name", + ) + @classmethod + def nonempty_optional_strings(cls, value: str | None) -> str | None: + if value is None: + return None + stripped = value.strip() + if not stripped: + raise ValueError("W&B string settings cannot be empty") + return stripped + + @field_validator("wandb_base_url") + @classmethod + def validate_base_url(cls, value: str | None) -> str | None: + if value is None: + return None + parsed = urlsplit(value) + if ( + parsed.scheme not in {"http", "https"} + or parsed.hostname is None + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + ): + raise ValueError("W&B base URL must be a credential-free HTTP(S) URL") + if parsed.scheme == "http" and not _is_loopback_host(parsed.hostname): + raise ValueError("W&B base URL must use HTTPS unless it targets loopback") + return value.rstrip("/") + + @field_validator("wandb_tags") + @classmethod + def validate_tags(cls, value: str) -> str: + tags = [tag for tag in (part.strip() for part in value.split(",")) if tag] + if any(not is_safe_wandb_tag(tag) for tag in tags): + raise ValueError("W&B tags must be safe identifiers between 1 and 64 characters") + return value + + @property + def enabled(self) -> bool: + return self.wandb_mode != WandbMode.disabled + + @property + def effective_wandb_project(self) -> str: + return self.wandb_project + + @property + def effective_wandb_tags(self) -> list[str]: + return [tag for tag in (part.strip() for part in self.wandb_tags.split(",")) if tag] + + @classmethod + def from_env_and_overrides( + cls, + *, + defaults: Mapping[str, Any] | None = None, + **overrides: Any, + ) -> ResolvedWandbConfig: + source = WandbInputs() + values = dict(defaults or {}) + values.update(source.model_dump(exclude_none=True)) + values.update({key: value for key, value in overrides.items() if value is not None}) + return cls.model_validate(values) + + def validated_update(self, **updates: Any) -> ResolvedWandbConfig: + values = self.model_dump() + values.update(updates) + return type(self).model_validate(values) + + +def _is_loopback_host(hostname: str) -> bool: + if hostname.lower() == "localhost": + return True + try: + return ip_address(hostname).is_loopback + except ValueError: + return False + + +def is_safe_wandb_tag(value: str) -> bool: + if not value or len(value) > WANDB_TAG_MAX_LENGTH or "://" in value or value.startswith("/"): + return False + return not wandb_tag_value_is_sensitive(value) + + +def generated_wandb_tag(namespace: str, value: str) -> str | None: + candidate = f"{namespace}:{value}" + if wandb_tag_value_is_sensitive(candidate) or "://" in candidate or candidate.startswith("/"): + return None + if len(candidate) <= WANDB_TAG_MAX_LENGTH: + return candidate + digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:_WANDB_TAG_DIGEST_LENGTH] + prefix_length = WANDB_TAG_MAX_LENGTH - len(namespace) - len(digest) - 2 + return f"{namespace}:{value[:prefix_length]}-{digest}" + + +def wandb_tag_value_is_sensitive(value: str) -> bool: + normalized = value.lower() + if normalized.startswith(_SENSITIVE_WANDB_TAG_VALUE_PREFIXES): + return True + parts = {part for part in re.split(r"[^a-z0-9]+", normalized) if part} + if "api" in parts and "key" in parts: + return True + return bool(parts & _SENSITIVE_WANDB_TAG_PARTS) From c0ba01384e3663dc7d0d4e8e24a2701dc58be769 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 19:12:31 +0000 Subject: [PATCH 03/17] refactor: split wandb metrics and publication Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 2 + tests/tools/test_measurement_wandb_logging.py | 4 +- .../measurement_tools/wandb_metrics.py | 292 ++++++++++++++ .../measurement_tools/wandb_models.py | 380 +++--------------- .../measurement_tools/wandb_publication.py | 78 ++++ 5 files changed, 419 insertions(+), 337 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_metrics.py create mode 100644 tools/measurement/measurement_tools/wandb_publication.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index d363c381..73c4dd8d 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -116,6 +116,8 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: "measurement_tools.wandb_policy", "measurement_tools.wandb_settings", "measurement_tools.wandb_metadata", + "measurement_tools.wandb_metrics", + "measurement_tools.wandb_publication", ], ) def test_wandb_model_facade_reexports_every_canonical_symbol(canonical_module: str) -> None: diff --git a/tests/tools/test_measurement_wandb_logging.py b/tests/tools/test_measurement_wandb_logging.py index 32fc13e0..aa13dcc1 100644 --- a/tests/tools/test_measurement_wandb_logging.py +++ b/tests/tools/test_measurement_wandb_logging.py @@ -3,10 +3,10 @@ from __future__ import annotations +import importlib import json import logging import os -import sys from pathlib import Path from types import ModuleType, SimpleNamespace from typing import Any @@ -608,7 +608,7 @@ def fail_publish(*_args: Any, **_kwargs: Any) -> Any: def test_wandb_outbound_models_structurally_exclude_sensitive_fields(wandb_setup_tool: ModuleType) -> None: - models = sys.modules[wandb_setup_tool.WandbPublishPayload.__module__] + models = importlib.import_module("measurement_tools.wandb_models") with pytest.raises(ValidationError, match="Extra inputs"): models.WandbHistoryPayload(metrics={}, text="Alice") diff --git a/tools/measurement/measurement_tools/wandb_metrics.py b/tools/measurement/measurement_tools/wandb_metrics.py new file mode 100644 index 00000000..87c935a0 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_metrics.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validated W&B metric and table projections.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Annotated, Literal + +from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, field_validator, model_validator + +from measurement_tools.validation import ( + MeasurementStrategy, + NonNegativeFloat, + NonNegativeInt, + Probability, + StrictFrozenModel, +) +from measurement_tools.wandb_metadata import SafeScalar +from measurement_tools.wandb_metric_schema import AGGREGATED_MEASUREMENT_FIELDS, BENCHMARK_METRIC_NAMES + +__all__ = [ + "EvaluationTableRow", + "ModelWorkflowTableRow", + "NddWorkflowTableRow", + "RecordTableRow", + "RunTableRow", + "StageTableRow", + "TraceCoverageTableRow", + "WANDB_TABLE_ROW_MODELS", + "WandbHistoryPayload", + "WandbSummaryPayload", + "WandbTablePayload", + "WandbTableRow", +] + +_PUBLICATION_COMPLETE_KEY = "publication/complete" +_PUBLICATION_SEAL_DIGEST_KEY = "publication/completion_seal_sha256" + + +class WandbHistoryPayload(StrictFrozenModel): + metrics: dict[StrictStr, SafeScalar] + + @field_validator("metrics") + @classmethod + def metric_keys_have_policy(cls, metrics: dict[str, SafeScalar]) -> dict[str, SafeScalar]: + unknown = [key for key in metrics if not _metric_has_aggregate_policy(key)] + if unknown: + raise ValueError(f"W&B metric has no aggregate exposure policy: {unknown!r}") + return metrics + + +class WandbSummaryPayload(StrictFrozenModel): + metrics: dict[StrictStr, SafeScalar] + + @field_validator("metrics") + @classmethod + def metric_keys_have_policy(cls, metrics: dict[str, SafeScalar]) -> dict[str, SafeScalar]: + return WandbHistoryPayload.metric_keys_have_policy(metrics) + + +class _MetricTableRow(StrictFrozenModel): + elapsed_sec: NonNegativeFloat | None = None + input_row_count: NonNegativeInt | None = None + seed_row_count: NonNegativeInt | None = None + output_row_count: NonNegativeInt | None = None + failed_record_count: NonNegativeInt | None = None + preview_num_records: NonNegativeInt | None = None + column_count: NonNegativeInt | None = None + observed_input_tokens: NonNegativeInt | None = None + observed_output_tokens: NonNegativeInt | None = None + observed_total_tokens: NonNegativeInt | None = None + observed_reasoning_tokens: NonNegativeInt | None = None + observed_successful_requests: NonNegativeInt | None = None + observed_failed_requests: NonNegativeInt | None = None + observed_total_requests: NonNegativeInt | None = None + text_length_chars: NonNegativeInt | None = None + text_length_tokens: NonNegativeInt | None = None + final_entity_count: NonNegativeInt | None = None + ground_truth_entity_count: NonNegativeInt | None = None + entity_true_positive_count: NonNegativeInt | None = None + entity_false_positive_count: NonNegativeInt | None = None + entity_false_negative_count: NonNegativeInt | None = None + entity_relaxed_gt_found_count: NonNegativeInt | None = None + entity_relaxed_detected_tp_count: NonNegativeInt | None = None + entity_relaxed_label_compatible_gt_found_count: NonNegativeInt | None = None + entity_relaxed_label_compatible_detected_tp_count: NonNegativeInt | None = None + replacement_count: NonNegativeInt | None = None + replacement_duplicate_value_count: NonNegativeInt | None = None + replacement_missing_final_entity_count: NonNegativeInt | None = None + replacement_missing_final_value_count: NonNegativeInt | None = None + replacement_synthetic_original_collision_count: NonNegativeInt | None = None + replacement_synthetic_original_collision_value_count: NonNegativeInt | None = None + original_value_leak_count: NonNegativeInt | None = None + detected_candidate_count: NonNegativeInt | None = None + validation_chunk_count: NonNegativeInt | None = None + llm_calls_estimated_total: NonNegativeInt | None = None + detection_invalid_entity_count: NonNegativeInt | None = None + type_fidelity_invalid_replacement_count: NonNegativeInt | None = None + relational_consistency_invalid_relation_count: NonNegativeInt | None = None + attribute_fidelity_invalid_entity_count: NonNegativeInt | None = None + observed_failed_request_rate: Probability | None = None + input_rows_per_sec: NonNegativeFloat | None = None + output_rows_per_sec: NonNegativeFloat | None = None + observed_tokens_per_sec: NonNegativeFloat | None = None + observed_requests_per_sec: NonNegativeFloat | None = None + observed_tokens_per_successful_request: NonNegativeFloat | None = None + entity_precision: Probability | None = None + entity_recall: Probability | None = None + entity_f1: Probability | None = None + utility_score: Probability | None = None + leakage_mass: NonNegativeFloat | None = None + weighted_leakage_rate: Probability | None = None + repair_iterations: NonNegativeInt | None = None + + +_METRIC_TABLE_POLICY_FIELDS = ( + "elapsed_sec", + "input_row_count", + "seed_row_count", + "output_row_count", + "failed_record_count", + "preview_num_records", + "column_count", + "observed_input_tokens", + "observed_output_tokens", + "observed_total_tokens", + "observed_reasoning_tokens", + "observed_successful_requests", + "observed_failed_requests", + "observed_total_requests", + "text_length_chars", + "text_length_tokens", + "final_entity_count", + "ground_truth_entity_count", + "entity_true_positive_count", + "entity_false_positive_count", + "entity_false_negative_count", + "entity_relaxed_gt_found_count", + "entity_relaxed_detected_tp_count", + "entity_relaxed_label_compatible_gt_found_count", + "entity_relaxed_label_compatible_detected_tp_count", + "replacement_count", + "replacement_duplicate_value_count", + "replacement_missing_final_entity_count", + "replacement_missing_final_value_count", + "replacement_synthetic_original_collision_count", + "replacement_synthetic_original_collision_value_count", + "original_value_leak_count", + "detected_candidate_count", + "validation_chunk_count", + "llm_calls_estimated_total", + "detection_invalid_entity_count", + "type_fidelity_invalid_replacement_count", + "relational_consistency_invalid_relation_count", + "attribute_fidelity_invalid_entity_count", + "observed_failed_request_rate", + "input_rows_per_sec", + "output_rows_per_sec", + "observed_tokens_per_sec", + "observed_requests_per_sec", + "observed_tokens_per_successful_request", + "entity_precision", + "entity_recall", + "entity_f1", + "utility_score", + "leakage_mass", + "weighted_leakage_rate", + "repair_iterations", +) + + +class RunTableRow(_MetricTableRow): + record_type: Literal["run"] + mode: Literal["replace", "rewrite"] + strategy: MeasurementStrategy + + +class StageTableRow(_MetricTableRow): + record_type: Literal["stage"] + stage: Literal[ + "Anonymizer._run_internal", + "EntityDetectionWorkflow.run", + "ReplacementWorkflow.run", + "RewriteWorkflow.run", + ] + status: Literal["completed", "error"] + + +class RecordTableRow(_MetricTableRow): + record_type: Literal["record"] + mode: Literal["replace", "rewrite"] + strategy: MeasurementStrategy + row_index: StrictInt | None + record_hash: Annotated[StrictStr, Field(pattern=r"^[0-9a-fA-F]{64}$")] + text_length_chars: None = None + text_length_tokens: None = None + text_length_chars_bucket: Literal["0", "1-127", "128-511", "512-2047", "2048-8191", "8192+"] + text_length_tokens_bucket: Literal["0", "1-127", "128-511", "512-2047", "2048-8191", "8192+"] + + +class EvaluationTableRow(RecordTableRow): + record_type: Literal["evaluation_record"] + detection_valid: StrictBool | None = None + type_fidelity_valid: StrictBool | None = None + relational_consistency_valid: StrictBool | None = None + attribute_fidelity_valid: StrictBool | None = None + + +class NddWorkflowTableRow(_MetricTableRow): + record_type: Literal["ndd_workflow"] + status: Literal["completed", "error"] + + +class ModelWorkflowTableRow(_MetricTableRow): + record_type: Literal["model_workflow"] + status: Literal["completed", "error"] + + +class TraceCoverageTableRow(StrictFrozenModel): + record_type: Literal["dd_trace_coverage"] + traced_column_count: NonNegativeInt + native_trace_column_count: NonNegativeInt + private_trace_column_count: NonNegativeInt + unsupported_column_count: NonNegativeInt + + +WandbTableRow = Annotated[ + RunTableRow + | StageTableRow + | RecordTableRow + | EvaluationTableRow + | NddWorkflowTableRow + | ModelWorkflowTableRow + | TraceCoverageTableRow, + Field(discriminator="record_type"), +] +WANDB_TABLE_ROW_MODELS: Mapping[str, type[BaseModel]] = MappingProxyType( + { + "run": RunTableRow, + "stage": StageTableRow, + "record": RecordTableRow, + "evaluation_record": EvaluationTableRow, + "ndd_workflow": NddWorkflowTableRow, + "model_workflow": ModelWorkflowTableRow, + "dd_trace_coverage": TraceCoverageTableRow, + } +) +_MEASUREMENT_RECORD_TYPES = frozenset(WANDB_TABLE_ROW_MODELS) + + +class WandbTablePayload(StrictFrozenModel): + record_type: StrictStr + rows: tuple[WandbTableRow, ...] + + @model_validator(mode="after") + def rows_match_record_type(self) -> WandbTablePayload: + if any(row.record_type != self.record_type for row in self.rows): + raise ValueError("W&B table rows must match the table record type") + return self + + @property + def name(self) -> str: + return f"measurement_table/{self.record_type}" + + @property + def columns(self) -> tuple[str, ...]: + present = {key for row in self.rows for key in row.model_dump(exclude_none=True)} + model_fields = type(self.rows[0]).model_fields if self.rows else {} + return tuple(name for name in model_fields if name in present) + + @property + def data(self) -> tuple[tuple[SafeScalar | None, ...], ...]: + columns = self.columns + return tuple(tuple(getattr(row, column) for column in columns) for row in self.rows) + + +def _metric_has_aggregate_policy(key: str) -> bool: + if key in BENCHMARK_METRIC_NAMES or key in { + "measurement/record_count", + _PUBLICATION_COMPLETE_KEY, + _PUBLICATION_SEAL_DIGEST_KEY, + }: + return True + parts = key.split("/") + return ( + len(parts) == 3 + and parts[0] == "measurement" + and parts[1] in _MEASUREMENT_RECORD_TYPES + and parts[2] in AGGREGATED_MEASUREMENT_FIELDS + ) diff --git a/tools/measurement/measurement_tools/wandb_models.py b/tools/measurement/measurement_tools/wandb_models.py index c72ea37e..f6571be6 100644 --- a/tools/measurement/measurement_tools/wandb_models.py +++ b/tools/measurement/measurement_tools/wandb_models.py @@ -4,30 +4,11 @@ from __future__ import annotations -from collections.abc import Mapping -from enum import StrEnum -from pathlib import Path -from types import MappingProxyType -from typing import Annotated, Literal, cast +from typing import cast -from pydantic import ( - BaseModel, - Field, - StrictBool, - StrictInt, - StrictStr, - field_validator, - model_validator, -) +from pydantic import BaseModel -from measurement_tools.validation import ( - MeasurementStrategy, - NonNegativeFloat, - NonNegativeInt, - Probability, - StrictFrozenModel, - VisibleIdentifier, -) +from measurement_tools.validation import StrictFrozenModel as StrictFrozenModel from measurement_tools.wandb_metadata import ( BenchmarkMetadata, ConfigMetadata, @@ -40,18 +21,55 @@ ReplaceMetadata, RewriteMetadata, RuntimeMetadata, - SafeScalar, SlurmJobMetadata, SlurmMetadata, SweepMetadata, WandbConfigPayload, WandbRunMetadata, - WandbTag, WorkloadMetadata, WorkloadSourceMetadata, ) -from measurement_tools.wandb_metric_schema import AGGREGATED_MEASUREMENT_FIELDS, BENCHMARK_METRIC_NAMES +from measurement_tools.wandb_metadata import ( + SafeScalar as SafeScalar, +) +from measurement_tools.wandb_metadata import ( + WandbTag as WandbTag, +) +from measurement_tools.wandb_metrics import ( + _METRIC_TABLE_POLICY_FIELDS, + EvaluationTableRow, + ModelWorkflowTableRow, + NddWorkflowTableRow, + RecordTableRow, + RunTableRow, + StageTableRow, + TraceCoverageTableRow, + WandbHistoryPayload, + WandbSummaryPayload, + WandbTablePayload, + _MetricTableRow, +) +from measurement_tools.wandb_metrics import ( + WANDB_TABLE_ROW_MODELS as WANDB_TABLE_ROW_MODELS, +) +from measurement_tools.wandb_metrics import ( + WandbTableRow as WandbTableRow, +) from measurement_tools.wandb_policy import DataClass, Exposure, FieldPolicy +from measurement_tools.wandb_publication import ( + PUBLICATION_COMPLETE_KEY as PUBLICATION_COMPLETE_KEY, +) +from measurement_tools.wandb_publication import ( + PUBLICATION_SEAL_DIGEST_KEY as PUBLICATION_SEAL_DIGEST_KEY, +) +from measurement_tools.wandb_publication import ( + WandbInitPayload, + WandbPublishPayload, + WandbPublishResult, +) +from measurement_tools.wandb_publication import ( + WandbPublicationState as WandbPublicationState, +) from measurement_tools.wandb_settings import ( DEFAULT_WANDB_PROJECT as DEFAULT_WANDB_PROJECT, ) @@ -65,7 +83,7 @@ WandbInputs as WandbInputs, ) from measurement_tools.wandb_settings import ( - WandbMode, + WandbMode as WandbMode, ) from measurement_tools.wandb_settings import ( generated_wandb_tag as generated_wandb_tag, @@ -73,299 +91,7 @@ from measurement_tools.wandb_settings import ( is_safe_wandb_tag as is_safe_wandb_tag, ) -from measurement_tools.wandb_settings import ( - wandb_tag_value_is_sensitive as wandb_tag_value_is_sensitive, -) - -PUBLICATION_COMPLETE_KEY = "publication/complete" -PUBLICATION_SEAL_DIGEST_KEY = "publication/completion_seal_sha256" - - -class WandbPublicationState(StrEnum): - created = "created" - resumed = "resumed" - already_complete = "already_complete" - - -class WandbInitPayload(StrictFrozenModel): - run_id: VisibleIdentifier - resume: Literal["allow", "never"] = "never" - project: StrictStr - name: StrictStr - mode: WandbMode - directory: Path - group: StrictStr - job_type: StrictStr - entity: StrictStr | None = None - tags: tuple[WandbTag, ...] = () - - -class WandbHistoryPayload(StrictFrozenModel): - metrics: dict[StrictStr, SafeScalar] - - @field_validator("metrics") - @classmethod - def metric_keys_have_policy(cls, metrics: dict[str, SafeScalar]) -> dict[str, SafeScalar]: - unknown = [key for key in metrics if not _metric_has_aggregate_policy(key)] - if unknown: - raise ValueError(f"W&B metric has no aggregate exposure policy: {unknown!r}") - return metrics - - -class WandbSummaryPayload(StrictFrozenModel): - metrics: dict[StrictStr, SafeScalar] - - @field_validator("metrics") - @classmethod - def metric_keys_have_policy(cls, metrics: dict[str, SafeScalar]) -> dict[str, SafeScalar]: - return WandbHistoryPayload.metric_keys_have_policy(metrics) - - -class _MetricTableRow(StrictFrozenModel): - elapsed_sec: NonNegativeFloat | None = None - input_row_count: NonNegativeInt | None = None - seed_row_count: NonNegativeInt | None = None - output_row_count: NonNegativeInt | None = None - failed_record_count: NonNegativeInt | None = None - preview_num_records: NonNegativeInt | None = None - column_count: NonNegativeInt | None = None - observed_input_tokens: NonNegativeInt | None = None - observed_output_tokens: NonNegativeInt | None = None - observed_total_tokens: NonNegativeInt | None = None - observed_reasoning_tokens: NonNegativeInt | None = None - observed_successful_requests: NonNegativeInt | None = None - observed_failed_requests: NonNegativeInt | None = None - observed_total_requests: NonNegativeInt | None = None - text_length_chars: NonNegativeInt | None = None - text_length_tokens: NonNegativeInt | None = None - final_entity_count: NonNegativeInt | None = None - ground_truth_entity_count: NonNegativeInt | None = None - entity_true_positive_count: NonNegativeInt | None = None - entity_false_positive_count: NonNegativeInt | None = None - entity_false_negative_count: NonNegativeInt | None = None - entity_relaxed_gt_found_count: NonNegativeInt | None = None - entity_relaxed_detected_tp_count: NonNegativeInt | None = None - entity_relaxed_label_compatible_gt_found_count: NonNegativeInt | None = None - entity_relaxed_label_compatible_detected_tp_count: NonNegativeInt | None = None - replacement_count: NonNegativeInt | None = None - replacement_duplicate_value_count: NonNegativeInt | None = None - replacement_missing_final_entity_count: NonNegativeInt | None = None - replacement_missing_final_value_count: NonNegativeInt | None = None - replacement_synthetic_original_collision_count: NonNegativeInt | None = None - replacement_synthetic_original_collision_value_count: NonNegativeInt | None = None - original_value_leak_count: NonNegativeInt | None = None - detected_candidate_count: NonNegativeInt | None = None - validation_chunk_count: NonNegativeInt | None = None - llm_calls_estimated_total: NonNegativeInt | None = None - detection_invalid_entity_count: NonNegativeInt | None = None - type_fidelity_invalid_replacement_count: NonNegativeInt | None = None - relational_consistency_invalid_relation_count: NonNegativeInt | None = None - attribute_fidelity_invalid_entity_count: NonNegativeInt | None = None - observed_failed_request_rate: Probability | None = None - input_rows_per_sec: NonNegativeFloat | None = None - output_rows_per_sec: NonNegativeFloat | None = None - observed_tokens_per_sec: NonNegativeFloat | None = None - observed_requests_per_sec: NonNegativeFloat | None = None - observed_tokens_per_successful_request: NonNegativeFloat | None = None - entity_precision: Probability | None = None - entity_recall: Probability | None = None - entity_f1: Probability | None = None - utility_score: Probability | None = None - leakage_mass: NonNegativeFloat | None = None - weighted_leakage_rate: Probability | None = None - repair_iterations: NonNegativeInt | None = None - - -_METRIC_TABLE_POLICY_FIELDS = ( - "elapsed_sec", - "input_row_count", - "seed_row_count", - "output_row_count", - "failed_record_count", - "preview_num_records", - "column_count", - "observed_input_tokens", - "observed_output_tokens", - "observed_total_tokens", - "observed_reasoning_tokens", - "observed_successful_requests", - "observed_failed_requests", - "observed_total_requests", - "text_length_chars", - "text_length_tokens", - "final_entity_count", - "ground_truth_entity_count", - "entity_true_positive_count", - "entity_false_positive_count", - "entity_false_negative_count", - "entity_relaxed_gt_found_count", - "entity_relaxed_detected_tp_count", - "entity_relaxed_label_compatible_gt_found_count", - "entity_relaxed_label_compatible_detected_tp_count", - "replacement_count", - "replacement_duplicate_value_count", - "replacement_missing_final_entity_count", - "replacement_missing_final_value_count", - "replacement_synthetic_original_collision_count", - "replacement_synthetic_original_collision_value_count", - "original_value_leak_count", - "detected_candidate_count", - "validation_chunk_count", - "llm_calls_estimated_total", - "detection_invalid_entity_count", - "type_fidelity_invalid_replacement_count", - "relational_consistency_invalid_relation_count", - "attribute_fidelity_invalid_entity_count", - "observed_failed_request_rate", - "input_rows_per_sec", - "output_rows_per_sec", - "observed_tokens_per_sec", - "observed_requests_per_sec", - "observed_tokens_per_successful_request", - "entity_precision", - "entity_recall", - "entity_f1", - "utility_score", - "leakage_mass", - "weighted_leakage_rate", - "repair_iterations", -) - - -class RunTableRow(_MetricTableRow): - record_type: Literal["run"] - mode: Literal["replace", "rewrite"] - strategy: MeasurementStrategy - - -class StageTableRow(_MetricTableRow): - record_type: Literal["stage"] - stage: Literal[ - "Anonymizer._run_internal", - "EntityDetectionWorkflow.run", - "ReplacementWorkflow.run", - "RewriteWorkflow.run", - ] - status: Literal["completed", "error"] - - -class RecordTableRow(_MetricTableRow): - record_type: Literal["record"] - mode: Literal["replace", "rewrite"] - strategy: MeasurementStrategy - row_index: StrictInt | None - record_hash: Annotated[StrictStr, Field(pattern=r"^[0-9a-fA-F]{64}$")] - text_length_chars: None = None - text_length_tokens: None = None - text_length_chars_bucket: Literal["0", "1-127", "128-511", "512-2047", "2048-8191", "8192+"] - text_length_tokens_bucket: Literal["0", "1-127", "128-511", "512-2047", "2048-8191", "8192+"] - - -class EvaluationTableRow(RecordTableRow): - record_type: Literal["evaluation_record"] - detection_valid: StrictBool | None = None - type_fidelity_valid: StrictBool | None = None - relational_consistency_valid: StrictBool | None = None - attribute_fidelity_valid: StrictBool | None = None - - -class NddWorkflowTableRow(_MetricTableRow): - record_type: Literal["ndd_workflow"] - status: Literal["completed", "error"] - - -class ModelWorkflowTableRow(_MetricTableRow): - record_type: Literal["model_workflow"] - status: Literal["completed", "error"] - - -class TraceCoverageTableRow(StrictFrozenModel): - record_type: Literal["dd_trace_coverage"] - traced_column_count: NonNegativeInt - native_trace_column_count: NonNegativeInt - private_trace_column_count: NonNegativeInt - unsupported_column_count: NonNegativeInt - - -WandbTableRow = Annotated[ - RunTableRow - | StageTableRow - | RecordTableRow - | EvaluationTableRow - | NddWorkflowTableRow - | ModelWorkflowTableRow - | TraceCoverageTableRow, - Field(discriminator="record_type"), -] -WANDB_TABLE_ROW_MODELS: Mapping[str, type[BaseModel]] = MappingProxyType( - { - "run": RunTableRow, - "stage": StageTableRow, - "record": RecordTableRow, - "evaluation_record": EvaluationTableRow, - "ndd_workflow": NddWorkflowTableRow, - "model_workflow": ModelWorkflowTableRow, - "dd_trace_coverage": TraceCoverageTableRow, - } -) -_MEASUREMENT_RECORD_TYPES = frozenset(WANDB_TABLE_ROW_MODELS) - - -class WandbTablePayload(StrictFrozenModel): - record_type: StrictStr - rows: tuple[WandbTableRow, ...] - - @model_validator(mode="after") - def rows_match_record_type(self) -> WandbTablePayload: - if any(row.record_type != self.record_type for row in self.rows): - raise ValueError("W&B table rows must match the table record type") - return self - - @property - def name(self) -> str: - return f"measurement_table/{self.record_type}" - - @property - def columns(self) -> tuple[str, ...]: - present = {key for row in self.rows for key in row.model_dump(exclude_none=True)} - model_fields = type(self.rows[0]).model_fields if self.rows else {} - return tuple(name for name in model_fields if name in present) - - @property - def data(self) -> tuple[tuple[SafeScalar | None, ...], ...]: - columns = self.columns - return tuple(tuple(getattr(row, column) for column in columns) for row in self.rows) - - -class WandbPublishPayload(StrictFrozenModel): - init: WandbInitPayload - config: WandbConfigPayload - history: WandbHistoryPayload - summary: WandbSummaryPayload - tables: tuple[WandbTablePayload, ...] = () - - @model_validator(mode="after") - def validate_resume_contract(self) -> WandbPublishPayload: - marker = self.summary.metrics.get(PUBLICATION_COMPLETE_KEY) - digest = self.summary.metrics.get(PUBLICATION_SEAL_DIGEST_KEY) - if self.init.resume == "allow": - imported = self.config.imported - if imported is None or marker is not True or digest != imported.completion_seal_sha256: - raise ValueError("resumable W&B payload requires matching imported publication markers") - elif marker is not None or digest is not None: - raise ValueError("fresh W&B payload cannot contain resume publication markers") - return self - - -class WandbPublishResult(StrictFrozenModel): - published: StrictBool - run_id: StrictStr | None = None - entity: StrictStr | None = None - project: StrictStr | None = None - run_url: StrictStr | None = None - publication_state: WandbPublicationState | None = None - measurement_sha256: Annotated[StrictStr | None, Field(pattern=r"^[0-9a-f]{64}$")] = None - record_count: StrictInt = 0 +from measurement_tools.wandb_settings import wandb_tag_value_is_sensitive as wandb_tag_value_is_sensitive def _aggregate_policies( @@ -602,19 +328,3 @@ def validate_outbound_field_policies() -> None: validate_outbound_field_policies() - - -def _metric_has_aggregate_policy(key: str) -> bool: - if key in BENCHMARK_METRIC_NAMES or key in { - "measurement/record_count", - PUBLICATION_COMPLETE_KEY, - PUBLICATION_SEAL_DIGEST_KEY, - }: - return True - parts = key.split("/") - return ( - len(parts) == 3 - and parts[0] == "measurement" - and parts[1] in _MEASUREMENT_RECORD_TYPES - and parts[2] in AGGREGATED_MEASUREMENT_FIELDS - ) diff --git a/tools/measurement/measurement_tools/wandb_publication.py b/tools/measurement/measurement_tools/wandb_publication.py new file mode 100644 index 00000000..110cf0ea --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_publication.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validated W&B publication envelopes and lifecycle results.""" + +from __future__ import annotations + +from enum import StrEnum +from pathlib import Path +from typing import Annotated, Literal + +from pydantic import Field, StrictBool, StrictInt, StrictStr, model_validator + +from measurement_tools.validation import StrictFrozenModel, VisibleIdentifier +from measurement_tools.wandb_metadata import WandbConfigPayload, WandbTag +from measurement_tools.wandb_metrics import WandbHistoryPayload, WandbSummaryPayload, WandbTablePayload +from measurement_tools.wandb_settings import WandbMode + +__all__ = [ + "PUBLICATION_COMPLETE_KEY", + "PUBLICATION_SEAL_DIGEST_KEY", + "WandbInitPayload", + "WandbPublicationState", + "WandbPublishPayload", + "WandbPublishResult", +] + +PUBLICATION_COMPLETE_KEY = "publication/complete" +PUBLICATION_SEAL_DIGEST_KEY = "publication/completion_seal_sha256" + + +class WandbPublicationState(StrEnum): + created = "created" + resumed = "resumed" + already_complete = "already_complete" + + +class WandbInitPayload(StrictFrozenModel): + run_id: VisibleIdentifier + resume: Literal["allow", "never"] = "never" + project: StrictStr + name: StrictStr + mode: WandbMode + directory: Path + group: StrictStr + job_type: StrictStr + entity: StrictStr | None = None + tags: tuple[WandbTag, ...] = () + + +class WandbPublishPayload(StrictFrozenModel): + init: WandbInitPayload + config: WandbConfigPayload + history: WandbHistoryPayload + summary: WandbSummaryPayload + tables: tuple[WandbTablePayload, ...] = () + + @model_validator(mode="after") + def validate_resume_contract(self) -> WandbPublishPayload: + marker = self.summary.metrics.get(PUBLICATION_COMPLETE_KEY) + digest = self.summary.metrics.get(PUBLICATION_SEAL_DIGEST_KEY) + if self.init.resume == "allow": + imported = self.config.imported + if imported is None or marker is not True or digest != imported.completion_seal_sha256: + raise ValueError("resumable W&B payload requires matching imported publication markers") + elif marker is not None or digest is not None: + raise ValueError("fresh W&B payload cannot contain resume publication markers") + return self + + +class WandbPublishResult(StrictFrozenModel): + published: StrictBool + run_id: StrictStr | None = None + entity: StrictStr | None = None + project: StrictStr | None = None + run_url: StrictStr | None = None + publication_state: WandbPublicationState | None = None + measurement_sha256: Annotated[StrictStr | None, Field(pattern=r"^[0-9a-f]{64}$")] = None + record_count: StrictInt = 0 From e8f9653b88c88f548bd41eeba7d7c86673cb9bff Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 19:27:35 +0000 Subject: [PATCH 04/17] refactor: isolate wandb field policies Signed-off-by: Aaron Gonzales --- .importlinter | 13 + .../test_measurement_module_compatibility.py | 1 + .../measurement_tools/wandb_field_policies.py | 288 +++++++++++++ .../measurement_tools/wandb_models.py | 380 +++--------------- 4 files changed, 359 insertions(+), 323 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_field_policies.py diff --git a/.importlinter b/.importlinter index 6bf5262d..c90bb5f3 100644 --- a/.importlinter +++ b/.importlinter @@ -47,3 +47,16 @@ source_modules = anonymizer.interface.run_telemetry forbidden_modules = anonymizer.interface.anonymizer + +[importlinter:contract:wandb-leaves-do-not-import-wandb-models] +name = W&B leaves do not import the compatibility facade +type = forbidden +source_modules = + measurement_tools.wandb_policy + measurement_tools.wandb_settings + measurement_tools.wandb_metadata + measurement_tools.wandb_metrics + measurement_tools.wandb_publication + measurement_tools.wandb_field_policies +forbidden_modules = + measurement_tools.wandb_models diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 73c4dd8d..470897ef 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -118,6 +118,7 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: "measurement_tools.wandb_metadata", "measurement_tools.wandb_metrics", "measurement_tools.wandb_publication", + "measurement_tools.wandb_field_policies", ], ) def test_wandb_model_facade_reexports_every_canonical_symbol(canonical_module: str) -> None: diff --git a/tools/measurement/measurement_tools/wandb_field_policies.py b/tools/measurement/measurement_tools/wandb_field_policies.py new file mode 100644 index 00000000..f8178af5 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_field_policies.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Outbound field-policy registry for W&B models.""" + +from __future__ import annotations + +from typing import cast + +from pydantic import BaseModel + +from measurement_tools.wandb_metadata import ( + BenchmarkMetadata, + ConfigMetadata, + DetectMetadata, + ExecutionMetadata, + GitMetadata, + ImportedRunMetadata, + MatrixMetadata, + ModelSourcesMetadata, + ReplaceMetadata, + RewriteMetadata, + RuntimeMetadata, + SlurmJobMetadata, + SlurmMetadata, + SweepMetadata, + WandbConfigPayload, + WandbRunMetadata, + WorkloadMetadata, + WorkloadSourceMetadata, +) +from measurement_tools.wandb_metrics import ( + _METRIC_TABLE_POLICY_FIELDS, + EvaluationTableRow, + ModelWorkflowTableRow, + NddWorkflowTableRow, + RecordTableRow, + RunTableRow, + StageTableRow, + TraceCoverageTableRow, + WandbHistoryPayload, + WandbSummaryPayload, + WandbTablePayload, + _MetricTableRow, +) +from measurement_tools.wandb_policy import DataClass, Exposure, FieldPolicy +from measurement_tools.wandb_publication import ( + WandbInitPayload, + WandbPublishPayload, + WandbPublishResult, +) + +__all__ = ["OUTBOUND_FIELD_POLICIES", "validate_outbound_field_policies"] + + +def _aggregate_policies( + *names: str, + data_class: DataClass = DataClass.operational, + exposure: Exposure = Exposure.aggregate, +) -> dict[str, FieldPolicy]: + return {name: FieldPolicy(data_class=data_class, exposure=exposure) for name in names} + + +OUTBOUND_FIELD_POLICIES: dict[type[BaseModel], dict[str, FieldPolicy]] = { + WandbInitPayload: { + "run_id": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.aggregate), + "resume": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "project": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "name": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "mode": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "directory": FieldPolicy(data_class=DataClass.sensitive, exposure=Exposure.local_only), + "group": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "job_type": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "entity": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "tags": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + }, + BenchmarkMetadata: _aggregate_policies( + "metadata_schema_version", + "suite_schema_version", + "wandb_sanitizer_version", + "measurement_schema_version", + "suite_id", + "workload_count", + "config_count", + "matrix_entry_count", + "case_count", + "case_retries", + "case_retry_backoff_sec", + "suite_file_hash", + ), + SlurmJobMetadata: _aggregate_policies("role", "job_id", data_class=DataClass.pseudonymous), + SlurmMetadata: _aggregate_policies( + "job_id", + "array_job_id", + "array_task_id", + "array_task_count", + "restart_count", + "ntasks", + "job_num_nodes", + "jobs", + data_class=DataClass.pseudonymous, + ), + ExecutionMetadata: _aggregate_policies("backend", "export", "fail_fast", "dd_trace", "dd_task_trace", "slurm"), + RuntimeMetadata: _aggregate_policies( + "anonymizer_version", + "datadesigner_version", + "wandb_version", + "python_version", + "platform_machine", + "platform_system", + ), + GitMetadata: _aggregate_policies("commit", "branch", "dirty", data_class=DataClass.pseudonymous), + ModelSourcesMetadata: _aggregate_policies("has_model_configs", "has_model_providers", "has_artifact_path"), + WorkloadSourceMetadata: _aggregate_policies("kind", "suffix"), + WorkloadMetadata: _aggregate_policies( + "id", + "text_column", + "has_id_column", + "has_data_summary", + "row_limit", + "row_offset", + "source", + data_class=DataClass.pseudonymous, + ), + DetectMetadata: _aggregate_policies( + "entity_label_source", + "entity_label_count", + "entity_label_set_hash", + "gliner_threshold", + "validation_max_entities_per_call", + "validation_excerpt_window_chars", + ), + ReplaceMetadata: _aggregate_policies( + "strategy", "normalize_label", "algorithm", "digest_length", "has_format_template", "has_instructions" + ), + RewriteMetadata: _aggregate_policies( + "risk_tolerance", + "max_repair_iterations", + "strict_entity_protection", + "has_privacy_goal", + "has_protect", + "has_preserve", + "has_instructions", + ), + ConfigMetadata: _aggregate_policies("id", "mode", "evaluate", "emit_telemetry", "detect", "replace", "rewrite"), + MatrixMetadata: _aggregate_policies("workload", "config", "repetitions", data_class=DataClass.pseudonymous), + SweepMetadata: _aggregate_policies("id", "arm_id", "params", data_class=DataClass.pseudonymous), + ImportedRunMetadata: _aggregate_policies( + "completion_seal_schema_version", + "completion_seal_sha256", + "producer_repository", + "producer_commit", + "phase", + "case_id", + data_class=DataClass.pseudonymous, + ), + WandbRunMetadata: _aggregate_policies( + "run_kind", + "benchmark", + "execution", + "runtime", + "git", + "model_sources", + "workloads", + "configs", + "matrix", + "sweep", + "imported", + data_class=DataClass.pseudonymous, + ), + WandbConfigPayload: _aggregate_policies( + "suite_id", + "wandb_mode", + "wandb_log_tables", + "benchmark_suite_id", + "benchmark_case_count", + "benchmark_workload_ids", + "benchmark_workload_row_limits", + "benchmark_workload_source_kinds", + "benchmark_workload_source_suffixes", + "benchmark_config_ids", + "benchmark_modes", + "benchmark_strategies", + "benchmark_gliner_thresholds", + "benchmark_entity_label_counts", + "benchmark_risk_tolerances", + "native_suite_id", + "sweep_id", + "sweep_arm_id", + "imported_config_id", + "sweep_params", + ), + WandbHistoryPayload: { + "metrics": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + }, + WandbSummaryPayload: { + "metrics": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + }, + _MetricTableRow: _aggregate_policies( + *_METRIC_TABLE_POLICY_FIELDS, + exposure=Exposure.table_opt_in, + ), + RunTableRow: _aggregate_policies( + "record_type", + "mode", + "strategy", + data_class=DataClass.pseudonymous, + exposure=Exposure.table_opt_in, + ), + StageTableRow: _aggregate_policies("record_type", "stage", "status", exposure=Exposure.table_opt_in), + RecordTableRow: _aggregate_policies( + "record_type", + "mode", + "strategy", + "row_index", + "record_hash", + "text_length_chars", + "text_length_tokens", + "text_length_chars_bucket", + "text_length_tokens_bucket", + data_class=DataClass.pseudonymous, + exposure=Exposure.table_opt_in, + ), + EvaluationTableRow: _aggregate_policies( + "record_type", + "detection_valid", + "type_fidelity_valid", + "relational_consistency_valid", + "attribute_fidelity_valid", + data_class=DataClass.pseudonymous, + exposure=Exposure.table_opt_in, + ), + NddWorkflowTableRow: _aggregate_policies( + "record_type", + "status", + exposure=Exposure.table_opt_in, + ), + ModelWorkflowTableRow: _aggregate_policies( + "record_type", + "status", + exposure=Exposure.table_opt_in, + ), + TraceCoverageTableRow: _aggregate_policies( + "record_type", + "traced_column_count", + "native_trace_column_count", + "private_trace_column_count", + "unsupported_column_count", + exposure=Exposure.table_opt_in, + ), + WandbTablePayload: { + "record_type": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.table_opt_in), + "rows": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.table_opt_in), + }, + WandbPublishPayload: { + "init": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "config": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "history": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "summary": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), + "tables": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.table_opt_in), + }, + WandbPublishResult: { + "published": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + "run_id": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.local_only), + "entity": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + "project": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + "run_url": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + "publication_state": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + "measurement_sha256": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.local_only), + "record_count": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), + }, +} + + +def validate_outbound_field_policies() -> None: + for model, policies in OUTBOUND_FIELD_POLICIES.items(): + inherited_policies: dict[str, FieldPolicy] = {} + for parent in reversed(model.mro()): + inherited_policies.update(OUTBOUND_FIELD_POLICIES.get(cast(type[BaseModel], parent), {})) + if set(model.model_fields) != set(inherited_policies): + raise RuntimeError(f"incomplete outbound field policy for {model.__name__}") + for name, policy in inherited_policies.items(): + if policy.data_class == DataClass.sensitive and policy.exposure != Exposure.local_only: + raise RuntimeError(f"sensitive outbound field {model.__name__}.{name} is not local-only") + if policy.exposure == Exposure.never: + raise RuntimeError(f"never-exposed field present in outbound model: {model.__name__}.{name}") + + +validate_outbound_field_policies() diff --git a/tools/measurement/measurement_tools/wandb_models.py b/tools/measurement/measurement_tools/wandb_models.py index f6571be6..d9c6e346 100644 --- a/tools/measurement/measurement_tools/wandb_models.py +++ b/tools/measurement/measurement_tools/wandb_models.py @@ -1,330 +1,64 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Typed configuration and outbound models for native W&B publication.""" +"""Compatibility facade for typed W&B configuration and publication models.""" from __future__ import annotations -from typing import cast - -from pydantic import BaseModel - from measurement_tools.validation import StrictFrozenModel as StrictFrozenModel -from measurement_tools.wandb_metadata import ( - BenchmarkMetadata, - ConfigMetadata, - DetectMetadata, - ExecutionMetadata, - GitMetadata, - ImportedRunMetadata, - MatrixMetadata, - ModelSourcesMetadata, - ReplaceMetadata, - RewriteMetadata, - RuntimeMetadata, - SlurmJobMetadata, - SlurmMetadata, - SweepMetadata, - WandbConfigPayload, - WandbRunMetadata, - WorkloadMetadata, - WorkloadSourceMetadata, -) -from measurement_tools.wandb_metadata import ( - SafeScalar as SafeScalar, -) -from measurement_tools.wandb_metadata import ( - WandbTag as WandbTag, -) -from measurement_tools.wandb_metrics import ( - _METRIC_TABLE_POLICY_FIELDS, - EvaluationTableRow, - ModelWorkflowTableRow, - NddWorkflowTableRow, - RecordTableRow, - RunTableRow, - StageTableRow, - TraceCoverageTableRow, - WandbHistoryPayload, - WandbSummaryPayload, - WandbTablePayload, - _MetricTableRow, -) -from measurement_tools.wandb_metrics import ( - WANDB_TABLE_ROW_MODELS as WANDB_TABLE_ROW_MODELS, -) -from measurement_tools.wandb_metrics import ( - WandbTableRow as WandbTableRow, -) -from measurement_tools.wandb_policy import DataClass, Exposure, FieldPolicy -from measurement_tools.wandb_publication import ( - PUBLICATION_COMPLETE_KEY as PUBLICATION_COMPLETE_KEY, -) -from measurement_tools.wandb_publication import ( - PUBLICATION_SEAL_DIGEST_KEY as PUBLICATION_SEAL_DIGEST_KEY, -) -from measurement_tools.wandb_publication import ( - WandbInitPayload, - WandbPublishPayload, - WandbPublishResult, -) -from measurement_tools.wandb_publication import ( - WandbPublicationState as WandbPublicationState, -) -from measurement_tools.wandb_settings import ( - DEFAULT_WANDB_PROJECT as DEFAULT_WANDB_PROJECT, -) -from measurement_tools.wandb_settings import ( - WANDB_TAG_MAX_LENGTH as WANDB_TAG_MAX_LENGTH, -) -from measurement_tools.wandb_settings import ( - ResolvedWandbConfig as ResolvedWandbConfig, -) -from measurement_tools.wandb_settings import ( - WandbInputs as WandbInputs, -) -from measurement_tools.wandb_settings import ( - WandbMode as WandbMode, -) -from measurement_tools.wandb_settings import ( - generated_wandb_tag as generated_wandb_tag, -) -from measurement_tools.wandb_settings import ( - is_safe_wandb_tag as is_safe_wandb_tag, -) +from measurement_tools.wandb_field_policies import ( + OUTBOUND_FIELD_POLICIES as OUTBOUND_FIELD_POLICIES, +) +from measurement_tools.wandb_field_policies import ( + validate_outbound_field_policies as validate_outbound_field_policies, +) +from measurement_tools.wandb_metadata import BenchmarkMetadata as BenchmarkMetadata +from measurement_tools.wandb_metadata import ConfigMetadata as ConfigMetadata +from measurement_tools.wandb_metadata import DetectMetadata as DetectMetadata +from measurement_tools.wandb_metadata import ExecutionMetadata as ExecutionMetadata +from measurement_tools.wandb_metadata import GitMetadata as GitMetadata +from measurement_tools.wandb_metadata import ImportedRunMetadata as ImportedRunMetadata +from measurement_tools.wandb_metadata import MatrixMetadata as MatrixMetadata +from measurement_tools.wandb_metadata import ModelSourcesMetadata as ModelSourcesMetadata +from measurement_tools.wandb_metadata import ReplaceMetadata as ReplaceMetadata +from measurement_tools.wandb_metadata import RewriteMetadata as RewriteMetadata +from measurement_tools.wandb_metadata import RuntimeMetadata as RuntimeMetadata +from measurement_tools.wandb_metadata import SafeScalar as SafeScalar +from measurement_tools.wandb_metadata import SlurmJobMetadata as SlurmJobMetadata +from measurement_tools.wandb_metadata import SlurmMetadata as SlurmMetadata +from measurement_tools.wandb_metadata import SweepMetadata as SweepMetadata +from measurement_tools.wandb_metadata import WandbConfigPayload as WandbConfigPayload +from measurement_tools.wandb_metadata import WandbRunMetadata as WandbRunMetadata +from measurement_tools.wandb_metadata import WandbTag as WandbTag +from measurement_tools.wandb_metadata import WorkloadMetadata as WorkloadMetadata +from measurement_tools.wandb_metadata import WorkloadSourceMetadata as WorkloadSourceMetadata +from measurement_tools.wandb_metrics import _METRIC_TABLE_POLICY_FIELDS as _METRIC_TABLE_POLICY_FIELDS +from measurement_tools.wandb_metrics import WANDB_TABLE_ROW_MODELS as WANDB_TABLE_ROW_MODELS +from measurement_tools.wandb_metrics import EvaluationTableRow as EvaluationTableRow +from measurement_tools.wandb_metrics import ModelWorkflowTableRow as ModelWorkflowTableRow +from measurement_tools.wandb_metrics import NddWorkflowTableRow as NddWorkflowTableRow +from measurement_tools.wandb_metrics import RecordTableRow as RecordTableRow +from measurement_tools.wandb_metrics import RunTableRow as RunTableRow +from measurement_tools.wandb_metrics import StageTableRow as StageTableRow +from measurement_tools.wandb_metrics import TraceCoverageTableRow as TraceCoverageTableRow +from measurement_tools.wandb_metrics import WandbHistoryPayload as WandbHistoryPayload +from measurement_tools.wandb_metrics import WandbSummaryPayload as WandbSummaryPayload +from measurement_tools.wandb_metrics import WandbTablePayload as WandbTablePayload +from measurement_tools.wandb_metrics import WandbTableRow as WandbTableRow +from measurement_tools.wandb_metrics import _MetricTableRow as _MetricTableRow +from measurement_tools.wandb_policy import DataClass as DataClass +from measurement_tools.wandb_policy import Exposure as Exposure +from measurement_tools.wandb_policy import FieldPolicy as FieldPolicy +from measurement_tools.wandb_publication import PUBLICATION_COMPLETE_KEY as PUBLICATION_COMPLETE_KEY +from measurement_tools.wandb_publication import PUBLICATION_SEAL_DIGEST_KEY as PUBLICATION_SEAL_DIGEST_KEY +from measurement_tools.wandb_publication import WandbInitPayload as WandbInitPayload +from measurement_tools.wandb_publication import WandbPublicationState as WandbPublicationState +from measurement_tools.wandb_publication import WandbPublishPayload as WandbPublishPayload +from measurement_tools.wandb_publication import WandbPublishResult as WandbPublishResult +from measurement_tools.wandb_settings import DEFAULT_WANDB_PROJECT as DEFAULT_WANDB_PROJECT +from measurement_tools.wandb_settings import WANDB_TAG_MAX_LENGTH as WANDB_TAG_MAX_LENGTH +from measurement_tools.wandb_settings import ResolvedWandbConfig as ResolvedWandbConfig +from measurement_tools.wandb_settings import WandbInputs as WandbInputs +from measurement_tools.wandb_settings import WandbMode as WandbMode +from measurement_tools.wandb_settings import generated_wandb_tag as generated_wandb_tag +from measurement_tools.wandb_settings import is_safe_wandb_tag as is_safe_wandb_tag from measurement_tools.wandb_settings import wandb_tag_value_is_sensitive as wandb_tag_value_is_sensitive - - -def _aggregate_policies( - *names: str, - data_class: DataClass = DataClass.operational, - exposure: Exposure = Exposure.aggregate, -) -> dict[str, FieldPolicy]: - return {name: FieldPolicy(data_class=data_class, exposure=exposure) for name in names} - - -OUTBOUND_FIELD_POLICIES: dict[type[BaseModel], dict[str, FieldPolicy]] = { - WandbInitPayload: { - "run_id": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.aggregate), - "resume": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "project": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "name": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "mode": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "directory": FieldPolicy(data_class=DataClass.sensitive, exposure=Exposure.local_only), - "group": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "job_type": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "entity": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "tags": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - }, - BenchmarkMetadata: _aggregate_policies( - "metadata_schema_version", - "suite_schema_version", - "wandb_sanitizer_version", - "measurement_schema_version", - "suite_id", - "workload_count", - "config_count", - "matrix_entry_count", - "case_count", - "case_retries", - "case_retry_backoff_sec", - "suite_file_hash", - ), - SlurmJobMetadata: _aggregate_policies("role", "job_id", data_class=DataClass.pseudonymous), - SlurmMetadata: _aggregate_policies( - "job_id", - "array_job_id", - "array_task_id", - "array_task_count", - "restart_count", - "ntasks", - "job_num_nodes", - "jobs", - data_class=DataClass.pseudonymous, - ), - ExecutionMetadata: _aggregate_policies("backend", "export", "fail_fast", "dd_trace", "dd_task_trace", "slurm"), - RuntimeMetadata: _aggregate_policies( - "anonymizer_version", - "datadesigner_version", - "wandb_version", - "python_version", - "platform_machine", - "platform_system", - ), - GitMetadata: _aggregate_policies("commit", "branch", "dirty", data_class=DataClass.pseudonymous), - ModelSourcesMetadata: _aggregate_policies("has_model_configs", "has_model_providers", "has_artifact_path"), - WorkloadSourceMetadata: _aggregate_policies("kind", "suffix"), - WorkloadMetadata: _aggregate_policies( - "id", - "text_column", - "has_id_column", - "has_data_summary", - "row_limit", - "row_offset", - "source", - data_class=DataClass.pseudonymous, - ), - DetectMetadata: _aggregate_policies( - "entity_label_source", - "entity_label_count", - "entity_label_set_hash", - "gliner_threshold", - "validation_max_entities_per_call", - "validation_excerpt_window_chars", - ), - ReplaceMetadata: _aggregate_policies( - "strategy", "normalize_label", "algorithm", "digest_length", "has_format_template", "has_instructions" - ), - RewriteMetadata: _aggregate_policies( - "risk_tolerance", - "max_repair_iterations", - "strict_entity_protection", - "has_privacy_goal", - "has_protect", - "has_preserve", - "has_instructions", - ), - ConfigMetadata: _aggregate_policies("id", "mode", "evaluate", "emit_telemetry", "detect", "replace", "rewrite"), - MatrixMetadata: _aggregate_policies("workload", "config", "repetitions", data_class=DataClass.pseudonymous), - SweepMetadata: _aggregate_policies("id", "arm_id", "params", data_class=DataClass.pseudonymous), - ImportedRunMetadata: _aggregate_policies( - "completion_seal_schema_version", - "completion_seal_sha256", - "producer_repository", - "producer_commit", - "phase", - "case_id", - data_class=DataClass.pseudonymous, - ), - WandbRunMetadata: _aggregate_policies( - "run_kind", - "benchmark", - "execution", - "runtime", - "git", - "model_sources", - "workloads", - "configs", - "matrix", - "sweep", - "imported", - data_class=DataClass.pseudonymous, - ), - WandbConfigPayload: _aggregate_policies( - "suite_id", - "wandb_mode", - "wandb_log_tables", - "benchmark_suite_id", - "benchmark_case_count", - "benchmark_workload_ids", - "benchmark_workload_row_limits", - "benchmark_workload_source_kinds", - "benchmark_workload_source_suffixes", - "benchmark_config_ids", - "benchmark_modes", - "benchmark_strategies", - "benchmark_gliner_thresholds", - "benchmark_entity_label_counts", - "benchmark_risk_tolerances", - "native_suite_id", - "sweep_id", - "sweep_arm_id", - "imported_config_id", - "sweep_params", - ), - WandbHistoryPayload: { - "metrics": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - }, - WandbSummaryPayload: { - "metrics": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - }, - _MetricTableRow: _aggregate_policies( - *_METRIC_TABLE_POLICY_FIELDS, - exposure=Exposure.table_opt_in, - ), - RunTableRow: _aggregate_policies( - "record_type", - "mode", - "strategy", - data_class=DataClass.pseudonymous, - exposure=Exposure.table_opt_in, - ), - StageTableRow: _aggregate_policies("record_type", "stage", "status", exposure=Exposure.table_opt_in), - RecordTableRow: _aggregate_policies( - "record_type", - "mode", - "strategy", - "row_index", - "record_hash", - "text_length_chars", - "text_length_tokens", - "text_length_chars_bucket", - "text_length_tokens_bucket", - data_class=DataClass.pseudonymous, - exposure=Exposure.table_opt_in, - ), - EvaluationTableRow: _aggregate_policies( - "record_type", - "detection_valid", - "type_fidelity_valid", - "relational_consistency_valid", - "attribute_fidelity_valid", - data_class=DataClass.pseudonymous, - exposure=Exposure.table_opt_in, - ), - NddWorkflowTableRow: _aggregate_policies( - "record_type", - "status", - exposure=Exposure.table_opt_in, - ), - ModelWorkflowTableRow: _aggregate_policies( - "record_type", - "status", - exposure=Exposure.table_opt_in, - ), - TraceCoverageTableRow: _aggregate_policies( - "record_type", - "traced_column_count", - "native_trace_column_count", - "private_trace_column_count", - "unsupported_column_count", - exposure=Exposure.table_opt_in, - ), - WandbTablePayload: { - "record_type": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.table_opt_in), - "rows": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.table_opt_in), - }, - WandbPublishPayload: { - "init": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "config": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "history": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "summary": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.aggregate), - "tables": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.table_opt_in), - }, - WandbPublishResult: { - "published": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - "run_id": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.local_only), - "entity": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - "project": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - "run_url": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - "publication_state": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - "measurement_sha256": FieldPolicy(data_class=DataClass.pseudonymous, exposure=Exposure.local_only), - "record_count": FieldPolicy(data_class=DataClass.operational, exposure=Exposure.local_only), - }, -} - - -def validate_outbound_field_policies() -> None: - for model, policies in OUTBOUND_FIELD_POLICIES.items(): - inherited_policies: dict[str, FieldPolicy] = {} - for parent in reversed(model.mro()): - inherited_policies.update(OUTBOUND_FIELD_POLICIES.get(cast(type[BaseModel], parent), {})) - if set(model.model_fields) != set(inherited_policies): - raise RuntimeError(f"incomplete outbound field policy for {model.__name__}") - for name, policy in inherited_policies.items(): - if policy.data_class == DataClass.sensitive and policy.exposure != Exposure.local_only: - raise RuntimeError(f"sensitive outbound field {model.__name__}.{name} is not local-only") - if policy.exposure == Exposure.never: - raise RuntimeError(f"never-exposed field present in outbound model: {model.__name__}.{name}") - - -validate_outbound_field_policies() From bf24957299a66d067b4b898b6ba57b20f2e50f53 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 19:42:08 +0000 Subject: [PATCH 05/17] refactor: split wandb staging and payload assembly Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 19 ++ .../measurement_tools/wandb_payload.py | 72 ++++++++ .../measurement_tools/wandb_run_identity.py | 53 ++++++ .../measurement_tools/wandb_setup.py | 168 +++--------------- .../measurement_tools/wandb_staging.py | 72 ++++++++ 5 files changed, 243 insertions(+), 141 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_payload.py create mode 100644 tools/measurement/measurement_tools/wandb_run_identity.py create mode 100644 tools/measurement/measurement_tools/wandb_staging.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 470897ef..666655a9 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -110,6 +110,25 @@ def test_wandb_facades_preserve_canonical_identity_and_reconstruction() -> None: sys.modules.pop(module_name, None) +def test_wandb_setup_facade_preserves_staging_identity_and_payload_contracts() -> None: + setup = importlib.import_module("measurement_tools.wandb_setup") + payload = importlib.import_module("measurement_tools.wandb_payload") + run_identity = importlib.import_module("measurement_tools.wandb_run_identity") + staging = importlib.import_module("measurement_tools.wandb_staging") + + assert payload.__all__ == ["BenchmarkWandbFinalization", "build_publish_payload"] + assert run_identity.__all__ == ["default_run_name", "effective_wandb_tags"] + assert staging.__all__ == [ + "open_directory_no_follow", + "prepare_wandb_staging_dir", + "validate_directory_metadata", + ] + assert setup.BenchmarkWandbFinalization is payload.BenchmarkWandbFinalization + assert setup._default_run_name is run_identity.default_run_name + assert setup._effective_wandb_tags is run_identity.effective_wandb_tags + assert setup.BenchmarkWandbFinalization.__module__ == "measurement_tools.wandb_payload" + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/measurement_tools/wandb_payload.py b/tools/measurement/measurement_tools/wandb_payload.py new file mode 100644 index 00000000..fe43a677 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_payload.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Typed W&B payload assembly from completed benchmark artifacts.""" + +from __future__ import annotations + +import secrets +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from measurement_tools.wandb_ingress import read_measurement_snapshot +from measurement_tools.wandb_metadata import BenchmarkMetadata, WandbConfigPayload, WandbRunMetadata +from measurement_tools.wandb_publication import WandbInitPayload, WandbPublishPayload +from measurement_tools.wandb_run_identity import default_run_name, effective_wandb_tags +from measurement_tools.wandb_settings import ResolvedWandbConfig +from measurement_tools.wandb_staging import prepare_wandb_staging_dir + +__all__ = ["BenchmarkWandbFinalization", "build_publish_payload"] + + +@dataclass(frozen=True) +class BenchmarkWandbFinalization: + """Artifacts needed to finalize benchmark W&B logging.""" + + measurement_path: Path + cases: Sequence[Any] + + +def build_publish_payload( + settings: ResolvedWandbConfig, + *, + suite_id: str, + output_dir: Path, + finalization: BenchmarkWandbFinalization, + metadata: WandbRunMetadata | None, +) -> tuple[WandbPublishPayload, str, int]: + from measurement_tools.wandb_logging import build_outbound_measurements # noqa: PLC0415 + + cases = list(finalization.cases) + expected_statuses = {str(case.case_id): str(case.status.value) for case in cases} + snapshot = read_measurement_snapshot(finalization.measurement_path, expected_statuses=expected_statuses) + history, summary, tables = build_outbound_measurements( + snapshot, + cases=cases, + log_tables=settings.wandb_log_tables, + ) + staging_dir = prepare_wandb_staging_dir(output_dir) + resolved_metadata = metadata or WandbRunMetadata( + benchmark=BenchmarkMetadata(suite_id=suite_id), + ) + init = WandbInitPayload( + run_id=secrets.token_hex(16), + project=settings.wandb_project, + name=settings.wandb_run_name or default_run_name(suite_id, resolved_metadata), + mode=settings.wandb_mode, + directory=staging_dir, + group=settings.wandb_group or suite_id, + job_type=settings.wandb_job_type or "benchmark", + entity=settings.wandb_entity, + tags=tuple(effective_wandb_tags(settings, suite_id=suite_id, metadata=resolved_metadata)), + ) + config = WandbConfigPayload.from_run_metadata(settings, suite_id=suite_id, metadata=resolved_metadata) + payload = WandbPublishPayload( + init=init, + config=config, + history=history, + summary=summary, + tables=tables, + ) + return payload, snapshot.sha256, len(snapshot.records) diff --git a/tools/measurement/measurement_tools/wandb_run_identity.py b/tools/measurement/measurement_tools/wandb_run_identity.py new file mode 100644 index 00000000..2bd61939 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_run_identity.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W&B run naming and tag identity policy.""" + +from __future__ import annotations + +from measurement_tools.wandb_metadata import WandbRunMetadata +from measurement_tools.wandb_settings import ( + ResolvedWandbConfig, + generated_wandb_tag, + is_safe_wandb_tag, +) + +__all__ = ["default_run_name", "effective_wandb_tags"] + + +def default_run_name(suite_id: str, metadata: WandbRunMetadata | None) -> str: + git = metadata.git if metadata is not None else None + if git is None: + return suite_id + commit = git.commit + branch = git.branch + if isinstance(commit, str) and commit: + suffix = commit[:7] + if isinstance(branch, str) and branch: + return f"{suite_id} {branch}@{suffix}" + return f"{suite_id} @{suffix}" + if isinstance(branch, str) and branch: + return f"{suite_id} {branch}" + return suite_id + + +def effective_wandb_tags( + settings: ResolvedWandbConfig, + *, + suite_id: str, + metadata: WandbRunMetadata | None, +) -> list[str]: + tags = [tag for tag in settings.effective_wandb_tags if is_safe_wandb_tag(tag)] + suite_tag = generated_wandb_tag("suite", suite_id) + if suite_tag is not None: + tags.append(suite_tag) + git = metadata.git if metadata is not None else None + if git is not None: + branch = git.branch + dirty = git.dirty + if isinstance(branch, str) and branch: + branch_tag = generated_wandb_tag("branch", branch) + if branch_tag is not None: + tags.append(branch_tag) + if isinstance(dirty, bool): + tags.append("dirty" if dirty else "clean") + return tags diff --git a/tools/measurement/measurement_tools/wandb_setup.py b/tools/measurement/measurement_tools/wandb_setup.py index 712c5dd9..804e4ad8 100644 --- a/tools/measurement/measurement_tools/wandb_setup.py +++ b/tools/measurement/measurement_tools/wandb_setup.py @@ -6,24 +6,18 @@ import logging import os -import secrets -import stat +import stat as stat import sys import threading -from collections.abc import Sequence -from dataclasses import dataclass from pathlib import Path from typing import Any, Callable from urllib.parse import quote -from measurement_tools.wandb_ingress import read_measurement_snapshot from measurement_tools.wandb_models import ( DEFAULT_WANDB_PROJECT, PUBLICATION_COMPLETE_KEY, PUBLICATION_SEAL_DIGEST_KEY, - BenchmarkMetadata, ResolvedWandbConfig, - WandbConfigPayload, WandbInitPayload, WandbInputs, WandbMode, @@ -31,9 +25,18 @@ WandbPublishPayload, WandbPublishResult, WandbRunMetadata, - generated_wandb_tag, - is_safe_wandb_tag, ) +from measurement_tools.wandb_models import ( + BenchmarkMetadata as BenchmarkMetadata, +) +from measurement_tools.wandb_models import ( + WandbConfigPayload as WandbConfigPayload, +) +from measurement_tools.wandb_payload import BenchmarkWandbFinalization as BenchmarkWandbFinalization +from measurement_tools.wandb_payload import build_publish_payload +from measurement_tools.wandb_run_identity import default_run_name, effective_wandb_tags +from measurement_tools.wandb_staging import open_directory_no_follow, validate_directory_metadata +from measurement_tools.wandb_staging import prepare_wandb_staging_dir as _prepare_wandb_staging_dir __all__ = [ "DEFAULT_WANDB_PROJECT", @@ -45,18 +48,13 @@ logger = logging.getLogger("measurement.wandb") +_default_run_name = default_run_name +_effective_wandb_tags = effective_wandb_tags + WANDB_SANITIZER_VERSION = 2 _WANDB_INSTALL_HINT = "Install the optional measurement dependency group: uv sync --group measurement" -@dataclass(frozen=True) -class BenchmarkWandbFinalization: - """Artifacts needed to finalize benchmark W&B logging.""" - - measurement_path: Path - cases: Sequence[Any] - - _PUBLISHER_WANDB_MODULE: Any | None = None @@ -296,40 +294,13 @@ def _build_publish_payload( finalization: BenchmarkWandbFinalization, metadata: WandbRunMetadata | None, ) -> tuple[WandbPublishPayload, str, int]: - from measurement_tools.wandb_logging import build_outbound_measurements # noqa: PLC0415 - - cases = list(finalization.cases) - expected_statuses = {str(case.case_id): str(case.status.value) for case in cases} - snapshot = read_measurement_snapshot(finalization.measurement_path, expected_statuses=expected_statuses) - history, summary, tables = build_outbound_measurements( - snapshot, - cases=cases, - log_tables=settings.wandb_log_tables, - ) - staging_dir = prepare_wandb_staging_dir(output_dir) - resolved_metadata = metadata or WandbRunMetadata( - benchmark=BenchmarkMetadata(suite_id=suite_id), - ) - init = WandbInitPayload( - run_id=secrets.token_hex(16), - project=settings.wandb_project, - name=settings.wandb_run_name or _default_run_name(suite_id, resolved_metadata), - mode=settings.wandb_mode, - directory=staging_dir, - group=settings.wandb_group or suite_id, - job_type=settings.wandb_job_type or "benchmark", - entity=settings.wandb_entity, - tags=tuple(_effective_wandb_tags(settings, suite_id=suite_id, metadata=resolved_metadata)), + return build_publish_payload( + settings, + suite_id=suite_id, + output_dir=output_dir, + finalization=finalization, + metadata=metadata, ) - config = WandbConfigPayload.from_run_metadata(settings, suite_id=suite_id, metadata=resolved_metadata) - payload = WandbPublishPayload( - init=init, - config=config, - history=history, - summary=summary, - tables=tables, - ) - return payload, snapshot.sha256, len(snapshot.records) def _publication_already_complete(run: Any, payload: WandbPublishPayload) -> bool: @@ -418,103 +389,18 @@ def _sdk_init_kwargs(wandb: Any, payload: WandbInitPayload) -> dict[str, Any]: def prepare_wandb_staging_dir(output_dir: Path) -> Path: - output_descriptor = _open_directory_no_follow(output_dir) - staging_dir = output_dir / ".wandb-private" - try: - try: - os.mkdir(".wandb-private", mode=0o700, dir_fd=output_descriptor) - except FileExistsError: - pass - flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_DIRECTORY", 0) - staging_descriptor = os.open( - ".wandb-private", - flags | getattr(os, "O_NOFOLLOW", 0), - dir_fd=output_descriptor, - ) - try: - metadata = os.fstat(staging_descriptor) - if not stat.S_ISDIR(metadata.st_mode) or metadata.st_uid != os.geteuid(): - raise ValueError("W&B staging directory must be an owned directory") - os.fchmod(staging_descriptor, 0o700) - finally: - os.close(staging_descriptor) - except OSError as exc: - raise ValueError("W&B staging directory cannot contain symlinks or special files") from exc - finally: - os.close(output_descriptor) - return staging_dir + """Prepare secure W&B staging through the canonical filesystem owner.""" + return _prepare_wandb_staging_dir(output_dir) def _open_directory_no_follow(path: Path) -> int: - absolute = path.absolute() - flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_DIRECTORY", 0) - no_follow = getattr(os, "O_NOFOLLOW", 0) - try: - descriptor = os.open(absolute.anchor, flags) - for index, part in enumerate(absolute.parts[1:], start=1): - child = os.open(part, flags | no_follow, dir_fd=descriptor) - os.close(descriptor) - descriptor = child - _validate_directory_metadata(os.fstat(descriptor), final=index == len(absolute.parts) - 1) - return descriptor - except (OSError, ValueError) as exc: - if "descriptor" in locals(): - os.close(descriptor) - raise ValueError("W&B output path cannot contain symlinks, untrusted directories, or special files") from exc + """Compatibility alias for the canonical no-follow traversal.""" + return open_directory_no_follow(path) def _validate_directory_metadata(metadata: os.stat_result, *, final: bool) -> None: - if not stat.S_ISDIR(metadata.st_mode): - raise ValueError("W&B output path must contain only directories") - if final and metadata.st_uid != os.geteuid(): - raise ValueError("W&B output directory must be owned by the current user") - sticky = metadata.st_mode & stat.S_ISVTX - world_writable = metadata.st_mode & stat.S_IWOTH - group_writable = metadata.st_mode & stat.S_IWGRP - root_owned = metadata.st_uid == 0 - if world_writable and not sticky: - raise ValueError("W&B output path contains an untrusted writable directory") - if group_writable and not sticky and not root_owned: - raise ValueError("W&B output path contains an untrusted writable directory") - - -def _default_run_name(suite_id: str, metadata: WandbRunMetadata | None) -> str: - git = metadata.git if metadata is not None else None - if git is None: - return suite_id - commit = git.commit - branch = git.branch - if isinstance(commit, str) and commit: - suffix = commit[:7] - if isinstance(branch, str) and branch: - return f"{suite_id} {branch}@{suffix}" - return f"{suite_id} @{suffix}" - if isinstance(branch, str) and branch: - return f"{suite_id} {branch}" - return suite_id - - -def _effective_wandb_tags( - settings: ResolvedWandbConfig, - *, - suite_id: str, - metadata: WandbRunMetadata | None, -) -> list[str]: - tags = [tag for tag in settings.effective_wandb_tags if is_safe_wandb_tag(tag)] - suite_tag = generated_wandb_tag("suite", suite_id) - if suite_tag is not None: - tags.append(suite_tag) - git = metadata.git if metadata is not None else None - if git is not None: - branch = git.branch - dirty = git.dirty - if isinstance(branch, str) and branch: - branch_tag = generated_wandb_tag("branch", branch) - if branch_tag is not None: - tags.append(branch_tag) - if isinstance(dirty, bool): - tags.append("dirty" if dirty else "clean") - return tags + """Compatibility alias for canonical directory metadata validation.""" + validate_directory_metadata(metadata, final=final) def _define_benchmark_metrics(run: Any) -> None: diff --git a/tools/measurement/measurement_tools/wandb_staging.py b/tools/measurement/measurement_tools/wandb_staging.py new file mode 100644 index 00000000..94eee185 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_staging.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Secure filesystem staging for W&B artifacts.""" + +from __future__ import annotations + +import os +import stat +from pathlib import Path + +__all__ = ["open_directory_no_follow", "prepare_wandb_staging_dir", "validate_directory_metadata"] + + +def prepare_wandb_staging_dir(output_dir: Path) -> Path: + output_descriptor = open_directory_no_follow(output_dir) + staging_dir = output_dir / ".wandb-private" + try: + try: + os.mkdir(".wandb-private", mode=0o700, dir_fd=output_descriptor) + except FileExistsError: + pass + flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_DIRECTORY", 0) + staging_descriptor = os.open( + ".wandb-private", + flags | getattr(os, "O_NOFOLLOW", 0), + dir_fd=output_descriptor, + ) + try: + metadata = os.fstat(staging_descriptor) + if not stat.S_ISDIR(metadata.st_mode) or metadata.st_uid != os.geteuid(): + raise ValueError("W&B staging directory must be an owned directory") + os.fchmod(staging_descriptor, 0o700) + finally: + os.close(staging_descriptor) + except OSError as exc: + raise ValueError("W&B staging directory cannot contain symlinks or special files") from exc + finally: + os.close(output_descriptor) + return staging_dir + + +def open_directory_no_follow(path: Path) -> int: + absolute = path.absolute() + flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_DIRECTORY", 0) + no_follow = getattr(os, "O_NOFOLLOW", 0) + try: + descriptor = os.open(absolute.anchor, flags) + for index, part in enumerate(absolute.parts[1:], start=1): + child = os.open(part, flags | no_follow, dir_fd=descriptor) + os.close(descriptor) + descriptor = child + validate_directory_metadata(os.fstat(descriptor), final=index == len(absolute.parts) - 1) + return descriptor + except (OSError, ValueError) as exc: + if "descriptor" in locals(): + os.close(descriptor) + raise ValueError("W&B output path cannot contain symlinks, untrusted directories, or special files") from exc + + +def validate_directory_metadata(metadata: os.stat_result, *, final: bool) -> None: + if not stat.S_ISDIR(metadata.st_mode): + raise ValueError("W&B output path must contain only directories") + if final and metadata.st_uid != os.geteuid(): + raise ValueError("W&B output directory must be owned by the current user") + sticky = metadata.st_mode & stat.S_ISVTX + world_writable = metadata.st_mode & stat.S_IWOTH + group_writable = metadata.st_mode & stat.S_IWGRP + root_owned = metadata.st_uid == 0 + if world_writable and not sticky: + raise ValueError("W&B output path contains an untrusted writable directory") + if group_writable and not sticky and not root_owned: + raise ValueError("W&B output path contains an untrusted writable directory") From 501ba1efc228741f7c65d8453ba920c4ea9e5346 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 19:57:26 +0000 Subject: [PATCH 06/17] refactor: split wandb sdk and publisher lifecycle Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 30 +- ...est_measurement_strict_import_publisher.py | 21 +- .../measurement_tools/wandb_publisher.py | 224 +++++++++++ .../wandb_sdk_environment.py | 124 ++++++ .../measurement_tools/wandb_setup.py | 361 ++---------------- 5 files changed, 431 insertions(+), 329 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_publisher.py create mode 100644 tools/measurement/measurement_tools/wandb_sdk_environment.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 666655a9..8a56337c 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -49,7 +49,12 @@ def _load_module(path: Path, name: str) -> ModuleType: "StrictFrozenModel", "measurement_tools.wandb_settings", ), - ("measurement_tools/wandb_setup.py", "WandbPublisher", "ResolvedWandbConfig", None), + ( + "measurement_tools/wandb_setup.py", + "WandbPublisher", + "ResolvedWandbConfig", + "measurement_tools.wandb_publisher", + ), ("create_wandb_report.py", "WandbReportResult", "WandbProjectPath", None), ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer", None), ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath", None), @@ -129,6 +134,27 @@ def test_wandb_setup_facade_preserves_staging_identity_and_payload_contracts() - assert setup.BenchmarkWandbFinalization.__module__ == "measurement_tools.wandb_payload" +def test_wandb_setup_facade_preserves_sdk_environment_and_publisher_contracts() -> None: + setup = importlib.import_module("measurement_tools.wandb_setup") + publisher = importlib.import_module("measurement_tools.wandb_publisher") + sdk_environment = importlib.import_module("measurement_tools.wandb_sdk_environment") + + assert publisher.__all__ == [ + "WandbPublisher", + "define_benchmark_metrics", + "publication_already_complete", + "publication_state", + "raise_lifecycle_failures", + "sdk_init_kwargs", + "wandb_run_url", + ] + assert sdk_environment.__all__ == ["WandbSdkEnvironment", "publisher_environment", "require_wandb"] + assert setup.WandbPublisher is publisher.WandbPublisher + assert setup.WandbSdkEnvironment is sdk_environment.WandbSdkEnvironment + assert setup._publisher_environment is sdk_environment.publisher_environment + assert setup.WandbPublisher.__module__ == "measurement_tools.wandb_publisher" + + @pytest.mark.parametrize( "canonical_module", [ @@ -267,7 +293,7 @@ def test_outbound_field_policy_validation_is_complete_at_import() -> None: ("tools/measurement/measurement_tools/wandb_logging.py", "from measurement_tools.wandb_models import"), ("tools/measurement/measurement_tools/wandb_report_models.py", "from measurement_tools.wandb_models import"), ("tools/measurement/sweep_benchmarks.py", "import run_benchmarks"), - ("tests/tools/test_measurement_strict_import_publisher.py", "publish_benchmark_wandb_best_effort.__module__"), + ("tests/tools/test_measurement_strict_import_publisher.py", "WandbPublisher.__module__"), (".github/workflows/benchmark-ci.yml", "uv run python tools/measurement/run_benchmarks.py"), ("tools/measurement/examples/run-repo-data-smoke-with-dd-traces.sh", "tools/measurement/run_benchmarks.py"), ("docs/development/observability.md", "tools/measurement/create_wandb_report.py"), diff --git a/tests/tools/test_measurement_strict_import_publisher.py b/tests/tools/test_measurement_strict_import_publisher.py index e2ee695a..9f805cb9 100644 --- a/tests/tools/test_measurement_strict_import_publisher.py +++ b/tests/tools/test_measurement_strict_import_publisher.py @@ -26,6 +26,11 @@ REPO_ROOT = Path(__file__).resolve().parents[2] +def _wandb_publisher_module() -> ModuleType: + publisher = importlib.import_module("measurement_tools.wandb_publisher") + return sys.modules[publisher.WandbPublisher.__module__] + + def _write_simple_suite(tmp_path: Path, *, suite_id: str = "base-suite") -> Path: return _write_yaml(tmp_path / "suite.yaml", _simple_suite_payload(suite_id=suite_id)) @@ -649,7 +654,7 @@ def test_wandb_publisher_uses_explicit_handle_and_finishes_once( case = _write_case_measurements(run_benchmarks_tool, tmp_path, _minimal_benchmark_case(run_benchmarks_tool)) state = _wandb_state() fake_wandb = _fake_wandb_module(state) - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) result = setup_module.WandbPublisher().publish( @@ -756,7 +761,7 @@ def test_wandb_table_opt_in_changes_only_table_payloads( run_benchmarks_tool: ModuleType, ) -> None: case = _write_case_measurements(run_benchmarks_tool, tmp_path, _minimal_benchmark_case(run_benchmarks_tool)) - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() states: list[SimpleNamespace] = [] for log_tables in (False, True): @@ -811,7 +816,7 @@ def failing_init(**kwargs: Any) -> Any: return run fake_wandb.init = failing_init - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(RuntimeError, match=f"{failure_stage} failed"): @@ -844,7 +849,7 @@ def init_with_two_failures(**kwargs: Any) -> Any: return run fake_wandb.init = init_with_two_failures - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(ExceptionGroup) as raised: @@ -876,7 +881,7 @@ def init_with_bad_finish(**kwargs: Any) -> Any: return run fake_wandb.init = init_with_bad_finish - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(RuntimeError, match="finish failed"): @@ -943,7 +948,7 @@ def init_with_changed_id(**kwargs: Any) -> Any: return run fake_wandb.init = init_with_changed_id - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(RuntimeError, match="different run identity"): @@ -967,7 +972,7 @@ def test_wandb_publisher_rejects_ambient_active_run( case = _write_case_measurements(run_benchmarks_tool, tmp_path, _minimal_benchmark_case(run_benchmarks_tool)) state = _wandb_state() fake_wandb = _fake_wandb_module(state, active_run=True) - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(RuntimeError, match="ambient W&B run"): @@ -1000,7 +1005,7 @@ def init_with_interrupt(**kwargs: Any) -> Any: return run fake_wandb.init = init_with_interrupt - setup_module = sys.modules[run_benchmarks_tool.publish_benchmark_wandb_best_effort.__module__] + setup_module = _wandb_publisher_module() monkeypatch.setattr(setup_module, "require_wandb", lambda: fake_wandb) with pytest.raises(KeyboardInterrupt): diff --git a/tools/measurement/measurement_tools/wandb_publisher.py b/tools/measurement/measurement_tools/wandb_publisher.py new file mode 100644 index 00000000..2767f213 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_publisher.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Strict W&B publication lifecycle and remote-state handling.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any +from urllib.parse import quote + +from measurement_tools.wandb_metadata import WandbRunMetadata +from measurement_tools.wandb_payload import BenchmarkWandbFinalization, build_publish_payload +from measurement_tools.wandb_publication import ( + PUBLICATION_COMPLETE_KEY, + PUBLICATION_SEAL_DIGEST_KEY, + WandbInitPayload, + WandbPublicationState, + WandbPublishPayload, + WandbPublishResult, +) +from measurement_tools.wandb_sdk_environment import WandbSdkEnvironment, require_wandb +from measurement_tools.wandb_settings import ResolvedWandbConfig, WandbMode + +__all__ = [ + "WandbPublisher", + "define_benchmark_metrics", + "publication_already_complete", + "publication_state", + "raise_lifecycle_failures", + "sdk_init_kwargs", + "wandb_run_url", +] + +logger = logging.getLogger("measurement.wandb") + + +class WandbPublisher: + """Own the complete strict native W&B publication lifecycle.""" + + def publish( + self, + settings: ResolvedWandbConfig, + *, + suite_id: str, + output_dir: Path, + finalization: BenchmarkWandbFinalization, + metadata: WandbRunMetadata | None = None, + ) -> WandbPublishResult: + if not settings.enabled: + return WandbPublishResult(published=False) + payload, snapshot_sha256, record_count = build_publish_payload( + settings, + suite_id=suite_id, + output_dir=output_dir, + finalization=finalization, + metadata=metadata, + ) + return self.publish_payload( + settings, + payload=payload, + measurement_sha256=snapshot_sha256, + record_count=record_count, + ) + + def publish_payload( + self, + settings: ResolvedWandbConfig, + *, + payload: WandbPublishPayload, + measurement_sha256: str, + record_count: int, + ) -> WandbPublishResult: + """Publish a complete typed payload without access to source artifacts.""" + if not settings.enabled: + return WandbPublishResult(published=False) + with WandbSdkEnvironment(settings): + wandb = require_wandb() + if getattr(wandb, "run", None) is not None: + raise RuntimeError("an ambient W&B run is active") + run: Any | None = None + result: WandbPublishResult | None = None + primary_error: BaseException | None = None + try: + run = wandb.init(**sdk_init_kwargs(wandb, payload.init)) + if run is None: + raise RuntimeError("wandb.init did not return an explicit run handle") + run_id = str(getattr(run, "id", "")) + if run_id != payload.init.run_id: + raise RuntimeError("wandb.init returned a different run identity") + already_complete = publication_already_complete(run, payload) + resolved_publication_state = publication_state(run, already_complete=already_complete) + if not already_complete: + define_benchmark_metrics(run) + run.config.update(payload.config.sdk_values(), allow_val_change=True) + tables = { + table.name: wandb.Table(columns=list(table.columns), data=[list(row) for row in table.data]) + for table in payload.tables + } + logged = {**payload.history.metrics, **tables} + if payload.init.resume == "allow": + run.log(logged, step=0) + else: + run.log(logged) + run.summary.update(payload.summary.metrics) + logger.info("W&B run id: %s", run_id) + result = WandbPublishResult( + published=True, + run_id=run_id, + entity=settings.wandb_entity, + project=settings.wandb_project, + run_url=wandb_run_url(settings, run_id=run_id), + publication_state=resolved_publication_state, + measurement_sha256=measurement_sha256, + record_count=record_count, + ) + except BaseException as exc: + primary_error = exc + finish_error: BaseException | None = None + if run is not None: + try: + run.finish() + except BaseException as exc: + finish_error = exc + raise_lifecycle_failures(primary_error, finish_error) + if result is None: + raise RuntimeError("W&B publisher completed without a result") + return result + + +def publication_already_complete(run: Any, payload: WandbPublishPayload) -> bool: + if payload.init.resume != "allow": + return False + expected_digest = payload.summary.metrics[PUBLICATION_SEAL_DIGEST_KEY] + summary = getattr(run, "summary", None) + get_value = getattr(summary, "get", None) + config = getattr(run, "config", None) + get_config = getattr(config, "get", None) + if not callable(get_value) or not callable(get_config): + raise RuntimeError("resumed W&B run does not expose readable publication state") + observed_complete = get_value(PUBLICATION_COMPLETE_KEY) + observed_digest = get_value(PUBLICATION_SEAL_DIGEST_KEY) + observed_imported = get_config("imported") + observed_config_digest = ( + observed_imported.get("completion_seal_sha256") if isinstance(observed_imported, dict) else None + ) + if observed_config_digest not in {None, expected_digest}: + raise RuntimeError("resumed W&B run config contains different sealed content") + if observed_complete is True: + if observed_digest != expected_digest or observed_config_digest != expected_digest: + raise RuntimeError("resumed W&B run is complete for different sealed content") + return True + if observed_complete not in {None, False}: + raise RuntimeError("resumed W&B run has an invalid publication marker") + if observed_digest not in {None, expected_digest}: + raise RuntimeError("resumed W&B run contains different sealed content") + return False + + +def publication_state(run: Any, *, already_complete: bool) -> WandbPublicationState: + if already_complete: + return WandbPublicationState.already_complete + return WandbPublicationState.resumed if bool(getattr(run, "resumed", False)) else WandbPublicationState.created + + +def wandb_run_url(settings: ResolvedWandbConfig, *, run_id: str) -> str | None: + if settings.wandb_mode != WandbMode.online or settings.wandb_entity is None: + return None + base_url = settings.wandb_base_url or "https://wandb.ai" + if base_url == "https://api.wandb.ai": + base_url = "https://wandb.ai" + segments = (settings.wandb_entity, settings.wandb_project, "runs", run_id) + return f"{base_url}/{'/'.join(quote(segment, safe='') for segment in segments)}" + + +def raise_lifecycle_failures(primary: BaseException | None, finish: BaseException | None) -> None: + if primary is not None and finish is not None: + if isinstance(primary, Exception) and isinstance(finish, Exception): + raise ExceptionGroup("W&B publication and finish both failed", [primary, finish]) from primary + raise BaseExceptionGroup("W&B publication and finish both failed", [primary, finish]) from primary + if primary is not None: + raise primary.with_traceback(primary.__traceback__) + if finish is not None: + raise finish.with_traceback(finish.__traceback__) + + +def sdk_init_kwargs(wandb: Any, payload: WandbInitPayload) -> dict[str, Any]: + values: dict[str, Any] = { + "project": payload.project, + "id": payload.run_id, + "resume": payload.resume, + "name": payload.name, + "mode": payload.mode.value, + "settings": wandb.Settings( + console="off", + disable_code=True, + disable_git=True, + host="redacted", + save_code=False, + x_disable_machine_info=True, + x_disable_meta=True, + x_disable_stats=True, + x_save_requirements=False, + ), + "dir": str(payload.directory), + "group": payload.group, + "job_type": payload.job_type, + } + if payload.entity is not None: + values["entity"] = payload.entity + if payload.tags: + values["tags"] = list(payload.tags) + return values + + +def define_benchmark_metrics(run: Any) -> None: + define_metric = getattr(run, "define_metric", None) + if not callable(define_metric): + return + for metric_name in ("benchmark/*", "measurement/*"): + try: + define_metric(metric_name, summary="last") + except Exception as exc: # noqa: BLE001 -- presentation polish is best-effort + logger.warning("Failed to define W&B metric %s (%s)", metric_name, type(exc).__name__) diff --git a/tools/measurement/measurement_tools/wandb_sdk_environment.py b/tools/measurement/measurement_tools/wandb_sdk_environment.py new file mode 100644 index 00000000..2503e0cc --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_sdk_environment.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Guarded process environment and SDK loading for W&B operations.""" + +from __future__ import annotations + +import os +import sys +import threading +from typing import Any + +from measurement_tools.wandb_settings import ResolvedWandbConfig + +__all__ = ["WandbSdkEnvironment", "publisher_environment", "require_wandb"] + +_WANDB_INSTALL_HINT = "Install the optional measurement dependency group: uv sync --group measurement" +_PUBLISHER_WANDB_MODULE: Any | None = None + + +def require_wandb() -> Any: + """Import wandb inside the publisher's guarded SDK environment.""" + global _PUBLISHER_WANDB_MODULE # noqa: PLW0603 + + if _WANDB_ENVIRONMENT_OWNER != threading.get_ident(): + raise RuntimeError("wandb must be imported inside the guarded SDK environment") + preloaded = sys.modules.get("wandb") + if _PUBLISHER_WANDB_MODULE is None and preloaded is not None: + raise RuntimeError("wandb must not be imported before the guarded publisher") + if _PUBLISHER_WANDB_MODULE is not None and preloaded is not _PUBLISHER_WANDB_MODULE: + raise RuntimeError("the loaded wandb module changed outside the guarded publisher") + try: + import wandb + except ImportError as exc: + raise ImportError( + f"W&B logging is enabled but the wandb package is not installed. {_WANDB_INSTALL_HINT}" + ) from exc + _PUBLISHER_WANDB_MODULE = wandb + return wandb + + +_WANDB_AMBIENT_ALLOWLIST = frozenset( + { + "WANDB_HTTP_TIMEOUT", + "WANDB_INIT_TIMEOUT", + "WANDB__SERVICE_WAIT", + } +) +_PROCESS_ENVIRONMENT_ALLOWLIST = frozenset( + { + "HOME", + "LANG", + "LC_ALL", + "LC_CTYPE", + "PATH", + "TEMP", + "TMP", + "TMPDIR", + "TZ", + } +) +_WANDB_ENVIRONMENT_LOCK = threading.Lock() +_WANDB_ENVIRONMENT_OWNER: int | None = None + + +class WandbSdkEnvironment: + """Process-wide exact environment transaction around W&B SDK use.""" + + def __init__(self, settings: ResolvedWandbConfig) -> None: + self._settings = settings + self._snapshot: dict[str, str] | None = None + + def __enter__(self) -> WandbSdkEnvironment: + global _WANDB_ENVIRONMENT_OWNER # noqa: PLW0603 + + if not _WANDB_ENVIRONMENT_LOCK.acquire(blocking=False): + raise RuntimeError("nested or concurrent W&B publisher use is not allowed") + _WANDB_ENVIRONMENT_OWNER = threading.get_ident() + self._snapshot = dict(os.environ) + try: + preserved = { + key: value + for key, value in self._snapshot.items() + if key in _PROCESS_ENVIRONMENT_ALLOWLIST or key in _WANDB_AMBIENT_ALLOWLIST + } + os.environ.clear() + os.environ.update(preserved) + os.environ.update(publisher_environment(self._settings)) + except BaseException: + self._restore() + raise + return self + + def __exit__(self, _exc_type: Any, _exc: Any, _traceback: Any) -> None: + self._restore() + + def _restore(self) -> None: + global _WANDB_ENVIRONMENT_OWNER # noqa: PLW0603 + + if self._snapshot is None: + return + os.environ.clear() + os.environ.update(self._snapshot) + self._snapshot = None + _WANDB_ENVIRONMENT_OWNER = None + _WANDB_ENVIRONMENT_LOCK.release() + + +def publisher_environment(settings: ResolvedWandbConfig) -> dict[str, str]: + environment = { + "WANDB_ERROR_REPORTING": "false", + "WANDB_MODE": settings.wandb_mode.value, + "WANDB_PROJECT": settings.wandb_project, + "WANDB_SILENT": "true", + } + optional = { + "WANDB_BASE_URL": settings.wandb_base_url, + "WANDB_ENTITY": settings.wandb_entity, + "WANDB_GROUP": settings.wandb_group, + "WANDB_JOB_TYPE": settings.wandb_job_type, + "WANDB_NAME": settings.wandb_run_name, + "WANDB_TAGS": settings.wandb_tags or None, + } + environment.update({key: value for key, value in optional.items() if value is not None}) + return environment diff --git a/tools/measurement/measurement_tools/wandb_setup.py b/tools/measurement/measurement_tools/wandb_setup.py index 804e4ad8..c5aa688e 100644 --- a/tools/measurement/measurement_tools/wandb_setup.py +++ b/tools/measurement/measurement_tools/wandb_setup.py @@ -1,40 +1,39 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Secure native W&B publishing for Anonymizer benchmark tooling.""" +"""Compatibility facade for secure native W&B publication.""" from __future__ import annotations import logging import os import stat as stat -import sys -import threading from pathlib import Path from typing import Any, Callable -from urllib.parse import quote -from measurement_tools.wandb_models import ( - DEFAULT_WANDB_PROJECT, - PUBLICATION_COMPLETE_KEY, - PUBLICATION_SEAL_DIGEST_KEY, - ResolvedWandbConfig, - WandbInitPayload, - WandbInputs, - WandbMode, - WandbPublicationState, - WandbPublishPayload, - WandbPublishResult, - WandbRunMetadata, -) -from measurement_tools.wandb_models import ( - BenchmarkMetadata as BenchmarkMetadata, -) -from measurement_tools.wandb_models import ( - WandbConfigPayload as WandbConfigPayload, -) +from measurement_tools import wandb_sdk_environment as _sdk_environment +from measurement_tools.wandb_models import DEFAULT_WANDB_PROJECT as DEFAULT_WANDB_PROJECT +from measurement_tools.wandb_models import BenchmarkMetadata as BenchmarkMetadata +from measurement_tools.wandb_models import ResolvedWandbConfig as ResolvedWandbConfig +from measurement_tools.wandb_models import WandbConfigPayload as WandbConfigPayload +from measurement_tools.wandb_models import WandbInitPayload as WandbInitPayload +from measurement_tools.wandb_models import WandbInputs as WandbInputs +from measurement_tools.wandb_models import WandbMode as WandbMode +from measurement_tools.wandb_models import WandbPublishPayload as WandbPublishPayload +from measurement_tools.wandb_models import WandbPublishResult as WandbPublishResult +from measurement_tools.wandb_models import WandbRunMetadata as WandbRunMetadata from measurement_tools.wandb_payload import BenchmarkWandbFinalization as BenchmarkWandbFinalization from measurement_tools.wandb_payload import build_publish_payload +from measurement_tools.wandb_publisher import WandbPublisher as WandbPublisher +from measurement_tools.wandb_publisher import ( + define_benchmark_metrics, + publication_already_complete, + publication_state, + raise_lifecycle_failures, + sdk_init_kwargs, + wandb_run_url, +) from measurement_tools.wandb_run_identity import default_run_name, effective_wandb_tags +from measurement_tools.wandb_sdk_environment import WandbSdkEnvironment as WandbSdkEnvironment from measurement_tools.wandb_staging import open_directory_no_follow, validate_directory_metadata from measurement_tools.wandb_staging import prepare_wandb_staging_dir as _prepare_wandb_staging_dir @@ -48,214 +47,34 @@ logger = logging.getLogger("measurement.wandb") -_default_run_name = default_run_name -_effective_wandb_tags = effective_wandb_tags - WANDB_SANITIZER_VERSION = 2 -_WANDB_INSTALL_HINT = "Install the optional measurement dependency group: uv sync --group measurement" - +_WANDB_INSTALL_HINT = _sdk_environment._WANDB_INSTALL_HINT +_PUBLISHER_WANDB_MODULE: Any | None = _sdk_environment._PUBLISHER_WANDB_MODULE +_WANDB_AMBIENT_ALLOWLIST = _sdk_environment._WANDB_AMBIENT_ALLOWLIST +_PROCESS_ENVIRONMENT_ALLOWLIST = _sdk_environment._PROCESS_ENVIRONMENT_ALLOWLIST +_WANDB_ENVIRONMENT_LOCK = _sdk_environment._WANDB_ENVIRONMENT_LOCK +_WANDB_ENVIRONMENT_OWNER = _sdk_environment._WANDB_ENVIRONMENT_OWNER -_PUBLISHER_WANDB_MODULE: Any | None = None +_default_run_name = default_run_name +_effective_wandb_tags = effective_wandb_tags +_publisher_environment = _sdk_environment.publisher_environment +_publication_already_complete = publication_already_complete +_publication_state = publication_state +_wandb_run_url = wandb_run_url +_raise_lifecycle_failures = raise_lifecycle_failures +_sdk_init_kwargs = sdk_init_kwargs +_define_benchmark_metrics = define_benchmark_metrics def require_wandb() -> Any: - """Import wandb inside the publisher's guarded SDK environment.""" + """Use the canonical SDK loader while preserving the legacy patch sentinel.""" global _PUBLISHER_WANDB_MODULE # noqa: PLW0603 - if _WANDB_ENVIRONMENT_OWNER != threading.get_ident(): - raise RuntimeError("wandb must be imported inside the guarded SDK environment") - preloaded = sys.modules.get("wandb") - if _PUBLISHER_WANDB_MODULE is None and preloaded is not None: - raise RuntimeError("wandb must not be imported before the guarded publisher") - if _PUBLISHER_WANDB_MODULE is not None and preloaded is not _PUBLISHER_WANDB_MODULE: - raise RuntimeError("the loaded wandb module changed outside the guarded publisher") + _sdk_environment._PUBLISHER_WANDB_MODULE = _PUBLISHER_WANDB_MODULE try: - import wandb - except ImportError as exc: - raise ImportError( - f"W&B logging is enabled but the wandb package is not installed. {_WANDB_INSTALL_HINT}" - ) from exc - _PUBLISHER_WANDB_MODULE = wandb - return wandb - - -_WANDB_AMBIENT_ALLOWLIST = frozenset( - { - "WANDB_HTTP_TIMEOUT", - "WANDB_INIT_TIMEOUT", - "WANDB__SERVICE_WAIT", - } -) -_PROCESS_ENVIRONMENT_ALLOWLIST = frozenset( - { - "HOME", - "LANG", - "LC_ALL", - "LC_CTYPE", - "PATH", - "TEMP", - "TMP", - "TMPDIR", - "TZ", - } -) -_WANDB_ENVIRONMENT_LOCK = threading.Lock() -_WANDB_ENVIRONMENT_OWNER: int | None = None - - -class WandbSdkEnvironment: - """Process-wide exact environment transaction around W&B SDK use.""" - - def __init__(self, settings: ResolvedWandbConfig) -> None: - self._settings = settings - self._snapshot: dict[str, str] | None = None - - def __enter__(self) -> WandbSdkEnvironment: - global _WANDB_ENVIRONMENT_OWNER # noqa: PLW0603 - - if not _WANDB_ENVIRONMENT_LOCK.acquire(blocking=False): - raise RuntimeError("nested or concurrent W&B publisher use is not allowed") - _WANDB_ENVIRONMENT_OWNER = threading.get_ident() - self._snapshot = dict(os.environ) - try: - preserved = { - key: value - for key, value in self._snapshot.items() - if key in _PROCESS_ENVIRONMENT_ALLOWLIST or key in _WANDB_AMBIENT_ALLOWLIST - } - os.environ.clear() - os.environ.update(preserved) - os.environ.update(_publisher_environment(self._settings)) - except BaseException: - self._restore() - raise - return self - - def __exit__(self, _exc_type: Any, _exc: Any, _traceback: Any) -> None: - self._restore() - - def _restore(self) -> None: - global _WANDB_ENVIRONMENT_OWNER # noqa: PLW0603 - - if self._snapshot is None: - return - os.environ.clear() - os.environ.update(self._snapshot) - self._snapshot = None - _WANDB_ENVIRONMENT_OWNER = None - _WANDB_ENVIRONMENT_LOCK.release() - - -def _publisher_environment(settings: ResolvedWandbConfig) -> dict[str, str]: - environment = { - "WANDB_ERROR_REPORTING": "false", - "WANDB_MODE": settings.wandb_mode.value, - "WANDB_PROJECT": settings.wandb_project, - "WANDB_SILENT": "true", - } - optional = { - "WANDB_BASE_URL": settings.wandb_base_url, - "WANDB_ENTITY": settings.wandb_entity, - "WANDB_GROUP": settings.wandb_group, - "WANDB_JOB_TYPE": settings.wandb_job_type, - "WANDB_NAME": settings.wandb_run_name, - "WANDB_TAGS": settings.wandb_tags or None, - } - environment.update({key: value for key, value in optional.items() if value is not None}) - return environment - - -class WandbPublisher: - """Own the complete strict native W&B publication lifecycle.""" - - def publish( - self, - settings: ResolvedWandbConfig, - *, - suite_id: str, - output_dir: Path, - finalization: BenchmarkWandbFinalization, - metadata: WandbRunMetadata | None = None, - ) -> WandbPublishResult: - if not settings.enabled: - return WandbPublishResult(published=False) - payload, snapshot_sha256, record_count = _build_publish_payload( - settings, - suite_id=suite_id, - output_dir=output_dir, - finalization=finalization, - metadata=metadata, - ) - return self.publish_payload( - settings, - payload=payload, - measurement_sha256=snapshot_sha256, - record_count=record_count, - ) - - def publish_payload( - self, - settings: ResolvedWandbConfig, - *, - payload: WandbPublishPayload, - measurement_sha256: str, - record_count: int, - ) -> WandbPublishResult: - """Publish a complete typed payload without access to source artifacts.""" - if not settings.enabled: - return WandbPublishResult(published=False) - with WandbSdkEnvironment(settings): - wandb = require_wandb() - if getattr(wandb, "run", None) is not None: - raise RuntimeError("an ambient W&B run is active") - run: Any | None = None - result: WandbPublishResult | None = None - primary_error: BaseException | None = None - try: - run = wandb.init(**_sdk_init_kwargs(wandb, payload.init)) - if run is None: - raise RuntimeError("wandb.init did not return an explicit run handle") - run_id = str(getattr(run, "id", "")) - if run_id != payload.init.run_id: - raise RuntimeError("wandb.init returned a different run identity") - already_complete = _publication_already_complete(run, payload) - publication_state = _publication_state(run, already_complete=already_complete) - if not already_complete: - _define_benchmark_metrics(run) - run.config.update(payload.config.sdk_values(), allow_val_change=True) - tables = { - table.name: wandb.Table(columns=list(table.columns), data=[list(row) for row in table.data]) - for table in payload.tables - } - logged = {**payload.history.metrics, **tables} - if payload.init.resume == "allow": - run.log(logged, step=0) - else: - run.log(logged) - run.summary.update(payload.summary.metrics) - logger.info("W&B run id: %s", run_id) - result = WandbPublishResult( - published=True, - run_id=run_id, - entity=settings.wandb_entity, - project=settings.wandb_project, - run_url=_wandb_run_url(settings, run_id=run_id), - publication_state=publication_state, - measurement_sha256=measurement_sha256, - record_count=record_count, - ) - except BaseException as exc: - primary_error = exc - finish_error: BaseException | None = None - if run is not None: - try: - run.finish() - except BaseException as exc: - finish_error = exc - _raise_lifecycle_failures(primary_error, finish_error) - if result is None: - raise RuntimeError("W&B publisher completed without a result") - return result + return _sdk_environment.require_wandb() + finally: + _PUBLISHER_WANDB_MODULE = _sdk_environment._PUBLISHER_WANDB_MODULE def publish_benchmark_wandb_best_effort( @@ -303,91 +122,6 @@ def _build_publish_payload( ) -def _publication_already_complete(run: Any, payload: WandbPublishPayload) -> bool: - if payload.init.resume != "allow": - return False - expected_digest = payload.summary.metrics[PUBLICATION_SEAL_DIGEST_KEY] - summary = getattr(run, "summary", None) - get_value = getattr(summary, "get", None) - config = getattr(run, "config", None) - get_config = getattr(config, "get", None) - if not callable(get_value) or not callable(get_config): - raise RuntimeError("resumed W&B run does not expose readable publication state") - observed_complete = get_value(PUBLICATION_COMPLETE_KEY) - observed_digest = get_value(PUBLICATION_SEAL_DIGEST_KEY) - observed_imported = get_config("imported") - observed_config_digest = ( - observed_imported.get("completion_seal_sha256") if isinstance(observed_imported, dict) else None - ) - if observed_config_digest not in {None, expected_digest}: - raise RuntimeError("resumed W&B run config contains different sealed content") - if observed_complete is True: - if observed_digest != expected_digest or observed_config_digest != expected_digest: - raise RuntimeError("resumed W&B run is complete for different sealed content") - return True - if observed_complete not in {None, False}: - raise RuntimeError("resumed W&B run has an invalid publication marker") - if observed_digest not in {None, expected_digest}: - raise RuntimeError("resumed W&B run contains different sealed content") - return False - - -def _publication_state(run: Any, *, already_complete: bool) -> WandbPublicationState: - if already_complete: - return WandbPublicationState.already_complete - return WandbPublicationState.resumed if bool(getattr(run, "resumed", False)) else WandbPublicationState.created - - -def _wandb_run_url(settings: ResolvedWandbConfig, *, run_id: str) -> str | None: - if settings.wandb_mode != WandbMode.online or settings.wandb_entity is None: - return None - base_url = settings.wandb_base_url or "https://wandb.ai" - if base_url == "https://api.wandb.ai": - base_url = "https://wandb.ai" - segments = (settings.wandb_entity, settings.wandb_project, "runs", run_id) - return f"{base_url}/{'/'.join(quote(segment, safe='') for segment in segments)}" - - -def _raise_lifecycle_failures(primary: BaseException | None, finish: BaseException | None) -> None: - if primary is not None and finish is not None: - if isinstance(primary, Exception) and isinstance(finish, Exception): - raise ExceptionGroup("W&B publication and finish both failed", [primary, finish]) from primary - raise BaseExceptionGroup("W&B publication and finish both failed", [primary, finish]) from primary - if primary is not None: - raise primary.with_traceback(primary.__traceback__) - if finish is not None: - raise finish.with_traceback(finish.__traceback__) - - -def _sdk_init_kwargs(wandb: Any, payload: WandbInitPayload) -> dict[str, Any]: - values: dict[str, Any] = { - "project": payload.project, - "id": payload.run_id, - "resume": payload.resume, - "name": payload.name, - "mode": payload.mode.value, - "settings": wandb.Settings( - console="off", - disable_code=True, - disable_git=True, - host="redacted", - save_code=False, - x_disable_machine_info=True, - x_disable_meta=True, - x_disable_stats=True, - x_save_requirements=False, - ), - "dir": str(payload.directory), - "group": payload.group, - "job_type": payload.job_type, - } - if payload.entity is not None: - values["entity"] = payload.entity - if payload.tags: - values["tags"] = list(payload.tags) - return values - - def prepare_wandb_staging_dir(output_dir: Path) -> Path: """Prepare secure W&B staging through the canonical filesystem owner.""" return _prepare_wandb_staging_dir(output_dir) @@ -401,14 +135,3 @@ def _open_directory_no_follow(path: Path) -> int: def _validate_directory_metadata(metadata: os.stat_result, *, final: bool) -> None: """Compatibility alias for canonical directory metadata validation.""" validate_directory_metadata(metadata, final=final) - - -def _define_benchmark_metrics(run: Any) -> None: - define_metric = getattr(run, "define_metric", None) - if not callable(define_metric): - return - for metric_name in ("benchmark/*", "measurement/*"): - try: - define_metric(metric_name, summary="last") - except Exception as exc: # noqa: BLE001 -- presentation polish is best-effort - logger.warning("Failed to define W&B metric %s (%s)", metric_name, type(exc).__name__) From 1087155a47170ca52ea89d372e041f41da34e2de Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 20:15:10 +0000 Subject: [PATCH 07/17] refactor: split W&B report vocabulary Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 26 +- tools/measurement/create_wandb_report.py | 435 +++--------------- .../measurement_tools/wandb_report_catalog.py | 195 ++++++++ .../wandb_report_contracts.py | 53 +++ .../measurement_tools/wandb_report_sdk.py | 102 ++++ .../measurement_tools/wandb_report_text.py | 103 +++++ 6 files changed, 540 insertions(+), 374 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_report_catalog.py create mode 100644 tools/measurement/measurement_tools/wandb_report_contracts.py create mode 100644 tools/measurement/measurement_tools/wandb_report_sdk.py create mode 100644 tools/measurement/measurement_tools/wandb_report_text.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 8a56337c..a99260cd 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -55,7 +55,12 @@ def _load_module(path: Path, name: str) -> ModuleType: "ResolvedWandbConfig", "measurement_tools.wandb_publisher", ), - ("create_wandb_report.py", "WandbReportResult", "WandbProjectPath", None), + ( + "create_wandb_report.py", + "WandbReportResult", + "WandbProjectPath", + "measurement_tools.wandb_report_contracts", + ), ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer", None), ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath", None), ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult", None), @@ -155,6 +160,25 @@ def test_wandb_setup_facade_preserves_sdk_environment_and_publisher_contracts() assert setup.WandbPublisher.__module__ == "measurement_tools.wandb_publisher" +def test_wandb_report_facade_preserves_leaf_contracts() -> None: + catalog = importlib.import_module("measurement_tools.wandb_report_catalog") + contracts = importlib.import_module("measurement_tools.wandb_report_contracts") + sdk = importlib.import_module("measurement_tools.wandb_report_sdk") + text = importlib.import_module("measurement_tools.wandb_report_text") + report = _load_module(MEASUREMENT_ROOT / "create_wandb_report.py", "compat_report_leaf_facade") + try: + assert report.WandbReportResult is contracts.WandbReportResult + assert report.WandbWorkspaceResult is contracts.WandbWorkspaceResult + assert report._all_report_metrics is catalog.all_report_metrics + assert report._group_visible_columns is catalog.group_visible_columns + assert report._read_group_views is sdk.read_group_views + assert report._report_settings is sdk.report_settings + assert report._plain_text is text.plain_text + assert report._validate_output_url is text.validate_output_url + finally: + sys.modules.pop("compat_report_leaf_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/create_wandb_report.py b/tools/measurement/create_wandb_report.py index b508d5af..1fa6bffb 100644 --- a/tools/measurement/create_wandb_report.py +++ b/tools/measurement/create_wandb_report.py @@ -13,19 +13,34 @@ from __future__ import annotations -import json import logging -import re import sys -import unicodedata -from html import escape as escape_html from typing import Annotated, Any, Literal -from urllib.parse import urlsplit import cyclopts from measurement_tools.cli import LogFormat, configure_logging, log_bad_input -from measurement_tools.wandb_metric_schema import metric_paths -from measurement_tools.wandb_models import ConfigMetadata, ResolvedWandbConfig, WandbMode, WorkloadMetadata +from measurement_tools.wandb_models import ( # noqa: F401 + ConfigMetadata, + ResolvedWandbConfig, + WandbMode, + WorkloadMetadata, +) +from measurement_tools.wandb_report_catalog import ( + _MEASUREMENT_TABLE_KEYS, + _WORKSPACE_BAR_SECTIONS, + _WORKSPACE_COMPARISON_COLUMNS, + _WORKSPACE_JOB_FILTERS, + _WORKSPACE_SUMMARY_SCALARS, + all_report_metrics, + bar_panel, + benchmark_panels, + group_visible_columns, + metric_panels, + metric_title, + single_run_visible_columns, + summary_columns, +) +from measurement_tools.wandb_report_contracts import WandbReportResult, WandbWorkspaceResult from measurement_tools.wandb_report_models import ( GroupComparison, WandbProjectPath, @@ -37,208 +52,51 @@ parse_wandb_run_view, validate_wandb_returned_url, ) +from measurement_tools.wandb_report_sdk import ( + read_group_views, + report_settings, + require_wandb_report_sdk, + require_wandb_workspace_sdk, + sweep_param_columns, +) +from measurement_tools.wandb_report_text import ( + code_span, + escape_heading, + escape_link_label, + escape_list_text, + escape_markdown_text, + plain_text, + scalar_text, + table_code, + validate_output_text, + validate_output_url, +) from measurement_tools.wandb_setup import WandbSdkEnvironment, require_wandb -from pydantic import BaseModel, ConfigDict, ValidationError, field_validator app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.wandb_report") -_WANDB_REPORT_INSTALL_HINT = "Install the optional measurement dependency group: uv sync --group measurement" - - -_ROW_FLOW_FIELDS = ("input_row_count", "output_row_count", "failed_record_count") -_ROW_THROUGHPUT_FIELDS = ("input_rows_per_sec_mean", "output_rows_per_sec_mean") -_CASE_HEALTH_METRICS = metric_paths( - "benchmark", - "case_total", - "case_completed", - "case_errored", - "case_success_rate", -) -_CASE_LATENCY_METRICS = [ - *metric_paths("benchmark", "case_elapsed_sec_mean", "case_elapsed_sec_sum"), - *metric_paths("measurement/stage", "elapsed_sec"), -] -_NDD_ROW_FLOW_METRICS = metric_paths( - "measurement/ndd_workflow", - "seed_row_count", - *_ROW_FLOW_FIELDS, - "column_count", -) -_NDD_REQUEST_HEALTH_METRICS = metric_paths( - "measurement/ndd_workflow", - "observed_total_requests", - "observed_successful_requests", - "observed_failed_requests", - "observed_failed_request_rate_mean", -) -_NDD_TOKEN_METRICS = metric_paths( - "measurement/ndd_workflow", - "observed_input_tokens", - "observed_output_tokens", - "observed_total_tokens", - "observed_tokens_per_successful_request_mean", -) -_NDD_THROUGHPUT_METRICS = metric_paths( - "measurement/ndd_workflow", - "elapsed_sec", - *_ROW_THROUGHPUT_FIELDS, - "observed_requests_per_sec_mean", - "observed_tokens_per_sec_mean", -) -_RECORD_PRIVACY_METRICS = metric_paths( - "measurement/record", - "final_entity_count", - "entity_precision_mean", - "entity_recall_mean", - "entity_f1_mean", - "original_value_leak_count", - "leakage_mass_mean", - "weighted_leakage_rate_mean", - "detected_candidate_count", - "validation_chunk_count", - "llm_calls_estimated_total", -) -_REWRITE_UTILITY_METRICS = metric_paths( - "measurement/record", - "utility_score_mean", - "repair_iterations_mean", -) -_REPLACEMENT_QUALITY_METRICS = metric_paths( - "measurement/record", - "replacement_count", - "replacement_duplicate_value_count", - "replacement_missing_final_entity_count", - "replacement_missing_final_value_count", - "replacement_synthetic_original_collision_count", - "replacement_synthetic_original_collision_value_count", -) -_STAGE_THROUGHPUT_METRICS = metric_paths( - "measurement/stage", - *_ROW_FLOW_FIELDS, - *_ROW_THROUGHPUT_FIELDS, -) -_METRIC_PANEL_GROUPS = [ - ("Case Health", _CASE_HEALTH_METRICS), - ("Case Latency", _CASE_LATENCY_METRICS), - ("NDD Row Flow", _NDD_ROW_FLOW_METRICS), - ("NDD Request Health", _NDD_REQUEST_HEALTH_METRICS), - ("NDD Token Usage", _NDD_TOKEN_METRICS), - ("NDD Throughput", _NDD_THROUGHPUT_METRICS), - ("Record Privacy", _RECORD_PRIVACY_METRICS), - ("Rewrite Utility", _REWRITE_UTILITY_METRICS), - ("Replacement Quality", _REPLACEMENT_QUALITY_METRICS), - ("Stage Throughput", _STAGE_THROUGHPUT_METRICS), -] -_MEASUREMENT_TABLE_KEYS = metric_paths( - "measurement_table", - "run", - "stage", - "ndd_workflow", - "model_workflow", - "record", - "evaluation_record", -) -_WORKSPACE_JOB_FILTERS = ("benchmark", "benchmark-import", "benchmark-sweep") -_WORKSPACE_SUMMARY_SCALARS = metric_paths( - "benchmark", - "case_success_rate", - "case_completed", - "case_errored", - "case_elapsed_sec_mean", -) -_WORKSPACE_COMPARISON_COLUMNS = [ - "config:sweep_arm_id", - "config:benchmark_strategies", - "config:benchmark_gliner_thresholds", - "config:benchmark_risk_tolerances", - "summary:benchmark/case_success_rate", - "summary:measurement/record/utility_score_mean", - "summary:measurement/record/weighted_leakage_rate_mean", - "summary:measurement/ndd_workflow/observed_total_tokens", -] -_WORKSPACE_BAR_SECTIONS = ( - ("Benchmark Summary", (("Case Health", _CASE_HEALTH_METRICS), ("Case Latency", _CASE_LATENCY_METRICS)), True), - ( - "Privacy", - (("Privacy Outcomes", _RECORD_PRIVACY_METRICS), ("Replacement Quality", _REPLACEMENT_QUALITY_METRICS)), - True, - ), - ("Utility", (("Rewrite Utility", _REWRITE_UTILITY_METRICS),), True), - ( - "Cost/Throughput", - ( - ("NDD Request Health", _NDD_REQUEST_HEALTH_METRICS), - ("NDD Token Usage", _NDD_TOKEN_METRICS), - ("NDD Throughput", _NDD_THROUGHPUT_METRICS), - ("Stage Throughput", _STAGE_THROUGHPUT_METRICS), - ), - True, - ), -) - - -class _WandbOutputValidation: - @field_validator("project_path", "group", "title") - @classmethod - def validate_output_text(cls, value: str | None) -> str | None: - return _validate_output_text(value) - - -class WandbReportResult(_WandbOutputValidation, BaseModel): - run_path: str | None = None - run_url: str | None = None - project_path: str - group: str | None = None - report_url: str - draft: bool - title: str - - model_config = ConfigDict(extra="forbid", frozen=True, strict=True) - - @field_validator("run_url", "report_url") - @classmethod - def validate_output_urls(cls, value: str | None) -> str | None: - return _validate_output_url(value) - - -class WandbWorkspaceResult(_WandbOutputValidation, BaseModel): - project_path: str - group: str | None = None - workspace_url: str - title: str - - model_config = ConfigDict(extra="forbid", frozen=True, strict=True) - - @field_validator("workspace_url") - @classmethod - def validate_output_urls(cls, value: str) -> str: - validated = _validate_output_url(value) - if validated is None: - raise ValueError("W&B workspace URL is required") - return validated - - -def require_wandb_report_sdk() -> tuple[Any, Any, Any]: - """Import W&B report APIs when report generation is requested.""" - wandb = require_wandb() - try: - from wandb.apis.reports import v2 as wr - from wandb_workspaces import expr - except ImportError as exc: - raise ImportError(f"W&B report generation requires wandb[workspaces]. {_WANDB_REPORT_INSTALL_HINT}") from exc - return wandb, wr, expr - - -def require_wandb_workspace_sdk() -> tuple[Any, Any]: - """Import W&B workspace APIs when workspace generation is requested.""" - require_wandb() - try: - import wandb_workspaces.reports.v2 as wr - import wandb_workspaces.workspaces as ws - except ImportError as exc: - raise ImportError(f"W&B workspace generation requires wandb[workspaces]. {_WANDB_REPORT_INSTALL_HINT}") from exc - return ws, wr +_all_report_metrics = all_report_metrics +_bar_panel = bar_panel +_benchmark_panels = benchmark_panels +_code_span = code_span +_escape_heading = escape_heading +_escape_link_label = escape_link_label +_escape_list_text = escape_list_text +_escape_markdown_text = escape_markdown_text +_group_visible_columns = group_visible_columns +_metric_panels = metric_panels +_metric_title = metric_title +_plain_text = plain_text +_read_group_views = read_group_views +_report_settings = report_settings +_scalar_text = scalar_text +_single_run_visible_columns = single_run_visible_columns +_summary_columns = summary_columns +_sweep_param_columns = sweep_param_columns +_table_code = table_code +_validate_output_text = validate_output_text +_validate_output_url = validate_output_url def create_benchmark_report( @@ -470,21 +328,6 @@ def build_benchmark_group_report( ) -def _single_run_visible_columns() -> list[str]: - return _summary_columns(_all_report_metrics()) - - -def _group_visible_columns(comparison: GroupComparison, parameter_columns: list[str] | None = None) -> list[str]: - columns = [ - f"config:{comparison.config_key}", - "config:run_kind", - *(["config:sweep_id", "config:sweep_params"] if comparison.run_kind == "sweep_arm" else []), - *(parameter_columns or []), - *_summary_columns(_all_report_metrics()), - ] - return list(dict.fromkeys(columns)) - - def _benchmark_workspace_settings(ws: Any) -> Any: return ws.WorkspaceSettings( sort_panels_alphabetically=False, @@ -578,83 +421,6 @@ def _workspace_metric_accessor(wr: Any, column: str) -> Any: raise ValueError(f"Unsupported workspace metric section: {section!r}") -def _metric_title(metric: str) -> str: - return metric.rsplit("/", maxsplit=1)[-1].replace("_", " ").title() - - -def _benchmark_panels(wr: Any, *, groupby: Any | None = None) -> list[Any]: - return [ - *_metric_panels(wr, groupby=groupby), - wr.MediaBrowser(title="Sanitized Measurement Tables", media_keys=_MEASUREMENT_TABLE_KEYS, mode="grid"), - ] - - -def _metric_panels(wr: Any, *, groupby: Any | None) -> list[Any]: - return [_bar_panel(wr, title, metrics, groupby=groupby) for title, metrics in _METRIC_PANEL_GROUPS] - - -def _bar_panel(wr: Any, title: str, metrics: list[str], *, groupby: Any | None = None) -> Any: - return wr.BarPlot(title=title, metrics=metrics, groupby=groupby) - - -def _all_report_metrics() -> list[str]: - return list(dict.fromkeys(metric for _, metrics in _METRIC_PANEL_GROUPS for metric in metrics)) - - -def _summary_columns(metrics: list[str]) -> list[str]: - return [f"summary:{metric}" for metric in metrics] - - -def _read_group_views(wandb: Any, *, project_path: WandbProjectPath, group: str) -> list[WandbRunView]: - runs = wandb.Api(timeout=60).runs(project_path.path, filters={"group": group}) - views: list[WandbRunView] = [] - for run in runs: - try: - run_id = getattr(run, "id", None) - if not isinstance(run_id, str): - raise ValueError("W&B group run is missing a string identity") - view = parse_wandb_run_view( - run, - run_path=WandbRunPath( - entity=project_path.entity, - project=project_path.project, - run_id=run_id, - base_url=project_path.base_url, - ), - allowed_metrics=frozenset(_all_report_metrics()), - ) - except (ValueError, ValidationError): - raise ValueError("W&B group contains invalid run metadata") from None - views.append(view) - return views - - -def _sweep_param_columns(views: list[WandbRunView]) -> list[str]: - keys = sorted({key for view in views if view.metadata.sweep is not None for key in view.metadata.sweep.params}) - return [f"config:sweep_param_{key}" for key in keys] - - -def _report_settings( - settings: ResolvedWandbConfig | None, - target: WandbRunPath | WandbProjectPath, -) -> ResolvedWandbConfig: - target_base_url = target.base_url if target.base_url != "https://wandb.ai" else None - if settings is None: - return ResolvedWandbConfig.from_env_and_overrides( - wandb_mode=WandbMode.online, - wandb_entity=target.entity, - wandb_project=target.project, - wandb_base_url=target_base_url, - ) - if target_base_url is not None and settings.wandb_base_url not in {None, target_base_url}: - raise ValueError("W&B target URL conflicts with resolved base URL") - return settings.validated_update( - wandb_entity=target.entity, - wandb_project=target.project, - wandb_base_url=target_base_url or settings.wandb_base_url, - ) - - def build_group_report_markdown(*, group: str, comparison: GroupComparison | None = None) -> str: resolved = comparison or GroupComparison(run_kind="sweep_arm", config_key="sweep_arm_id", label="Sweep Arm") return f"""### Group Summary @@ -816,83 +582,6 @@ def _number_metric(summary: dict[str, Any], key: str) -> str: return f"{float(value):.4g}" -def _scalar_text(value: Any) -> str: - if value is None: - return "n/a" - if isinstance(value, bool): - return "true" if value else "false" - if isinstance(value, int | float | str): - return str(value) - return json.dumps(value, ensure_ascii=True, sort_keys=True) - - -def _plain_text(value: str) -> str: - if any( - unicodedata.category(character) == "Cf" - or (unicodedata.category(character) == "Cc" and character not in "\t\n\r") - for character in value - ): - raise ValueError("W&B display text contains unsafe control characters") - return " ".join(value.split()) - - -def _validate_output_url(value: str | None) -> str | None: - if value is None: - return None - if any(unicodedata.category(character) in {"Cc", "Cf"} for character in value): - raise ValueError("W&B output URL contains unsafe control characters") - parsed = urlsplit(value) - if ( - parsed.scheme not in {"http", "https"} - or parsed.hostname is None - or parsed.username is not None - or parsed.password is not None - ): - raise ValueError("W&B output URL must be a credential-free HTTP(S) URL") - if parsed.scheme == "http" and parsed.hostname not in {"localhost", "127.0.0.1", "::1"}: - raise ValueError("W&B output URL must use HTTPS unless it targets loopback") - return value - - -def _validate_output_text(value: str | None) -> str | None: - if value is not None: - _plain_text(value) - return value - - -def _escape_link_label(value: str) -> str: - text = escape_html(_plain_text(value), quote=False) - return text.replace("\\", "\\\\").replace("[", "\\[").replace("]", "\\]") - - -def _escape_markdown_text(value: str) -> str: - text = escape_html(_plain_text(value), quote=False) - for character in ("\\", "`", "*", "_", "{", "}", "[", "]", "(", ")", "#", "+", "-", ".", "!", "|"): - text = text.replace(character, f"\\{character}") - return text - - -def _escape_heading(value: str) -> str: - return _escape_markdown_text(value) - - -def _code_span(value: Any) -> str: - text = escape_html(_scalar_text(value), quote=False).replace("|", "|") - text = _plain_text(text) - fence = "`" * (max((len(match) for match in re.findall(r"`+", text)), default=0) + 1) - if "`" in text: - return f"{fence} {text} {fence}" - return f"`{text}`" - - -def _table_code(value: Any) -> str: - return _code_span(value) - - -def _escape_list_text(value: Any) -> str: - return escape_html(_plain_text(_scalar_text(value)), quote=False).replace("|", "|") - - def render_result(result: WandbReportResult | WandbWorkspaceResult, *, json_output: bool) -> str: if json_output: return result.model_dump_json(indent=2) diff --git a/tools/measurement/measurement_tools/wandb_report_catalog.py b/tools/measurement/measurement_tools/wandb_report_catalog.py new file mode 100644 index 00000000..54d6094e --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_report_catalog.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Metric, column, and panel catalogs shared by W&B report surfaces.""" + +from __future__ import annotations + +from typing import Any + +from measurement_tools.wandb_metric_schema import metric_paths +from measurement_tools.wandb_report_models import GroupComparison + +_ROW_FLOW_FIELDS = ("input_row_count", "output_row_count", "failed_record_count") +_ROW_THROUGHPUT_FIELDS = ("input_rows_per_sec_mean", "output_rows_per_sec_mean") +_CASE_HEALTH_METRICS = metric_paths( + "benchmark", + "case_total", + "case_completed", + "case_errored", + "case_success_rate", +) +_CASE_LATENCY_METRICS = [ + *metric_paths("benchmark", "case_elapsed_sec_mean", "case_elapsed_sec_sum"), + *metric_paths("measurement/stage", "elapsed_sec"), +] +_NDD_ROW_FLOW_METRICS = metric_paths( + "measurement/ndd_workflow", + "seed_row_count", + *_ROW_FLOW_FIELDS, + "column_count", +) +_NDD_REQUEST_HEALTH_METRICS = metric_paths( + "measurement/ndd_workflow", + "observed_total_requests", + "observed_successful_requests", + "observed_failed_requests", + "observed_failed_request_rate_mean", +) +_NDD_TOKEN_METRICS = metric_paths( + "measurement/ndd_workflow", + "observed_input_tokens", + "observed_output_tokens", + "observed_total_tokens", + "observed_tokens_per_successful_request_mean", +) +_NDD_THROUGHPUT_METRICS = metric_paths( + "measurement/ndd_workflow", + "elapsed_sec", + *_ROW_THROUGHPUT_FIELDS, + "observed_requests_per_sec_mean", + "observed_tokens_per_sec_mean", +) +_RECORD_PRIVACY_METRICS = metric_paths( + "measurement/record", + "final_entity_count", + "entity_precision_mean", + "entity_recall_mean", + "entity_f1_mean", + "original_value_leak_count", + "leakage_mass_mean", + "weighted_leakage_rate_mean", + "detected_candidate_count", + "validation_chunk_count", + "llm_calls_estimated_total", +) +_REWRITE_UTILITY_METRICS = metric_paths( + "measurement/record", + "utility_score_mean", + "repair_iterations_mean", +) +_REPLACEMENT_QUALITY_METRICS = metric_paths( + "measurement/record", + "replacement_count", + "replacement_duplicate_value_count", + "replacement_missing_final_entity_count", + "replacement_missing_final_value_count", + "replacement_synthetic_original_collision_count", + "replacement_synthetic_original_collision_value_count", +) +_STAGE_THROUGHPUT_METRICS = metric_paths( + "measurement/stage", + *_ROW_FLOW_FIELDS, + *_ROW_THROUGHPUT_FIELDS, +) +_METRIC_PANEL_GROUPS = [ + ("Case Health", _CASE_HEALTH_METRICS), + ("Case Latency", _CASE_LATENCY_METRICS), + ("NDD Row Flow", _NDD_ROW_FLOW_METRICS), + ("NDD Request Health", _NDD_REQUEST_HEALTH_METRICS), + ("NDD Token Usage", _NDD_TOKEN_METRICS), + ("NDD Throughput", _NDD_THROUGHPUT_METRICS), + ("Record Privacy", _RECORD_PRIVACY_METRICS), + ("Rewrite Utility", _REWRITE_UTILITY_METRICS), + ("Replacement Quality", _REPLACEMENT_QUALITY_METRICS), + ("Stage Throughput", _STAGE_THROUGHPUT_METRICS), +] +_MEASUREMENT_TABLE_KEYS = metric_paths( + "measurement_table", + "run", + "stage", + "ndd_workflow", + "model_workflow", + "record", + "evaluation_record", +) +_WORKSPACE_JOB_FILTERS = ("benchmark", "benchmark-import", "benchmark-sweep") +_WORKSPACE_SUMMARY_SCALARS = metric_paths( + "benchmark", + "case_success_rate", + "case_completed", + "case_errored", + "case_elapsed_sec_mean", +) +_WORKSPACE_COMPARISON_COLUMNS = [ + "config:sweep_arm_id", + "config:benchmark_strategies", + "config:benchmark_gliner_thresholds", + "config:benchmark_risk_tolerances", + "summary:benchmark/case_success_rate", + "summary:measurement/record/utility_score_mean", + "summary:measurement/record/weighted_leakage_rate_mean", + "summary:measurement/ndd_workflow/observed_total_tokens", +] +_WORKSPACE_BAR_SECTIONS = ( + ("Benchmark Summary", (("Case Health", _CASE_HEALTH_METRICS), ("Case Latency", _CASE_LATENCY_METRICS)), True), + ( + "Privacy", + (("Privacy Outcomes", _RECORD_PRIVACY_METRICS), ("Replacement Quality", _REPLACEMENT_QUALITY_METRICS)), + True, + ), + ("Utility", (("Rewrite Utility", _REWRITE_UTILITY_METRICS),), True), + ( + "Cost/Throughput", + ( + ("NDD Request Health", _NDD_REQUEST_HEALTH_METRICS), + ("NDD Token Usage", _NDD_TOKEN_METRICS), + ("NDD Throughput", _NDD_THROUGHPUT_METRICS), + ("Stage Throughput", _STAGE_THROUGHPUT_METRICS), + ), + True, + ), +) + + +def single_run_visible_columns() -> list[str]: + return summary_columns(all_report_metrics()) + + +def group_visible_columns(comparison: GroupComparison, parameter_columns: list[str] | None = None) -> list[str]: + columns = [ + f"config:{comparison.config_key}", + "config:run_kind", + *(["config:sweep_id", "config:sweep_params"] if comparison.run_kind == "sweep_arm" else []), + *(parameter_columns or []), + *summary_columns(all_report_metrics()), + ] + return list(dict.fromkeys(columns)) + + +def metric_title(metric: str) -> str: + return metric.rsplit("/", maxsplit=1)[-1].replace("_", " ").title() + + +def benchmark_panels(wr: Any, *, groupby: Any | None = None) -> list[Any]: + return [ + *metric_panels(wr, groupby=groupby), + wr.MediaBrowser(title="Sanitized Measurement Tables", media_keys=_MEASUREMENT_TABLE_KEYS, mode="grid"), + ] + + +def metric_panels(wr: Any, *, groupby: Any | None) -> list[Any]: + return [bar_panel(wr, title, metrics, groupby=groupby) for title, metrics in _METRIC_PANEL_GROUPS] + + +def bar_panel(wr: Any, title: str, metrics: list[str], *, groupby: Any | None = None) -> Any: + return wr.BarPlot(title=title, metrics=metrics, groupby=groupby) + + +def all_report_metrics() -> list[str]: + return list(dict.fromkeys(metric for _, metrics in _METRIC_PANEL_GROUPS for metric in metrics)) + + +def summary_columns(metrics: list[str]) -> list[str]: + return [f"summary:{metric}" for metric in metrics] + + +__all__ = [ + "all_report_metrics", + "bar_panel", + "benchmark_panels", + "group_visible_columns", + "metric_panels", + "metric_title", + "single_run_visible_columns", + "summary_columns", +] diff --git a/tools/measurement/measurement_tools/wandb_report_contracts.py b/tools/measurement/measurement_tools/wandb_report_contracts.py new file mode 100644 index 00000000..81d70fff --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_report_contracts.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validated result contracts returned by W&B report tooling.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, field_validator + +from measurement_tools.wandb_report_text import validate_output_text, validate_output_url + + +class _WandbOutputValidation: + @field_validator("project_path", "group", "title") + @classmethod + def validate_output_text(cls, value: str | None) -> str | None: + return validate_output_text(value) + + +class WandbReportResult(_WandbOutputValidation, BaseModel): + run_path: str | None = None + run_url: str | None = None + project_path: str + group: str | None = None + report_url: str + draft: bool + title: str + + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + @field_validator("run_url", "report_url") + @classmethod + def validate_output_urls(cls, value: str | None) -> str | None: + return validate_output_url(value) + + +class WandbWorkspaceResult(_WandbOutputValidation, BaseModel): + project_path: str + group: str | None = None + workspace_url: str + title: str + + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + @field_validator("workspace_url") + @classmethod + def validate_output_urls(cls, value: str) -> str: + validated = validate_output_url(value) + if validated is None: + raise ValueError("W&B workspace URL is required") + return validated + + +__all__ = ["WandbReportResult", "WandbWorkspaceResult"] diff --git a/tools/measurement/measurement_tools/wandb_report_sdk.py b/tools/measurement/measurement_tools/wandb_report_sdk.py new file mode 100644 index 00000000..fa0a7f59 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_report_sdk.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W&B SDK loading and remote read plumbing for reports and workspaces.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import ValidationError + +from measurement_tools.wandb_models import ResolvedWandbConfig, WandbMode +from measurement_tools.wandb_report_catalog import all_report_metrics +from measurement_tools.wandb_report_models import ( + WandbProjectPath, + WandbRunPath, + WandbRunView, + parse_wandb_run_view, +) +from measurement_tools.wandb_setup import require_wandb + +_WANDB_REPORT_INSTALL_HINT = "Install the optional measurement dependency group: uv sync --group measurement" + + +def require_wandb_report_sdk() -> tuple[Any, Any, Any]: + """Import W&B report APIs when report generation is requested.""" + wandb = require_wandb() + try: + from wandb.apis.reports import v2 as wr + from wandb_workspaces import expr + except ImportError as exc: + raise ImportError(f"W&B report generation requires wandb[workspaces]. {_WANDB_REPORT_INSTALL_HINT}") from exc + return wandb, wr, expr + + +def require_wandb_workspace_sdk() -> tuple[Any, Any]: + """Import W&B workspace APIs when workspace generation is requested.""" + require_wandb() + try: + import wandb_workspaces.reports.v2 as wr + import wandb_workspaces.workspaces as ws + except ImportError as exc: + raise ImportError(f"W&B workspace generation requires wandb[workspaces]. {_WANDB_REPORT_INSTALL_HINT}") from exc + return ws, wr + + +def read_group_views(wandb: Any, *, project_path: WandbProjectPath, group: str) -> list[WandbRunView]: + runs = wandb.Api(timeout=60).runs(project_path.path, filters={"group": group}) + views: list[WandbRunView] = [] + for run in runs: + try: + run_id = getattr(run, "id", None) + if not isinstance(run_id, str): + raise ValueError("W&B group run is missing a string identity") + view = parse_wandb_run_view( + run, + run_path=WandbRunPath( + entity=project_path.entity, + project=project_path.project, + run_id=run_id, + base_url=project_path.base_url, + ), + allowed_metrics=frozenset(all_report_metrics()), + ) + except (ValueError, ValidationError): + raise ValueError("W&B group contains invalid run metadata") from None + views.append(view) + return views + + +def sweep_param_columns(views: list[WandbRunView]) -> list[str]: + keys = sorted({key for view in views if view.metadata.sweep is not None for key in view.metadata.sweep.params}) + return [f"config:sweep_param_{key}" for key in keys] + + +def report_settings( + settings: ResolvedWandbConfig | None, + target: WandbRunPath | WandbProjectPath, +) -> ResolvedWandbConfig: + target_base_url = target.base_url if target.base_url != "https://wandb.ai" else None + if settings is None: + return ResolvedWandbConfig.from_env_and_overrides( + wandb_mode=WandbMode.online, + wandb_entity=target.entity, + wandb_project=target.project, + wandb_base_url=target_base_url, + ) + if target_base_url is not None and settings.wandb_base_url not in {None, target_base_url}: + raise ValueError("W&B target URL conflicts with resolved base URL") + return settings.validated_update( + wandb_entity=target.entity, + wandb_project=target.project, + wandb_base_url=target_base_url or settings.wandb_base_url, + ) + + +__all__ = [ + "read_group_views", + "report_settings", + "require_wandb_report_sdk", + "require_wandb_workspace_sdk", + "sweep_param_columns", +] diff --git a/tools/measurement/measurement_tools/wandb_report_text.py b/tools/measurement/measurement_tools/wandb_report_text.py new file mode 100644 index 00000000..f7695c1f --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_report_text.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validated text and Markdown rendering primitives for W&B reports.""" + +from __future__ import annotations + +import json +import re +import unicodedata +from html import escape as escape_html +from typing import Any +from urllib.parse import urlsplit + + +def scalar_text(value: Any) -> str: + if value is None: + return "n/a" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int | float | str): + return str(value) + return json.dumps(value, ensure_ascii=True, sort_keys=True) + + +def plain_text(value: str) -> str: + if any( + unicodedata.category(character) == "Cf" + or (unicodedata.category(character) == "Cc" and character not in "\t\n\r") + for character in value + ): + raise ValueError("W&B display text contains unsafe control characters") + return " ".join(value.split()) + + +def validate_output_url(value: str | None) -> str | None: + if value is None: + return None + if any(unicodedata.category(character) in {"Cc", "Cf"} for character in value): + raise ValueError("W&B output URL contains unsafe control characters") + parsed = urlsplit(value) + if ( + parsed.scheme not in {"http", "https"} + or parsed.hostname is None + or parsed.username is not None + or parsed.password is not None + ): + raise ValueError("W&B output URL must be a credential-free HTTP(S) URL") + if parsed.scheme == "http" and parsed.hostname not in {"localhost", "127.0.0.1", "::1"}: + raise ValueError("W&B output URL must use HTTPS unless it targets loopback") + return value + + +def validate_output_text(value: str | None) -> str | None: + if value is not None: + plain_text(value) + return value + + +def escape_link_label(value: str) -> str: + text = escape_html(plain_text(value), quote=False) + return text.replace("\\", "\\\\").replace("[", "\\[").replace("]", "\\]") + + +def escape_markdown_text(value: str) -> str: + text = escape_html(plain_text(value), quote=False) + for character in ("\\", "`", "*", "_", "{", "}", "[", "]", "(", ")", "#", "+", "-", ".", "!", "|"): + text = text.replace(character, f"\\{character}") + return text + + +def escape_heading(value: str) -> str: + return escape_markdown_text(value) + + +def code_span(value: Any) -> str: + text = escape_html(scalar_text(value), quote=False).replace("|", "|") + text = plain_text(text) + fence = "`" * (max((len(match) for match in re.findall(r"`+", text)), default=0) + 1) + if "`" in text: + return f"{fence} {text} {fence}" + return f"`{text}`" + + +def table_code(value: Any) -> str: + return code_span(value) + + +def escape_list_text(value: Any) -> str: + return escape_html(plain_text(scalar_text(value)), quote=False).replace("|", "|") + + +__all__ = [ + "code_span", + "escape_heading", + "escape_link_label", + "escape_list_text", + "escape_markdown_text", + "plain_text", + "scalar_text", + "table_code", + "validate_output_text", + "validate_output_url", +] From 17d048de68f3c8005d0c5a37dcd693975cbd2e85 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 20:30:14 +0000 Subject: [PATCH 08/17] refactor: split W&B report construction Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 31 ++ tests/tools/test_measurement_wandb_reports.py | 26 ++ tools/measurement/create_wandb_report.py | 351 +++------------- .../measurement_tools/wandb_reports.py | 384 ++++++++++++++++++ 4 files changed, 501 insertions(+), 291 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_reports.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index a99260cd..a45d48b0 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -179,6 +179,37 @@ def test_wandb_report_facade_preserves_leaf_contracts() -> None: sys.modules.pop("compat_report_leaf_facade", None) +def test_wandb_report_facade_preserves_report_construction_contracts() -> None: + reports = importlib.import_module("measurement_tools.wandb_reports") + report = _load_module(MEASUREMENT_ROOT / "create_wandb_report.py", "compat_report_construction_facade") + try: + assert reports.__all__ == [ + "build_benchmark_group_report", + "build_benchmark_report", + "build_group_report_markdown", + "build_report_markdown", + "config_line", + "create_benchmark_group_report", + "create_benchmark_report", + "default_report_title", + "int_metric", + "metric", + "number_metric", + "save_report", + "workload_line", + ] + assert report.build_benchmark_report is reports.build_benchmark_report + assert report.build_benchmark_group_report is reports.build_benchmark_group_report + assert report.build_report_markdown is reports.build_report_markdown + assert report._save_report is reports.save_report + assert report._default_report_title is reports.default_report_title + assert report.create_benchmark_report is not reports.create_benchmark_report + assert report.create_benchmark_group_report is not reports.create_benchmark_group_report + assert report.create_benchmark_report.__module__ == "compat_report_construction_facade" + finally: + sys.modules.pop("compat_report_construction_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tests/tools/test_measurement_wandb_reports.py b/tests/tools/test_measurement_wandb_reports.py index 5006e4b2..ceee0c53 100644 --- a/tests/tools/test_measurement_wandb_reports.py +++ b/tests/tools/test_measurement_wandb_reports.py @@ -3,6 +3,7 @@ from __future__ import annotations +import importlib import json import logging import os @@ -514,6 +515,31 @@ def test_wandb_report_markdown_uses_sanitized_fields(wandb_report_tool: ModuleTy assert "sk-secret-token" not in markdown +def test_wandb_report_save_uses_auth_preflight_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + wandb_reports = importlib.import_module("measurement_tools.wandb_reports") + + report = SimpleNamespace(save=lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("relogin required"))) + fallback_calls: list[tuple[Any, bool]] = [] + + def fallback(value: Any, *, draft: bool) -> str: + fallback_calls.append((value, draft)) + return "saved" + + monkeypatch.setattr(wandb_reports, "_save_report_without_project_preflight", fallback) + + assert wandb_reports.save_report(report, draft=True) == "saved" + assert fallback_calls == [(report, True)] + + +def test_wandb_report_save_preserves_unrelated_sdk_errors() -> None: + wandb_reports = importlib.import_module("measurement_tools.wandb_reports") + + report = SimpleNamespace(save=lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("permission denied"))) + + with pytest.raises(RuntimeError, match="permission denied"): + wandb_reports.save_report(report, draft=False) + + def test_wandb_report_views_parse_v1_and_v2_with_explicit_axes(wandb_report_tool: ModuleType) -> None: summary, v1_config = _wandb_report_fixture() path = wandb_report_tool.WandbRunPath(entity="entity", project="project", run_id="run-a") diff --git a/tools/measurement/create_wandb_report.py b/tools/measurement/create_wandb_report.py index 1fa6bffb..17bac52a 100644 --- a/tools/measurement/create_wandb_report.py +++ b/tools/measurement/create_wandb_report.py @@ -19,12 +19,10 @@ import cyclopts from measurement_tools.cli import LogFormat, configure_logging, log_bad_input -from measurement_tools.wandb_models import ( # noqa: F401 - ConfigMetadata, - ResolvedWandbConfig, - WandbMode, - WorkloadMetadata, -) +from measurement_tools.wandb_models import ConfigMetadata as ConfigMetadata +from measurement_tools.wandb_models import ResolvedWandbConfig as ResolvedWandbConfig +from measurement_tools.wandb_models import WandbMode as WandbMode +from measurement_tools.wandb_models import WorkloadMetadata as WorkloadMetadata from measurement_tools.wandb_report_catalog import ( _MEASUREMENT_TABLE_KEYS, _WORKSPACE_BAR_SECTIONS, @@ -45,13 +43,17 @@ GroupComparison, WandbProjectPath, WandbRunPath, - WandbRunView, group_comparison, parse_wandb_project_path, parse_wandb_run_path, - parse_wandb_run_view, validate_wandb_returned_url, ) +from measurement_tools.wandb_report_models import ( + WandbRunView as WandbRunView, +) +from measurement_tools.wandb_report_models import ( + parse_wandb_run_view as parse_wandb_run_view, +) from measurement_tools.wandb_report_sdk import ( read_group_views, report_settings, @@ -71,7 +73,31 @@ validate_output_text, validate_output_url, ) -from measurement_tools.wandb_setup import WandbSdkEnvironment, require_wandb +from measurement_tools.wandb_reports import ( + build_benchmark_group_report, + build_benchmark_report, + config_line, + default_report_title, + int_metric, + metric, + number_metric, + save_report, + workload_line, +) +from measurement_tools.wandb_reports import ( + build_group_report_markdown as build_group_report_markdown, +) +from measurement_tools.wandb_reports import ( + build_report_markdown as build_report_markdown, +) +from measurement_tools.wandb_reports import ( + create_benchmark_group_report as _create_benchmark_group_report, +) +from measurement_tools.wandb_reports import ( + create_benchmark_report as _create_benchmark_report, +) +from measurement_tools.wandb_setup import WandbSdkEnvironment +from measurement_tools.wandb_setup import require_wandb as require_wandb app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.wandb_report") @@ -97,6 +123,13 @@ _table_code = table_code _validate_output_text = validate_output_text _validate_output_url = validate_output_url +_config_line = config_line +_default_report_title = default_report_title +_int_metric = int_metric +_metric = metric +_number_metric = number_metric +_save_report = save_report +_workload_line = workload_line def create_benchmark_report( @@ -109,36 +142,16 @@ def create_benchmark_report( timeout: int = 60, ) -> WandbReportResult: """Create a W&B report for one benchmark run.""" - resolved = _report_settings(settings, run_path) - effective_run_path = run_path.with_base_url(resolved.wandb_base_url or run_path.base_url) - with WandbSdkEnvironment(resolved): - wandb, wr, expr = require_wandb_report_sdk() - run = wandb.Api(timeout=timeout).run(run_path.path) - view = parse_wandb_run_view( - run, - run_path=effective_run_path, - allowed_metrics=frozenset(_all_report_metrics()), - ) - report_title = _plain_text(title) if title is not None else _default_report_title(view) - report_description = _plain_text( - description or "SDK-generated benchmark report for a NeMo Anonymizer measurement run." - ) - report = build_benchmark_report( - view, - title=report_title, - description=report_description, - wr=wr, - expr=expr, - ) - _save_report(report, draft=draft) - report_url = validate_wandb_returned_url(report.url, expected_base_url=effective_run_path.base_url) - return WandbReportResult( - run_path=run_path.path, - run_url=effective_run_path.url, - project_path=f"{run_path.entity}/{run_path.project}", - report_url=report_url, + return _create_benchmark_report( + run_path, + settings=settings, + title=title, + description=description, draft=draft, - title=report_title, + timeout=timeout, + sdk_loader=require_wandb_report_sdk, + report_builder=build_benchmark_report, + report_saver=_save_report, ) @@ -190,36 +203,17 @@ def create_benchmark_group_report( expected_run_kind: Literal["native_suite", "sweep_arm", "imported_case"] | None = None, ) -> WandbReportResult: """Create a W&B report for a benchmark run group.""" - resolved = _report_settings(settings, project_path) - effective_project_path = project_path.with_base_url(resolved.wandb_base_url or project_path.base_url) - with WandbSdkEnvironment(resolved): - wandb, wr, expr = require_wandb_report_sdk() - views = _read_group_views(wandb, project_path=effective_project_path, group=group) - comparison = group_comparison(views, expected_run_kind=expected_run_kind) - report_title = ( - _plain_text(title) if title is not None else f"NeMo Anonymizer Benchmark Group: {_plain_text(group)}" - ) - report_description = _plain_text( - description or "SDK-generated benchmark sweep report for NeMo Anonymizer measurement runs." - ) - report = build_benchmark_group_report( - effective_project_path, - group=group, - title=report_title, - description=report_description, - sweep_param_columns=_sweep_param_columns(views), - comparison=comparison, - wr=wr, - expr=expr, - ) - _save_report(report, draft=draft) - report_url = validate_wandb_returned_url(report.url, expected_base_url=effective_project_path.base_url) - return WandbReportResult( - project_path=project_path.path, + return _create_benchmark_group_report( + project_path, + settings=settings, group=group, - report_url=report_url, + title=title, + description=description, draft=draft, - title=report_title, + expected_run_kind=expected_run_kind, + sdk_loader=require_wandb_report_sdk, + report_builder=build_benchmark_group_report, + report_saver=_save_report, ) @@ -254,80 +248,6 @@ def build_benchmark_workspace( ) -def build_benchmark_report(view: WandbRunView, *, title: str, description: str, wr: Any, expr: Any) -> Any: - """Build an unsaved report object from a W&B run.""" - run_path = view.path - run_id = run_path.run_id - runset = wr.Runset( - entity=run_path.entity, - project=run_path.project, - name="Benchmark run", - filters=[expr.Metric("ID") == run_id], - pinned_columns=["Name", "State", "CreatedTimestamp", "Tags"], - visible_columns=_single_run_visible_columns(), - run_settings={run_id: wr.RunSettings(color="#2F80ED", disabled=False)}, - ) - safe_title = _escape_heading(title) - return wr.Report( - entity=run_path.entity, - project=run_path.project, - title=safe_title, - description=_escape_markdown_text(description), - width="fluid", - blocks=[ - wr.H1(safe_title), - wr.MarkdownBlock(build_report_markdown(view)), - wr.H2("Run Panels"), - wr.PanelGrid(runsets=[runset], panels=_benchmark_panels(wr)), - ], - ) - - -def build_benchmark_group_report( - project_path: WandbProjectPath, - *, - group: str, - title: str, - description: str, - sweep_param_columns: list[str] | None = None, - comparison: GroupComparison | None = None, - wr: Any, - expr: Any, -) -> Any: - """Build an unsaved report object for a W&B run group.""" - resolved_comparison = comparison or GroupComparison( - run_kind="sweep_arm", - config_key="sweep_arm_id", - label="Sweep Arm", - ) - runset = wr.Runset( - entity=project_path.entity, - project=project_path.project, - name="Benchmark sweep", - filters=[expr.Metric("Group") == group], - groupby=[f"config.{resolved_comparison.config_key}"], - pinned_columns=["Name", "State", "CreatedTimestamp", "Group", "Tags"], - visible_columns=_group_visible_columns(resolved_comparison, sweep_param_columns), - ) - safe_title = _escape_heading(title) - return wr.Report( - entity=project_path.entity, - project=project_path.project, - title=safe_title, - description=_escape_markdown_text(description), - width="fluid", - blocks=[ - wr.H1(safe_title), - wr.MarkdownBlock(build_group_report_markdown(group=group, comparison=resolved_comparison)), - wr.H2("Sweep Panels"), - wr.PanelGrid( - runsets=[runset], - panels=_benchmark_panels(wr, groupby=wr.Config(resolved_comparison.config_key)), - ), - ], - ) - - def _benchmark_workspace_settings(ws: Any) -> Any: return ws.WorkspaceSettings( sort_panels_alphabetically=False, @@ -421,167 +341,16 @@ def _workspace_metric_accessor(wr: Any, column: str) -> Any: raise ValueError(f"Unsupported workspace metric section: {section!r}") -def build_group_report_markdown(*, group: str, comparison: GroupComparison | None = None) -> str: - resolved = comparison or GroupComparison(run_kind="sweep_arm", config_key="sweep_arm_id", label="Sweep Arm") - return f"""### Group Summary - -This report compares benchmark runs in W&B group {_code_span(group)}. - -Run kind: {_code_span(resolved.run_kind)}. Comparison axis: {_code_span(resolved.label)}. - -### Privacy Boundary - -This report is built from benchmark W&B summary/config fields. The benchmark runner sanitizes these fields before upload and excludes raw text, prompts, model responses, replacement maps, entity payloads, trace records, paths, URLs, provider config payloads, and sensitive-looking run tags. -""" - - -def build_report_markdown(view: WandbRunView) -> str: - """Render report Markdown exclusively from a closed typed run view.""" - summary = view.summary.metrics - metadata = view.metadata - benchmark = metadata.benchmark - runtime = metadata.runtime - git = metadata.git - workload_lines = [_workload_line(item) for item in metadata.workloads] - config_lines = [_config_line(item) for item in metadata.configs] - return f"""### Run Summary - -[{_escape_link_label(view.name)}]({view.path.url}) finished with **{_int_metric(summary, "benchmark/case_completed")}/{_int_metric(summary, "benchmark/case_total")} cases completed** and **{_int_metric(summary, "benchmark/case_errored")} errors**. - -| Metric | Value | -| --- | ---: | -| Case success rate | {_number_metric(summary, "benchmark/case_success_rate")} | -| Mean case elapsed seconds | {_number_metric(summary, "benchmark/case_elapsed_sec_mean")} | -| Final entities | {_int_metric(summary, "measurement/record/final_entity_count")} | -| Original value leaks | {_int_metric(summary, "measurement/record/original_value_leak_count")} | -| Model requests | {_int_metric(summary, "measurement/ndd_workflow/observed_total_requests")} | -| Failed model requests | {_int_metric(summary, "measurement/ndd_workflow/observed_failed_requests")} | -| Total model tokens | {_int_metric(summary, "measurement/ndd_workflow/observed_total_tokens")} | - -### Benchmark Metadata - -| Field | Value | -| --- | --- | -| Run kind | {_table_code(metadata.run_kind)} | -| Suite | {_table_code(benchmark.suite_id if benchmark else None)} | -| Suite file hash | {_table_code(benchmark.suite_file_hash if benchmark else None)} | -| Git branch | {_table_code(git.branch if git else None)} | -| Git commit | {_table_code(git.commit if git else None)} | -| Git dirty | {_table_code(git.dirty if git else None)} | -| Anonymizer | {_table_code(runtime.anonymizer_version if runtime else None)} | -| DataDesigner | {_table_code(runtime.datadesigner_version if runtime else None)} | -| W&B | {_table_code(runtime.wandb_version if runtime else None)} | - -### Workloads - -{chr(10).join(workload_lines) or "- none"} - -### Configs - -{chr(10).join(config_lines) or "- none"} - -### Privacy Boundary - -This report is built from benchmark W&B summary/config fields. The benchmark runner sanitizes these fields before upload and excludes raw text, prompts, model responses, replacement maps, entity payloads, trace records, paths, URLs, provider config payloads, and sensitive-looking run tags. -""" - - -def _save_report(report: Any, *, draft: bool) -> Any: - """Save a report, falling back around a W&B SDK project-list auth edge.""" - try: - return report.save(draft=draft) - except Exception as exc: # noqa: BLE001 -- preserve the SDK error as fallback context - if "relogin required" not in str(exc).lower(): - raise - logger.info("Falling back to direct W&B report upsert after project preflight auth failure.") - return _save_report_without_project_preflight(report, draft=draft) - - -def _save_report_without_project_preflight(report: Any, *, draft: bool) -> Any: - """Call the same report upsert mutation without the project-list preflight.""" - wandb = require_wandb() - from wandb_workspaces._graphql import execute_graphql - from wandb_workspaces.reports.v2 import gql, internal - - api = wandb.Api() - model = report._to_model() - result = execute_graphql( - api, - gql.upsert_view, - { - "id": None if not model.id else model.id, - "name": internal._generate_name() if not model.name else model.name, - "entityName": model.project.entity_name, - "projectName": model.project.name, - "description": model.description, - "displayName": model.display_name, - "type": "runs/draft" if draft else "runs", - "spec": model.spec.model_dump_json(by_alias=True, exclude_none=True), - }, - ) - report.id = result["upsertView"]["view"]["id"] - return report - - def _save_workspace(workspace: Any) -> Any: """Save a workspace through the W&B Workspaces API.""" return workspace.save() -def _default_report_title(view: WandbRunView) -> str: - return f"NeMo Anonymizer Benchmark: {_plain_text(view.name)}"[:128] - - def _default_workspace_title(*, group: str | None) -> str: suffix = f": {_plain_text(group)}" if group else "" return f"NeMo Anonymizer Benchmark Workspace{suffix}"[:128] -def _workload_line(item: WorkloadMetadata) -> str: - source_suffix = item.source.suffix if item.source else None - source_kind = item.source.kind if item.source else None - return ( - f"- {_code_span(item.id)}: source_kind={_code_span(source_kind)}, " - f"source_suffix={_code_span(source_suffix)}, row_limit={_escape_list_text(item.row_limit)}, " - f"text_column={_code_span(item.text_column)}" - ) - - -def _config_line(item: ConfigMetadata) -> str: - detect = item.detect - replace = item.replace - rewrite = item.rewrite - parts = [ - f"strategy={_code_span(replace.strategy if replace else ('rewrite' if rewrite else None))}", - f"entity_label_count={_escape_list_text(detect.entity_label_count if detect else None)}", - f"gliner_threshold={_escape_list_text(detect.gliner_threshold if detect else None)}", - ] - if rewrite: - parts.append(f"risk_tolerance={_code_span(rewrite.risk_tolerance)}") - return f"- {_code_span(item.id)}: " + ", ".join(parts) - - -def _metric(summary: dict[str, Any], key: str) -> int | float | None: - value = summary.get(key) - if isinstance(value, bool): - return int(value) - if isinstance(value, int | float): - return value - return None - - -def _int_metric(summary: dict[str, Any], key: str) -> int: - value = _metric(summary, key) - return int(value) if value is not None else 0 - - -def _number_metric(summary: dict[str, Any], key: str) -> str: - value = _metric(summary, key) - if value is None: - return "n/a" - return f"{float(value):.4g}" - - def render_result(result: WandbReportResult | WandbWorkspaceResult, *, json_output: bool) -> str: if json_output: return result.model_dump_json(indent=2) diff --git a/tools/measurement/measurement_tools/wandb_reports.py b/tools/measurement/measurement_tools/wandb_reports.py new file mode 100644 index 00000000..57511581 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_reports.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W&B benchmark report orchestration, construction, and persistence.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any, Literal + +from measurement_tools.wandb_models import ConfigMetadata, ResolvedWandbConfig, WorkloadMetadata +from measurement_tools.wandb_report_catalog import ( + all_report_metrics, + benchmark_panels, + group_visible_columns, + single_run_visible_columns, +) +from measurement_tools.wandb_report_contracts import WandbReportResult +from measurement_tools.wandb_report_models import ( + GroupComparison, + WandbProjectPath, + WandbRunPath, + WandbRunView, + group_comparison, + parse_wandb_run_view, + validate_wandb_returned_url, +) +from measurement_tools.wandb_report_sdk import ( + read_group_views, + report_settings, + require_wandb_report_sdk, + sweep_param_columns, +) +from measurement_tools.wandb_report_text import ( + code_span, + escape_heading, + escape_link_label, + escape_list_text, + escape_markdown_text, + plain_text, + table_code, +) +from measurement_tools.wandb_setup import WandbSdkEnvironment, require_wandb + +logger = logging.getLogger("measurement.wandb_report") + + +def create_benchmark_report( + run_path: WandbRunPath, + *, + settings: ResolvedWandbConfig | None = None, + title: str | None = None, + description: str | None = None, + draft: bool = True, + timeout: int = 60, + sdk_loader: Callable[[], tuple[Any, Any, Any]] = require_wandb_report_sdk, + report_builder: Callable[..., Any] | None = None, + report_saver: Callable[..., Any] | None = None, +) -> WandbReportResult: + """Create a W&B report for one benchmark run.""" + builder = report_builder or build_benchmark_report + saver = report_saver or save_report + resolved = report_settings(settings, run_path) + effective_run_path = run_path.with_base_url(resolved.wandb_base_url or run_path.base_url) + with WandbSdkEnvironment(resolved): + wandb, wr, expr = sdk_loader() + run = wandb.Api(timeout=timeout).run(run_path.path) + view = parse_wandb_run_view( + run, + run_path=effective_run_path, + allowed_metrics=frozenset(all_report_metrics()), + ) + report_title = plain_text(title) if title is not None else default_report_title(view) + report_description = plain_text( + description or "SDK-generated benchmark report for a NeMo Anonymizer measurement run." + ) + report = builder( + view, + title=report_title, + description=report_description, + wr=wr, + expr=expr, + ) + saver(report, draft=draft) + report_url = validate_wandb_returned_url(report.url, expected_base_url=effective_run_path.base_url) + return WandbReportResult( + run_path=run_path.path, + run_url=effective_run_path.url, + project_path=f"{run_path.entity}/{run_path.project}", + report_url=report_url, + draft=draft, + title=report_title, + ) + + +def create_benchmark_group_report( + project_path: WandbProjectPath, + *, + settings: ResolvedWandbConfig | None = None, + group: str, + title: str | None = None, + description: str | None = None, + draft: bool = True, + expected_run_kind: Literal["native_suite", "sweep_arm", "imported_case"] | None = None, + sdk_loader: Callable[[], tuple[Any, Any, Any]] = require_wandb_report_sdk, + report_builder: Callable[..., Any] | None = None, + report_saver: Callable[..., Any] | None = None, +) -> WandbReportResult: + """Create a W&B report for a benchmark run group.""" + builder = report_builder or build_benchmark_group_report + saver = report_saver or save_report + resolved = report_settings(settings, project_path) + effective_project_path = project_path.with_base_url(resolved.wandb_base_url or project_path.base_url) + with WandbSdkEnvironment(resolved): + wandb, wr, expr = sdk_loader() + views = read_group_views(wandb, project_path=effective_project_path, group=group) + comparison = group_comparison(views, expected_run_kind=expected_run_kind) + report_title = ( + plain_text(title) if title is not None else f"NeMo Anonymizer Benchmark Group: {plain_text(group)}" + ) + report_description = plain_text( + description or "SDK-generated benchmark sweep report for NeMo Anonymizer measurement runs." + ) + report = builder( + effective_project_path, + group=group, + title=report_title, + description=report_description, + sweep_param_columns=sweep_param_columns(views), + comparison=comparison, + wr=wr, + expr=expr, + ) + saver(report, draft=draft) + report_url = validate_wandb_returned_url(report.url, expected_base_url=effective_project_path.base_url) + return WandbReportResult( + project_path=project_path.path, + group=group, + report_url=report_url, + draft=draft, + title=report_title, + ) + + +def build_benchmark_report(view: WandbRunView, *, title: str, description: str, wr: Any, expr: Any) -> Any: + """Build an unsaved report object from a W&B run.""" + run_path = view.path + run_id = run_path.run_id + runset = wr.Runset( + entity=run_path.entity, + project=run_path.project, + name="Benchmark run", + filters=[expr.Metric("ID") == run_id], + pinned_columns=["Name", "State", "CreatedTimestamp", "Tags"], + visible_columns=single_run_visible_columns(), + run_settings={run_id: wr.RunSettings(color="#2F80ED", disabled=False)}, + ) + safe_title = escape_heading(title) + return wr.Report( + entity=run_path.entity, + project=run_path.project, + title=safe_title, + description=escape_markdown_text(description), + width="fluid", + blocks=[ + wr.H1(safe_title), + wr.MarkdownBlock(build_report_markdown(view)), + wr.H2("Run Panels"), + wr.PanelGrid(runsets=[runset], panels=benchmark_panels(wr)), + ], + ) + + +def build_benchmark_group_report( + project_path: WandbProjectPath, + *, + group: str, + title: str, + description: str, + sweep_param_columns: list[str] | None = None, + comparison: GroupComparison | None = None, + wr: Any, + expr: Any, +) -> Any: + """Build an unsaved report object for a W&B run group.""" + resolved_comparison = comparison or GroupComparison( + run_kind="sweep_arm", + config_key="sweep_arm_id", + label="Sweep Arm", + ) + runset = wr.Runset( + entity=project_path.entity, + project=project_path.project, + name="Benchmark sweep", + filters=[expr.Metric("Group") == group], + groupby=[f"config.{resolved_comparison.config_key}"], + pinned_columns=["Name", "State", "CreatedTimestamp", "Group", "Tags"], + visible_columns=group_visible_columns(resolved_comparison, sweep_param_columns), + ) + safe_title = escape_heading(title) + return wr.Report( + entity=project_path.entity, + project=project_path.project, + title=safe_title, + description=escape_markdown_text(description), + width="fluid", + blocks=[ + wr.H1(safe_title), + wr.MarkdownBlock(build_group_report_markdown(group=group, comparison=resolved_comparison)), + wr.H2("Sweep Panels"), + wr.PanelGrid( + runsets=[runset], + panels=benchmark_panels(wr, groupby=wr.Config(resolved_comparison.config_key)), + ), + ], + ) + + +def build_group_report_markdown(*, group: str, comparison: GroupComparison | None = None) -> str: + resolved = comparison or GroupComparison(run_kind="sweep_arm", config_key="sweep_arm_id", label="Sweep Arm") + return f"""### Group Summary + +This report compares benchmark runs in W&B group {code_span(group)}. + +Run kind: {code_span(resolved.run_kind)}. Comparison axis: {code_span(resolved.label)}. + +### Privacy Boundary + +This report is built from benchmark W&B summary/config fields. The benchmark runner sanitizes these fields before upload and excludes raw text, prompts, model responses, replacement maps, entity payloads, trace records, paths, URLs, provider config payloads, and sensitive-looking run tags. +""" + + +def build_report_markdown(view: WandbRunView) -> str: + """Render report Markdown exclusively from a closed typed run view.""" + summary = view.summary.metrics + metadata = view.metadata + benchmark = metadata.benchmark + runtime = metadata.runtime + git = metadata.git + workload_lines = [workload_line(item) for item in metadata.workloads] + config_lines = [config_line(item) for item in metadata.configs] + return f"""### Run Summary + +[{escape_link_label(view.name)}]({view.path.url}) finished with **{int_metric(summary, "benchmark/case_completed")}/{int_metric(summary, "benchmark/case_total")} cases completed** and **{int_metric(summary, "benchmark/case_errored")} errors**. + +| Metric | Value | +| --- | ---: | +| Case success rate | {number_metric(summary, "benchmark/case_success_rate")} | +| Mean case elapsed seconds | {number_metric(summary, "benchmark/case_elapsed_sec_mean")} | +| Final entities | {int_metric(summary, "measurement/record/final_entity_count")} | +| Original value leaks | {int_metric(summary, "measurement/record/original_value_leak_count")} | +| Model requests | {int_metric(summary, "measurement/ndd_workflow/observed_total_requests")} | +| Failed model requests | {int_metric(summary, "measurement/ndd_workflow/observed_failed_requests")} | +| Total model tokens | {int_metric(summary, "measurement/ndd_workflow/observed_total_tokens")} | + +### Benchmark Metadata + +| Field | Value | +| --- | --- | +| Run kind | {table_code(metadata.run_kind)} | +| Suite | {table_code(benchmark.suite_id if benchmark else None)} | +| Suite file hash | {table_code(benchmark.suite_file_hash if benchmark else None)} | +| Git branch | {table_code(git.branch if git else None)} | +| Git commit | {table_code(git.commit if git else None)} | +| Git dirty | {table_code(git.dirty if git else None)} | +| Anonymizer | {table_code(runtime.anonymizer_version if runtime else None)} | +| DataDesigner | {table_code(runtime.datadesigner_version if runtime else None)} | +| W&B | {table_code(runtime.wandb_version if runtime else None)} | + +### Workloads + +{chr(10).join(workload_lines) or "- none"} + +### Configs + +{chr(10).join(config_lines) or "- none"} + +### Privacy Boundary + +This report is built from benchmark W&B summary/config fields. The benchmark runner sanitizes these fields before upload and excludes raw text, prompts, model responses, replacement maps, entity payloads, trace records, paths, URLs, provider config payloads, and sensitive-looking run tags. +""" + + +def save_report(report: Any, *, draft: bool) -> Any: + """Save a report, falling back around a W&B SDK project-list auth edge.""" + try: + return report.save(draft=draft) + except Exception as exc: # noqa: BLE001 -- preserve the SDK error as fallback context + if "relogin required" not in str(exc).lower(): + raise + logger.info("Falling back to direct W&B report upsert after project preflight auth failure.") + return _save_report_without_project_preflight(report, draft=draft) + + +def _save_report_without_project_preflight(report: Any, *, draft: bool) -> Any: + """Call the same report upsert mutation without the project-list preflight.""" + wandb = require_wandb() + from wandb_workspaces._graphql import execute_graphql + from wandb_workspaces.reports.v2 import gql, internal + + api = wandb.Api() + model = report._to_model() + result = execute_graphql( + api, + gql.upsert_view, + { + "id": None if not model.id else model.id, + "name": internal._generate_name() if not model.name else model.name, + "entityName": model.project.entity_name, + "projectName": model.project.name, + "description": model.description, + "displayName": model.display_name, + "type": "runs/draft" if draft else "runs", + "spec": model.spec.model_dump_json(by_alias=True, exclude_none=True), + }, + ) + report.id = result["upsertView"]["view"]["id"] + return report + + +def default_report_title(view: WandbRunView) -> str: + return f"NeMo Anonymizer Benchmark: {plain_text(view.name)}"[:128] + + +def workload_line(item: WorkloadMetadata) -> str: + source_suffix = item.source.suffix if item.source else None + source_kind = item.source.kind if item.source else None + return ( + f"- {code_span(item.id)}: source_kind={code_span(source_kind)}, " + f"source_suffix={code_span(source_suffix)}, row_limit={escape_list_text(item.row_limit)}, " + f"text_column={code_span(item.text_column)}" + ) + + +def config_line(item: ConfigMetadata) -> str: + detect = item.detect + replace = item.replace + rewrite = item.rewrite + parts = [ + f"strategy={code_span(replace.strategy if replace else ('rewrite' if rewrite else None))}", + f"entity_label_count={escape_list_text(detect.entity_label_count if detect else None)}", + f"gliner_threshold={escape_list_text(detect.gliner_threshold if detect else None)}", + ] + if rewrite: + parts.append(f"risk_tolerance={code_span(rewrite.risk_tolerance)}") + return f"- {code_span(item.id)}: " + ", ".join(parts) + + +def metric(summary: dict[str, Any], key: str) -> int | float | None: + value = summary.get(key) + if isinstance(value, bool): + return int(value) + if isinstance(value, int | float): + return value + return None + + +def int_metric(summary: dict[str, Any], key: str) -> int: + value = metric(summary, key) + return int(value) if value is not None else 0 + + +def number_metric(summary: dict[str, Any], key: str) -> str: + value = metric(summary, key) + if value is None: + return "n/a" + return f"{float(value):.4g}" + + +__all__ = [ + "build_benchmark_group_report", + "build_benchmark_report", + "build_group_report_markdown", + "build_report_markdown", + "config_line", + "create_benchmark_group_report", + "create_benchmark_report", + "default_report_title", + "int_metric", + "metric", + "number_metric", + "save_report", + "workload_line", +] From 3a09d6464883de1d6083a8bd66f0bd292f06fa0f Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 20:40:31 +0000 Subject: [PATCH 09/17] refactor(measurement): split W&B workspace construction Signed-off-by: Aaron Gonzales --- tests/conftest.py | 5 +- .../test_measurement_module_compatibility.py | 37 +++ tools/measurement/create_wandb_report.py | 216 ++++------------- .../measurement_tools/wandb_workspaces.py | 228 ++++++++++++++++++ 4 files changed, 317 insertions(+), 169 deletions(-) create mode 100644 tools/measurement/measurement_tools/wandb_workspaces.py diff --git a/tests/conftest.py b/tests/conftest.py index 69f2c6c2..4b23610b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -42,7 +42,10 @@ def loader( assert spec.loader is not None module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - for import_path in (*additional_paths, path.parent): + import_paths = list(additional_paths) + if not any(path.is_relative_to(import_path) for import_path in additional_paths): + import_paths.append(path.parent) + for import_path in import_paths: sys.path.insert(0, str(import_path)) spec.loader.exec_module(module) return module diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index a45d48b0..5d66e766 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -210,6 +210,43 @@ def test_wandb_report_facade_preserves_report_construction_contracts() -> None: sys.modules.pop("compat_report_construction_facade", None) +def test_wandb_report_facade_preserves_workspace_construction_contracts() -> None: + workspaces = importlib.import_module("measurement_tools.wandb_workspaces") + report = _load_module(MEASUREMENT_ROOT / "create_wandb_report.py", "compat_workspace_construction_facade") + try: + assert workspaces.__all__ == [ + "benchmark_run_filters", + "benchmark_run_groupby", + "benchmark_runset_settings", + "benchmark_workspace_sections", + "benchmark_workspace_settings", + "build_benchmark_workspace", + "comparison_workspace_panels", + "create_benchmark_workspace", + "default_workspace_title", + "save_workspace", + "table_workspace_panels", + "workspace_bar_panels", + "workspace_metric_accessor", + ] + assert report.build_benchmark_workspace is workspaces.build_benchmark_workspace + assert report._benchmark_workspace_settings is workspaces.benchmark_workspace_settings + assert report._benchmark_runset_settings is workspaces.benchmark_runset_settings + assert report._benchmark_run_filters is workspaces.benchmark_run_filters + assert report._benchmark_run_groupby is workspaces.benchmark_run_groupby + assert report._benchmark_workspace_sections is workspaces.benchmark_workspace_sections + assert report._workspace_bar_panels is workspaces.workspace_bar_panels + assert report._comparison_workspace_panels is workspaces.comparison_workspace_panels + assert report._table_workspace_panels is workspaces.table_workspace_panels + assert report._workspace_metric_accessor is workspaces.workspace_metric_accessor + assert report._save_workspace is workspaces.save_workspace + assert report._default_workspace_title is workspaces.default_workspace_title + assert report.create_benchmark_workspace is not workspaces.create_benchmark_workspace + assert report.create_benchmark_workspace.__module__ == "compat_workspace_construction_facade" + finally: + sys.modules.pop("compat_workspace_construction_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/create_wandb_report.py b/tools/measurement/create_wandb_report.py index 17bac52a..08fcea09 100644 --- a/tools/measurement/create_wandb_report.py +++ b/tools/measurement/create_wandb_report.py @@ -15,7 +15,7 @@ import logging import sys -from typing import Annotated, Any, Literal +from typing import Annotated, Literal import cyclopts from measurement_tools.cli import LogFormat, configure_logging, log_bad_input @@ -24,11 +24,6 @@ from measurement_tools.wandb_models import WandbMode as WandbMode from measurement_tools.wandb_models import WorkloadMetadata as WorkloadMetadata from measurement_tools.wandb_report_catalog import ( - _MEASUREMENT_TABLE_KEYS, - _WORKSPACE_BAR_SECTIONS, - _WORKSPACE_COMPARISON_COLUMNS, - _WORKSPACE_JOB_FILTERS, - _WORKSPACE_SUMMARY_SCALARS, all_report_metrics, bar_panel, benchmark_panels, @@ -40,20 +35,26 @@ ) from measurement_tools.wandb_report_contracts import WandbReportResult, WandbWorkspaceResult from measurement_tools.wandb_report_models import ( - GroupComparison, + GroupComparison as GroupComparison, +) +from measurement_tools.wandb_report_models import ( WandbProjectPath, WandbRunPath, - group_comparison, parse_wandb_project_path, parse_wandb_run_path, - validate_wandb_returned_url, ) from measurement_tools.wandb_report_models import ( WandbRunView as WandbRunView, ) +from measurement_tools.wandb_report_models import ( + group_comparison as group_comparison, +) from measurement_tools.wandb_report_models import ( parse_wandb_run_view as parse_wandb_run_view, ) +from measurement_tools.wandb_report_models import ( + validate_wandb_returned_url as validate_wandb_returned_url, +) from measurement_tools.wandb_report_sdk import ( read_group_views, report_settings, @@ -96,8 +97,25 @@ from measurement_tools.wandb_reports import ( create_benchmark_report as _create_benchmark_report, ) -from measurement_tools.wandb_setup import WandbSdkEnvironment +from measurement_tools.wandb_setup import WandbSdkEnvironment as WandbSdkEnvironment from measurement_tools.wandb_setup import require_wandb as require_wandb +from measurement_tools.wandb_workspaces import ( + benchmark_run_filters, + benchmark_run_groupby, + benchmark_runset_settings, + benchmark_workspace_sections, + benchmark_workspace_settings, + build_benchmark_workspace, + comparison_workspace_panels, + default_workspace_title, + save_workspace, + table_workspace_panels, + workspace_bar_panels, + workspace_metric_accessor, +) +from measurement_tools.wandb_workspaces import ( + create_benchmark_workspace as _create_benchmark_workspace, +) app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.wandb_report") @@ -130,6 +148,17 @@ _number_metric = number_metric _save_report = save_report _workload_line = workload_line +_benchmark_run_filters = benchmark_run_filters +_benchmark_run_groupby = benchmark_run_groupby +_benchmark_runset_settings = benchmark_runset_settings +_benchmark_workspace_sections = benchmark_workspace_sections +_benchmark_workspace_settings = benchmark_workspace_settings +_comparison_workspace_panels = comparison_workspace_panels +_default_workspace_title = default_workspace_title +_save_workspace = save_workspace +_table_workspace_panels = table_workspace_panels +_workspace_bar_panels = workspace_bar_panels +_workspace_metric_accessor = workspace_metric_accessor def create_benchmark_report( @@ -164,31 +193,16 @@ def create_benchmark_workspace( expected_run_kind: Literal["native_suite", "sweep_arm", "imported_case"] | None = None, ) -> WandbWorkspaceResult: """Create a W&B workspace for benchmark runs in a project.""" - resolved = _report_settings(settings, project_path) - effective_project_path = project_path.with_base_url(resolved.wandb_base_url or project_path.base_url) - with WandbSdkEnvironment(resolved): - comparison: GroupComparison | None = None - if group is not None: - wandb, _wr, _expr = require_wandb_report_sdk() - views = _read_group_views(wandb, project_path=effective_project_path, group=group) - comparison = group_comparison(views, expected_run_kind=expected_run_kind) - ws, workspace_wr = require_wandb_workspace_sdk() - workspace_title = _plain_text(title) if title is not None else _default_workspace_title(group=group) - workspace = build_benchmark_workspace( - effective_project_path, - group=group, - title=workspace_title, - comparison_config_key=comparison.config_key if comparison is not None else "sweep_arm_id", - ws=ws, - wr=workspace_wr, - ) - saved = _save_workspace(workspace) - workspace_url = validate_wandb_returned_url(saved.url, expected_base_url=effective_project_path.base_url) - return WandbWorkspaceResult( - project_path=project_path.path, + return _create_benchmark_workspace( + project_path, + settings=settings, group=group, - workspace_url=workspace_url, - title=workspace_title, + title=title, + expected_run_kind=expected_run_kind, + report_sdk_loader=require_wandb_report_sdk, + workspace_sdk_loader=require_wandb_workspace_sdk, + workspace_builder=build_benchmark_workspace, + workspace_saver=_save_workspace, ) @@ -217,140 +231,6 @@ def create_benchmark_group_report( ) -def build_benchmark_workspace( - project_path: WandbProjectPath, - *, - group: str | None, - title: str, - comparison_config_key: str = "sweep_arm_id", - ws: Any, - wr: Any, -) -> Any: - """Build an unsaved manual W&B workspace for benchmark runs.""" - groupby = wr.Config(comparison_config_key) if group else None - return ws.Workspace( - entity=project_path.entity, - project=project_path.project, - name=title, - auto_generate_panels=False, - settings=_benchmark_workspace_settings(ws), - runset_settings=_benchmark_runset_settings( - ws, - group=group, - comparison_config_key=comparison_config_key, - ), - sections=_benchmark_workspace_sections( - ws, - wr, - groupby=groupby, - comparison_config_key=comparison_config_key, - ), - ) - - -def _benchmark_workspace_settings(ws: Any) -> Any: - return ws.WorkspaceSettings( - sort_panels_alphabetically=False, - group_by_prefix="first", - max_runs=10, - ) - - -def _benchmark_runset_settings(ws: Any, *, group: str | None, comparison_config_key: str) -> Any: - return ws.RunsetSettings( - filters=_benchmark_run_filters(ws, group=group), - groupby=_benchmark_run_groupby(ws, group=group, comparison_config_key=comparison_config_key), - order=[ws.Ordering(ws.Metric("CreatedTimestamp"), ascending=False)], - pinned_columns=["Name", "State", "CreatedTimestamp", "Group", "JobType", "Tags"], - ) - - -def _benchmark_run_filters(ws: Any, *, group: str | None) -> Any: - job_filter = ws.Or(*(ws.Metric("JobType") == job_type for job_type in _WORKSPACE_JOB_FILTERS)) - return ws.And(job_filter, ws.Metric("Group") == group) if group else job_filter - - -def _benchmark_run_groupby(ws: Any, *, group: str | None, comparison_config_key: str) -> list[Any]: - return [ws.Config(comparison_config_key)] if group else [ws.Metric("Group")] - - -def _benchmark_workspace_sections( - ws: Any, - wr: Any, - *, - groupby: Any | None, - comparison_config_key: str, -) -> list[Any]: - sections = [] - for name, panel_defs, is_open in _WORKSPACE_BAR_SECTIONS: - panels = _workspace_bar_panels(wr, panel_defs, groupby=groupby) - if name == "Benchmark Summary": - panels = [ - *( - wr.ScalarChart(title=_metric_title(metric), metric=metric, groupby_aggfunc="mean") - for metric in _WORKSPACE_SUMMARY_SCALARS - ), - *panels, - ] - sections.append(ws.Section(name=name, panels=panels, is_open=is_open)) - if groupby is not None: - sections.append( - ws.Section( - name="Sweep Comparison", - panels=_comparison_workspace_panels(wr, comparison_config_key=comparison_config_key), - is_open=True, - ) - ) - sections.append(ws.Section(name="Tables", panels=_table_workspace_panels(wr), is_open=False)) - return sections - - -def _workspace_bar_panels( - wr: Any, - panel_defs: tuple[tuple[str, list[str]], ...], - *, - groupby: Any | None, -) -> list[Any]: - return [_bar_panel(wr, title, metrics, groupby=groupby) for title, metrics in panel_defs] - - -def _comparison_workspace_panels(wr: Any, *, comparison_config_key: str = "sweep_arm_id") -> list[Any]: - columns = [f"config:{comparison_config_key}", *_WORKSPACE_COMPARISON_COLUMNS[1:]] - return [ - wr.RunComparer(diff_only=True), - wr.ParallelCoordinatesPlot( - title="Sweep Parameter Tradeoffs", - columns=[ - wr.ParallelCoordinatesPlotColumn(metric=_workspace_metric_accessor(wr, column)) for column in columns - ], - ), - wr.ParameterImportancePlot(with_respect_to="measurement/record/weighted_leakage_rate_mean"), - ] - - -def _table_workspace_panels(wr: Any) -> list[Any]: - return [wr.MediaBrowser(title="Sanitized Measurement Tables", media_keys=_MEASUREMENT_TABLE_KEYS, mode="grid")] - - -def _workspace_metric_accessor(wr: Any, column: str) -> Any: - section, name = column.split(":", maxsplit=1) - if section == "config": - return wr.Config(name) - if section == "summary": - return wr.SummaryMetric(name) - raise ValueError(f"Unsupported workspace metric section: {section!r}") - - -def _save_workspace(workspace: Any) -> Any: - """Save a workspace through the W&B Workspaces API.""" - return workspace.save() - - -def _default_workspace_title(*, group: str | None) -> str: - suffix = f": {_plain_text(group)}" if group else "" - return f"NeMo Anonymizer Benchmark Workspace{suffix}"[:128] - - def render_result(result: WandbReportResult | WandbWorkspaceResult, *, json_output: bool) -> str: if json_output: return result.model_dump_json(indent=2) diff --git a/tools/measurement/measurement_tools/wandb_workspaces.py b/tools/measurement/measurement_tools/wandb_workspaces.py new file mode 100644 index 00000000..7f04aef5 --- /dev/null +++ b/tools/measurement/measurement_tools/wandb_workspaces.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W&B benchmark workspace orchestration, panels, and persistence.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Literal + +from measurement_tools.wandb_models import ResolvedWandbConfig +from measurement_tools.wandb_report_catalog import ( + _MEASUREMENT_TABLE_KEYS, + _WORKSPACE_BAR_SECTIONS, + _WORKSPACE_COMPARISON_COLUMNS, + _WORKSPACE_JOB_FILTERS, + _WORKSPACE_SUMMARY_SCALARS, + bar_panel, + metric_title, +) +from measurement_tools.wandb_report_contracts import WandbWorkspaceResult +from measurement_tools.wandb_report_models import ( + GroupComparison, + WandbProjectPath, + group_comparison, + validate_wandb_returned_url, +) +from measurement_tools.wandb_report_sdk import ( + read_group_views, + report_settings, + require_wandb_report_sdk, + require_wandb_workspace_sdk, +) +from measurement_tools.wandb_report_text import plain_text +from measurement_tools.wandb_setup import WandbSdkEnvironment + + +def create_benchmark_workspace( + project_path: WandbProjectPath, + *, + settings: ResolvedWandbConfig | None = None, + group: str | None = None, + title: str | None = None, + expected_run_kind: Literal["native_suite", "sweep_arm", "imported_case"] | None = None, + report_sdk_loader: Callable[[], tuple[Any, Any, Any]] = require_wandb_report_sdk, + workspace_sdk_loader: Callable[[], tuple[Any, Any]] = require_wandb_workspace_sdk, + workspace_builder: Callable[..., Any] | None = None, + workspace_saver: Callable[[Any], Any] | None = None, +) -> WandbWorkspaceResult: + """Create a W&B workspace for benchmark runs in a project.""" + builder = workspace_builder or build_benchmark_workspace + saver = workspace_saver or save_workspace + resolved = report_settings(settings, project_path) + effective_project_path = project_path.with_base_url(resolved.wandb_base_url or project_path.base_url) + with WandbSdkEnvironment(resolved): + comparison: GroupComparison | None = None + if group is not None: + wandb, _wr, _expr = report_sdk_loader() + views = read_group_views(wandb, project_path=effective_project_path, group=group) + comparison = group_comparison(views, expected_run_kind=expected_run_kind) + ws, workspace_wr = workspace_sdk_loader() + workspace_title = plain_text(title) if title is not None else default_workspace_title(group=group) + workspace = builder( + effective_project_path, + group=group, + title=workspace_title, + comparison_config_key=comparison.config_key if comparison is not None else "sweep_arm_id", + ws=ws, + wr=workspace_wr, + ) + saved = saver(workspace) + workspace_url = validate_wandb_returned_url(saved.url, expected_base_url=effective_project_path.base_url) + return WandbWorkspaceResult( + project_path=project_path.path, + group=group, + workspace_url=workspace_url, + title=workspace_title, + ) + + +def build_benchmark_workspace( + project_path: WandbProjectPath, + *, + group: str | None, + title: str, + comparison_config_key: str = "sweep_arm_id", + ws: Any, + wr: Any, +) -> Any: + """Build an unsaved manual W&B workspace for benchmark runs.""" + groupby = wr.Config(comparison_config_key) if group else None + return ws.Workspace( + entity=project_path.entity, + project=project_path.project, + name=title, + auto_generate_panels=False, + settings=benchmark_workspace_settings(ws), + runset_settings=benchmark_runset_settings( + ws, + group=group, + comparison_config_key=comparison_config_key, + ), + sections=benchmark_workspace_sections( + ws, + wr, + groupby=groupby, + comparison_config_key=comparison_config_key, + ), + ) + + +def benchmark_workspace_settings(ws: Any) -> Any: + return ws.WorkspaceSettings( + sort_panels_alphabetically=False, + group_by_prefix="first", + max_runs=10, + ) + + +def benchmark_runset_settings(ws: Any, *, group: str | None, comparison_config_key: str) -> Any: + return ws.RunsetSettings( + filters=benchmark_run_filters(ws, group=group), + groupby=benchmark_run_groupby(ws, group=group, comparison_config_key=comparison_config_key), + order=[ws.Ordering(ws.Metric("CreatedTimestamp"), ascending=False)], + pinned_columns=["Name", "State", "CreatedTimestamp", "Group", "JobType", "Tags"], + ) + + +def benchmark_run_filters(ws: Any, *, group: str | None) -> Any: + job_filter = ws.Or(*(ws.Metric("JobType") == job_type for job_type in _WORKSPACE_JOB_FILTERS)) + return ws.And(job_filter, ws.Metric("Group") == group) if group else job_filter + + +def benchmark_run_groupby(ws: Any, *, group: str | None, comparison_config_key: str) -> list[Any]: + return [ws.Config(comparison_config_key)] if group else [ws.Metric("Group")] + + +def benchmark_workspace_sections( + ws: Any, + wr: Any, + *, + groupby: Any | None, + comparison_config_key: str, +) -> list[Any]: + sections = [] + for name, panel_defs, is_open in _WORKSPACE_BAR_SECTIONS: + panels = workspace_bar_panels(wr, panel_defs, groupby=groupby) + if name == "Benchmark Summary": + panels = [ + *( + wr.ScalarChart(title=metric_title(metric), metric=metric, groupby_aggfunc="mean") + for metric in _WORKSPACE_SUMMARY_SCALARS + ), + *panels, + ] + sections.append(ws.Section(name=name, panels=panels, is_open=is_open)) + if groupby is not None: + sections.append( + ws.Section( + name="Sweep Comparison", + panels=comparison_workspace_panels(wr, comparison_config_key=comparison_config_key), + is_open=True, + ) + ) + sections.append(ws.Section(name="Tables", panels=table_workspace_panels(wr), is_open=False)) + return sections + + +def workspace_bar_panels( + wr: Any, + panel_defs: tuple[tuple[str, list[str]], ...], + *, + groupby: Any | None, +) -> list[Any]: + return [bar_panel(wr, title, metrics, groupby=groupby) for title, metrics in panel_defs] + + +def comparison_workspace_panels(wr: Any, *, comparison_config_key: str = "sweep_arm_id") -> list[Any]: + columns = [f"config:{comparison_config_key}", *_WORKSPACE_COMPARISON_COLUMNS[1:]] + return [ + wr.RunComparer(diff_only=True), + wr.ParallelCoordinatesPlot( + title="Sweep Parameter Tradeoffs", + columns=[ + wr.ParallelCoordinatesPlotColumn(metric=workspace_metric_accessor(wr, column)) for column in columns + ], + ), + wr.ParameterImportancePlot(with_respect_to="measurement/record/weighted_leakage_rate_mean"), + ] + + +def table_workspace_panels(wr: Any) -> list[Any]: + return [wr.MediaBrowser(title="Sanitized Measurement Tables", media_keys=_MEASUREMENT_TABLE_KEYS, mode="grid")] + + +def workspace_metric_accessor(wr: Any, column: str) -> Any: + section, name = column.split(":", maxsplit=1) + if section == "config": + return wr.Config(name) + if section == "summary": + return wr.SummaryMetric(name) + raise ValueError(f"Unsupported workspace metric section: {section!r}") + + +def save_workspace(workspace: Any) -> Any: + """Save a workspace through the W&B Workspaces API.""" + return workspace.save() + + +def default_workspace_title(*, group: str | None) -> str: + suffix = f": {plain_text(group)}" if group else "" + return f"NeMo Anonymizer Benchmark Workspace{suffix}"[:128] + + +__all__ = [ + "benchmark_run_filters", + "benchmark_run_groupby", + "benchmark_runset_settings", + "benchmark_workspace_sections", + "benchmark_workspace_settings", + "build_benchmark_workspace", + "comparison_workspace_panels", + "create_benchmark_workspace", + "default_workspace_title", + "save_workspace", + "table_workspace_panels", + "workspace_bar_panels", + "workspace_metric_accessor", +] From b66f49373772f6c5017d0dfc8ea754ca1a753c67 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 20:54:39 +0000 Subject: [PATCH 10/17] refactor(measurement): split benchmark inputs and planning Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 80 ++- .../measurement_tools/benchmark_inputs.py | 203 ++++++ .../measurement_tools/benchmark_models.py | 223 ++++++ .../measurement_tools/benchmark_planning.py | 84 +++ .../measurement_tools/benchmark_specs.py | 225 ++++++ tools/measurement/run_benchmarks.py | 679 +++--------------- 6 files changed, 919 insertions(+), 575 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_inputs.py create mode 100644 tools/measurement/measurement_tools/benchmark_models.py create mode 100644 tools/measurement/measurement_tools/benchmark_planning.py create mode 100644 tools/measurement/measurement_tools/benchmark_specs.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 5d66e766..f71b27c0 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -61,7 +61,12 @@ def _load_module(path: Path, name: str) -> ModuleType: "WandbProjectPath", "measurement_tools.wandb_report_contracts", ), - ("run_benchmarks.py", "BenchmarkSpec", "Anonymizer", None), + ( + "run_benchmarks.py", + "BenchmarkSpec", + "Anonymizer", + "measurement_tools.benchmark_models", + ), ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath", None), ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult", None), ], @@ -247,6 +252,79 @@ def test_wandb_report_facade_preserves_workspace_construction_contracts() -> Non sys.modules.pop("compat_workspace_construction_facade", None) +def test_benchmark_facade_preserves_models_specs_inputs_and_planning_contracts() -> None: + inputs = importlib.import_module("measurement_tools.benchmark_inputs") + models = importlib.import_module("measurement_tools.benchmark_models") + planning = importlib.import_module("measurement_tools.benchmark_planning") + specs = importlib.import_module("measurement_tools.benchmark_specs") + runner = _load_module(MEASUREMENT_ROOT / "run_benchmarks.py", "compat_benchmark_b1_facade") + try: + assert models.__all__ == [ + "BenchmarkCase", + "BenchmarkResult", + "BenchmarkSpec", + "CaseRunPaths", + "CaseStatus", + "ConfigSpec", + "DDTraceMode", + "MatrixEntry", + "ReplaceKind", + "ReplaceSpec", + "RESERVED_RUN_TAG_KEYS", + "RewriteSpec", + "WorkloadSpec", + "duplicate_matrix_entries", + "duplicates", + ] + assert specs.__all__ == [ + "active_config_ids", + "build_cases", + "cross_product_matrix", + "input_columns", + "load_spec", + "preflight_config_errors", + "preflight_model_configs", + "preflight_model_providers", + "preflight_model_providers_with_errors", + "preflight_suite", + "preflight_workload", + "preflight_workload_errors", + "prepare_output_dir", + ] + assert inputs.__all__ == [ + "build_anonymizer_config", + "build_input", + "build_replace", + "build_rewrite", + "is_local_input_source", + "materialize_sliced_source", + "present", + "privacy_goal", + "read_local_input_dataframe", + "resolve_config_source", + "resolve_input_source", + "resolve_optional_path", + "resolve_path", + "safe_case_filename", + "slice_bounds", + "workload_has_row_slice", + "write_local_input_dataframe", + ] + assert planning.__all__ == ["dry_run_result", "plan_suite"] + assert runner.BenchmarkSpec is models.BenchmarkSpec + assert runner._CaseRunPaths is models.CaseRunPaths + assert runner.load_spec is specs.load_spec + assert runner.build_cases is specs.build_cases + assert runner.build_input is inputs.build_input + assert runner.build_anonymizer_config is inputs.build_anonymizer_config + assert runner._resolve_input_source is inputs.resolve_input_source + assert runner.dry_run_result is not planning.dry_run_result + assert runner.plan_suite is not planning.plan_suite + assert runner.plan_suite.__module__ == "compat_benchmark_b1_facade" + finally: + sys.modules.pop("compat_benchmark_b1_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/measurement_tools/benchmark_inputs.py b/tools/measurement/measurement_tools/benchmark_inputs.py new file mode 100644 index 00000000..5fb6f237 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_inputs.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark input resolution, slicing, and Anonymizer configuration.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pandas as pd + +from anonymizer.config.anonymizer_config import ( + AnonymizerConfig, + AnonymizerInput, + Detect, + Rewrite, + infer_input_source_suffix, +) +from anonymizer.config.replace_strategies import Annotate, Hash, Redact, Substitute +from anonymizer.config.rewrite import DEFAULT_PRESERVE_TEXT, DEFAULT_PROTECT_TEXT, PrivacyGoal +from anonymizer.engine.io.constants import SUPPORTED_IO_FORMATS +from measurement_tools.benchmark_models import ConfigSpec, ReplaceKind, ReplaceSpec, RewriteSpec, WorkloadSpec + + +def build_input( + workload: WorkloadSpec, + base_dir: Path, + *, + slice_dir: Path | None = None, + case_id: str | None = None, +) -> AnonymizerInput: + resolved_source = resolve_input_source(workload.source, base_dir) + source = ( + materialize_sliced_source(workload, resolved_source, slice_dir=slice_dir, case_id=case_id) + if workload_has_row_slice(workload) + else resolved_source + ) + return AnonymizerInput( + source=str(source), + text_column=workload.text_column, + id_column=workload.id_column, + data_summary=workload.data_summary, + ) + + +def workload_has_row_slice(workload: WorkloadSpec) -> bool: + return workload.row_limit is not None or workload.row_offset > 0 + + +def is_local_input_source(source: str) -> bool: + return "://" not in source + + +def materialize_sliced_source( + workload: WorkloadSpec, + source: str | Path, + *, + slice_dir: Path | None, + case_id: str | None, +) -> Path: + if not is_local_input_source(str(source)): + raise ValueError(f"workload '{workload.id}' row slicing requires a local workload source") + if slice_dir is None or case_id is None: + raise ValueError("row slicing requires slice_dir and case_id") + source_path = Path(source) + suffix = infer_input_source_suffix(str(source_path)) + dataframe = read_local_input_dataframe(source_path, suffix=suffix) + sliced = dataframe.iloc[slice_bounds(workload)] + slice_dir.mkdir(parents=True, exist_ok=True) + destination = slice_dir / f"{safe_case_filename(case_id)}{suffix}" + write_local_input_dataframe(sliced, destination, suffix=suffix) + return destination + + +def slice_bounds(workload: WorkloadSpec) -> slice: + start = workload.row_offset + stop = start + workload.row_limit if workload.row_limit is not None else None + return slice(start, stop) + + +def read_local_input_dataframe(source: Path, *, suffix: str) -> pd.DataFrame: + if suffix == ".csv": + return pd.read_csv(source) + if suffix == ".parquet": + return pd.read_parquet(source) + supported_formats = " or ".join(SUPPORTED_IO_FORMATS) + raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") + + +def write_local_input_dataframe(dataframe: pd.DataFrame, destination: Path, *, suffix: str) -> None: + if suffix == ".csv": + dataframe.to_csv(destination, index=False) + return + if suffix == ".parquet": + dataframe.to_parquet(destination, index=False) + return + supported_formats = " or ".join(SUPPORTED_IO_FORMATS) + raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") + + +def safe_case_filename(case_id: str) -> str: + return "".join(char if char.isalnum() or char in "._-" else "_" for char in case_id) + + +def build_anonymizer_config(config: ConfigSpec) -> AnonymizerConfig: + detect = Detect.model_validate(config.detect) + if config.replace is not None: + return AnonymizerConfig( + detect=detect, replace=build_replace(config.replace), emit_telemetry=config.emit_telemetry + ) + return AnonymizerConfig(detect=detect, rewrite=build_rewrite(config.rewrite), emit_telemetry=config.emit_telemetry) + + +def build_replace(raw: str | ReplaceSpec) -> Redact | Hash | Annotate | Substitute: + spec = ReplaceSpec(strategy=ReplaceKind(raw)) if isinstance(raw, str) else raw + if spec.strategy == ReplaceKind.redact: + return Redact(**present({"format_template": spec.format_template, "normalize_label": spec.normalize_label})) + if spec.strategy == ReplaceKind.hash: + return Hash( + **present( + { + "format_template": spec.format_template, + "algorithm": spec.algorithm, + "digest_length": spec.digest_length, + } + ) + ) + if spec.strategy == ReplaceKind.annotate: + return Annotate(**present({"format_template": spec.format_template})) + return Substitute(**present({"instructions": spec.instructions})) + + +def build_rewrite(spec: RewriteSpec | None) -> Rewrite: + if spec is None: + raise ValueError("rewrite config is missing") + privacy_goal_value = privacy_goal(spec) + return Rewrite( + privacy_goal=privacy_goal_value, + instructions=spec.instructions, + risk_tolerance=spec.risk_tolerance, + max_repair_iterations=spec.max_repair_iterations, + strict_entity_protection=spec.strict_entity_protection, + ) + + +def privacy_goal(spec: RewriteSpec) -> PrivacyGoal | None: + if spec.protect is None and spec.preserve is None: + return None + return PrivacyGoal( + protect=spec.protect or DEFAULT_PROTECT_TEXT, + preserve=spec.preserve or DEFAULT_PRESERVE_TEXT, + ) + + +def present(values: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in values.items() if value is not None} + + +def resolve_input_source(source: str, base_dir: Path) -> str | Path: + if "://" in source: + return source + return resolve_path(source, base_dir) + + +def resolve_optional_path(raw: str | None, base_dir: Path) -> Path | None: + if raw is None: + return None + return resolve_path(raw, base_dir) + + +def resolve_config_source(raw: str | None, base_dir: Path) -> str | None: + if raw is None or "\n" in raw: + return raw + candidate = Path(raw).expanduser() + if candidate.suffix in {".yaml", ".yml"}: + return str(resolve_path(raw, base_dir)) + return raw + + +def resolve_path(raw: str, base_dir: Path) -> Path: + path = Path(raw).expanduser() + return path if path.is_absolute() else base_dir / path + + +__all__ = [ + "build_anonymizer_config", + "build_input", + "build_replace", + "build_rewrite", + "is_local_input_source", + "materialize_sliced_source", + "present", + "privacy_goal", + "read_local_input_dataframe", + "resolve_config_source", + "resolve_input_source", + "resolve_optional_path", + "resolve_path", + "safe_case_filename", + "slice_bounds", + "workload_has_row_slice", + "write_local_input_dataframe", +] diff --git a/tools/measurement/measurement_tools/benchmark_models.py b/tools/measurement/measurement_tools/benchmark_models.py new file mode 100644 index 00000000..8f17a0e4 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_models.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark suite models and their validation policy.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from anonymizer.config.rewrite import RiskTolerance + +RESERVED_RUN_TAG_KEYS = frozenset({"suite_id", "workload_id", "config_id", "repetition", "case_id"}) + + +class CaseStatus(StrEnum): + planned = "planned" + completed = "completed" + error = "error" + + +class DDTraceMode(StrEnum): + none = "none" + last_message = "last_message" + all_messages = "all_messages" + + +class ReplaceKind(StrEnum): + redact = "redact" + hash = "hash" + annotate = "annotate" + substitute = "substitute" + + +class WorkloadSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + source: str + text_column: str = "text" + id_column: str | None = None + data_summary: str | None = None + row_limit: int | None = Field(default=None, ge=1) + row_offset: int = Field(default=0, ge=0) + + +class ReplaceSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + strategy: ReplaceKind + format_template: str | None = None + normalize_label: bool | None = None + algorithm: str | None = None + digest_length: int | None = None + instructions: str | None = None + + +class RewriteSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + protect: str | None = None + preserve: str | None = None + instructions: str | None = None + risk_tolerance: RiskTolerance = RiskTolerance.low + max_repair_iterations: int = 3 + strict_entity_protection: bool = False + + +class ConfigSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + id: str + detect: dict[str, Any] = Field(default_factory=dict) + replace: str | ReplaceSpec | None = None + rewrite: RewriteSpec | None = None + evaluate: bool = False + emit_telemetry: bool = False + + @model_validator(mode="after") + def validate_mode(self) -> "ConfigSpec": + if self.replace is None and self.rewrite is None: + raise ValueError("config must define replace or rewrite") + if self.replace is not None and self.rewrite is not None: + raise ValueError("config cannot define both replace and rewrite") + if self.evaluate and self.rewrite is not None: + raise ValueError("evaluate is only supported for replace configs") + return self + + +class MatrixEntry(BaseModel): + model_config = ConfigDict(extra="forbid") + + workload: str + config: str + repetitions: int = Field(default=1, ge=1) + + +def duplicates(values: list[str]) -> list[str]: + seen: set[str] = set() + duplicates: set[str] = set() + for value in values: + if value in seen: + duplicates.add(value) + seen.add(value) + return sorted(duplicates) + + +def duplicate_matrix_entries(matrix: list[MatrixEntry]) -> list[tuple[str, str]]: + seen: set[tuple[str, str]] = set() + duplicates: set[tuple[str, str]] = set() + for entry in matrix: + key = (entry.workload, entry.config) + if key in seen: + duplicates.add(key) + seen.add(key) + return sorted(duplicates) + + +class BenchmarkSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + suite_id: str + model_configs: str | None = None + model_providers: str | None = None + artifact_path: str | None = None + run_tags: dict[str, Any] = Field(default_factory=dict) + case_retries: int = Field(default=0, ge=0) + case_retry_backoff_sec: float = Field(default=0.0, ge=0.0) + workloads: list[WorkloadSpec] = Field(min_length=1) + configs: list[ConfigSpec] = Field(min_length=1) + matrix: list[MatrixEntry] | None = Field(default=None, min_length=1) + + @model_validator(mode="after") + def validate_ids(self) -> "BenchmarkSpec": + workload_ids = [workload.id for workload in self.workloads] + config_ids = [config.id for config in self.configs] + if duplicate_workloads := duplicates(workload_ids): + raise ValueError(f"duplicate workload id(s): {', '.join(duplicate_workloads)}") + if duplicate_configs := duplicates(config_ids): + raise ValueError(f"duplicate config id(s): {', '.join(duplicate_configs)}") + self._validate_matrix_references(set(workload_ids), set(config_ids)) + self._validate_run_tags() + return self + + def _validate_matrix_references(self, workload_ids: set[str], config_ids: set[str]) -> None: + if self.matrix is None: + return + missing_workloads = sorted({entry.workload for entry in self.matrix} - workload_ids) + missing_configs = sorted({entry.config for entry in self.matrix} - config_ids) + if missing_workloads: + raise ValueError(f"matrix references unknown workload id(s): {', '.join(missing_workloads)}") + if missing_configs: + raise ValueError(f"matrix references unknown config id(s): {', '.join(missing_configs)}") + duplicate_entries = duplicate_matrix_entries(self.matrix) + if duplicate_entries: + formatted = ", ".join(f"{workload}/{config}" for workload, config in duplicate_entries) + raise ValueError(f"duplicate matrix workload/config entry(s): {formatted}; use repetitions for repeats") + + def _validate_run_tags(self) -> None: + reserved_tags = sorted(set(self.run_tags) & RESERVED_RUN_TAG_KEYS) + if reserved_tags: + formatted = ", ".join(reserved_tags) + raise ValueError(f"run_tags cannot define reserved benchmark tag(s): {formatted}") + + +class BenchmarkCase(BaseModel): + suite_id: str + workload_id: str + config_id: str + repetition: int + case_id: str + status: CaseStatus = CaseStatus.planned + elapsed_sec: float | None = None + measurement_path: str | None = None + detection_artifact_path: str | None = None + trace_path: str | None = None + task_trace_path: str | None = None + error: str | None = None + attempt_count: int = 0 + attempt_errors: list[str] = Field(default_factory=list) + + +class BenchmarkResult(BaseModel): + suite_id: str + output_dir: str + measurement_path: str + summary_path: str + table_dir: str | None + detection_artifact_analysis_path: str | None = None + cases: list[BenchmarkCase] + execution: dict[str, Any] = Field(default_factory=dict) + + +@dataclass(frozen=True) +class CaseRunPaths: + raw_path: Path + artifact_output_path: Path + trace_path: Path | None + task_trace_path: Path | None + artifact_snapshot: dict[str, int] | None + export_detection_artifacts: bool + + +__all__ = [ + "BenchmarkCase", + "BenchmarkResult", + "BenchmarkSpec", + "CaseRunPaths", + "CaseStatus", + "ConfigSpec", + "DDTraceMode", + "MatrixEntry", + "ReplaceKind", + "ReplaceSpec", + "RESERVED_RUN_TAG_KEYS", + "RewriteSpec", + "WorkloadSpec", + "duplicate_matrix_entries", + "duplicates", +] diff --git a/tools/measurement/measurement_tools/benchmark_planning.py b/tools/measurement/measurement_tools/benchmark_planning.py new file mode 100644 index 00000000..df44fccf --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_planning.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark dry-run result construction and suite planning.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from measurement_tools.benchmark_models import BenchmarkResult, BenchmarkSpec, DDTraceMode +from measurement_tools.benchmark_specs import build_cases, preflight_suite + + +def dry_run_result( + spec: BenchmarkSpec, + *, + output_dir: Path, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode, + trace_dir: Path | None, + dd_task_trace: bool = False, + task_trace_dir: Path | None = None, + execution_metadata: Callable[..., dict[str, Any]], +) -> BenchmarkResult: + cases = build_cases(spec) + if dd_trace != DDTraceMode.none: + resolved_trace_dir = trace_dir or output_dir / "traces" + cases = [ + case.model_copy(update={"trace_path": str(resolved_trace_dir / f"{case.case_id}.jsonl")}) for case in cases + ] + if dd_task_trace: + resolved_task_trace_dir = task_trace_dir or output_dir / "task-traces" + cases = [ + case.model_copy(update={"task_trace_path": str(resolved_task_trace_dir / f"{case.case_id}.jsonl")}) + for case in cases + ] + return BenchmarkResult( + suite_id=spec.suite_id, + output_dir=str(output_dir), + measurement_path=str(output_dir / "measurements.jsonl"), + summary_path=str(output_dir / "summary.json"), + table_dir=str(output_dir / "tables") if export else None, + detection_artifact_analysis_path=str(output_dir / "detection-artifacts.jsonl") if export else None, + cases=cases, + execution=execution_metadata( + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + dd_task_trace=dd_task_trace, + ), + ) + + +def plan_suite( + spec: BenchmarkSpec, + *, + spec_path: Path, + output_dir: Path, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode = DDTraceMode.none, + trace_dir: Path | None = None, + dd_task_trace: bool = False, + task_trace_dir: Path | None = None, + execution_metadata: Callable[..., dict[str, Any]], +) -> BenchmarkResult: + preflight_suite(spec, spec_path=spec_path) + return dry_run_result( + spec, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + trace_dir=trace_dir, + dd_task_trace=dd_task_trace, + task_trace_dir=task_trace_dir, + execution_metadata=execution_metadata, + ) + + +__all__ = ["dry_run_result", "plan_suite"] diff --git a/tools/measurement/measurement_tools/benchmark_specs.py b/tools/measurement/measurement_tools/benchmark_specs.py new file mode 100644 index 00000000..15654dba --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_specs.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark suite loading, expansion, output preparation, and preflight.""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Any + +import pandas as pd +import pyarrow.parquet as pq +import yaml +from data_designer.config.models import ModelProvider +from data_designer.config.utils.io_helpers import load_config_file + +from anonymizer.config.anonymizer_config import AnonymizerInput, infer_input_source_suffix, is_remote_input_source +from anonymizer.config.replace_strategies import Substitute +from anonymizer.engine.io.constants import SUPPORTED_IO_FORMATS +from anonymizer.engine.ndd.model_loader import parse_model_configs, validate_model_alias_references +from measurement_tools.benchmark_inputs import ( + build_anonymizer_config, + is_local_input_source, + resolve_config_source, + resolve_input_source, + workload_has_row_slice, +) +from measurement_tools.benchmark_models import BenchmarkCase, BenchmarkSpec, MatrixEntry, WorkloadSpec + + +def load_spec(path: Path) -> BenchmarkSpec: + if not path.exists() or path.is_dir(): + raise ValueError(f"spec path is not a file: {path}") + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError("benchmark spec must be a YAML mapping") + return BenchmarkSpec.model_validate(raw) + + +def build_cases(spec: BenchmarkSpec) -> list[BenchmarkCase]: + matrix = spec.matrix or cross_product_matrix(spec) + return [ + BenchmarkCase( + suite_id=spec.suite_id, + workload_id=entry.workload, + config_id=entry.config, + repetition=repetition, + case_id=f"{entry.workload}__{entry.config}__r{repetition:03d}", + ) + for entry in matrix + for repetition in range(entry.repetitions) + ] + + +def cross_product_matrix(spec: BenchmarkSpec) -> list[MatrixEntry]: + return [ + MatrixEntry(workload=workload.id, config=config.id, repetitions=1) + for workload in spec.workloads + for config in spec.configs + ] + + +def prepare_output_dir(output_dir: Path, *, overwrite: bool, dry_run: bool) -> None: + if dry_run: + return + if output_dir.exists() and not output_dir.is_dir(): + raise ValueError(f"output path exists and is not a directory: {output_dir}") + if output_dir.exists(): + if overwrite: + shutil.rmtree(output_dir) + elif any(output_dir.iterdir()): + raise ValueError(f"output directory is not empty: {output_dir}; pass --overwrite to replace it") + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "raw").mkdir(exist_ok=True) + + +def preflight_suite(spec: BenchmarkSpec, *, spec_path: Path) -> None: + """Validate cheap suite inputs before any benchmark case consumes model time.""" + base_dir = spec_path.parent + errors: list[str] = [] + parsed_models = preflight_model_configs(spec, base_dir=base_dir, errors=errors) + + preflight_model_providers_with_errors(spec, base_dir=base_dir, errors=errors) + errors.extend(preflight_workload_errors(spec, base_dir=base_dir)) + errors.extend(preflight_config_errors(spec, parsed_models=parsed_models)) + if errors: + raise ValueError("Benchmark preflight failed:\n- " + "\n- ".join(errors)) + + +def preflight_model_configs(spec: BenchmarkSpec, *, base_dir: Path, errors: list[str]) -> Any | None: + try: + return parse_model_configs(resolve_config_source(spec.model_configs, base_dir)) + except Exception as exc: + errors.append(f"model_configs invalid: {exc}") + return None + + +def preflight_model_providers_with_errors( + spec: BenchmarkSpec, + *, + base_dir: Path, + errors: list[str], +) -> None: + try: + preflight_model_providers(spec, base_dir=base_dir) + except Exception as exc: + errors.append(f"model_providers invalid: {exc}") + + +def preflight_workload_errors(spec: BenchmarkSpec, *, base_dir: Path) -> list[str]: + errors: list[str] = [] + for workload in spec.workloads: + try: + preflight_workload(workload, base_dir=base_dir) + except Exception as exc: + errors.append(str(exc)) + return errors + + +def preflight_config_errors(spec: BenchmarkSpec, *, parsed_models: Any | None) -> list[str]: + errors: list[str] = [] + active_ids = active_config_ids(spec) + for config in spec.configs: + if config.id not in active_ids: + continue + try: + anonymizer_config = build_anonymizer_config(config) + except Exception as exc: + errors.append(f"config '{config.id}' invalid: {exc}") + continue + if parsed_models is None: + continue + try: + validate_model_alias_references( + parsed_models.model_configs, + parsed_models.selected_models, + check_substitute=isinstance(anonymizer_config.replace, Substitute) + or anonymizer_config.rewrite is not None, + check_rewrite=anonymizer_config.rewrite is not None, + check_evaluate=config.evaluate, + ) + except ValueError as exc: + errors.append(f"config '{config.id}' model aliases invalid: {exc}") + return errors + + +def active_config_ids(spec: BenchmarkSpec) -> set[str]: + if spec.matrix is None: + return {config.id for config in spec.configs} + return {entry.config for entry in spec.matrix} + + +def preflight_model_providers(spec: BenchmarkSpec, *, base_dir: Path) -> None: + raw = resolve_config_source(spec.model_providers, base_dir) + if raw is None: + return + if "\n" in raw: + config_dict = yaml.safe_load(raw) + else: + candidate = Path(raw.strip()).expanduser() + if candidate.suffix in (".yaml", ".yml"): + if not candidate.is_file(): + raise FileNotFoundError(f"Providers config file not found: {candidate}") + config_dict = load_config_file(candidate) + else: + config_dict = yaml.safe_load(raw) + raw_providers = config_dict.get("providers") if isinstance(config_dict, dict) else None + if not isinstance(raw_providers, list): + raise ValueError("model_providers YAML must contain a top-level 'providers' list.") + for provider in raw_providers: + ModelProvider.model_validate(provider) + + +def preflight_workload(workload: WorkloadSpec, *, base_dir: Path) -> None: + resolved_source = resolve_input_source(workload.source, base_dir) + if workload_has_row_slice(workload) and not is_local_input_source(str(resolved_source)): + raise ValueError(f"workload '{workload.id}' row slicing requires a local workload source") + input_data = AnonymizerInput( + source=str(resolved_source), + text_column=workload.text_column, + id_column=workload.id_column, + data_summary=workload.data_summary, + ) + columns = input_columns(input_data.source) + if columns is None: + return + if workload.text_column not in columns: + raise ValueError( + f"workload '{workload.id}' text_column '{workload.text_column}' not found in {input_data.source}; " + f"available columns: {sorted(columns)}" + ) + if workload.id_column is not None and workload.id_column not in columns: + raise ValueError( + f"workload '{workload.id}' id_column '{workload.id_column}' not found in {input_data.source}; " + f"available columns: {sorted(columns)}" + ) + + +def input_columns(source: str) -> set[str] | None: + suffix = infer_input_source_suffix(source) + if suffix not in SUPPORTED_IO_FORMATS: + supported_formats = " or ".join(SUPPORTED_IO_FORMATS) + raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") + if is_remote_input_source(source): + return None + if suffix == ".csv": + return set(pd.read_csv(source, nrows=0).columns) + return set(pq.ParquetFile(source).schema_arrow.names) + + +__all__ = [ + "active_config_ids", + "build_cases", + "cross_product_matrix", + "input_columns", + "load_spec", + "preflight_config_errors", + "preflight_model_configs", + "preflight_model_providers", + "preflight_model_providers_with_errors", + "preflight_suite", + "preflight_workload", + "preflight_workload_errors", + "prepare_output_dir", +] diff --git a/tools/measurement/run_benchmarks.py b/tools/measurement/run_benchmarks.py index 89b5408c..2d920769 100755 --- a/tools/measurement/run_benchmarks.py +++ b/tools/measurement/run_benchmarks.py @@ -10,27 +10,83 @@ import logging import platform -import shutil import subprocess import sys import time -from dataclasses import dataclass -from enum import StrEnum from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Annotated, Any import cyclopts import pandas as pd -import pyarrow.parquet as pq -import yaml from analyze_detection_artifacts import ( analyze_artifacts, iter_detection_parquet_files, ) -from data_designer.config.models import ModelProvider -from data_designer.config.utils.io_helpers import load_config_file from export_measurements import export_tables, read_measurements +from measurement_tools.benchmark_inputs import ( + build_anonymizer_config, + build_input, + is_local_input_source, + materialize_sliced_source, + present, + privacy_goal, + read_local_input_dataframe, + resolve_config_source, + resolve_input_source, + resolve_optional_path, + resolve_path, + safe_case_filename, + slice_bounds, + workload_has_row_slice, + write_local_input_dataframe, +) +from measurement_tools.benchmark_inputs import ( + build_replace as build_replace, +) +from measurement_tools.benchmark_inputs import ( + build_rewrite as build_rewrite, +) +from measurement_tools.benchmark_models import ( + RESERVED_RUN_TAG_KEYS as RESERVED_RUN_TAG_KEYS, +) +from measurement_tools.benchmark_models import ( + BenchmarkCase, + BenchmarkResult, + BenchmarkSpec, + CaseRunPaths, + CaseStatus, + ConfigSpec, + DDTraceMode, + ReplaceSpec, + RewriteSpec, + WorkloadSpec, + duplicate_matrix_entries, + duplicates, +) +from measurement_tools.benchmark_models import ( + MatrixEntry as MatrixEntry, +) +from measurement_tools.benchmark_models import ( + ReplaceKind as ReplaceKind, +) +from measurement_tools.benchmark_planning import dry_run_result as _canonical_dry_run_result +from measurement_tools.benchmark_planning import plan_suite as _canonical_plan_suite +from measurement_tools.benchmark_specs import ( + active_config_ids, + build_cases, + cross_product_matrix, + input_columns, + load_spec, + preflight_config_errors, + preflight_model_configs, + preflight_model_providers, + preflight_model_providers_with_errors, + preflight_suite, + preflight_workload, + preflight_workload_errors, + prepare_output_dir, +) from measurement_tools.cli import LogFormat, configure_logging, log_bad_input, summarize_validation_error from measurement_tools.execution import benchmark_execution_metadata, stable_metadata_hash from measurement_tools.tables import ExportFormat @@ -42,20 +98,15 @@ WandbRunMetadata, publish_benchmark_wandb_best_effort, ) -from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator +from pydantic import ValidationError +from anonymizer.config.anonymizer_config import Rewrite as Rewrite from anonymizer.config.anonymizer_config import ( - AnonymizerConfig, - AnonymizerInput, - Detect, - Rewrite, infer_input_source_suffix, is_remote_input_source, ) -from anonymizer.config.replace_strategies import Annotate, Hash, Redact, Substitute -from anonymizer.config.rewrite import DEFAULT_PRESERVE_TEXT, DEFAULT_PROTECT_TEXT, PrivacyGoal, RiskTolerance -from anonymizer.engine.io.constants import SUPPORTED_IO_FORMATS -from anonymizer.engine.ndd.model_loader import parse_model_configs, validate_model_alias_references +from anonymizer.config.replace_strategies import Redact as Redact +from anonymizer.config.rewrite import RiskTolerance as RiskTolerance from anonymizer.interface.anonymizer import Anonymizer from anonymizer.measurement import ( MEASUREMENT_SCHEMA_VERSION, @@ -67,378 +118,34 @@ app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.benchmark") - -class CaseStatus(StrEnum): - planned = "planned" - completed = "completed" - error = "error" - - -class DDTraceMode(StrEnum): - none = "none" - last_message = "last_message" - all_messages = "all_messages" - - -class ReplaceKind(StrEnum): - redact = "redact" - hash = "hash" - annotate = "annotate" - substitute = "substitute" - - -class WorkloadSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - id: str - source: str - text_column: str = "text" - id_column: str | None = None - data_summary: str | None = None - row_limit: int | None = Field(default=None, ge=1) - row_offset: int = Field(default=0, ge=0) - - -class ReplaceSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - strategy: ReplaceKind - format_template: str | None = None - normalize_label: bool | None = None - algorithm: str | None = None - digest_length: int | None = None - instructions: str | None = None - - -class RewriteSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - protect: str | None = None - preserve: str | None = None - instructions: str | None = None - risk_tolerance: RiskTolerance = RiskTolerance.low - max_repair_iterations: int = 3 - strict_entity_protection: bool = False - - -class ConfigSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - id: str - detect: dict[str, Any] = Field(default_factory=dict) - replace: str | ReplaceSpec | None = None - rewrite: RewriteSpec | None = None - evaluate: bool = False - emit_telemetry: bool = False - - @model_validator(mode="after") - def validate_mode(self) -> "ConfigSpec": - if self.replace is None and self.rewrite is None: - raise ValueError("config must define replace or rewrite") - if self.replace is not None and self.rewrite is not None: - raise ValueError("config cannot define both replace and rewrite") - if self.evaluate and self.rewrite is not None: - raise ValueError("evaluate is only supported for replace configs") - return self - - -class MatrixEntry(BaseModel): - model_config = ConfigDict(extra="forbid") - - workload: str - config: str - repetitions: int = Field(default=1, ge=1) - - -RESERVED_RUN_TAG_KEYS = frozenset({"suite_id", "workload_id", "config_id", "repetition", "case_id"}) BENCHMARK_SUITE_SCHEMA_VERSION = 1 WANDB_METADATA_SCHEMA_VERSION = 2 - -def _duplicates(values: list[str]) -> list[str]: - seen: set[str] = set() - duplicates: set[str] = set() - for value in values: - if value in seen: - duplicates.add(value) - seen.add(value) - return sorted(duplicates) - - -class BenchmarkSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - suite_id: str - model_configs: str | None = None - model_providers: str | None = None - artifact_path: str | None = None - run_tags: dict[str, Any] = Field(default_factory=dict) - case_retries: int = Field(default=0, ge=0) - case_retry_backoff_sec: float = Field(default=0.0, ge=0.0) - workloads: list[WorkloadSpec] = Field(min_length=1) - configs: list[ConfigSpec] = Field(min_length=1) - matrix: list[MatrixEntry] | None = Field(default=None, min_length=1) - - @model_validator(mode="after") - def validate_ids(self) -> "BenchmarkSpec": - workload_ids = [workload.id for workload in self.workloads] - config_ids = [config.id for config in self.configs] - if duplicate_workloads := _duplicates(workload_ids): - raise ValueError(f"duplicate workload id(s): {', '.join(duplicate_workloads)}") - if duplicate_configs := _duplicates(config_ids): - raise ValueError(f"duplicate config id(s): {', '.join(duplicate_configs)}") - self._validate_matrix_references(set(workload_ids), set(config_ids)) - self._validate_run_tags() - return self - - def _validate_matrix_references(self, workload_ids: set[str], config_ids: set[str]) -> None: - if self.matrix is None: - return - missing_workloads = sorted({entry.workload for entry in self.matrix} - workload_ids) - missing_configs = sorted({entry.config for entry in self.matrix} - config_ids) - if missing_workloads: - raise ValueError(f"matrix references unknown workload id(s): {', '.join(missing_workloads)}") - if missing_configs: - raise ValueError(f"matrix references unknown config id(s): {', '.join(missing_configs)}") - duplicate_entries = _duplicate_matrix_entries(self.matrix) - if duplicate_entries: - formatted = ", ".join(f"{workload}/{config}" for workload, config in duplicate_entries) - raise ValueError(f"duplicate matrix workload/config entry(s): {formatted}; use repetitions for repeats") - - def _validate_run_tags(self) -> None: - reserved_tags = sorted(set(self.run_tags) & RESERVED_RUN_TAG_KEYS) - if reserved_tags: - formatted = ", ".join(reserved_tags) - raise ValueError(f"run_tags cannot define reserved benchmark tag(s): {formatted}") - - -class BenchmarkCase(BaseModel): - suite_id: str - workload_id: str - config_id: str - repetition: int - case_id: str - status: CaseStatus = CaseStatus.planned - elapsed_sec: float | None = None - measurement_path: str | None = None - detection_artifact_path: str | None = None - trace_path: str | None = None - task_trace_path: str | None = None - error: str | None = None - attempt_count: int = 0 - attempt_errors: list[str] = Field(default_factory=list) - - -class BenchmarkResult(BaseModel): - suite_id: str - output_dir: str - measurement_path: str - summary_path: str - table_dir: str | None - detection_artifact_analysis_path: str | None = None - cases: list[BenchmarkCase] - execution: dict[str, Any] = Field(default_factory=dict) - - -@dataclass(frozen=True) -class _CaseRunPaths: - raw_path: Path - artifact_output_path: Path - trace_path: Path | None - task_trace_path: Path | None - artifact_snapshot: dict[str, int] | None - export_detection_artifacts: bool - - -def load_spec(path: Path) -> BenchmarkSpec: - if not path.exists() or path.is_dir(): - raise ValueError(f"spec path is not a file: {path}") - raw = yaml.safe_load(path.read_text(encoding="utf-8")) - if not isinstance(raw, dict): - raise ValueError("benchmark spec must be a YAML mapping") - return BenchmarkSpec.model_validate(raw) - - -def build_cases(spec: BenchmarkSpec) -> list[BenchmarkCase]: - matrix = spec.matrix or _cross_product_matrix(spec) - return [ - BenchmarkCase( - suite_id=spec.suite_id, - workload_id=entry.workload, - config_id=entry.config, - repetition=repetition, - case_id=f"{entry.workload}__{entry.config}__r{repetition:03d}", - ) - for entry in matrix - for repetition in range(entry.repetitions) - ] - - -def _cross_product_matrix(spec: BenchmarkSpec) -> list[MatrixEntry]: - return [ - MatrixEntry(workload=workload.id, config=config.id, repetitions=1) - for workload in spec.workloads - for config in spec.configs - ] - - -def _duplicate_matrix_entries(matrix: list[MatrixEntry]) -> list[tuple[str, str]]: - seen: set[tuple[str, str]] = set() - duplicates: set[tuple[str, str]] = set() - for entry in matrix: - key = (entry.workload, entry.config) - if key in seen: - duplicates.add(key) - seen.add(key) - return sorted(duplicates) - - -def prepare_output_dir(output_dir: Path, *, overwrite: bool, dry_run: bool) -> None: - if dry_run: - return - if output_dir.exists() and not output_dir.is_dir(): - raise ValueError(f"output path exists and is not a directory: {output_dir}") - if output_dir.exists(): - if overwrite: - shutil.rmtree(output_dir) - elif any(output_dir.iterdir()): - raise ValueError(f"output directory is not empty: {output_dir}; pass --overwrite to replace it") - output_dir.mkdir(parents=True, exist_ok=True) - (output_dir / "raw").mkdir(exist_ok=True) - - -def preflight_suite(spec: BenchmarkSpec, *, spec_path: Path) -> None: - """Validate cheap suite inputs before any benchmark case consumes model time.""" - base_dir = spec_path.parent - errors: list[str] = [] - parsed_models = _preflight_model_configs(spec, base_dir=base_dir, errors=errors) - - _preflight_model_providers_with_errors(spec, base_dir=base_dir, errors=errors) - errors.extend(_preflight_workload_errors(spec, base_dir=base_dir)) - errors.extend(_preflight_config_errors(spec, parsed_models=parsed_models)) - if errors: - raise ValueError("Benchmark preflight failed:\n- " + "\n- ".join(errors)) - - -def _preflight_model_configs(spec: BenchmarkSpec, *, base_dir: Path, errors: list[str]) -> Any | None: - try: - return parse_model_configs(_resolve_config_source(spec.model_configs, base_dir)) - except Exception as exc: - errors.append(f"model_configs invalid: {exc}") - return None - - -def _preflight_model_providers_with_errors( - spec: BenchmarkSpec, - *, - base_dir: Path, - errors: list[str], -) -> None: - try: - _preflight_model_providers(spec, base_dir=base_dir) - except Exception as exc: - errors.append(f"model_providers invalid: {exc}") - - -def _preflight_workload_errors(spec: BenchmarkSpec, *, base_dir: Path) -> list[str]: - errors: list[str] = [] - for workload in spec.workloads: - try: - _preflight_workload(workload, base_dir=base_dir) - except Exception as exc: - errors.append(str(exc)) - return errors - - -def _preflight_config_errors(spec: BenchmarkSpec, *, parsed_models: Any | None) -> list[str]: - errors: list[str] = [] - active_config_ids = _active_config_ids(spec) - for config in spec.configs: - if config.id not in active_config_ids: - continue - try: - anonymizer_config = build_anonymizer_config(config) - except Exception as exc: - errors.append(f"config '{config.id}' invalid: {exc}") - continue - if parsed_models is None: - continue - try: - validate_model_alias_references( - parsed_models.model_configs, - parsed_models.selected_models, - check_substitute=isinstance(anonymizer_config.replace, Substitute) - or anonymizer_config.rewrite is not None, - check_rewrite=anonymizer_config.rewrite is not None, - check_evaluate=config.evaluate, - ) - except ValueError as exc: - errors.append(f"config '{config.id}' model aliases invalid: {exc}") - return errors - - -def _active_config_ids(spec: BenchmarkSpec) -> set[str]: - if spec.matrix is None: - return {config.id for config in spec.configs} - return {entry.config for entry in spec.matrix} - - -def _preflight_model_providers(spec: BenchmarkSpec, *, base_dir: Path) -> None: - raw = _resolve_config_source(spec.model_providers, base_dir) - if raw is None: - return - if "\n" in raw: - config_dict = yaml.safe_load(raw) - else: - candidate = Path(raw.strip()).expanduser() - if candidate.suffix in (".yaml", ".yml"): - if not candidate.is_file(): - raise FileNotFoundError(f"Providers config file not found: {candidate}") - config_dict = load_config_file(candidate) - else: - config_dict = yaml.safe_load(raw) - raw_providers = config_dict.get("providers") if isinstance(config_dict, dict) else None - if not isinstance(raw_providers, list): - raise ValueError("model_providers YAML must contain a top-level 'providers' list.") - for provider in raw_providers: - ModelProvider.model_validate(provider) - - -def _preflight_workload(workload: WorkloadSpec, *, base_dir: Path) -> None: - resolved_source = _resolve_input_source(workload.source, base_dir) - if _workload_has_row_slice(workload) and not _is_local_input_source(str(resolved_source)): - raise ValueError(f"workload '{workload.id}' row slicing requires a local workload source") - input_data = AnonymizerInput( - source=str(resolved_source), - text_column=workload.text_column, - id_column=workload.id_column, - data_summary=workload.data_summary, - ) - columns = _input_columns(input_data.source) - if columns is None: - return - if workload.text_column not in columns: - raise ValueError( - f"workload '{workload.id}' text_column '{workload.text_column}' not found in {input_data.source}; " - f"available columns: {sorted(columns)}" - ) - if workload.id_column is not None and workload.id_column not in columns: - raise ValueError( - f"workload '{workload.id}' id_column '{workload.id_column}' not found in {input_data.source}; " - f"available columns: {sorted(columns)}" - ) - - -def _input_columns(source: str) -> set[str] | None: - suffix = infer_input_source_suffix(source) - if suffix not in SUPPORTED_IO_FORMATS: - supported_formats = " or ".join(SUPPORTED_IO_FORMATS) - raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") - if is_remote_input_source(source): - return None - if suffix == ".csv": - return set(pd.read_csv(source, nrows=0).columns) - return set(pq.ParquetFile(source).schema_arrow.names) +_CaseRunPaths = CaseRunPaths +_active_config_ids = active_config_ids +_cross_product_matrix = cross_product_matrix +_duplicate_matrix_entries = duplicate_matrix_entries +_duplicates = duplicates +_input_columns = input_columns +_is_local_input_source = is_local_input_source +_materialize_sliced_source = materialize_sliced_source +_preflight_config_errors = preflight_config_errors +_preflight_model_configs = preflight_model_configs +_preflight_model_providers = preflight_model_providers +_preflight_model_providers_with_errors = preflight_model_providers_with_errors +_preflight_workload = preflight_workload +_preflight_workload_errors = preflight_workload_errors +_present = present +_privacy_goal = privacy_goal +_read_local_input_dataframe = read_local_input_dataframe +_resolve_config_source = resolve_config_source +_resolve_input_source = resolve_input_source +_resolve_optional_path = resolve_optional_path +_resolve_path = resolve_path +_safe_case_filename = safe_case_filename +_slice_bounds = slice_bounds +_workload_has_row_slice = workload_has_row_slice +_write_local_input_dataframe = write_local_input_dataframe def run_suite( @@ -856,136 +563,6 @@ def _execute_case( ) -def build_input( - workload: WorkloadSpec, - base_dir: Path, - *, - slice_dir: Path | None = None, - case_id: str | None = None, -) -> AnonymizerInput: - resolved_source = _resolve_input_source(workload.source, base_dir) - source = ( - _materialize_sliced_source(workload, resolved_source, slice_dir=slice_dir, case_id=case_id) - if _workload_has_row_slice(workload) - else resolved_source - ) - return AnonymizerInput( - source=str(source), - text_column=workload.text_column, - id_column=workload.id_column, - data_summary=workload.data_summary, - ) - - -def _workload_has_row_slice(workload: WorkloadSpec) -> bool: - return workload.row_limit is not None or workload.row_offset > 0 - - -def _is_local_input_source(source: str) -> bool: - return "://" not in source - - -def _materialize_sliced_source( - workload: WorkloadSpec, - source: str | Path, - *, - slice_dir: Path | None, - case_id: str | None, -) -> Path: - if not _is_local_input_source(str(source)): - raise ValueError(f"workload '{workload.id}' row slicing requires a local workload source") - if slice_dir is None or case_id is None: - raise ValueError("row slicing requires slice_dir and case_id") - source_path = Path(source) - suffix = infer_input_source_suffix(str(source_path)) - dataframe = _read_local_input_dataframe(source_path, suffix=suffix) - sliced = dataframe.iloc[_slice_bounds(workload)] - slice_dir.mkdir(parents=True, exist_ok=True) - destination = slice_dir / f"{_safe_case_filename(case_id)}{suffix}" - _write_local_input_dataframe(sliced, destination, suffix=suffix) - return destination - - -def _slice_bounds(workload: WorkloadSpec) -> slice: - start = workload.row_offset - stop = start + workload.row_limit if workload.row_limit is not None else None - return slice(start, stop) - - -def _read_local_input_dataframe(source: Path, *, suffix: str) -> pd.DataFrame: - if suffix == ".csv": - return pd.read_csv(source) - if suffix == ".parquet": - return pd.read_parquet(source) - supported_formats = " or ".join(SUPPORTED_IO_FORMATS) - raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") - - -def _write_local_input_dataframe(dataframe: pd.DataFrame, destination: Path, *, suffix: str) -> None: - if suffix == ".csv": - dataframe.to_csv(destination, index=False) - return - if suffix == ".parquet": - dataframe.to_parquet(destination, index=False) - return - supported_formats = " or ".join(SUPPORTED_IO_FORMATS) - raise ValueError(f"Unsupported input format: {suffix}. Use {supported_formats}.") - - -def _safe_case_filename(case_id: str) -> str: - return "".join(char if char.isalnum() or char in "._-" else "_" for char in case_id) - - -def build_anonymizer_config(config: ConfigSpec) -> AnonymizerConfig: - detect = Detect.model_validate(config.detect) - if config.replace is not None: - return AnonymizerConfig( - detect=detect, replace=build_replace(config.replace), emit_telemetry=config.emit_telemetry - ) - return AnonymizerConfig(detect=detect, rewrite=build_rewrite(config.rewrite), emit_telemetry=config.emit_telemetry) - - -def build_replace(raw: str | ReplaceSpec) -> Redact | Hash | Annotate | Substitute: - spec = ReplaceSpec(strategy=ReplaceKind(raw)) if isinstance(raw, str) else raw - if spec.strategy == ReplaceKind.redact: - return Redact(**_present({"format_template": spec.format_template, "normalize_label": spec.normalize_label})) - if spec.strategy == ReplaceKind.hash: - return Hash( - **_present( - { - "format_template": spec.format_template, - "algorithm": spec.algorithm, - "digest_length": spec.digest_length, - } - ) - ) - if spec.strategy == ReplaceKind.annotate: - return Annotate(**_present({"format_template": spec.format_template})) - return Substitute(**_present({"instructions": spec.instructions})) - - -def build_rewrite(spec: RewriteSpec | None) -> Rewrite: - if spec is None: - raise ValueError("rewrite config is missing") - privacy_goal = _privacy_goal(spec) - return Rewrite( - privacy_goal=privacy_goal, - instructions=spec.instructions, - risk_tolerance=spec.risk_tolerance, - max_repair_iterations=spec.max_repair_iterations, - strict_entity_protection=spec.strict_entity_protection, - ) - - -def _privacy_goal(spec: RewriteSpec) -> PrivacyGoal | None: - if spec.protect is None and spec.preserve is None: - return None - return PrivacyGoal( - protect=spec.protect or DEFAULT_PROTECT_TEXT, - preserve=spec.preserve or DEFAULT_PRESERVE_TEXT, - ) - - def combine_measurements(cases: list[BenchmarkCase], destination: Path) -> Path: with destination.open("w", encoding="utf-8") as output: for case in cases: @@ -1126,42 +703,12 @@ def _run_tags(case: BenchmarkCase, spec: BenchmarkSpec) -> dict[str, Any]: } -def _present(values: dict[str, Any]) -> dict[str, Any]: - return {key: value for key, value in values.items() if value is not None} - - def _get_item(items: dict[str, Any], item_id: str, item_type: str) -> Any: if item_id not in items: raise ValueError(f"unknown {item_type}: {item_id}") return items[item_id] -def _resolve_input_source(source: str, base_dir: Path) -> str | Path: - if "://" in source: - return source - return _resolve_path(source, base_dir) - - -def _resolve_optional_path(raw: str | None, base_dir: Path) -> Path | None: - if raw is None: - return None - return _resolve_path(raw, base_dir) - - -def _resolve_config_source(raw: str | None, base_dir: Path) -> str | None: - if raw is None or "\n" in raw: - return raw - candidate = Path(raw).expanduser() - if candidate.suffix in {".yaml", ".yml"}: - return str(_resolve_path(raw, base_dir)) - return raw - - -def _resolve_path(raw: str, base_dir: Path) -> Path: - path = Path(raw).expanduser() - return path if path.is_absolute() else base_dir / path - - def dry_run_result( spec: BenchmarkSpec, *, @@ -1173,33 +720,16 @@ def dry_run_result( dd_task_trace: bool = False, task_trace_dir: Path | None = None, ) -> BenchmarkResult: - cases = build_cases(spec) - if dd_trace != DDTraceMode.none: - resolved_trace_dir = trace_dir or output_dir / "traces" - cases = [ - case.model_copy(update={"trace_path": str(resolved_trace_dir / f"{case.case_id}.jsonl")}) for case in cases - ] - if dd_task_trace: - resolved_task_trace_dir = task_trace_dir or output_dir / "task-traces" - cases = [ - case.model_copy(update={"task_trace_path": str(resolved_task_trace_dir / f"{case.case_id}.jsonl")}) - for case in cases - ] - return BenchmarkResult( - suite_id=spec.suite_id, - output_dir=str(output_dir), - measurement_path=str(output_dir / "measurements.jsonl"), - summary_path=str(output_dir / "summary.json"), - table_dir=str(output_dir / "tables") if export else None, - detection_artifact_analysis_path=str(output_dir / "detection-artifacts.jsonl") if export else None, - cases=cases, - execution=_execution_metadata( - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - dd_trace=dd_trace, - dd_task_trace=dd_task_trace, - ), + return _canonical_dry_run_result( + spec, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + trace_dir=trace_dir, + dd_task_trace=dd_task_trace, + task_trace_dir=task_trace_dir, + execution_metadata=_execution_metadata, ) @@ -1215,9 +745,9 @@ def plan_suite( dd_task_trace: bool = False, task_trace_dir: Path | None = None, ) -> BenchmarkResult: - preflight_suite(spec, spec_path=spec_path) - return dry_run_result( + return _canonical_plan_suite( spec, + spec_path=spec_path, output_dir=output_dir, export=export, fail_fast=fail_fast, @@ -1225,6 +755,7 @@ def plan_suite( trace_dir=trace_dir, dd_task_trace=dd_task_trace, task_trace_dir=task_trace_dir, + execution_metadata=_execution_metadata, ) From a11f7c8e0cd0eb0a509fd739147605661889773b Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 21:02:54 +0000 Subject: [PATCH 11/17] refactor(measurement): split benchmark artifacts and metadata Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 55 +++ .../measurement_tools/benchmark_artifacts.py | 160 +++++++ .../benchmark_wandb_metadata.py | 299 ++++++++++++ tools/measurement/run_benchmarks.py | 445 +++--------------- 4 files changed, 590 insertions(+), 369 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_artifacts.py create mode 100644 tools/measurement/measurement_tools/benchmark_wandb_metadata.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index f71b27c0..ed685147 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -325,6 +325,61 @@ def test_benchmark_facade_preserves_models_specs_inputs_and_planning_contracts() sys.modules.pop("compat_benchmark_b1_facade", None) +def test_benchmark_facade_preserves_artifact_and_metadata_contracts() -> None: + artifacts = importlib.import_module("measurement_tools.benchmark_artifacts") + metadata = importlib.import_module("measurement_tools.benchmark_wandb_metadata") + runner = _load_module(MEASUREMENT_ROOT / "run_benchmarks.py", "compat_benchmark_b2_facade") + try: + assert artifacts.__all__ == [ + "changed_detection_artifact_files", + "combine_detection_artifact_analysis", + "combine_measurements", + "export_case_detection_artifact_analysis", + "export_detection_artifact_analysis", + "export_measurement_tables", + "jsonl_chunk", + "render_result", + "snapshot_detection_artifacts", + "with_case_metadata", + "write_detection_artifact_payloads", + "write_summary", + ] + assert metadata.__all__ == [ + "BENCHMARK_SUITE_SCHEMA_VERSION", + "WANDB_METADATA_SCHEMA_VERSION", + "benchmark_metadata", + "build_wandb_metadata", + "config_metadata", + "detect_metadata", + "execution_metadata", + "file_hash", + "git_metadata", + "git_output", + "model_source_metadata", + "package_version", + "package_versions", + "replace_metadata", + "rewrite_metadata", + "run_tags", + "runtime_metadata", + "source_metadata", + "source_suffix", + "stable_hash", + "sweep_metadata", + "workload_metadata", + ] + assert runner.combine_measurements is artifacts.combine_measurements + assert runner.export_measurement_tables is artifacts.export_measurement_tables + assert runner._with_case_metadata is artifacts.with_case_metadata + assert runner._git_metadata is metadata.git_metadata + assert runner._package_versions is metadata.package_versions + assert runner._execution_metadata is metadata.execution_metadata + assert runner.build_wandb_metadata is not metadata.build_wandb_metadata + assert runner.build_wandb_metadata.__module__ == "compat_benchmark_b2_facade" + finally: + sys.modules.pop("compat_benchmark_b2_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/measurement_tools/benchmark_artifacts.py b/tools/measurement/measurement_tools/benchmark_artifacts.py new file mode 100644 index 00000000..c58bb617 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_artifacts.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark measurement and detection-artifact exports.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pandas as pd +from analyze_detection_artifacts import analyze_artifacts, iter_detection_parquet_files +from export_measurements import export_tables, read_measurements + +from measurement_tools.benchmark_models import BenchmarkCase, BenchmarkResult, CaseStatus +from measurement_tools.tables import ExportFormat + + +def combine_measurements(cases: list[BenchmarkCase], destination: Path) -> Path: + with destination.open("w", encoding="utf-8") as output: + for case in cases: + if case.measurement_path is None: + continue + source = Path(case.measurement_path) + if source.exists(): + output.write(source.read_text(encoding="utf-8")) + return destination + + +def combine_detection_artifact_analysis(cases: list[BenchmarkCase], destination: Path) -> Path | None: + chunks: list[str] = [] + for case in cases: + if case.detection_artifact_path is None: + continue + source = Path(case.detection_artifact_path) + if source.exists(): + chunks.append(jsonl_chunk(source.read_text(encoding="utf-8"))) + if not chunks: + return None + destination.write_text("".join(chunks), encoding="utf-8") + return destination + + +def jsonl_chunk(text: str) -> str: + if not text or text.endswith("\n"): + return text + return text + "\n" + + +def export_measurement_tables(measurement_path: Path, table_dir: Path) -> Path: + dataframe = read_measurements(measurement_path) + export_tables( + dataframe, input_path=measurement_path, output_dir=table_dir, export_format=ExportFormat.parquet, overwrite=True + ) + return table_dir + + +def snapshot_detection_artifacts(artifact_path: Path) -> dict[str, int]: + if not artifact_path.exists(): + return {} + return { + str(parquet_file.relative_to(artifact_path)): parquet_file.stat().st_mtime_ns + for parquet_file in iter_detection_parquet_files(artifact_path) + } + + +def changed_detection_artifact_files(artifact_path: Path, snapshot: dict[str, int]) -> list[Path]: + if not artifact_path.exists(): + return [] + changed: list[Path] = [] + for parquet_file in iter_detection_parquet_files(artifact_path): + key = str(parquet_file.relative_to(artifact_path)) + if snapshot.get(key) != parquet_file.stat().st_mtime_ns: + changed.append(parquet_file) + return changed + + +def export_detection_artifact_analysis( + artifact_path: Path, + output_path: Path, + *, + artifact_snapshot: dict[str, int] | None = None, +) -> Path | None: + if not artifact_path.exists(): + return None + parquet_files = ( + changed_detection_artifact_files(artifact_path, artifact_snapshot) if artifact_snapshot is not None else None + ) + analysis = analyze_artifacts(artifact_path, parquet_files=parquet_files) + if not analysis.rows: + return None + write_detection_artifact_payloads([row.model_dump() for row in analysis.rows], output_path) + return output_path + + +def export_case_detection_artifact_analysis( + artifact_path: Path, + output_path: Path, + *, + case: BenchmarkCase, + artifact_snapshot: dict[str, int], +) -> Path | None: + if not artifact_path.exists(): + return None + parquet_files = changed_detection_artifact_files(artifact_path, artifact_snapshot) + analysis = analyze_artifacts(artifact_path, parquet_files=parquet_files) + if not analysis.rows: + return None + write_detection_artifact_payloads( + [with_case_metadata(row.model_dump(), case=case) for row in analysis.rows], + output_path, + ) + return output_path + + +def with_case_metadata(row: dict[str, Any], *, case: BenchmarkCase) -> dict[str, Any]: + return { + "suite_id": case.suite_id, + "workload_id": case.workload_id, + "config_id": case.config_id, + "repetition": case.repetition, + "case_id": case.case_id, + "run_id": case.case_id, + **row, + } + + +def write_detection_artifact_payloads(rows: list[dict[str, Any]], output_path: Path) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + pd.json_normalize(rows, sep=".").to_json(output_path, orient="records", lines=True) + + +def write_summary(result: BenchmarkResult) -> None: + Path(result.summary_path).write_text(result.model_dump_json(indent=2) + "\n", encoding="utf-8") + + +def render_result(result: BenchmarkResult, *, json_output: bool) -> str: + if json_output: + return result.model_dump_json(indent=2) + completed = sum(case.status == CaseStatus.completed for case in result.cases) + errored = sum(case.status == CaseStatus.error for case in result.cases) + planned = sum(case.status == CaseStatus.planned for case in result.cases) + if planned and completed == 0 and errored == 0: + return f"Planned {planned} case(s); output={result.output_dir}" + return f"Ran {completed}/{len(result.cases)} case(s); errors={errored}; output={result.output_dir}" + + +__all__ = [ + "changed_detection_artifact_files", + "combine_detection_artifact_analysis", + "combine_measurements", + "export_case_detection_artifact_analysis", + "export_detection_artifact_analysis", + "export_measurement_tables", + "jsonl_chunk", + "render_result", + "snapshot_detection_artifacts", + "with_case_metadata", + "write_detection_artifact_payloads", + "write_summary", +] diff --git a/tools/measurement/measurement_tools/benchmark_wandb_metadata.py b/tools/measurement/measurement_tools/benchmark_wandb_metadata.py new file mode 100644 index 00000000..7c9e37c8 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_wandb_metadata.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sanitized W&B metadata for benchmark suites and cases.""" + +from __future__ import annotations + +import platform +import subprocess +from collections.abc import Callable +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +from anonymizer.config.anonymizer_config import infer_input_source_suffix, is_remote_input_source +from anonymizer.measurement import MEASUREMENT_SCHEMA_VERSION +from measurement_tools.benchmark_models import ( + BenchmarkCase, + BenchmarkSpec, + ConfigSpec, + DDTraceMode, + ReplaceSpec, + RewriteSpec, + WorkloadSpec, +) +from measurement_tools.benchmark_specs import cross_product_matrix +from measurement_tools.execution import benchmark_execution_metadata, stable_metadata_hash +from measurement_tools.wandb_setup import WANDB_SANITIZER_VERSION, WandbRunMetadata + +BENCHMARK_SUITE_SCHEMA_VERSION = 1 +WANDB_METADATA_SCHEMA_VERSION = 2 + + +def run_tags(case: BenchmarkCase, spec: BenchmarkSpec) -> dict[str, Any]: + return { + **spec.run_tags, + "suite_id": spec.suite_id, + "workload_id": case.workload_id, + "config_id": case.config_id, + "repetition": case.repetition, + "case_id": case.case_id, + } + + +def build_wandb_metadata( + spec: BenchmarkSpec, + *, + spec_path: Path, + output_dir: Path, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode, + dd_task_trace: bool, + execution_metadata_builder: Callable[..., dict[str, Any]] | None = None, + git_metadata_builder: Callable[[Path], dict[str, str | bool | None]] | None = None, + package_versions_builder: Callable[[], dict[str, str | None]] | None = None, +) -> WandbRunMetadata: + execution_builder = execution_metadata_builder or execution_metadata + git_builder = git_metadata_builder or git_metadata + packages_builder = package_versions_builder or package_versions + sweep = sweep_metadata(spec.run_tags) + metadata = { + "run_kind": "sweep_arm" if sweep is not None else "native_suite", + "benchmark": benchmark_metadata(spec, spec_path=spec_path), + "execution": execution_builder( + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + dd_task_trace=dd_task_trace, + ), + "runtime": runtime_metadata(package_versions_builder=packages_builder), + "git": git_builder(Path.cwd()), + "model_sources": model_source_metadata(spec), + "workloads": [workload_metadata(workload) for workload in spec.workloads], + "configs": [config_metadata(config) for config in spec.configs], + "matrix": [entry.model_dump(mode="json") for entry in (spec.matrix or cross_product_matrix(spec))], + "sweep": sweep, + } + return WandbRunMetadata.model_validate(metadata, strict=True) + + +def benchmark_metadata(spec: BenchmarkSpec, *, spec_path: Path) -> dict[str, Any]: + matrix = spec.matrix or cross_product_matrix(spec) + metadata = { + "metadata_schema_version": WANDB_METADATA_SCHEMA_VERSION, + "suite_schema_version": BENCHMARK_SUITE_SCHEMA_VERSION, + "wandb_sanitizer_version": WANDB_SANITIZER_VERSION, + "measurement_schema_version": MEASUREMENT_SCHEMA_VERSION, + "suite_id": spec.suite_id, + "workload_count": len(spec.workloads), + "config_count": len(spec.configs), + "matrix_entry_count": len(matrix), + "case_count": sum(entry.repetitions for entry in matrix), + "case_retries": spec.case_retries, + "case_retry_backoff_sec": spec.case_retry_backoff_sec, + } + if spec_hash := file_hash(spec_path): + metadata["suite_file_hash"] = spec_hash + return metadata + + +def execution_metadata( + *, + output_dir: Path, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode, + dd_task_trace: bool, +) -> dict[str, Any]: + metadata = benchmark_execution_metadata( + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace.value, + dd_task_trace=dd_task_trace, + ) + metadata.pop("output_dir_hash", None) + return metadata + + +def runtime_metadata(*, package_versions_builder: Callable[[], dict[str, str | None]] | None = None) -> dict[str, Any]: + versions_builder = package_versions_builder or package_versions + return { + **versions_builder(), + "python_version": platform.python_version(), + "platform_machine": platform.machine(), + "platform_system": platform.system(), + } + + +def package_versions() -> dict[str, str | None]: + return { + "anonymizer_version": package_version("nemo-anonymizer"), + "datadesigner_version": package_version("data-designer"), + "wandb_version": package_version("wandb"), + } + + +def package_version(package: str) -> str | None: + try: + return version(package) + except PackageNotFoundError: + return None + + +def model_source_metadata(spec: BenchmarkSpec) -> dict[str, bool]: + return { + "has_model_configs": spec.model_configs is not None, + "has_model_providers": spec.model_providers is not None, + "has_artifact_path": spec.artifact_path is not None, + } + + +def workload_metadata(workload: WorkloadSpec) -> dict[str, Any]: + return { + "id": workload.id, + "source": source_metadata(workload.source), + "text_column": workload.text_column, + "has_id_column": workload.id_column is not None, + "has_data_summary": workload.data_summary is not None, + "row_limit": workload.row_limit, + "row_offset": workload.row_offset, + } + + +def source_metadata(source: str) -> dict[str, str | None]: + return { + "kind": "remote_file" if is_remote_input_source(source) else "local_file", + "suffix": source_suffix(source), + } + + +def source_suffix(source: str) -> str | None: + try: + return infer_input_source_suffix(source) + except ValueError: + return None + + +def config_metadata(config: ConfigSpec) -> dict[str, Any]: + return { + "id": config.id, + "mode": "rewrite" if config.rewrite is not None else "replace", + "detect": detect_metadata(config.detect), + "replace": replace_metadata(config.replace), + "rewrite": rewrite_metadata(config.rewrite), + "evaluate": config.evaluate, + "emit_telemetry": config.emit_telemetry, + } + + +def sweep_metadata(run_tags: dict[str, Any]) -> dict[str, Any] | None: + sweep = run_tags.get("wandb_sweep") + return sweep if isinstance(sweep, dict) else None + + +def detect_metadata(detect: dict[str, Any]) -> dict[str, Any]: + entity_labels = detect.get("entity_labels") + metadata = { + "entity_label_source": "custom" if isinstance(entity_labels, list) else "default", + "entity_label_count": len(entity_labels) if isinstance(entity_labels, list) else None, + } + if isinstance(entity_labels, list): + metadata["entity_label_set_hash"] = stable_hash(",".join(sorted(map(str, entity_labels)))) + for key in ("gliner_threshold", "validation_max_entities_per_call", "validation_excerpt_window_chars"): + if key in detect: + metadata[key] = detect[key] + return metadata + + +def replace_metadata(replace: str | ReplaceSpec | None) -> dict[str, Any] | None: + if replace is None: + return None + if isinstance(replace, str): + return {"strategy": replace} + metadata: dict[str, Any] = {"strategy": replace.strategy.value} + for key in ("normalize_label", "algorithm", "digest_length"): + value = getattr(replace, key) + if value is not None: + metadata[key] = value + if replace.format_template is not None: + metadata["has_format_template"] = True + if replace.instructions is not None: + metadata["has_instructions"] = True + return metadata + + +def rewrite_metadata(rewrite: RewriteSpec | None) -> dict[str, Any] | None: + if rewrite is None: + return None + return { + "risk_tolerance": rewrite.risk_tolerance.value, + "max_repair_iterations": rewrite.max_repair_iterations, + "strict_entity_protection": rewrite.strict_entity_protection, + "has_privacy_goal": rewrite.protect is not None or rewrite.preserve is not None, + "has_protect": rewrite.protect is not None, + "has_preserve": rewrite.preserve is not None, + "has_instructions": rewrite.instructions is not None, + } + + +def git_metadata(cwd: Path) -> dict[str, str | bool | None]: + commit = git_output(cwd, "rev-parse", "HEAD") + branch = git_output(cwd, "rev-parse", "--abbrev-ref", "HEAD") + status = git_output(cwd, "status", "--short") + return {"commit": commit, "branch": branch, "dirty": bool(status) if status is not None else None} + + +def git_output(cwd: Path, *args: str) -> str | None: + try: + result = subprocess.run( + ["git", *args], + cwd=cwd, + check=False, + capture_output=True, + text=True, + timeout=2, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + return result.stdout.strip() or None + + +def file_hash(path: Path) -> str | None: + if not path.is_file(): + return None + return stable_hash(path.read_text(encoding="utf-8")) + + +def stable_hash(value: str) -> str: + return stable_metadata_hash(value) + + +__all__ = [ + "BENCHMARK_SUITE_SCHEMA_VERSION", + "WANDB_METADATA_SCHEMA_VERSION", + "benchmark_metadata", + "build_wandb_metadata", + "config_metadata", + "detect_metadata", + "execution_metadata", + "file_hash", + "git_metadata", + "git_output", + "model_source_metadata", + "package_version", + "package_versions", + "replace_metadata", + "rewrite_metadata", + "run_tags", + "runtime_metadata", + "source_metadata", + "source_suffix", + "stable_hash", + "sweep_metadata", + "workload_metadata", +] diff --git a/tools/measurement/run_benchmarks.py b/tools/measurement/run_benchmarks.py index 2d920769..6c8269a6 100755 --- a/tools/measurement/run_benchmarks.py +++ b/tools/measurement/run_benchmarks.py @@ -9,21 +9,32 @@ """ import logging -import platform -import subprocess import sys import time -from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Annotated, Any import cyclopts -import pandas as pd -from analyze_detection_artifacts import ( - analyze_artifacts, - iter_detection_parquet_files, +from measurement_tools.benchmark_artifacts import ( + changed_detection_artifact_files as changed_detection_artifact_files, +) +from measurement_tools.benchmark_artifacts import ( + combine_detection_artifact_analysis, + combine_measurements, + export_case_detection_artifact_analysis, + export_measurement_tables, + jsonl_chunk, + render_result, + snapshot_detection_artifacts, + with_case_metadata, + write_summary, +) +from measurement_tools.benchmark_artifacts import ( + export_detection_artifact_analysis as export_detection_artifact_analysis, +) +from measurement_tools.benchmark_artifacts import ( + write_detection_artifact_payloads as write_detection_artifact_payloads, ) -from export_measurements import export_tables, read_measurements from measurement_tools.benchmark_inputs import ( build_anonymizer_config, build_input, @@ -58,8 +69,6 @@ CaseStatus, ConfigSpec, DDTraceMode, - ReplaceSpec, - RewriteSpec, WorkloadSpec, duplicate_matrix_entries, duplicates, @@ -70,6 +79,8 @@ from measurement_tools.benchmark_models import ( ReplaceKind as ReplaceKind, ) +from measurement_tools.benchmark_models import ReplaceSpec as ReplaceSpec +from measurement_tools.benchmark_models import RewriteSpec as RewriteSpec from measurement_tools.benchmark_planning import dry_run_result as _canonical_dry_run_result from measurement_tools.benchmark_planning import plan_suite as _canonical_plan_suite from measurement_tools.benchmark_specs import ( @@ -87,11 +98,35 @@ preflight_workload_errors, prepare_output_dir, ) +from measurement_tools.benchmark_wandb_metadata import ( + BENCHMARK_SUITE_SCHEMA_VERSION as BENCHMARK_SUITE_SCHEMA_VERSION, +) +from measurement_tools.benchmark_wandb_metadata import ( + WANDB_METADATA_SCHEMA_VERSION as WANDB_METADATA_SCHEMA_VERSION, +) +from measurement_tools.benchmark_wandb_metadata import ( + benchmark_metadata, + config_metadata, + detect_metadata, + execution_metadata, + file_hash, + git_metadata, + git_output, + model_source_metadata, + package_version, + package_versions, + replace_metadata, + rewrite_metadata, + run_tags, + runtime_metadata, + source_metadata, + source_suffix, + stable_hash, + sweep_metadata, +) +from measurement_tools.benchmark_wandb_metadata import build_wandb_metadata as _canonical_build_wandb_metadata from measurement_tools.cli import LogFormat, configure_logging, log_bad_input, summarize_validation_error -from measurement_tools.execution import benchmark_execution_metadata, stable_metadata_hash -from measurement_tools.tables import ExportFormat from measurement_tools.wandb_setup import ( - WANDB_SANITIZER_VERSION, BenchmarkWandbFinalization, ResolvedWandbConfig, WandbMode, @@ -101,15 +136,11 @@ from pydantic import ValidationError from anonymizer.config.anonymizer_config import Rewrite as Rewrite -from anonymizer.config.anonymizer_config import ( - infer_input_source_suffix, - is_remote_input_source, -) +from anonymizer.config.anonymizer_config import is_remote_input_source as is_remote_input_source from anonymizer.config.replace_strategies import Redact as Redact from anonymizer.config.rewrite import RiskTolerance as RiskTolerance from anonymizer.interface.anonymizer import Anonymizer from anonymizer.measurement import ( - MEASUREMENT_SCHEMA_VERSION, MeasurementConfig, configured_measurement_session, record_evaluation_metrics, @@ -118,9 +149,6 @@ app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.benchmark") -BENCHMARK_SUITE_SCHEMA_VERSION = 1 -WANDB_METADATA_SCHEMA_VERSION = 2 - _CaseRunPaths = CaseRunPaths _active_config_ids = active_config_ids _cross_product_matrix = cross_product_matrix @@ -146,6 +174,26 @@ _slice_bounds = slice_bounds _workload_has_row_slice = workload_has_row_slice _write_local_input_dataframe = write_local_input_dataframe +_benchmark_metadata = benchmark_metadata +_config_metadata = config_metadata +_detect_metadata = detect_metadata +_execution_metadata = execution_metadata +_file_hash = file_hash +_git_metadata = git_metadata +_git_output = git_output +_jsonl_chunk = jsonl_chunk +_model_source_metadata = model_source_metadata +_package_version = package_version +_package_versions = package_versions +_replace_metadata = replace_metadata +_rewrite_metadata = rewrite_metadata +_run_tags = run_tags +_runtime_metadata = runtime_metadata +_source_metadata = source_metadata +_source_suffix = source_suffix +_stable_hash = stable_hash +_sweep_metadata = sweep_metadata +_with_case_metadata = with_case_metadata def run_suite( @@ -563,146 +611,6 @@ def _execute_case( ) -def combine_measurements(cases: list[BenchmarkCase], destination: Path) -> Path: - with destination.open("w", encoding="utf-8") as output: - for case in cases: - if case.measurement_path is None: - continue - source = Path(case.measurement_path) - if source.exists(): - output.write(source.read_text(encoding="utf-8")) - return destination - - -def combine_detection_artifact_analysis(cases: list[BenchmarkCase], destination: Path) -> Path | None: - chunks: list[str] = [] - for case in cases: - if case.detection_artifact_path is None: - continue - source = Path(case.detection_artifact_path) - if source.exists(): - chunks.append(_jsonl_chunk(source.read_text(encoding="utf-8"))) - if not chunks: - return None - destination.write_text("".join(chunks), encoding="utf-8") - return destination - - -def _jsonl_chunk(text: str) -> str: - if not text or text.endswith("\n"): - return text - return text + "\n" - - -def export_measurement_tables(measurement_path: Path, table_dir: Path) -> Path: - dataframe = read_measurements(measurement_path) - export_tables( - dataframe, input_path=measurement_path, output_dir=table_dir, export_format=ExportFormat.parquet, overwrite=True - ) - return table_dir - - -def snapshot_detection_artifacts(artifact_path: Path) -> dict[str, int]: - if not artifact_path.exists(): - return {} - return { - str(parquet_file.relative_to(artifact_path)): parquet_file.stat().st_mtime_ns - for parquet_file in iter_detection_parquet_files(artifact_path) - } - - -def changed_detection_artifact_files(artifact_path: Path, snapshot: dict[str, int]) -> list[Path]: - if not artifact_path.exists(): - return [] - changed: list[Path] = [] - for parquet_file in iter_detection_parquet_files(artifact_path): - key = str(parquet_file.relative_to(artifact_path)) - if snapshot.get(key) != parquet_file.stat().st_mtime_ns: - changed.append(parquet_file) - return changed - - -def export_detection_artifact_analysis( - artifact_path: Path, - output_path: Path, - *, - artifact_snapshot: dict[str, int] | None = None, -) -> Path | None: - if not artifact_path.exists(): - return None - parquet_files = ( - changed_detection_artifact_files(artifact_path, artifact_snapshot) if artifact_snapshot is not None else None - ) - analysis = analyze_artifacts(artifact_path, parquet_files=parquet_files) - if not analysis.rows: - return None - write_detection_artifact_payloads([row.model_dump() for row in analysis.rows], output_path) - return output_path - - -def export_case_detection_artifact_analysis( - artifact_path: Path, - output_path: Path, - *, - case: BenchmarkCase, - artifact_snapshot: dict[str, int], -) -> Path | None: - if not artifact_path.exists(): - return None - parquet_files = changed_detection_artifact_files(artifact_path, artifact_snapshot) - analysis = analyze_artifacts(artifact_path, parquet_files=parquet_files) - if not analysis.rows: - return None - write_detection_artifact_payloads( - [_with_case_metadata(row.model_dump(), case=case) for row in analysis.rows], - output_path, - ) - return output_path - - -def _with_case_metadata(row: dict[str, Any], *, case: BenchmarkCase) -> dict[str, Any]: - return { - "suite_id": case.suite_id, - "workload_id": case.workload_id, - "config_id": case.config_id, - "repetition": case.repetition, - "case_id": case.case_id, - "run_id": case.case_id, - **row, - } - - -def write_detection_artifact_payloads(rows: list[dict[str, Any]], output_path: Path) -> None: - output_path.parent.mkdir(parents=True, exist_ok=True) - pd.json_normalize(rows, sep=".").to_json(output_path, orient="records", lines=True) - - -def write_summary(result: BenchmarkResult) -> None: - Path(result.summary_path).write_text(result.model_dump_json(indent=2) + "\n", encoding="utf-8") - - -def render_result(result: BenchmarkResult, *, json_output: bool) -> str: - if json_output: - return result.model_dump_json(indent=2) - completed = sum(case.status == CaseStatus.completed for case in result.cases) - errored = sum(case.status == CaseStatus.error for case in result.cases) - planned = sum(case.status == CaseStatus.planned for case in result.cases) - if planned and completed == 0 and errored == 0: - return f"Planned {planned} case(s); output={result.output_dir}" - return f"Ran {completed}/{len(result.cases)} case(s); errors={errored}; output={result.output_dir}" - - -def _run_tags(case: BenchmarkCase, spec: BenchmarkSpec) -> dict[str, Any]: - return { - **spec.run_tags, - "suite_id": spec.suite_id, - "workload_id": case.workload_id, - "config_id": case.config_id, - "repetition": case.repetition, - "case_id": case.case_id, - } - - def _get_item(items: dict[str, Any], item_id: str, item_type: str) -> Any: if item_id not in items: raise ValueError(f"unknown {item_type}: {item_id}") @@ -769,219 +677,18 @@ def build_wandb_metadata( dd_trace: DDTraceMode, dd_task_trace: bool, ) -> WandbRunMetadata: - sweep = _sweep_metadata(spec.run_tags) - metadata = { - "run_kind": "sweep_arm" if sweep is not None else "native_suite", - "benchmark": _benchmark_metadata(spec, spec_path=spec_path), - "execution": _execution_metadata( - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - dd_trace=dd_trace, - dd_task_trace=dd_task_trace, - ), - "runtime": _runtime_metadata(), - "git": _git_metadata(Path.cwd()), - "model_sources": _model_source_metadata(spec), - "workloads": [_workload_metadata(workload) for workload in spec.workloads], - "configs": [_config_metadata(config) for config in spec.configs], - "matrix": [entry.model_dump(mode="json") for entry in (spec.matrix or _cross_product_matrix(spec))], - "sweep": sweep, - } - return WandbRunMetadata.model_validate(metadata, strict=True) - - -def _benchmark_metadata(spec: BenchmarkSpec, *, spec_path: Path) -> dict[str, Any]: - matrix = spec.matrix or _cross_product_matrix(spec) - metadata = { - "metadata_schema_version": WANDB_METADATA_SCHEMA_VERSION, - "suite_schema_version": BENCHMARK_SUITE_SCHEMA_VERSION, - "wandb_sanitizer_version": WANDB_SANITIZER_VERSION, - "measurement_schema_version": MEASUREMENT_SCHEMA_VERSION, - "suite_id": spec.suite_id, - "workload_count": len(spec.workloads), - "config_count": len(spec.configs), - "matrix_entry_count": len(matrix), - "case_count": sum(entry.repetitions for entry in matrix), - "case_retries": spec.case_retries, - "case_retry_backoff_sec": spec.case_retry_backoff_sec, - } - if spec_hash := _file_hash(spec_path): - metadata["suite_file_hash"] = spec_hash - return metadata - - -def _execution_metadata( - *, - output_dir: Path, - export: bool, - fail_fast: bool, - dd_trace: DDTraceMode, - dd_task_trace: bool, -) -> dict[str, Any]: - metadata = benchmark_execution_metadata( + return _canonical_build_wandb_metadata( + spec, + spec_path=spec_path, output_dir=output_dir, export=export, fail_fast=fail_fast, - dd_trace=dd_trace.value, + dd_trace=dd_trace, dd_task_trace=dd_task_trace, + execution_metadata_builder=_execution_metadata, + git_metadata_builder=_git_metadata, + package_versions_builder=_package_versions, ) - metadata.pop("output_dir_hash", None) - return metadata - - -def _runtime_metadata() -> dict[str, Any]: - return { - **_package_versions(), - "python_version": platform.python_version(), - "platform_machine": platform.machine(), - "platform_system": platform.system(), - } - - -def _package_versions() -> dict[str, str | None]: - return { - "anonymizer_version": _package_version("nemo-anonymizer"), - "datadesigner_version": _package_version("data-designer"), - "wandb_version": _package_version("wandb"), - } - - -def _package_version(package: str) -> str | None: - try: - return version(package) - except PackageNotFoundError: - return None - - -def _model_source_metadata(spec: BenchmarkSpec) -> dict[str, bool]: - return { - "has_model_configs": spec.model_configs is not None, - "has_model_providers": spec.model_providers is not None, - "has_artifact_path": spec.artifact_path is not None, - } - - -def _workload_metadata(workload: WorkloadSpec) -> dict[str, Any]: - return { - "id": workload.id, - "source": _source_metadata(workload.source), - "text_column": workload.text_column, - "has_id_column": workload.id_column is not None, - "has_data_summary": workload.data_summary is not None, - "row_limit": workload.row_limit, - "row_offset": workload.row_offset, - } - - -def _source_metadata(source: str) -> dict[str, str | None]: - return { - "kind": "remote_file" if is_remote_input_source(source) else "local_file", - "suffix": _source_suffix(source), - } - - -def _source_suffix(source: str) -> str | None: - try: - return infer_input_source_suffix(source) - except ValueError: - return None - - -def _config_metadata(config: ConfigSpec) -> dict[str, Any]: - return { - "id": config.id, - "mode": "rewrite" if config.rewrite is not None else "replace", - "detect": _detect_metadata(config.detect), - "replace": _replace_metadata(config.replace), - "rewrite": _rewrite_metadata(config.rewrite), - "evaluate": config.evaluate, - "emit_telemetry": config.emit_telemetry, - } - - -def _sweep_metadata(run_tags: dict[str, Any]) -> dict[str, Any] | None: - sweep = run_tags.get("wandb_sweep") - return sweep if isinstance(sweep, dict) else None - - -def _detect_metadata(detect: dict[str, Any]) -> dict[str, Any]: - entity_labels = detect.get("entity_labels") - metadata = { - "entity_label_source": "custom" if isinstance(entity_labels, list) else "default", - "entity_label_count": len(entity_labels) if isinstance(entity_labels, list) else None, - } - if isinstance(entity_labels, list): - metadata["entity_label_set_hash"] = _stable_hash(",".join(sorted(map(str, entity_labels)))) - for key in ("gliner_threshold", "validation_max_entities_per_call", "validation_excerpt_window_chars"): - if key in detect: - metadata[key] = detect[key] - return metadata - - -def _replace_metadata(replace: str | ReplaceSpec | None) -> dict[str, Any] | None: - if replace is None: - return None - if isinstance(replace, str): - return {"strategy": replace} - metadata: dict[str, Any] = {"strategy": replace.strategy.value} - for key in ("normalize_label", "algorithm", "digest_length"): - value = getattr(replace, key) - if value is not None: - metadata[key] = value - if replace.format_template is not None: - metadata["has_format_template"] = True - if replace.instructions is not None: - metadata["has_instructions"] = True - return metadata - - -def _rewrite_metadata(rewrite: RewriteSpec | None) -> dict[str, Any] | None: - if rewrite is None: - return None - return { - "risk_tolerance": rewrite.risk_tolerance.value, - "max_repair_iterations": rewrite.max_repair_iterations, - "strict_entity_protection": rewrite.strict_entity_protection, - "has_privacy_goal": rewrite.protect is not None or rewrite.preserve is not None, - "has_protect": rewrite.protect is not None, - "has_preserve": rewrite.preserve is not None, - "has_instructions": rewrite.instructions is not None, - } - - -def _git_metadata(cwd: Path) -> dict[str, str | bool | None]: - commit = _git_output(cwd, "rev-parse", "HEAD") - branch = _git_output(cwd, "rev-parse", "--abbrev-ref", "HEAD") - status = _git_output(cwd, "status", "--short") - return {"commit": commit, "branch": branch, "dirty": bool(status) if status is not None else None} - - -def _git_output(cwd: Path, *args: str) -> str | None: - try: - result = subprocess.run( - ["git", *args], - cwd=cwd, - check=False, - capture_output=True, - text=True, - timeout=2, - ) - except (OSError, subprocess.TimeoutExpired): - return None - if result.returncode != 0: - return None - return result.stdout.strip() or None - - -def _file_hash(path: Path) -> str | None: - if not path.is_file(): - return None - return _stable_hash(path.read_text(encoding="utf-8")) - - -def _stable_hash(value: str) -> str: - return stable_metadata_hash(value) @app.default From bedf1841fbc48464f460d76b7b03ba7f5a8ee59e Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 21:12:46 +0000 Subject: [PATCH 12/17] refactor(measurement): split benchmark execution Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 40 ++ .../measurement_tools/benchmark_execution.py | 608 ++++++++++++++++++ tools/measurement/run_benchmarks.py | 492 +++----------- 3 files changed, 739 insertions(+), 401 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_execution.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index ed685147..acb006b9 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -380,6 +380,46 @@ def test_benchmark_facade_preserves_artifact_and_metadata_contracts() -> None: sys.modules.pop("compat_benchmark_b2_facade", None) +def test_benchmark_facade_preserves_execution_contracts() -> None: + execution = importlib.import_module("measurement_tools.benchmark_execution") + runner = _load_module(MEASUREMENT_ROOT / "run_benchmarks.py", "compat_benchmark_b3_facade") + try: + assert execution.__all__ == [ + "benchmark_result", + "build_contexts", + "case_detection_artifact_path", + "case_run_paths", + "case_task_trace_path", + "case_trace_path", + "case_with_result", + "combine_suite_detection_artifacts", + "execute_case", + "export_case_detection_artifacts_if_requested", + "export_suite_tables", + "get_item", + "run_case", + "run_case_error", + "run_case_success", + "run_cases", + "run_or_plan", + "run_suite", + "should_export_measurements", + "sleep_before_case_retry", + ] + assert runner._build_contexts is execution.build_contexts + assert runner._run_cases is execution.run_cases + assert runner._run_case_success is execution.run_case_success + assert runner._case_run_paths is execution.case_run_paths + assert runner.run_suite is not execution.run_suite + assert runner._run_case is not execution.run_case + assert runner._run_case_error is not execution.run_case_error + assert runner._execute_case is not execution.execute_case + assert runner.run_or_plan is not execution.run_or_plan + assert runner.run_or_plan.__module__ == "compat_benchmark_b3_facade" + finally: + sys.modules.pop("compat_benchmark_b3_facade", None) + + @pytest.mark.parametrize( "canonical_module", [ diff --git a/tools/measurement/measurement_tools/benchmark_execution.py b/tools/measurement/measurement_tools/benchmark_execution.py new file mode 100644 index 00000000..5f01174c --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_execution.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark suite and case execution orchestration.""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from anonymizer.interface.anonymizer import Anonymizer +from anonymizer.measurement import MeasurementConfig, configured_measurement_session, record_evaluation_metrics +from measurement_tools.benchmark_artifacts import ( + combine_detection_artifact_analysis, + combine_measurements, + export_case_detection_artifact_analysis, + export_measurement_tables, + snapshot_detection_artifacts, + write_summary, +) +from measurement_tools.benchmark_inputs import ( + build_anonymizer_config, + build_input, + resolve_config_source, + resolve_optional_path, +) +from measurement_tools.benchmark_models import ( + BenchmarkCase, + BenchmarkResult, + BenchmarkSpec, + CaseRunPaths, + CaseStatus, + ConfigSpec, + DDTraceMode, + WorkloadSpec, +) +from measurement_tools.benchmark_planning import plan_suite +from measurement_tools.benchmark_specs import build_cases, load_spec, preflight_suite, prepare_output_dir +from measurement_tools.benchmark_wandb_metadata import build_wandb_metadata, execution_metadata, run_tags +from measurement_tools.wandb_setup import ( + BenchmarkWandbFinalization, + ResolvedWandbConfig, + publish_benchmark_wandb_best_effort, +) + +logger = logging.getLogger("measurement.benchmark") + + +def run_suite( + spec: BenchmarkSpec, + *, + spec_path: Path, + output_dir: Path, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode, + trace_dir: Path | None, + dd_task_trace: bool = False, + task_trace_dir: Path | None = None, + anonymizer_factory: Callable[..., Anonymizer] = Anonymizer, + run_case_operation: Callable[..., BenchmarkCase] | None = None, + combine_measurements_operation: Callable[[list[BenchmarkCase], Path], Path] = combine_measurements, + export_tables_operation: Callable[[Path, Path], Path] = export_measurement_tables, + combine_artifacts_operation: Callable[ + [list[BenchmarkCase], Path], Path | None + ] = combine_detection_artifact_analysis, + execution_metadata_builder: Callable[..., dict[str, Any]] = execution_metadata, + write_summary_operation: Callable[[BenchmarkResult], None] = write_summary, +) -> BenchmarkResult: + case_operation = run_case_operation or run_case + contexts = build_contexts( + spec, + spec_path=spec_path, + output_dir=output_dir, + dd_trace=dd_trace, + trace_dir=trace_dir, + dd_task_trace=dd_task_trace, + task_trace_dir=task_trace_dir, + ) + anonymizer = anonymizer_factory(**contexts["anonymizer_kwargs"]) + cases = run_cases( + spec, + contexts=contexts, + anonymizer=anonymizer, + fail_fast=fail_fast, + export=export, + run_case_operation=case_operation, + ) + measurement_path = combine_measurements_operation(cases, output_dir / "measurements.jsonl") + should_export = should_export_measurements(export=export, measurement_path=measurement_path) + table_dir = export_suite_tables( + measurement_path, + output_dir=output_dir, + should_export=should_export, + export_tables_operation=export_tables_operation, + ) + artifact_analysis_path = combine_suite_detection_artifacts( + cases, + output_dir=output_dir, + should_export=should_export, + combine_artifacts_operation=combine_artifacts_operation, + ) + result = benchmark_result( + spec, + output_dir=output_dir, + measurement_path=measurement_path, + table_dir=table_dir, + artifact_analysis_path=artifact_analysis_path, + cases=cases, + execution=execution_metadata_builder( + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + dd_task_trace=dd_task_trace, + ), + ) + write_summary_operation(result) + return result + + +def run_cases( + spec: BenchmarkSpec, + *, + contexts: dict[str, Any], + anonymizer: Anonymizer, + fail_fast: bool, + export: bool, + run_case_operation: Callable[..., BenchmarkCase] | None = None, +) -> list[BenchmarkCase]: + case_operation = run_case_operation or run_case + return [ + case_operation( + case, + spec, + contexts=contexts, + anonymizer=anonymizer, + fail_fast=fail_fast, + export_detection_artifacts=export, + ) + for case in build_cases(spec) + ] + + +def should_export_measurements(*, export: bool, measurement_path: Path) -> bool: + return export and measurement_path.stat().st_size > 0 + + +def export_suite_tables( + measurement_path: Path, + *, + output_dir: Path, + should_export: bool, + export_tables_operation: Callable[[Path, Path], Path] = export_measurement_tables, +) -> Path | None: + if not should_export: + return None + return export_tables_operation(measurement_path, output_dir / "tables") + + +def combine_suite_detection_artifacts( + cases: list[BenchmarkCase], + *, + output_dir: Path, + should_export: bool, + combine_artifacts_operation: Callable[ + [list[BenchmarkCase], Path], Path | None + ] = combine_detection_artifact_analysis, +) -> Path | None: + if not should_export: + return None + return combine_artifacts_operation(cases, output_dir / "detection-artifacts.jsonl") + + +def benchmark_result( + spec: BenchmarkSpec, + *, + output_dir: Path, + measurement_path: Path, + table_dir: Path | None, + artifact_analysis_path: Path | None, + cases: list[BenchmarkCase], + execution: dict[str, Any] | None = None, +) -> BenchmarkResult: + return BenchmarkResult( + suite_id=spec.suite_id, + output_dir=str(output_dir), + measurement_path=str(measurement_path), + summary_path=str(output_dir / "summary.json"), + table_dir=str(table_dir) if table_dir is not None else None, + detection_artifact_analysis_path=str(artifact_analysis_path) if artifact_analysis_path is not None else None, + cases=cases, + execution=execution or {}, + ) + + +def build_contexts( + spec: BenchmarkSpec, + *, + spec_path: Path, + output_dir: Path, + dd_trace: DDTraceMode, + trace_dir: Path | None, + dd_task_trace: bool = False, + task_trace_dir: Path | None = None, +) -> dict[str, Any]: + base_dir = spec_path.parent + artifact_path = resolve_optional_path(spec.artifact_path, base_dir) or output_dir / "artifacts" + return { + "base_dir": base_dir, + "workloads": {workload.id: workload for workload in spec.workloads}, + "configs": {config.id: config for config in spec.configs}, + "raw_dir": output_dir / "raw", + "dd_trace": dd_trace, + "trace_dir": trace_dir or output_dir / "traces", + "dd_task_trace": dd_task_trace, + "task_trace_dir": task_trace_dir or output_dir / "task-traces", + "artifact_path": artifact_path, + "anonymizer_kwargs": { + "model_configs": resolve_config_source(spec.model_configs, base_dir), + "model_providers": resolve_config_source(spec.model_providers, base_dir), + "artifact_path": artifact_path, + }, + } + + +def run_case( + case: BenchmarkCase, + spec: BenchmarkSpec, + *, + contexts: dict[str, Any], + anonymizer: Anonymizer, + fail_fast: bool, + export_detection_artifacts: bool, + execute_case_operation: Callable[..., None] | None = None, + run_case_error_operation: Callable[..., BenchmarkCase] | None = None, +) -> BenchmarkCase: + execute_operation = execute_case_operation or execute_case + error_operation = run_case_error_operation or run_case_error + started = time.perf_counter() + attempt_errors: list[str] = [] + max_attempts = 1 if fail_fast else spec.case_retries + 1 + for attempt_number in range(1, max_attempts + 1): + paths = case_run_paths(case, contexts=contexts, export_detection_artifacts=export_detection_artifacts) + try: + return run_case_success( + case, + spec, + contexts=contexts, + anonymizer=anonymizer, + paths=paths, + started=started, + attempt_count=attempt_number, + attempt_errors=attempt_errors, + execute_case_operation=execute_operation, + ) + except Exception as exc: + if fail_fast: + raise + attempt_errors.append(str(exc)) + if attempt_number >= max_attempts: + return error_operation( + case, + contexts=contexts, + paths=paths, + started=started, + error=exc, + attempt_count=attempt_number, + attempt_errors=attempt_errors, + ) + sleep_before_case_retry(spec, case=case, attempt_number=attempt_number, error=exc) + + raise RuntimeError("unreachable benchmark retry state") + + +def sleep_before_case_retry( + spec: BenchmarkSpec, + *, + case: BenchmarkCase, + attempt_number: int, + error: Exception, +) -> None: + logger.warning( + "case %s attempt %d failed; retrying after %.2fs: %s", + case.case_id, + attempt_number, + spec.case_retry_backoff_sec, + error, + ) + if spec.case_retry_backoff_sec > 0: + time.sleep(spec.case_retry_backoff_sec) + + +def case_run_paths( + case: BenchmarkCase, + *, + contexts: dict[str, Any], + export_detection_artifacts: bool, +) -> CaseRunPaths: + return CaseRunPaths( + raw_path=contexts["raw_dir"] / f"{case.case_id}.jsonl", + artifact_output_path=contexts["raw_dir"] / f"{case.case_id}.detection-artifacts.jsonl", + trace_path=case_trace_path(case, contexts=contexts), + task_trace_path=case_task_trace_path(case, contexts=contexts), + artifact_snapshot=snapshot_detection_artifacts(contexts["artifact_path"]) + if export_detection_artifacts + else None, + export_detection_artifacts=export_detection_artifacts, + ) + + +def run_case_success( + case: BenchmarkCase, + spec: BenchmarkSpec, + *, + contexts: dict[str, Any], + anonymizer: Anonymizer, + paths: CaseRunPaths, + started: float, + attempt_count: int, + attempt_errors: list[str], + execute_case_operation: Callable[..., None] | None = None, +) -> BenchmarkCase: + execute_operation = execute_case_operation or execute_case + workload = get_item(contexts["workloads"], case.workload_id, "workload") + config = get_item(contexts["configs"], case.config_id, "config") + execute_operation( + anonymizer, + workload, + config, + raw_path=paths.raw_path, + trace_path=paths.trace_path, + task_trace_path=paths.task_trace_path, + case=case, + spec=spec, + base_dir=contexts["base_dir"], + dd_trace=contexts["dd_trace"], + ) + detection_artifact_path = case_detection_artifact_path(contexts, paths, case=case) + return case_with_result( + case, + status=CaseStatus.completed, + started=started, + raw_path=paths.raw_path, + detection_artifact_path=detection_artifact_path, + trace_path=paths.trace_path, + task_trace_path=paths.task_trace_path, + attempt_count=attempt_count, + attempt_errors=attempt_errors, + ) + + +def run_case_error( + case: BenchmarkCase, + *, + contexts: dict[str, Any], + paths: CaseRunPaths, + started: float, + error: Exception, + attempt_count: int, + attempt_errors: list[str], + artifact_export_operation: Callable[..., Path | None] = export_case_detection_artifact_analysis, +) -> BenchmarkCase: + detection_artifact_path = export_case_detection_artifacts_if_requested( + contexts, + paths.artifact_output_path, + case=case, + artifact_snapshot=paths.artifact_snapshot, + artifact_export_operation=artifact_export_operation, + ) + return case_with_result( + case, + status=CaseStatus.error, + started=started, + raw_path=paths.raw_path, + detection_artifact_path=detection_artifact_path, + trace_path=paths.trace_path, + task_trace_path=paths.task_trace_path, + error=str(error), + attempt_count=attempt_count, + attempt_errors=attempt_errors, + ) + + +def case_detection_artifact_path( + contexts: dict[str, Any], + paths: CaseRunPaths, + *, + case: BenchmarkCase, +) -> Path | None: + detection_artifact_path = export_case_detection_artifacts_if_requested( + contexts, + paths.artifact_output_path, + case=case, + artifact_snapshot=paths.artifact_snapshot, + ) + if detection_artifact_path is not None or paths.artifact_snapshot is None: + return detection_artifact_path + return None + + +def case_with_result( + case: BenchmarkCase, + *, + status: CaseStatus, + started: float, + raw_path: Path, + detection_artifact_path: Path | None, + trace_path: Path | None, + task_trace_path: Path | None, + attempt_count: int, + attempt_errors: list[str], + error: str | None = None, +) -> BenchmarkCase: + return case.model_copy( + update={ + "status": status, + "elapsed_sec": time.perf_counter() - started, + "measurement_path": str(raw_path), + "detection_artifact_path": (str(detection_artifact_path) if detection_artifact_path is not None else None), + "trace_path": str(trace_path) if trace_path is not None else None, + "task_trace_path": str(task_trace_path) if task_trace_path is not None else None, + "error": error, + "attempt_count": attempt_count, + "attempt_errors": list(attempt_errors), + } + ) + + +def export_case_detection_artifacts_if_requested( + contexts: dict[str, Any], + output_path: Path, + *, + case: BenchmarkCase, + artifact_snapshot: dict[str, int] | None, + artifact_export_operation: Callable[..., Path | None] = export_case_detection_artifact_analysis, +) -> Path | None: + if artifact_snapshot is None: + return None + return artifact_export_operation( + contexts["artifact_path"], + output_path, + case=case, + artifact_snapshot=artifact_snapshot, + ) + + +def case_trace_path(case: BenchmarkCase, *, contexts: dict[str, Any]) -> Path | None: + if contexts["dd_trace"] == DDTraceMode.none: + return None + return contexts["trace_dir"] / f"{case.case_id}.jsonl" + + +def case_task_trace_path(case: BenchmarkCase, *, contexts: dict[str, Any]) -> Path | None: + if not contexts["dd_task_trace"]: + return None + return contexts["task_trace_dir"] / f"{case.case_id}.jsonl" + + +def execute_case( + anonymizer: Anonymizer, + workload: WorkloadSpec, + config: ConfigSpec, + *, + raw_path: Path, + trace_path: Path | None, + task_trace_path: Path | None, + case: BenchmarkCase, + spec: BenchmarkSpec, + base_dir: Path, + dd_trace: DDTraceMode, + build_config_operation: Callable[[ConfigSpec], Any] = build_anonymizer_config, + build_input_operation: Callable[..., Any] = build_input, + run_tags_operation: Callable[[BenchmarkCase, BenchmarkSpec], dict[str, Any]] = run_tags, + measurement_session: Callable[..., Any] = configured_measurement_session, + record_metrics_operation: Callable[..., None] = record_evaluation_metrics, +) -> None: + anonymizer_config = build_config_operation(config) + input_data = build_input_operation( + workload, + base_dir, + slice_dir=raw_path.parent / "inputs", + case_id=case.case_id, + ) + measurement = MeasurementConfig( + output_path=raw_path, + run_id=case.case_id, + run_tags=run_tags_operation(case, spec), + streaming=True, + keep_records=False, + dd_trace=dd_trace.value, + dd_trace_path=trace_path, + dd_task_trace_path=task_trace_path, + fail_on_write_error=True, + ) + with measurement_session(measurement): + result = anonymizer.run(config=anonymizer_config, data=input_data) + if config.evaluate: + evaluated = anonymizer.evaluate(result) + record_metrics_operation( + evaluated.trace_dataframe, + mode="replace", + strategy=type(anonymizer_config.replace).__name__, + text_column=evaluated.resolved_text_column, + ) + + +def get_item(items: dict[str, Any], item_id: str, item_type: str) -> Any: + if item_id not in items: + raise ValueError(f"unknown {item_type}: {item_id}") + return items[item_id] + + +def run_or_plan( + spec_path: Path, + *, + output: Path | None, + overwrite: bool, + dry_run: bool, + export: bool, + fail_fast: bool, + dd_trace: DDTraceMode = DDTraceMode.none, + trace_dir: Path | None = None, + dd_task_trace: bool = False, + task_trace_dir: Path | None = None, + wandb_settings: ResolvedWandbConfig | None = None, + load_spec_operation: Callable[[Path], BenchmarkSpec] = load_spec, + plan_suite_operation: Callable[..., BenchmarkResult] = plan_suite, + preflight_operation: Callable[..., None] = preflight_suite, + prepare_output_operation: Callable[..., None] = prepare_output_dir, + run_suite_operation: Callable[..., BenchmarkResult] = run_suite, + publisher_operation: Callable[..., Any] = publish_benchmark_wandb_best_effort, + metadata_builder: Callable[..., Any] = build_wandb_metadata, +) -> BenchmarkResult: + benchmark_spec = load_spec_operation(spec_path) + output_dir = output or Path("benchmark-runs") / benchmark_spec.suite_id + resolved_wandb = wandb_settings or ResolvedWandbConfig() + if trace_dir is not None and dd_trace == DDTraceMode.none: + raise ValueError("--trace-dir requires --dd-trace") + if task_trace_dir is not None and not dd_task_trace: + raise ValueError("--task-trace-dir requires --dd-task-trace") + if dry_run: + return plan_suite_operation( + benchmark_spec, + spec_path=spec_path, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + trace_dir=trace_dir, + dd_task_trace=dd_task_trace, + task_trace_dir=task_trace_dir, + ) + preflight_operation(benchmark_spec, spec_path=spec_path) + prepare_output_operation(output_dir, overwrite=overwrite, dry_run=dry_run) + result = run_suite_operation( + benchmark_spec, + spec_path=spec_path, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + trace_dir=trace_dir, + dd_task_trace=dd_task_trace, + task_trace_dir=task_trace_dir, + ) + publisher_operation( + resolved_wandb, + suite_id=benchmark_spec.suite_id, + output_dir=output_dir, + finalization=BenchmarkWandbFinalization(measurement_path=Path(result.measurement_path), cases=result.cases), + metadata_factory=lambda: metadata_builder( + benchmark_spec, + spec_path=spec_path, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + dd_trace=dd_trace, + dd_task_trace=dd_task_trace, + ), + ) + return result + + +__all__ = [ + "benchmark_result", + "build_contexts", + "case_detection_artifact_path", + "case_run_paths", + "case_task_trace_path", + "case_trace_path", + "case_with_result", + "combine_suite_detection_artifacts", + "execute_case", + "export_case_detection_artifacts_if_requested", + "export_suite_tables", + "get_item", + "run_case", + "run_case_error", + "run_case_success", + "run_cases", + "run_or_plan", + "run_suite", + "should_export_measurements", + "sleep_before_case_retry", +] diff --git a/tools/measurement/run_benchmarks.py b/tools/measurement/run_benchmarks.py index 6c8269a6..0521d798 100755 --- a/tools/measurement/run_benchmarks.py +++ b/tools/measurement/run_benchmarks.py @@ -10,7 +10,6 @@ import logging import sys -import time from pathlib import Path from typing import Annotated, Any @@ -25,16 +24,38 @@ export_measurement_tables, jsonl_chunk, render_result, - snapshot_detection_artifacts, with_case_metadata, write_summary, ) from measurement_tools.benchmark_artifacts import ( export_detection_artifact_analysis as export_detection_artifact_analysis, ) +from measurement_tools.benchmark_artifacts import snapshot_detection_artifacts as snapshot_detection_artifacts from measurement_tools.benchmark_artifacts import ( write_detection_artifact_payloads as write_detection_artifact_payloads, ) +from measurement_tools.benchmark_execution import ( + benchmark_result, + build_contexts, + case_detection_artifact_path, + case_run_paths, + case_task_trace_path, + case_trace_path, + case_with_result, + combine_suite_detection_artifacts, + export_case_detection_artifacts_if_requested, + export_suite_tables, + get_item, + run_case_success, + run_cases, + should_export_measurements, + sleep_before_case_retry, +) +from measurement_tools.benchmark_execution import execute_case as _canonical_execute_case +from measurement_tools.benchmark_execution import run_case as _canonical_run_case +from measurement_tools.benchmark_execution import run_case_error as _canonical_run_case_error +from measurement_tools.benchmark_execution import run_or_plan as _canonical_run_or_plan +from measurement_tools.benchmark_execution import run_suite as _canonical_run_suite from measurement_tools.benchmark_inputs import ( build_anonymizer_config, build_input, @@ -85,7 +106,6 @@ from measurement_tools.benchmark_planning import plan_suite as _canonical_plan_suite from measurement_tools.benchmark_specs import ( active_config_ids, - build_cases, cross_product_matrix, input_columns, load_spec, @@ -98,6 +118,7 @@ preflight_workload_errors, prepare_output_dir, ) +from measurement_tools.benchmark_specs import build_cases as build_cases from measurement_tools.benchmark_wandb_metadata import ( BENCHMARK_SUITE_SCHEMA_VERSION as BENCHMARK_SUITE_SCHEMA_VERSION, ) @@ -126,8 +147,8 @@ ) from measurement_tools.benchmark_wandb_metadata import build_wandb_metadata as _canonical_build_wandb_metadata from measurement_tools.cli import LogFormat, configure_logging, log_bad_input, summarize_validation_error +from measurement_tools.wandb_setup import BenchmarkWandbFinalization as BenchmarkWandbFinalization from measurement_tools.wandb_setup import ( - BenchmarkWandbFinalization, ResolvedWandbConfig, WandbMode, WandbRunMetadata, @@ -141,7 +162,6 @@ from anonymizer.config.rewrite import RiskTolerance as RiskTolerance from anonymizer.interface.anonymizer import Anonymizer from anonymizer.measurement import ( - MeasurementConfig, configured_measurement_session, record_evaluation_metrics, ) @@ -194,6 +214,21 @@ _stable_hash = stable_hash _sweep_metadata = sweep_metadata _with_case_metadata = with_case_metadata +_benchmark_result = benchmark_result +_build_contexts = build_contexts +_case_detection_artifact_path = case_detection_artifact_path +_case_run_paths = case_run_paths +_case_task_trace_path = case_task_trace_path +_case_trace_path = case_trace_path +_case_with_result = case_with_result +_combine_suite_detection_artifacts = combine_suite_detection_artifacts +_export_case_detection_artifacts_if_requested = export_case_detection_artifacts_if_requested +_export_suite_tables = export_suite_tables +_get_item = get_item +_run_case_success = run_case_success +_run_cases = run_cases +_should_export_measurements = should_export_measurements +_sleep_before_case_retry = sleep_before_case_retry def run_suite( @@ -208,134 +243,24 @@ def run_suite( dd_task_trace: bool = False, task_trace_dir: Path | None = None, ) -> BenchmarkResult: - contexts = _build_contexts( + return _canonical_run_suite( spec, spec_path=spec_path, output_dir=output_dir, + export=export, + fail_fast=fail_fast, dd_trace=dd_trace, trace_dir=trace_dir, dd_task_trace=dd_task_trace, task_trace_dir=task_trace_dir, + anonymizer_factory=Anonymizer, + run_case_operation=_run_case, + combine_measurements_operation=combine_measurements, + export_tables_operation=export_measurement_tables, + combine_artifacts_operation=combine_detection_artifact_analysis, + execution_metadata_builder=_execution_metadata, + write_summary_operation=write_summary, ) - anonymizer = Anonymizer(**contexts["anonymizer_kwargs"]) - cases = _run_cases(spec, contexts=contexts, anonymizer=anonymizer, fail_fast=fail_fast, export=export) - measurement_path = combine_measurements(cases, output_dir / "measurements.jsonl") - should_export = _should_export_measurements(export=export, measurement_path=measurement_path) - table_dir = _export_suite_tables(measurement_path, output_dir=output_dir, should_export=should_export) - artifact_analysis_path = _combine_suite_detection_artifacts( - cases, output_dir=output_dir, should_export=should_export - ) - result = _benchmark_result( - spec, - output_dir=output_dir, - measurement_path=measurement_path, - table_dir=table_dir, - artifact_analysis_path=artifact_analysis_path, - cases=cases, - execution=_execution_metadata( - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - dd_trace=dd_trace, - dd_task_trace=dd_task_trace, - ), - ) - write_summary(result) - return result - - -def _run_cases( - spec: BenchmarkSpec, - *, - contexts: dict[str, Any], - anonymizer: Anonymizer, - fail_fast: bool, - export: bool, -) -> list[BenchmarkCase]: - return [ - _run_case( - case, - spec, - contexts=contexts, - anonymizer=anonymizer, - fail_fast=fail_fast, - export_detection_artifacts=export, - ) - for case in build_cases(spec) - ] - - -def _should_export_measurements(*, export: bool, measurement_path: Path) -> bool: - return export and measurement_path.stat().st_size > 0 - - -def _export_suite_tables(measurement_path: Path, *, output_dir: Path, should_export: bool) -> Path | None: - if not should_export: - return None - return export_measurement_tables(measurement_path, output_dir / "tables") - - -def _combine_suite_detection_artifacts( - cases: list[BenchmarkCase], - *, - output_dir: Path, - should_export: bool, -) -> Path | None: - if not should_export: - return None - return combine_detection_artifact_analysis(cases, output_dir / "detection-artifacts.jsonl") - - -def _benchmark_result( - spec: BenchmarkSpec, - *, - output_dir: Path, - measurement_path: Path, - table_dir: Path | None, - artifact_analysis_path: Path | None, - cases: list[BenchmarkCase], - execution: dict[str, Any] | None = None, -) -> BenchmarkResult: - return BenchmarkResult( - suite_id=spec.suite_id, - output_dir=str(output_dir), - measurement_path=str(measurement_path), - summary_path=str(output_dir / "summary.json"), - table_dir=str(table_dir) if table_dir is not None else None, - detection_artifact_analysis_path=str(artifact_analysis_path) if artifact_analysis_path is not None else None, - cases=cases, - execution=execution or {}, - ) - - -def _build_contexts( - spec: BenchmarkSpec, - *, - spec_path: Path, - output_dir: Path, - dd_trace: DDTraceMode, - trace_dir: Path | None, - dd_task_trace: bool = False, - task_trace_dir: Path | None = None, -) -> dict[str, Any]: - base_dir = spec_path.parent - artifact_path = _resolve_optional_path(spec.artifact_path, base_dir) or output_dir / "artifacts" - return { - "base_dir": base_dir, - "workloads": {workload.id: workload for workload in spec.workloads}, - "configs": {config.id: config for config in spec.configs}, - "raw_dir": output_dir / "raw", - "dd_trace": dd_trace, - "trace_dir": trace_dir or output_dir / "traces", - "dd_task_trace": dd_task_trace, - "task_trace_dir": task_trace_dir or output_dir / "task-traces", - "artifact_path": artifact_path, - "anonymizer_kwargs": { - "model_configs": _resolve_config_source(spec.model_configs, base_dir), - "model_providers": _resolve_config_source(spec.model_providers, base_dir), - "artifact_path": artifact_path, - }, - } def _run_case( @@ -347,117 +272,15 @@ def _run_case( fail_fast: bool, export_detection_artifacts: bool, ) -> BenchmarkCase: - started = time.perf_counter() - attempt_errors: list[str] = [] - max_attempts = 1 if fail_fast else spec.case_retries + 1 - for attempt_number in range(1, max_attempts + 1): - paths = _case_run_paths(case, contexts=contexts, export_detection_artifacts=export_detection_artifacts) - try: - return _run_case_success( - case, - spec, - contexts=contexts, - anonymizer=anonymizer, - paths=paths, - started=started, - attempt_count=attempt_number, - attempt_errors=attempt_errors, - ) - except Exception as exc: - if fail_fast: - raise - attempt_errors.append(str(exc)) - if attempt_number >= max_attempts: - return _run_case_error( - case, - contexts=contexts, - paths=paths, - started=started, - error=exc, - attempt_count=attempt_number, - attempt_errors=attempt_errors, - ) - _sleep_before_case_retry(spec, case=case, attempt_number=attempt_number, error=exc) - - raise RuntimeError("unreachable benchmark retry state") - - -def _sleep_before_case_retry( - spec: BenchmarkSpec, - *, - case: BenchmarkCase, - attempt_number: int, - error: Exception, -) -> None: - logger.warning( - "case %s attempt %d failed; retrying after %.2fs: %s", - case.case_id, - attempt_number, - spec.case_retry_backoff_sec, - error, - ) - if spec.case_retry_backoff_sec > 0: - time.sleep(spec.case_retry_backoff_sec) - - -def _case_run_paths( - case: BenchmarkCase, - *, - contexts: dict[str, Any], - export_detection_artifacts: bool, -) -> _CaseRunPaths: - return _CaseRunPaths( - raw_path=contexts["raw_dir"] / f"{case.case_id}.jsonl", - artifact_output_path=contexts["raw_dir"] / f"{case.case_id}.detection-artifacts.jsonl", - trace_path=_case_trace_path(case, contexts=contexts), - task_trace_path=_case_task_trace_path(case, contexts=contexts), - artifact_snapshot=snapshot_detection_artifacts(contexts["artifact_path"]) - if export_detection_artifacts - else None, - export_detection_artifacts=export_detection_artifacts, - ) - - -def _run_case_success( - case: BenchmarkCase, - spec: BenchmarkSpec, - *, - contexts: dict[str, Any], - anonymizer: Anonymizer, - paths: _CaseRunPaths, - started: float, - attempt_count: int, - attempt_errors: list[str], -) -> BenchmarkCase: - workload = _get_item(contexts["workloads"], case.workload_id, "workload") - config = _get_item(contexts["configs"], case.config_id, "config") - _execute_case( - anonymizer, - workload, - config, - raw_path=paths.raw_path, - trace_path=paths.trace_path, - task_trace_path=paths.task_trace_path, - case=case, - spec=spec, - base_dir=contexts["base_dir"], - dd_trace=contexts["dd_trace"], - ) - detection_artifact_path = _case_detection_artifact_path( - contexts, - paths, - case=case, - ) - return _case_with_result( + return _canonical_run_case( case, - status=CaseStatus.completed, - started=started, - raw_path=paths.raw_path, - detection_artifact_path=detection_artifact_path, - trace_path=paths.trace_path, - task_trace_path=paths.task_trace_path, - attempt_count=attempt_count, - attempt_errors=attempt_errors, + spec, + contexts=contexts, + anonymizer=anonymizer, + fail_fast=fail_fast, + export_detection_artifacts=export_detection_artifacts, + execute_case_operation=_execute_case, + run_case_error_operation=_run_case_error, ) @@ -471,100 +294,18 @@ def _run_case_error( attempt_count: int, attempt_errors: list[str], ) -> BenchmarkCase: - detection_artifact_path = _export_case_detection_artifacts_if_requested( - contexts, - paths.artifact_output_path, - case=case, - artifact_snapshot=paths.artifact_snapshot, - ) - return _case_with_result( + return _canonical_run_case_error( case, - status=CaseStatus.error, + contexts=contexts, + paths=paths, started=started, - raw_path=paths.raw_path, - detection_artifact_path=detection_artifact_path, - trace_path=paths.trace_path, - task_trace_path=paths.task_trace_path, - error=str(error), + error=error, attempt_count=attempt_count, attempt_errors=attempt_errors, + artifact_export_operation=export_case_detection_artifact_analysis, ) -def _case_detection_artifact_path( - contexts: dict[str, Any], - paths: _CaseRunPaths, - *, - case: BenchmarkCase, -) -> Path | None: - detection_artifact_path = _export_case_detection_artifacts_if_requested( - contexts, - paths.artifact_output_path, - case=case, - artifact_snapshot=paths.artifact_snapshot, - ) - if detection_artifact_path is not None or paths.artifact_snapshot is None: - return detection_artifact_path - return None - - -def _case_with_result( - case: BenchmarkCase, - *, - status: CaseStatus, - started: float, - raw_path: Path, - detection_artifact_path: Path | None, - trace_path: Path | None, - task_trace_path: Path | None, - attempt_count: int, - attempt_errors: list[str], - error: str | None = None, -) -> BenchmarkCase: - return case.model_copy( - update={ - "status": status, - "elapsed_sec": time.perf_counter() - started, - "measurement_path": str(raw_path), - "detection_artifact_path": (str(detection_artifact_path) if detection_artifact_path is not None else None), - "trace_path": str(trace_path) if trace_path is not None else None, - "task_trace_path": str(task_trace_path) if task_trace_path is not None else None, - "error": error, - "attempt_count": attempt_count, - "attempt_errors": list(attempt_errors), - } - ) - - -def _export_case_detection_artifacts_if_requested( - contexts: dict[str, Any], - output_path: Path, - *, - case: BenchmarkCase, - artifact_snapshot: dict[str, int] | None, -) -> Path | None: - if artifact_snapshot is None: - return None - return export_case_detection_artifact_analysis( - contexts["artifact_path"], - output_path, - case=case, - artifact_snapshot=artifact_snapshot, - ) - - -def _case_trace_path(case: BenchmarkCase, *, contexts: dict[str, Any]) -> Path | None: - if contexts["dd_trace"] == DDTraceMode.none: - return None - return contexts["trace_dir"] / f"{case.case_id}.jsonl" - - -def _case_task_trace_path(case: BenchmarkCase, *, contexts: dict[str, Any]) -> Path | None: - if not contexts["dd_task_trace"]: - return None - return contexts["task_trace_dir"] / f"{case.case_id}.jsonl" - - def _execute_case( anonymizer: Anonymizer, workload: WorkloadSpec, @@ -578,43 +319,23 @@ def _execute_case( base_dir: Path, dd_trace: DDTraceMode, ) -> None: - anonymizer_config = build_anonymizer_config(config) - input_data = build_input( + return _canonical_execute_case( + anonymizer, workload, - base_dir, - slice_dir=raw_path.parent / "inputs", - case_id=case.case_id, - ) - measurement = MeasurementConfig( - output_path=raw_path, - run_id=case.case_id, - run_tags=_run_tags(case, spec), - streaming=True, - keep_records=False, - dd_trace=dd_trace.value, - dd_trace_path=trace_path, - dd_task_trace_path=task_trace_path, - fail_on_write_error=True, + config, + raw_path=raw_path, + trace_path=trace_path, + task_trace_path=task_trace_path, + case=case, + spec=spec, + base_dir=base_dir, + dd_trace=dd_trace, + build_config_operation=build_anonymizer_config, + build_input_operation=build_input, + run_tags_operation=_run_tags, + measurement_session=configured_measurement_session, + record_metrics_operation=record_evaluation_metrics, ) - with configured_measurement_session(measurement): - result = anonymizer.run( - config=anonymizer_config, - data=input_data, - ) - if config.evaluate: - evaluated = anonymizer.evaluate(result) - record_evaluation_metrics( - evaluated.trace_dataframe, - mode="replace", - strategy=type(anonymizer_config.replace).__name__, - text_column=evaluated.resolved_text_column, - ) - - -def _get_item(items: dict[str, Any], item_id: str, item_type: str) -> Any: - if item_id not in items: - raise ValueError(f"unknown {item_type}: {item_id}") - return items[item_id] def dry_run_result( @@ -792,57 +513,26 @@ def run_or_plan( task_trace_dir: Path | None = None, wandb_settings: ResolvedWandbConfig | None = None, ) -> BenchmarkResult: - benchmark_spec = load_spec(spec_path) - output_dir = output or Path("benchmark-runs") / benchmark_spec.suite_id - resolved_wandb = wandb_settings or ResolvedWandbConfig() - if trace_dir is not None and dd_trace == DDTraceMode.none: - raise ValueError("--trace-dir requires --dd-trace") - if task_trace_dir is not None and not dd_task_trace: - raise ValueError("--task-trace-dir requires --dd-task-trace") - if dry_run: - return plan_suite( - benchmark_spec, - spec_path=spec_path, - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - dd_trace=dd_trace, - trace_dir=trace_dir, - dd_task_trace=dd_task_trace, - task_trace_dir=task_trace_dir, - ) - preflight_suite(benchmark_spec, spec_path=spec_path) - prepare_output_dir(output_dir, overwrite=overwrite, dry_run=dry_run) - result = run_suite( - benchmark_spec, - spec_path=spec_path, - output_dir=output_dir, + return _canonical_run_or_plan( + spec_path, + output=output, + overwrite=overwrite, + dry_run=dry_run, export=export, fail_fast=fail_fast, dd_trace=dd_trace, trace_dir=trace_dir, dd_task_trace=dd_task_trace, task_trace_dir=task_trace_dir, + wandb_settings=wandb_settings, + load_spec_operation=load_spec, + plan_suite_operation=plan_suite, + preflight_operation=preflight_suite, + prepare_output_operation=prepare_output_dir, + run_suite_operation=run_suite, + publisher_operation=publish_benchmark_wandb_best_effort, + metadata_builder=build_wandb_metadata, ) - publish_benchmark_wandb_best_effort( - resolved_wandb, - suite_id=benchmark_spec.suite_id, - output_dir=output_dir, - finalization=BenchmarkWandbFinalization( - measurement_path=Path(result.measurement_path), - cases=result.cases, - ), - metadata_factory=lambda: build_wandb_metadata( - benchmark_spec, - spec_path=spec_path, - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - dd_trace=dd_trace, - dd_task_trace=dd_task_trace, - ), - ) - return result if __name__ == "__main__": From c6cd6ae78136caecc764fd47221a988a9b159219 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 21:21:53 +0000 Subject: [PATCH 13/17] refactor(measurement): split sweep contracts and compiler Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 7 +- .../benchmark_sweep_compiler.py | 167 +++++++++++++ .../benchmark_sweep_models.py | 88 +++++++ tools/measurement/sweep_benchmarks.py | 222 ++---------------- 4 files changed, 276 insertions(+), 208 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_sweep_compiler.py create mode 100644 tools/measurement/measurement_tools/benchmark_sweep_models.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index acb006b9..56a45af1 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -67,7 +67,12 @@ def _load_module(path: Path, name: str) -> ModuleType: "Anonymizer", "measurement_tools.benchmark_models", ), - ("sweep_benchmarks.py", "SweepSpec", "WandbProjectPath", None), + ( + "sweep_benchmarks.py", + "SweepSpec", + "WandbProjectPath", + "measurement_tools.benchmark_sweep_models", + ), ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult", None), ], ) diff --git a/tools/measurement/measurement_tools/benchmark_sweep_compiler.py b/tools/measurement/measurement_tools/benchmark_sweep_compiler.py new file mode 100644 index 00000000..80c30167 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_sweep_compiler.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compile benchmark sweep specifications into materialized benchmark suites.""" + +from __future__ import annotations + +import copy +import itertools +from pathlib import Path +from typing import Any, cast + +import yaml + +from anonymizer.config.anonymizer_config import is_remote_input_source +from measurement_tools.benchmark_sweep_models import SweepArm, SweepSpec +from measurement_tools.wandb_models import SweepMetadata + + +def load_sweep_spec(path: Path) -> SweepSpec: + raw = read_yaml_mapping(path) + spec = SweepSpec.model_validate(raw) + values = spec.model_dump() + values["base_suite"] = str(resolve_path(spec.base_suite, path.parent)) + return SweepSpec.model_validate(values) + + +def expand_sweep_arms(spec: SweepSpec) -> list[SweepArm]: + names = list(spec.parameters) + value_grid = itertools.product(*(spec.parameters[name] for name in names)) + return [ + SweepArm(arm_id=f"arm-{index:03d}", parameters=dict(zip(names, values, strict=True))) + for index, values in enumerate(value_grid) + ] + + +def materialize_arm_suite(spec: SweepSpec, arm: SweepArm, *, output_root: Path, overwrite: bool) -> Path: + patched = arm_suite_payload(spec, arm) + arm_dir = output_root / arm.arm_id + arm_dir.mkdir(parents=True, exist_ok=True) + suite_path = arm_dir / "suite.yaml" + if suite_path.exists() and not overwrite: + raise ValueError(f"sweep arm suite already exists: {suite_path}") + suite_path.write_text(yaml.safe_dump(patched, sort_keys=False), encoding="utf-8") + return suite_path + + +def arm_suite_payload(spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: + base_path = Path(spec.base_suite) + suite = rebase_suite_paths(read_yaml_mapping(base_path), base_path.parent) + return patched_suite(suite, spec=spec, arm=arm) + + +def patched_suite(suite: dict[str, Any], *, spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: + patched = copy.deepcopy(suite) + run_tags = dict(patched.get("run_tags") or {}) + run_tags.update(sweep_run_tags(spec, arm)) + patched["run_tags"] = run_tags + for path, value in arm.parameters.items(): + apply_parameter(patched, path, value) + return patched + + +def sweep_run_tags(spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: + sweep = SweepMetadata.from_arm( + sweep_id=spec.sweep_id, + arm_id=arm.arm_id, + params={safe_param_name(path): value for path, value in arm.parameters.items()}, + ) + return {"wandb_sweep": sweep.model_dump(mode="json")} + + +def apply_parameter(suite: dict[str, Any], path: str, value: Any) -> None: + parts = path.split(".") + if len(parts) >= 3 and parts[0] == "configs": + apply_config_parameter(suite, parts[1], parts[2:], value) + return + set_nested(suite, parts, value) + + +def apply_config_parameter(suite: dict[str, Any], config_selector: str, parts: list[str], value: Any) -> None: + configs = suite.get("configs") + if not isinstance(configs, list): + raise ValueError("base suite must define a configs list") + matched = [config for config in configs if isinstance(config, dict) and config_matches(config, config_selector)] + if not matched: + raise ValueError(f"sweep parameter references unknown config selector: {config_selector}") + for config in matched: + set_nested(config, parts, value) + + +def config_matches(config: dict[str, Any], selector: str) -> bool: + return selector == "*" or config.get("id") == selector + + +def set_nested(target: dict[str, Any], parts: list[str], value: Any) -> None: + current = target + for part in parts[:-1]: + existing = current.get(part) + if isinstance(existing, str): + existing = {"strategy": existing} + current[part] = existing + if existing is None: + existing = {} + current[part] = existing + if not isinstance(existing, dict): + raise ValueError(f"cannot set nested sweep parameter through non-mapping field: {part}") + current = cast(dict[str, Any], existing) + current[parts[-1]] = value + + +def resolve_path(value: str, base_dir: Path) -> Path: + path = Path(value).expanduser() + return path.resolve() if path.is_absolute() else (base_dir / path).resolve() + + +def read_yaml_mapping(path: Path) -> dict[str, Any]: + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError(f"YAML file must contain a mapping: {path}") + return raw + + +def safe_param_name(path: str) -> str: + return path.replace("*", "all").replace(".", "_").replace("-", "_") + + +def rebase_suite_paths(suite: dict[str, Any], base_dir: Path) -> dict[str, Any]: + rebased = copy.deepcopy(suite) + for key in ("model_configs", "model_providers"): + if isinstance(rebased.get(key), str) and is_yaml_file_reference(rebased[key]): + rebased[key] = str(resolve_path(rebased[key], base_dir)) + if isinstance(rebased.get("artifact_path"), str): + rebased["artifact_path"] = str(resolve_path(rebased["artifact_path"], base_dir)) + for workload in rebased.get("workloads", []): + if isinstance(workload, dict) and isinstance(workload.get("source"), str): + workload["source"] = rebase_source(workload["source"], base_dir) + return rebased + + +def is_yaml_file_reference(value: str) -> bool: + return "\n" not in value and Path(value).suffix.lower() in {".yaml", ".yml"} + + +def rebase_source(source: str, base_dir: Path) -> str: + if is_remote_input_source(source): + return source + return str(resolve_path(source, base_dir)) + + +__all__ = [ + "apply_config_parameter", + "apply_parameter", + "arm_suite_payload", + "config_matches", + "expand_sweep_arms", + "is_yaml_file_reference", + "load_sweep_spec", + "materialize_arm_suite", + "patched_suite", + "read_yaml_mapping", + "rebase_source", + "rebase_suite_paths", + "resolve_path", + "safe_param_name", + "set_nested", + "sweep_run_tags", +] diff --git a/tools/measurement/measurement_tools/benchmark_sweep_models.py b/tools/measurement/measurement_tools/benchmark_sweep_models.py new file mode 100644 index 00000000..b1e15f50 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_sweep_models.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contracts for benchmark sweep specifications, arms, results, and CLI inputs.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from measurement_tools.wandb_settings import WandbMode + + +class SweepSpec(BaseModel): + model_config = ConfigDict(extra="forbid") + + sweep_id: str + base_suite: str + output_root: str | None = None + parameters: dict[str, list[Any]] = Field(min_length=1) + + @model_validator(mode="after") + def validate_parameters(self) -> "SweepSpec": + empty = [name for name, values in self.parameters.items() if not values] + if empty: + raise ValueError(f"sweep parameter(s) must have at least one value: {', '.join(sorted(empty))}") + return self + + +class SweepArm(BaseModel): + arm_id: str + parameters: dict[str, Any] + + +class SweepArmResult(BaseModel): + arm_id: str + parameters: dict[str, Any] + suite_path: str + output_dir: str + wandb_run_name: str + status: str + completed_cases: int = 0 + errored_cases: int = 0 + total_cases: int = 0 + error: str | None = None + + +class SweepResult(BaseModel): + sweep_id: str + output_root: str + arms: list[SweepArmResult] + report_url: str | None = None + workspace_url: str | None = None + report_error: str | None = None + workspace_error: str | None = None + + @property + def completed_arms(self) -> int: + return sum(1 for arm in self.arms if arm.status == "completed") + + @property + def errored_arms(self) -> int: + return sum(1 for arm in self.arms if arm.status == "error") + + +@dataclass(frozen=True) +class SweepCliOptions: + spec: Path + output_root: Path | None + overwrite: bool + dry_run: bool + export: bool + fail_fast: bool + create_report: bool + create_workspace: bool + wandb_mode: WandbMode | None + wandb_entity: str | None + wandb_project: str | None + wandb_base_url: str | None + wandb_group: str | None + wandb_job_type: str | None + wandb_tags: str | None + wandb_log_tables: bool | None + + +__all__ = ["SweepArm", "SweepArmResult", "SweepCliOptions", "SweepResult", "SweepSpec"] diff --git a/tools/measurement/sweep_benchmarks.py b/tools/measurement/sweep_benchmarks.py index be6c7c11..66daba0e 100644 --- a/tools/measurement/sweep_benchmarks.py +++ b/tools/measurement/sweep_benchmarks.py @@ -11,98 +11,32 @@ from __future__ import annotations -import copy -import itertools import logging import sys from dataclasses import dataclass from pathlib import Path -from typing import Annotated, Any, Callable, cast +from typing import Annotated, Any, Callable import cyclopts import run_benchmarks -import yaml from create_wandb_report import WandbProjectPath, create_benchmark_group_report, create_benchmark_workspace +from measurement_tools.benchmark_sweep_compiler import ( + arm_suite_payload, + expand_sweep_arms, + load_sweep_spec, + materialize_arm_suite, + rebase_suite_paths, + resolve_path, +) +from measurement_tools.benchmark_sweep_models import SweepArm, SweepArmResult, SweepCliOptions, SweepResult, SweepSpec from measurement_tools.cli import LogFormat, configure_logging, log_bad_input, summarize_validation_error -from measurement_tools.wandb_models import SweepMetadata, generated_wandb_tag -from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator +from measurement_tools.wandb_models import generated_wandb_tag +from pydantic import ValidationError app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.sweep") -class SweepSpec(BaseModel): - model_config = ConfigDict(extra="forbid") - - sweep_id: str - base_suite: str - output_root: str | None = None - parameters: dict[str, list[Any]] = Field(min_length=1) - - @model_validator(mode="after") - def validate_parameters(self) -> "SweepSpec": - empty = [name for name, values in self.parameters.items() if not values] - if empty: - raise ValueError(f"sweep parameter(s) must have at least one value: {', '.join(sorted(empty))}") - return self - - -class SweepArm(BaseModel): - arm_id: str - parameters: dict[str, Any] - - -class SweepArmResult(BaseModel): - arm_id: str - parameters: dict[str, Any] - suite_path: str - output_dir: str - wandb_run_name: str - status: str - completed_cases: int = 0 - errored_cases: int = 0 - total_cases: int = 0 - error: str | None = None - - -class SweepResult(BaseModel): - sweep_id: str - output_root: str - arms: list[SweepArmResult] - report_url: str | None = None - workspace_url: str | None = None - report_error: str | None = None - workspace_error: str | None = None - - @property - def completed_arms(self) -> int: - return sum(1 for arm in self.arms if arm.status == "completed") - - @property - def errored_arms(self) -> int: - return sum(1 for arm in self.arms if arm.status == "error") - - -@dataclass(frozen=True) -class SweepCliOptions: - spec: Path - output_root: Path | None - overwrite: bool - dry_run: bool - export: bool - fail_fast: bool - create_report: bool - create_workspace: bool - wandb_mode: run_benchmarks.WandbMode | None - wandb_entity: str | None - wandb_project: str | None - wandb_base_url: str | None - wandb_group: str | None - wandb_job_type: str | None - wandb_tags: str | None - wandb_log_tables: bool | None - - class WandbViewCreationError(RuntimeError): """A post-benchmark W&B presentation operation failed.""" @@ -115,96 +49,9 @@ class _WandbViewOperation: create: Callable[..., Any] -def load_sweep_spec(path: Path) -> SweepSpec: - raw = _read_yaml_mapping(path) - spec = SweepSpec.model_validate(raw) - values = spec.model_dump() - values["base_suite"] = str(_resolve_path(spec.base_suite, path.parent)) - return SweepSpec.model_validate(values) - - -def expand_sweep_arms(spec: SweepSpec) -> list[SweepArm]: - names = list(spec.parameters) - value_grid = itertools.product(*(spec.parameters[name] for name in names)) - return [ - SweepArm(arm_id=f"arm-{index:03d}", parameters=dict(zip(names, values, strict=True))) - for index, values in enumerate(value_grid) - ] - - -def materialize_arm_suite(spec: SweepSpec, arm: SweepArm, *, output_root: Path, overwrite: bool) -> Path: - patched = _arm_suite_payload(spec, arm) - arm_dir = output_root / arm.arm_id - arm_dir.mkdir(parents=True, exist_ok=True) - suite_path = arm_dir / "suite.yaml" - if suite_path.exists() and not overwrite: - raise ValueError(f"sweep arm suite already exists: {suite_path}") - suite_path.write_text(yaml.safe_dump(patched, sort_keys=False), encoding="utf-8") - return suite_path - - -def _arm_suite_payload(spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: - base_path = Path(spec.base_suite) - suite = _rebase_suite_paths(_read_yaml_mapping(base_path), base_path.parent) - return _patched_suite(suite, spec=spec, arm=arm) - - -def _patched_suite(suite: dict[str, Any], *, spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: - patched = copy.deepcopy(suite) - run_tags = dict(patched.get("run_tags") or {}) - run_tags.update(_sweep_run_tags(spec, arm)) - patched["run_tags"] = run_tags - for path, value in arm.parameters.items(): - _apply_parameter(patched, path, value) - return patched - - -def _sweep_run_tags(spec: SweepSpec, arm: SweepArm) -> dict[str, Any]: - sweep = SweepMetadata.from_arm( - sweep_id=spec.sweep_id, - arm_id=arm.arm_id, - params={_safe_param_name(path): value for path, value in arm.parameters.items()}, - ) - return {"wandb_sweep": sweep.model_dump(mode="json")} - - -def _apply_parameter(suite: dict[str, Any], path: str, value: Any) -> None: - parts = path.split(".") - if len(parts) >= 3 and parts[0] == "configs": - _apply_config_parameter(suite, parts[1], parts[2:], value) - return - _set_nested(suite, parts, value) - - -def _apply_config_parameter(suite: dict[str, Any], config_selector: str, parts: list[str], value: Any) -> None: - configs = suite.get("configs") - if not isinstance(configs, list): - raise ValueError("base suite must define a configs list") - matched = [config for config in configs if isinstance(config, dict) and _config_matches(config, config_selector)] - if not matched: - raise ValueError(f"sweep parameter references unknown config selector: {config_selector}") - for config in matched: - _set_nested(config, parts, value) - - -def _config_matches(config: dict[str, Any], selector: str) -> bool: - return selector == "*" or config.get("id") == selector - - -def _set_nested(target: dict[str, Any], parts: list[str], value: Any) -> None: - current = target - for part in parts[:-1]: - existing = current.get(part) - if isinstance(existing, str): - existing = {"strategy": existing} - current[part] = existing - if existing is None: - existing = {} - current[part] = existing - if not isinstance(existing, dict): - raise ValueError(f"cannot set nested sweep parameter through non-mapping field: {part}") - current = cast(dict[str, Any], existing) - current[parts[-1]] = value +_arm_suite_payload = arm_suite_payload +_rebase_suite_paths = rebase_suite_paths +_resolve_path = resolve_path def run_sweep( @@ -481,45 +328,6 @@ def _default_output_root(spec: SweepSpec, sweep_path: Path) -> Path: return sweep_path.with_suffix("").with_name(spec.sweep_id) -def _resolve_path(value: str, base_dir: Path) -> Path: - path = Path(value).expanduser() - return path.resolve() if path.is_absolute() else (base_dir / path).resolve() - - -def _read_yaml_mapping(path: Path) -> dict[str, Any]: - raw = yaml.safe_load(path.read_text(encoding="utf-8")) - if not isinstance(raw, dict): - raise ValueError(f"YAML file must contain a mapping: {path}") - return raw - - -def _safe_param_name(path: str) -> str: - return path.replace("*", "all").replace(".", "_").replace("-", "_") - - -def _rebase_suite_paths(suite: dict[str, Any], base_dir: Path) -> dict[str, Any]: - rebased = copy.deepcopy(suite) - for key in ("model_configs", "model_providers"): - if isinstance(rebased.get(key), str) and _is_yaml_file_reference(rebased[key]): - rebased[key] = str(_resolve_path(rebased[key], base_dir)) - if isinstance(rebased.get("artifact_path"), str): - rebased["artifact_path"] = str(_resolve_path(rebased["artifact_path"], base_dir)) - for workload in rebased.get("workloads", []): - if isinstance(workload, dict) and isinstance(workload.get("source"), str): - workload["source"] = _rebase_source(workload["source"], base_dir) - return rebased - - -def _is_yaml_file_reference(value: str) -> bool: - return "\n" not in value and Path(value).suffix.lower() in {".yaml", ".yml"} - - -def _rebase_source(source: str, base_dir: Path) -> str: - if run_benchmarks.is_remote_input_source(source): - return source - return str(_resolve_path(source, base_dir)) - - def render_result(result: SweepResult, *, json_output: bool) -> str: if json_output: return result.model_dump_json(indent=2) From 53405ee3b1ffa59444b8142b0c5fe6174f8a2a83 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 21:29:23 +0000 Subject: [PATCH 14/17] refactor(measurement): split sweep execution and views Signed-off-by: Aaron Gonzales --- .../benchmark_sweep_execution.py | 342 ++++++++++++++++++ .../benchmark_sweep_views.py | 52 +++ tools/measurement/sweep_benchmarks.py | 324 +++-------------- 3 files changed, 443 insertions(+), 275 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_sweep_execution.py create mode 100644 tools/measurement/measurement_tools/benchmark_sweep_views.py diff --git a/tools/measurement/measurement_tools/benchmark_sweep_execution.py b/tools/measurement/measurement_tools/benchmark_sweep_execution.py new file mode 100644 index 00000000..91e74310 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_sweep_execution.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plan and execute benchmark sweep arms and assemble their results.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Callable + +from measurement_tools.benchmark_models import BenchmarkResult, BenchmarkSpec, CaseStatus +from measurement_tools.benchmark_sweep_compiler import ( + arm_suite_payload, + expand_sweep_arms, + load_sweep_spec, + materialize_arm_suite, + resolve_path, +) +from measurement_tools.benchmark_sweep_models import SweepArm, SweepArmResult, SweepResult, SweepSpec +from measurement_tools.benchmark_sweep_views import ( + WandbViewCreationError, + WandbViewOperation, + create_view, +) +from measurement_tools.wandb_models import generated_wandb_tag +from measurement_tools.wandb_settings import ResolvedWandbConfig + + +def run_sweep( + sweep_path: Path, + *, + output_root: Path | None, + overwrite: bool, + dry_run: bool, + export: bool, + fail_fast: bool, + wandb_settings: ResolvedWandbConfig, + create_report: bool, + create_workspace: bool = False, + run_or_plan_operation: Callable[..., BenchmarkResult], + create_report_operation: Callable[..., Any], + create_workspace_operation: Callable[..., Any], + benchmark_spec_type: type[BenchmarkSpec], + build_cases_operation: Callable[..., list[Any]], + load_spec_operation: Callable[..., BenchmarkSpec], + plan_suite_operation: Callable[..., BenchmarkResult], +) -> SweepResult: + spec = load_sweep_spec(sweep_path) + resolved_output_root = output_root or default_output_root(spec, sweep_path) + arms = expand_sweep_arms(spec) + results = run_sweep_arms( + spec, + arms, + output_root=resolved_output_root, + overwrite=overwrite, + dry_run=dry_run, + export=export, + fail_fast=fail_fast, + wandb_settings=wandb_settings, + run_or_plan_operation=run_or_plan_operation, + benchmark_spec_type=benchmark_spec_type, + build_cases_operation=build_cases_operation, + load_spec_operation=load_spec_operation, + plan_suite_operation=plan_suite_operation, + ) + report_url = None + workspace_url = None + report_error = None + workspace_error = None + for operation in ( + WandbViewOperation(create_report, "report", "report_url", create_report_operation), + WandbViewOperation(create_workspace, "workspace", "workspace_url", create_workspace_operation), + ): + try: + url = create_view(spec, wandb_settings=wandb_settings, dry_run=dry_run, operation=operation) + except WandbViewCreationError as exc: + if operation.label == "report": + report_error = str(exc) + else: + workspace_error = str(exc) + else: + if operation.label == "report": + report_url = url + else: + workspace_url = url + return SweepResult( + sweep_id=spec.sweep_id, + output_root=str(resolved_output_root), + arms=results, + report_url=report_url, + workspace_url=workspace_url, + report_error=report_error, + workspace_error=workspace_error, + ) + + +def run_sweep_arms( + spec: SweepSpec, + arms: list[SweepArm], + *, + output_root: Path, + overwrite: bool, + dry_run: bool, + export: bool, + fail_fast: bool, + wandb_settings: ResolvedWandbConfig, + run_or_plan_operation: Callable[..., BenchmarkResult], + benchmark_spec_type: type[BenchmarkSpec], + build_cases_operation: Callable[..., list[Any]], + load_spec_operation: Callable[..., BenchmarkSpec], + plan_suite_operation: Callable[..., BenchmarkResult], +) -> list[SweepArmResult]: + results: list[SweepArmResult] = [] + for arm in arms: + result = run_arm( + spec, + arm, + output_root=output_root, + overwrite=overwrite, + dry_run=dry_run, + export=export, + fail_fast=fail_fast, + wandb_settings=wandb_settings, + run_or_plan_operation=run_or_plan_operation, + benchmark_spec_type=benchmark_spec_type, + build_cases_operation=build_cases_operation, + load_spec_operation=load_spec_operation, + plan_suite_operation=plan_suite_operation, + ) + results.append(result) + if fail_fast and result.status == "error": + break + return results + + +def run_arm( + spec: SweepSpec, + arm: SweepArm, + *, + output_root: Path, + overwrite: bool, + dry_run: bool, + export: bool, + fail_fast: bool, + wandb_settings: ResolvedWandbConfig, + run_or_plan_operation: Callable[..., BenchmarkResult], + benchmark_spec_type: type[BenchmarkSpec], + build_cases_operation: Callable[..., list[Any]], + load_spec_operation: Callable[..., BenchmarkSpec], + plan_suite_operation: Callable[..., BenchmarkResult], +) -> SweepArmResult: + suite_path = output_root / arm.arm_id / "suite.yaml" + output_dir = output_root / arm.arm_id / "output" + run_name = f"{spec.sweep_id}-{arm.arm_id}" + total_cases = 0 + try: + arm_wandb_settings_value = arm_wandb_settings(wandb_settings, spec=spec, arm=arm, run_name=run_name) + if dry_run: + total_cases, result = plan_arm( + spec, + arm, + suite_path=suite_path, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + benchmark_spec_type=benchmark_spec_type, + build_cases_operation=build_cases_operation, + plan_suite_operation=plan_suite_operation, + ) + else: + suite_path = materialize_arm_suite(spec, arm, output_root=output_root, overwrite=overwrite) + total_cases = planned_case_count( + suite_path, + build_cases_operation=build_cases_operation, + load_spec_operation=load_spec_operation, + ) + result = run_or_plan_operation( + suite_path, + output=output_dir, + overwrite=overwrite, + dry_run=False, + export=export, + fail_fast=fail_fast, + wandb_settings=arm_wandb_settings_value, + ) + except Exception as exc: # noqa: BLE001 -- keep sweeping other arms and report failure + return arm_error( + arm, + suite_path=suite_path, + output_dir=output_dir, + run_name=run_name, + total_cases=total_cases + or planned_case_count_if_readable( + suite_path, + build_cases_operation=build_cases_operation, + load_spec_operation=load_spec_operation, + ), + error=str(exc), + ) + return arm_result(arm, suite_path=suite_path, output_dir=output_dir, run_name=run_name, result=result) + + +def plan_arm( + spec: SweepSpec, + arm: SweepArm, + *, + suite_path: Path, + output_dir: Path, + export: bool, + fail_fast: bool, + benchmark_spec_type: type[BenchmarkSpec], + build_cases_operation: Callable[..., list[Any]], + plan_suite_operation: Callable[..., BenchmarkResult], +) -> tuple[int, BenchmarkResult]: + benchmark_spec = benchmark_spec_type.model_validate(arm_suite_payload(spec, arm)) + total_cases = len(build_cases_operation(benchmark_spec)) + result = plan_suite_operation( + benchmark_spec, + spec_path=suite_path, + output_dir=output_dir, + export=export, + fail_fast=fail_fast, + ) + return total_cases, result + + +def arm_wandb_settings( + settings: ResolvedWandbConfig, + *, + spec: SweepSpec, + arm: SweepArm, + run_name: str, +) -> ResolvedWandbConfig: + if not settings.enabled: + return settings + generated_tags = [ + tag + for tag in (generated_wandb_tag("sweep", spec.sweep_id), generated_wandb_tag("arm", arm.arm_id)) + if tag is not None + ] + return settings.validated_update( + wandb_group=settings.wandb_group or spec.sweep_id, + wandb_job_type=settings.wandb_job_type or "benchmark-sweep", + wandb_run_name=settings.wandb_run_name or run_name, + wandb_tags=joined_tags(settings.effective_wandb_tags, ["sweep", *generated_tags]), + ) + + +def arm_result( + arm: SweepArm, + *, + suite_path: Path, + output_dir: Path, + run_name: str, + result: BenchmarkResult, +) -> SweepArmResult: + errored = sum(1 for case in result.cases if case.status == CaseStatus.error) + completed = sum(1 for case in result.cases if case.status == CaseStatus.completed) + status = "error" if errored else "completed" + return SweepArmResult( + arm_id=arm.arm_id, + parameters=arm.parameters, + suite_path=str(suite_path), + output_dir=str(output_dir), + wandb_run_name=run_name, + status=status, + completed_cases=completed, + errored_cases=errored, + total_cases=len(result.cases), + ) + + +def arm_error( + arm: SweepArm, + *, + suite_path: Path, + output_dir: Path, + run_name: str, + total_cases: int, + error: str, +) -> SweepArmResult: + return SweepArmResult( + arm_id=arm.arm_id, + parameters=arm.parameters, + suite_path=str(suite_path), + output_dir=str(output_dir), + wandb_run_name=run_name, + status="error", + total_cases=total_cases, + error=error, + ) + + +def planned_case_count( + suite_path: Path, + *, + build_cases_operation: Callable[..., list[Any]], + load_spec_operation: Callable[..., BenchmarkSpec], +) -> int: + return len(build_cases_operation(load_spec_operation(suite_path))) + + +def planned_case_count_if_readable( + suite_path: Path, + *, + build_cases_operation: Callable[..., list[Any]], + load_spec_operation: Callable[..., BenchmarkSpec], +) -> int: + if not suite_path.exists(): + return 0 + try: + return planned_case_count( + suite_path, + build_cases_operation=build_cases_operation, + load_spec_operation=load_spec_operation, + ) + except Exception: # noqa: BLE001 -- keep the original per-arm failure as the reported error + return 0 + + +def joined_tags(first: list[str], second: list[str]) -> str: + return ",".join(dict.fromkeys([*first, *second])) + + +def default_output_root(spec: SweepSpec, sweep_path: Path) -> Path: + if spec.output_root: + return resolve_path(spec.output_root, sweep_path.parent) + return sweep_path.with_suffix("").with_name(spec.sweep_id) + + +__all__ = [ + "arm_error", + "arm_result", + "arm_wandb_settings", + "default_output_root", + "joined_tags", + "plan_arm", + "planned_case_count", + "planned_case_count_if_readable", + "run_arm", + "run_sweep", + "run_sweep_arms", +] diff --git a/tools/measurement/measurement_tools/benchmark_sweep_views.py b/tools/measurement/measurement_tools/benchmark_sweep_views.py new file mode 100644 index 00000000..564bac71 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_sweep_views.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W&B presentation operations for completed benchmark sweeps.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +from measurement_tools.benchmark_sweep_models import SweepSpec +from measurement_tools.wandb_report_models import WandbProjectPath +from measurement_tools.wandb_settings import ResolvedWandbConfig + + +class WandbViewCreationError(RuntimeError): + """A post-benchmark W&B presentation operation failed.""" + + +@dataclass(frozen=True) +class WandbViewOperation: + enabled: bool + label: str + url_attribute: str + create: Callable[..., Any] + + +def create_view( + spec: SweepSpec, + *, + wandb_settings: ResolvedWandbConfig, + dry_run: bool, + operation: WandbViewOperation, +) -> str | None: + if not operation.enabled or dry_run or not wandb_settings.enabled: + return None + entity = wandb_settings.wandb_entity + if not entity: + return None + try: + project = WandbProjectPath(entity=entity, project=wandb_settings.effective_wandb_project) + result = operation.create( + project, + settings=wandb_settings, + group=wandb_settings.wandb_group or spec.sweep_id, + expected_run_kind="sweep_arm", + ) + except Exception as exc: # noqa: BLE001 -- remote SDK errors must not expose response contents + raise WandbViewCreationError(f"W&B {operation.label} creation failed ({type(exc).__name__})") from None + return getattr(result, operation.url_attribute) + + +__all__ = ["WandbViewCreationError", "WandbViewOperation", "create_view"] diff --git a/tools/measurement/sweep_benchmarks.py b/tools/measurement/sweep_benchmarks.py index 66daba0e..14931b54 100644 --- a/tools/measurement/sweep_benchmarks.py +++ b/tools/measurement/sweep_benchmarks.py @@ -13,45 +13,61 @@ import logging import sys -from dataclasses import dataclass from pathlib import Path -from typing import Annotated, Any, Callable +from typing import Annotated, Any import cyclopts import run_benchmarks -from create_wandb_report import WandbProjectPath, create_benchmark_group_report, create_benchmark_workspace +from create_wandb_report import create_benchmark_group_report, create_benchmark_workspace +from measurement_tools import benchmark_sweep_execution +from measurement_tools.benchmark_sweep_compiler import ( + arm_suite_payload as arm_suite_payload, +) +from measurement_tools.benchmark_sweep_compiler import ( + expand_sweep_arms as expand_sweep_arms, +) +from measurement_tools.benchmark_sweep_compiler import ( + load_sweep_spec as load_sweep_spec, +) +from measurement_tools.benchmark_sweep_compiler import ( + materialize_arm_suite as materialize_arm_suite, +) from measurement_tools.benchmark_sweep_compiler import ( - arm_suite_payload, - expand_sweep_arms, - load_sweep_spec, - materialize_arm_suite, rebase_suite_paths, resolve_path, ) -from measurement_tools.benchmark_sweep_models import SweepArm, SweepArmResult, SweepCliOptions, SweepResult, SweepSpec +from measurement_tools.benchmark_sweep_models import ( + SweepArm as SweepArm, +) +from measurement_tools.benchmark_sweep_models import ( + SweepArmResult as SweepArmResult, +) +from measurement_tools.benchmark_sweep_models import ( + SweepCliOptions, + SweepResult, +) +from measurement_tools.benchmark_sweep_models import ( + SweepSpec as SweepSpec, +) +from measurement_tools.benchmark_sweep_views import ( + WandbViewCreationError as WandbViewCreationError, +) +from measurement_tools.benchmark_sweep_views import ( + WandbViewOperation, +) from measurement_tools.cli import LogFormat, configure_logging, log_bad_input, summarize_validation_error -from measurement_tools.wandb_models import generated_wandb_tag +from measurement_tools.wandb_report_models import WandbProjectPath as WandbProjectPath from pydantic import ValidationError app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.sweep") -class WandbViewCreationError(RuntimeError): - """A post-benchmark W&B presentation operation failed.""" - - -@dataclass(frozen=True) -class _WandbViewOperation: - enabled: bool - label: str - url_attribute: str - create: Callable[..., Any] - - _arm_suite_payload = arm_suite_payload _rebase_suite_paths = rebase_suite_paths _resolve_path = resolve_path +_WandbViewOperation = WandbViewOperation +_arm_wandb_settings = benchmark_sweep_execution.arm_wandb_settings def run_sweep( @@ -66,266 +82,24 @@ def run_sweep( create_report: bool, create_workspace: bool = False, ) -> SweepResult: - spec = load_sweep_spec(sweep_path) - resolved_output_root = output_root or _default_output_root(spec, sweep_path) - arms = expand_sweep_arms(spec) - results = _run_sweep_arms( - spec, - arms, - output_root=resolved_output_root, + return benchmark_sweep_execution.run_sweep( + sweep_path, + output_root=output_root, overwrite=overwrite, dry_run=dry_run, export=export, fail_fast=fail_fast, wandb_settings=wandb_settings, + create_report=create_report, + create_workspace=create_workspace, + run_or_plan_operation=run_benchmarks.run_or_plan, + create_report_operation=create_benchmark_group_report, + create_workspace_operation=create_benchmark_workspace, + benchmark_spec_type=run_benchmarks.BenchmarkSpec, + build_cases_operation=run_benchmarks.build_cases, + load_spec_operation=run_benchmarks.load_spec, + plan_suite_operation=run_benchmarks.plan_suite, ) - report_url = None - workspace_url = None - report_error = None - workspace_error = None - for operation in ( - _WandbViewOperation(create_report, "report", "report_url", create_benchmark_group_report), - _WandbViewOperation(create_workspace, "workspace", "workspace_url", create_benchmark_workspace), - ): - try: - url = _maybe_create_view(spec, wandb_settings=wandb_settings, dry_run=dry_run, operation=operation) - except WandbViewCreationError as exc: - if operation.label == "report": - report_error = str(exc) - else: - workspace_error = str(exc) - else: - if operation.label == "report": - report_url = url - else: - workspace_url = url - return SweepResult( - sweep_id=spec.sweep_id, - output_root=str(resolved_output_root), - arms=results, - report_url=report_url, - workspace_url=workspace_url, - report_error=report_error, - workspace_error=workspace_error, - ) - - -def _run_sweep_arms( - spec: SweepSpec, - arms: list[SweepArm], - *, - output_root: Path, - overwrite: bool, - dry_run: bool, - export: bool, - fail_fast: bool, - wandb_settings: run_benchmarks.ResolvedWandbConfig, -) -> list[SweepArmResult]: - results: list[SweepArmResult] = [] - for arm in arms: - result = _run_arm( - spec, - arm, - output_root=output_root, - overwrite=overwrite, - dry_run=dry_run, - export=export, - fail_fast=fail_fast, - wandb_settings=wandb_settings, - ) - results.append(result) - if fail_fast and result.status == "error": - break - return results - - -def _run_arm( - spec: SweepSpec, - arm: SweepArm, - *, - output_root: Path, - overwrite: bool, - dry_run: bool, - export: bool, - fail_fast: bool, - wandb_settings: run_benchmarks.ResolvedWandbConfig, -) -> SweepArmResult: - suite_path = output_root / arm.arm_id / "suite.yaml" - output_dir = output_root / arm.arm_id / "output" - run_name = f"{spec.sweep_id}-{arm.arm_id}" - total_cases = 0 - try: - arm_wandb_settings = _arm_wandb_settings(wandb_settings, spec=spec, arm=arm, run_name=run_name) - if dry_run: - total_cases, result = _plan_arm( - spec, - arm, - suite_path=suite_path, - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - ) - else: - suite_path = materialize_arm_suite(spec, arm, output_root=output_root, overwrite=overwrite) - total_cases = _planned_case_count(suite_path) - result = run_benchmarks.run_or_plan( - suite_path, - output=output_dir, - overwrite=overwrite, - dry_run=False, - export=export, - fail_fast=fail_fast, - wandb_settings=arm_wandb_settings, - ) - except Exception as exc: # noqa: BLE001 -- keep sweeping other arms and report failure - return _arm_error( - arm, - suite_path=suite_path, - output_dir=output_dir, - run_name=run_name, - total_cases=total_cases or _planned_case_count_if_readable(suite_path), - error=str(exc), - ) - return _arm_result(arm, suite_path=suite_path, output_dir=output_dir, run_name=run_name, result=result) - - -def _plan_arm( - spec: SweepSpec, - arm: SweepArm, - *, - suite_path: Path, - output_dir: Path, - export: bool, - fail_fast: bool, -) -> tuple[int, run_benchmarks.BenchmarkResult]: - benchmark_spec = run_benchmarks.BenchmarkSpec.model_validate(_arm_suite_payload(spec, arm)) - total_cases = len(run_benchmarks.build_cases(benchmark_spec)) - result = run_benchmarks.plan_suite( - benchmark_spec, - spec_path=suite_path, - output_dir=output_dir, - export=export, - fail_fast=fail_fast, - ) - return total_cases, result - - -def _arm_wandb_settings( - settings: run_benchmarks.ResolvedWandbConfig, - *, - spec: SweepSpec, - arm: SweepArm, - run_name: str, -) -> run_benchmarks.ResolvedWandbConfig: - if not settings.enabled: - return settings - generated_tags = [ - tag - for tag in (generated_wandb_tag("sweep", spec.sweep_id), generated_wandb_tag("arm", arm.arm_id)) - if tag is not None - ] - return settings.validated_update( - wandb_group=settings.wandb_group or spec.sweep_id, - wandb_job_type=settings.wandb_job_type or "benchmark-sweep", - wandb_run_name=settings.wandb_run_name or run_name, - wandb_tags=_joined_tags(settings.effective_wandb_tags, ["sweep", *generated_tags]), - ) - - -def _arm_result( - arm: SweepArm, - *, - suite_path: Path, - output_dir: Path, - run_name: str, - result: run_benchmarks.BenchmarkResult, -) -> SweepArmResult: - errored = sum(1 for case in result.cases if case.status == run_benchmarks.CaseStatus.error) - completed = sum(1 for case in result.cases if case.status == run_benchmarks.CaseStatus.completed) - status = "error" if errored else "completed" - return SweepArmResult( - arm_id=arm.arm_id, - parameters=arm.parameters, - suite_path=str(suite_path), - output_dir=str(output_dir), - wandb_run_name=run_name, - status=status, - completed_cases=completed, - errored_cases=errored, - total_cases=len(result.cases), - ) - - -def _arm_error( - arm: SweepArm, - *, - suite_path: Path, - output_dir: Path, - run_name: str, - total_cases: int, - error: str, -) -> SweepArmResult: - return SweepArmResult( - arm_id=arm.arm_id, - parameters=arm.parameters, - suite_path=str(suite_path), - output_dir=str(output_dir), - wandb_run_name=run_name, - status="error", - total_cases=total_cases, - error=error, - ) - - -def _planned_case_count(suite_path: Path) -> int: - return len(run_benchmarks.build_cases(run_benchmarks.load_spec(suite_path))) - - -def _planned_case_count_if_readable(suite_path: Path) -> int: - if not suite_path.exists(): - return 0 - try: - return _planned_case_count(suite_path) - except Exception: # noqa: BLE001 -- keep the original per-arm failure as the reported error - return 0 - - -def _maybe_create_view( - spec: SweepSpec, - *, - wandb_settings: run_benchmarks.ResolvedWandbConfig, - dry_run: bool, - operation: _WandbViewOperation, -) -> str | None: - if not operation.enabled or dry_run or not wandb_settings.enabled: - return None - entity = wandb_settings.wandb_entity - if not entity: - return None - try: - project = WandbProjectPath( - entity=entity, - project=wandb_settings.effective_wandb_project, - ) - result = operation.create( - project, - settings=wandb_settings, - group=wandb_settings.wandb_group or spec.sweep_id, - expected_run_kind="sweep_arm", - ) - except Exception as exc: # noqa: BLE001 -- remote SDK errors must not expose response contents - raise WandbViewCreationError(f"W&B {operation.label} creation failed ({type(exc).__name__})") from None - return getattr(result, operation.url_attribute) - - -def _joined_tags(first: list[str], second: list[str]) -> str: - return ",".join(dict.fromkeys([*first, *second])) - - -def _default_output_root(spec: SweepSpec, sweep_path: Path) -> Path: - if spec.output_root: - return _resolve_path(spec.output_root, sweep_path.parent) - return sweep_path.with_suffix("").with_name(spec.sweep_id) def render_result(result: SweepResult, *, json_output: bool) -> str: From 831a4e05f4cf817a23abc96fccc357f39fe91f52 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 21:55:10 +0000 Subject: [PATCH 15/17] refactor(measurement): split analysis contracts and math Signed-off-by: Aaron Gonzales --- .../test_measurement_module_compatibility.py | 7 +- tools/measurement/analyze_benchmark_output.py | 557 +++--------------- .../benchmark_analysis_math.py | 159 +++++ .../benchmark_analysis_models.py | 326 ++++++++++ 4 files changed, 569 insertions(+), 480 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_analysis_math.py create mode 100644 tools/measurement/measurement_tools/benchmark_analysis_models.py diff --git a/tests/tools/test_measurement_module_compatibility.py b/tests/tools/test_measurement_module_compatibility.py index 56a45af1..fda5a027 100644 --- a/tests/tools/test_measurement_module_compatibility.py +++ b/tests/tools/test_measurement_module_compatibility.py @@ -73,7 +73,12 @@ def _load_module(path: Path, name: str) -> ModuleType: "WandbProjectPath", "measurement_tools.benchmark_sweep_models", ), - ("analyze_benchmark_output.py", "BenchmarkOutputAnalysis", "AnalysisExportResult", None), + ( + "analyze_benchmark_output.py", + "BenchmarkOutputAnalysis", + "AnalysisExportResult", + "measurement_tools.benchmark_analysis_models", + ), ], ) def test_exact_tool_paths_support_repeated_dynamic_loading( diff --git a/tools/measurement/analyze_benchmark_output.py b/tools/measurement/analyze_benchmark_output.py index 434cef6f..6735c675 100644 --- a/tools/measurement/analyze_benchmark_output.py +++ b/tools/measurement/analyze_benchmark_output.py @@ -16,12 +16,89 @@ import logging import math import sys -from dataclasses import dataclass from pathlib import Path from typing import Annotated, Any, cast import cyclopts import pandas as pd +from measurement_tools.benchmark_analysis_math import ( + coalesce_number as _coalesce_number, +) +from measurement_tools.benchmark_analysis_math import ( + f1 as _f1, +) +from measurement_tools.benchmark_analysis_math import ( + first_float as _first_float, +) +from measurement_tools.benchmark_analysis_math import ( + first_int as _first_int, +) +from measurement_tools.benchmark_analysis_math import ( + first_value as _first_value, +) +from measurement_tools.benchmark_analysis_math import ( + model_workflow_rows as _model_workflow_rows, +) +from measurement_tools.benchmark_analysis_math import ( + non_null_count as _non_null_count, +) +from measurement_tools.benchmark_analysis_math import ( + pipeline_stage_rows as _pipeline_stage_rows, +) +from measurement_tools.benchmark_analysis_math import ( + positive_count as _positive_count, +) +from measurement_tools.benchmark_analysis_math import ( + records_of_type as _records_of_type, +) +from measurement_tools.benchmark_analysis_math import ( + request_failure_rate as _request_failure_rate, +) +from measurement_tools.benchmark_analysis_math import ( + safe_rate as _safe_rate, +) +from measurement_tools.benchmark_analysis_math import ( + safe_ratio as _safe_ratio, +) +from measurement_tools.benchmark_analysis_math import ( + sum_bool_or_zero as _sum_bool_or_zero, +) +from measurement_tools.benchmark_analysis_math import ( + sum_int_or_none as _sum_int_or_none, +) +from measurement_tools.benchmark_analysis_math import ( + sum_optional_numbers as _sum_optional_numbers, +) +from measurement_tools.benchmark_analysis_math import ( + sum_prefixed_ints as _sum_prefixed_ints, +) +from measurement_tools.benchmark_analysis_math import ( + zero_count as _zero_count, +) +from measurement_tools.benchmark_analysis_math import ( + zero_with_positive_count as _zero_with_positive_count, +) +from measurement_tools.benchmark_analysis_models import ( + _EVALUATION_ROLLUPS, +) +from measurement_tools.benchmark_analysis_models import ( + BenchmarkOutputAnalysis as BenchmarkOutputAnalysis, +) +from measurement_tools.benchmark_analysis_models import ( + CaseAnalysisRow as CaseAnalysisRow, +) +from measurement_tools.benchmark_analysis_models import ( + GroupAnalysisRow as GroupAnalysisRow, +) +from measurement_tools.benchmark_analysis_models import ( + ModelUsageAnalysisRow as ModelUsageAnalysisRow, +) +from measurement_tools.benchmark_analysis_models import ( + ModelUsageGroupAnalysisRow as ModelUsageGroupAnalysisRow, +) +from measurement_tools.benchmark_analysis_models import ( + _EvaluationRollup as _EvaluationRollup, +) from measurement_tools.cli import LogFormat, configure_logging, log_bad_input from measurement_tools.stats import median_or_none as _median_or_none from measurement_tools.stats import none_if_nan as _none_if_nan @@ -30,7 +107,6 @@ from measurement_tools.stats import sum_or_zero as _sum_or_zero from measurement_tools.tables import AnalysisExportResult, ExportFormat, ModelTableSpec from measurement_tools.tables import write_analysis_tables as _write_analysis_table_specs -from pydantic import BaseModel, Field, computed_field app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.benchmark_output") @@ -46,319 +122,6 @@ } -@dataclass(frozen=True) -class _EvaluationRollup: - prefix: str - valid_column: str - invalid_count_column: str - - -_EVALUATION_ROLLUPS = ( - _EvaluationRollup("detection", "detection_valid", "detection_invalid_entity_count"), - _EvaluationRollup("type_fidelity", "type_fidelity_valid", "type_fidelity_invalid_replacement_count"), - _EvaluationRollup( - "relational_consistency", - "relational_consistency_valid", - "relational_consistency_invalid_relation_count", - ), - _EvaluationRollup("attribute_fidelity", "attribute_fidelity_valid", "attribute_fidelity_invalid_entity_count"), -) - - -class CaseAnalysisRow(BaseModel): - suite_id: str | None = None - workload_id: str | None = None - workload_category: str | None = None - config_id: str | None = None - experimental_detection_strategy: str | None = None - experimental_replacement_strategy: str | None = None - dd_parser_compat: str | None = None - entity_label_set_id: str | None = None - entity_label_count: int | None = None - gliner_threshold: float | None = None - repetition: int | None = None - case_id: str - run_id: str - case_failed: bool = False - error_stage_count: int = 0 - error_ndd_workflow_count: int = 0 - error_model_workflow_count: int = 0 - pipeline_elapsed_sec: float | None = None - ndd_workflow_count: int = 0 - ndd_elapsed_sec_total: float = 0.0 - observed_total_requests: int = 0 - observed_successful_requests: int = 0 - observed_input_tokens: int = 0 - observed_output_tokens: int = 0 - observed_total_tokens: int = 0 - observed_failed_requests: int = 0 - observed_failed_request_rate: float | None = None - dd_trace_record_count: int = 0 - dd_trace_error_count: int = 0 - dd_trace_sync_client_unavailable_count: int = 0 - observed_bridge_fallback_requests: int | None = None - observed_non_bridge_total_requests: int | None = None - observed_non_bridge_failed_requests: int | None = None - observed_non_bridge_failed_request_rate: float | None = None - record_count: int = 0 - input_text_tokens_total: int | None = None - records_per_pipeline_sec: float | None = None - records_per_ndd_sec: float | None = None - input_text_tokens_per_pipeline_sec: float | None = None - input_text_tokens_per_ndd_sec: float | None = None - topology_endpoint_count: float | None = None - topology_gpu_count: float | None = None - topology_tensor_parallelism: float | None = None - topology_shard_count: float | None = None - input_text_tokens_per_endpoint_sec: float | None = None - input_text_tokens_per_gpu_sec: float | None = None - final_entity_count: float | None = None - empty_detection_count: int = 0 - empty_detection_rate: float | None = None - empty_detection_with_ground_truth_count: int = 0 - empty_detection_with_ground_truth_rate: float | None = None - ground_truth_record_count: int = 0 - ground_truth_entity_count: float | None = None - entity_true_positive_count: float | None = None - entity_false_positive_count: float | None = None - entity_false_negative_count: float | None = None - entity_precision: float | None = None - entity_recall: float | None = None - entity_f1: float | None = None - entity_relaxed_gt_found_count: float | None = None - entity_relaxed_detected_tp_count: float | None = None - entity_relaxed_label_compatible_gt_found_count: float | None = None - entity_relaxed_label_compatible_detected_tp_count: float | None = None - entity_relaxed_precision: float | None = None - entity_relaxed_recall: float | None = None - entity_relaxed_f1: float | None = None - entity_relaxed_label_compatible_precision: float | None = None - entity_relaxed_label_compatible_recall: float | None = None - entity_relaxed_label_compatible_f1: float | None = None - replacement_count: float | None = None - replacement_missing_final_entity_count: float | None = None - replacement_missing_final_entity_label_counts: dict[str, int] = Field(default_factory=dict) - replacement_missing_final_value_count: float | None = None - replacement_synthetic_original_collision_count: float | None = None - replacement_synthetic_original_collision_label_counts: dict[str, int] = Field(default_factory=dict) - replacement_synthetic_original_collision_value_count: float | None = None - original_value_leak_count: float | None = None - original_value_leak_record_count: int = 0 - original_value_leak_label_counts: dict[str, int] = Field(default_factory=dict) - detection_judged_record_count: int = 0 - detection_valid_record_count: int = 0 - detection_valid_rate: float | None = None - detection_invalid_entity_count: int = 0 - type_fidelity_judged_record_count: int = 0 - type_fidelity_valid_record_count: int = 0 - type_fidelity_valid_rate: float | None = None - type_fidelity_invalid_replacement_count: int = 0 - relational_consistency_judged_record_count: int = 0 - relational_consistency_valid_record_count: int = 0 - relational_consistency_valid_rate: float | None = None - relational_consistency_invalid_relation_count: int = 0 - attribute_fidelity_judged_record_count: int = 0 - attribute_fidelity_valid_record_count: int = 0 - attribute_fidelity_valid_rate: float | None = None - attribute_fidelity_invalid_entity_count: int = 0 - validation_max_entities_per_call: int | None = None - detection_artifact_rows: int = 0 - seed_entity_count: float | None = None - seed_validation_candidate_count: float | None = None - estimated_seed_validation_chunk_count: float | None = None - augmented_entity_count: float | None = None - augmented_new_final_value_count: float | None = None - artifact_final_entity_count: float | None = None - artifact_final_detector_entity_count: float | None = None - artifact_final_augmenter_entity_count: float | None = None - artifact_final_entity_signature_count: float | None = None - artifact_final_entity_signature_hashes: list[str] = Field(default_factory=list) - artifact_final_entity_signature_labels: dict[str, str] = Field(default_factory=dict) - artifact_final_entity_signature_details: dict[str, dict[str, Any]] = Field(default_factory=dict) - - -class GroupAnalysisRow(BaseModel): - workload_id: str | None = None - workload_category: str | None = None - config_id: str | None = None - experimental_detection_strategy: str | None = None - experimental_replacement_strategy: str | None = None - entity_label_set_id: str | None = None - entity_label_count: int | None = None - gliner_threshold: float | None = None - case_count: int - failed_case_count: int = 0 - failed_case_rate: float | None = None - error_stage_count: int = 0 - error_ndd_workflow_count: int = 0 - error_model_workflow_count: int = 0 - median_pipeline_elapsed_sec: float | None = None - median_ndd_elapsed_sec_total: float | None = None - median_observed_total_requests: float | None = None - median_observed_successful_requests: float | None = None - median_observed_input_tokens: float | None = None - median_observed_output_tokens: float | None = None - median_observed_total_tokens: float | None = None - median_observed_failed_requests: float | None = None - median_observed_failed_request_rate: float | None = None - median_observed_bridge_fallback_requests: float | None = None - median_observed_non_bridge_total_requests: float | None = None - median_observed_non_bridge_failed_requests: float | None = None - median_observed_non_bridge_failed_request_rate: float | None = None - total_record_count: int = 0 - median_record_count: float | None = None - total_input_text_tokens: int | None = None - median_input_text_tokens_total: float | None = None - median_records_per_pipeline_sec: float | None = None - median_records_per_ndd_sec: float | None = None - median_input_text_tokens_per_pipeline_sec: float | None = None - median_input_text_tokens_per_ndd_sec: float | None = None - median_topology_endpoint_count: float | None = None - median_topology_gpu_count: float | None = None - median_topology_tensor_parallelism: float | None = None - median_topology_shard_count: float | None = None - median_input_text_tokens_per_endpoint_sec: float | None = None - median_input_text_tokens_per_gpu_sec: float | None = None - median_final_entity_count: float | None = None - total_empty_detection_count: int = 0 - empty_detection_rate: float | None = None - total_empty_detection_with_ground_truth_count: int = 0 - empty_detection_with_ground_truth_rate: float | None = None - total_ground_truth_record_count: int = 0 - sum_ground_truth_entity_count: float | None = None - sum_entity_true_positive_count: float | None = None - sum_entity_false_positive_count: float | None = None - sum_entity_false_negative_count: float | None = None - micro_entity_precision: float | None = None - micro_entity_recall: float | None = None - micro_entity_f1: float | None = None - sum_entity_relaxed_gt_found_count: float | None = None - sum_entity_relaxed_detected_tp_count: float | None = None - sum_entity_relaxed_label_compatible_gt_found_count: float | None = None - sum_entity_relaxed_label_compatible_detected_tp_count: float | None = None - micro_entity_relaxed_precision: float | None = None - micro_entity_relaxed_recall: float | None = None - micro_entity_relaxed_f1: float | None = None - micro_entity_relaxed_label_compatible_precision: float | None = None - micro_entity_relaxed_label_compatible_recall: float | None = None - micro_entity_relaxed_label_compatible_f1: float | None = None - median_entity_relaxed_f1: float | None = None - median_entity_relaxed_label_compatible_f1: float | None = None - median_replacement_missing_final_entity_count: float | None = None - median_replacement_missing_final_value_count: float | None = None - replacement_missing_final_entity_label_counts: dict[str, int] = Field(default_factory=dict) - median_replacement_synthetic_original_collision_count: float | None = None - median_replacement_synthetic_original_collision_value_count: float | None = None - replacement_synthetic_original_collision_label_counts: dict[str, int] = Field(default_factory=dict) - sum_original_value_leak_count: float | None = None - leaking_case_count: int = 0 - median_original_value_leak_count: float | None = None - sum_detection_judged_record_count: int = 0 - sum_detection_valid_record_count: int = 0 - micro_detection_valid_rate: float | None = None - sum_detection_invalid_entity_count: int = 0 - sum_type_fidelity_judged_record_count: int = 0 - sum_type_fidelity_valid_record_count: int = 0 - micro_type_fidelity_valid_rate: float | None = None - sum_type_fidelity_invalid_replacement_count: int = 0 - sum_relational_consistency_judged_record_count: int = 0 - sum_relational_consistency_valid_record_count: int = 0 - micro_relational_consistency_valid_rate: float | None = None - sum_relational_consistency_invalid_relation_count: int = 0 - sum_attribute_fidelity_judged_record_count: int = 0 - sum_attribute_fidelity_valid_record_count: int = 0 - micro_attribute_fidelity_valid_rate: float | None = None - sum_attribute_fidelity_invalid_entity_count: int = 0 - median_seed_entity_count: float | None = None - median_seed_validation_candidate_count: float | None = None - median_estimated_seed_validation_chunk_count: float | None = None - median_augmented_entity_count: float | None = None - median_augmented_new_final_value_count: float | None = None - median_artifact_final_entity_count: float | None = None - median_artifact_final_detector_entity_count: float | None = None - median_artifact_final_augmenter_entity_count: float | None = None - median_artifact_final_entity_signature_count: float | None = None - - -class ModelUsageAnalysisRow(BaseModel): - suite_id: str | None = None - workload_id: str | None = None - config_id: str | None = None - experimental_detection_strategy: str | None = None - experimental_replacement_strategy: str | None = None - dd_parser_compat: str | None = None - repetition: int | None = None - case_id: str - run_id: str - workflow_name: str | None = None - model_alias: str | None = None - model_name: str - model_provider_name: str | None = None - ndd_elapsed_sec: float | None = None - observed_total_requests: int = 0 - observed_successful_requests: int = 0 - observed_failed_requests: int = 0 - observed_input_tokens: int = 0 - observed_output_tokens: int = 0 - observed_total_tokens: int = 0 - observed_reasoning_tokens: int | None = None - observed_failed_request_rate: float | None = None - - -class ModelUsageGroupAnalysisRow(BaseModel): - workload_id: str | None = None - config_id: str | None = None - experimental_detection_strategy: str | None = None - experimental_replacement_strategy: str | None = None - dd_parser_compat: str | None = None - workflow_name: str | None = None - model_alias: str | None = None - model_name: str - model_provider_name: str | None = None - case_count: int - workflow_count: int - sum_observed_total_requests: int = 0 - sum_observed_successful_requests: int = 0 - sum_observed_failed_requests: int = 0 - sum_observed_input_tokens: int = 0 - sum_observed_output_tokens: int = 0 - sum_observed_total_tokens: int = 0 - sum_observed_reasoning_tokens: int | None = None - observed_failed_request_rate: float | None = None - median_observed_total_requests: float | None = None - median_observed_failed_requests: float | None = None - median_observed_total_tokens: float | None = None - - -class BenchmarkOutputAnalysis(BaseModel): - benchmark_dir: str - detection_artifacts_path: str | None = None - cases: list[CaseAnalysisRow] = Field(default_factory=list) - groups: list[GroupAnalysisRow] = Field(default_factory=list) - model_usage: list[ModelUsageAnalysisRow] = Field(default_factory=list) - model_usage_groups: list[ModelUsageGroupAnalysisRow] = Field(default_factory=list) - - @computed_field - @property - def case_count(self) -> int: - return len(self.cases) - - @computed_field - @property - def group_count(self) -> int: - return len(self.groups) - - @computed_field - @property - def model_usage_count(self) -> int: - return len(self.model_usage) - - @computed_field - @property - def model_usage_group_count(self) -> int: - return len(self.model_usage_groups) - - def analyze_benchmark_output( benchmark_dir: Path, *, @@ -823,29 +586,6 @@ def _rows_for_case(dataframe: pd.DataFrame, case_id: str) -> pd.DataFrame: return dataframe[mask] -def _records_of_type(dataframe: pd.DataFrame, record_type: str) -> pd.DataFrame: - if "record_type" not in dataframe.columns: - return dataframe.iloc[0:0] - return dataframe[dataframe["record_type"] == record_type] - - -def _records_of_types(dataframe: pd.DataFrame, record_types: set[str]) -> pd.DataFrame: - if "record_type" not in dataframe.columns: - return dataframe.iloc[0:0] - return dataframe[dataframe["record_type"].isin(record_types)] - - -def _model_workflow_rows(dataframe: pd.DataFrame) -> pd.DataFrame: - return _records_of_types(dataframe, {"ndd_workflow", "model_workflow"}) - - -def _pipeline_stage_rows(dataframe: pd.DataFrame) -> pd.DataFrame: - stages = _records_of_type(dataframe, "stage") - if "stage" not in stages.columns: - return stages.iloc[0:0] - return stages[stages["stage"] == "Anonymizer._run_internal"] - - _MODEL_USAGE_SUFFIXES = { ".request_usage.total_requests": "observed_total_requests", ".request_usage.successful_requests": "observed_successful_requests", @@ -988,34 +728,6 @@ def _coerce_int(value: Any) -> int: return int(float(value)) -def _first_value(frames: list[pd.DataFrame], columns: list[str]) -> str | None: - for frame in frames: - for column in columns: - if column not in frame.columns: - continue - values = frame[column].dropna() - if not values.empty: - return str(values.iloc[0]) - return None - - -def _first_int(frames: list[pd.DataFrame], columns: list[str]) -> int | None: - value = _first_value(frames, columns) - return int(float(value)) if value is not None else None - - -def _first_float(frames: list[pd.DataFrame], columns: list[str]) -> float | None: - value = _first_value(frames, columns) - return float(value) if value is not None else None - - -def _coalesce_number(*values: float | None) -> float | None: - for value in values: - if value is not None: - return value - return None - - def _artifact_signature_hashes(artifact_rows: pd.DataFrame) -> list[str]: if "final_entity_signature_hashes" not in artifact_rows.columns: return [] @@ -1104,65 +816,6 @@ def _signature_count(artifact_rows: pd.DataFrame, *, signature_hashes: list[str] return _sum_or_none(artifact_rows, "final_entity_signature_count") -def _positive_count(dataframe: pd.DataFrame, column: str) -> int: - if column not in dataframe.columns: - return 0 - values = pd.to_numeric(dataframe[column], errors="coerce").fillna(0) - return int((values > 0).sum()) - - -def _zero_count(dataframe: pd.DataFrame, column: str) -> int: - if column not in dataframe.columns: - return 0 - values = pd.to_numeric(dataframe[column], errors="coerce").dropna() - return int((values == 0).sum()) - - -def _non_null_count(dataframe: pd.DataFrame, column: str) -> int: - if column not in dataframe.columns: - return 0 - return int(pd.to_numeric(dataframe[column], errors="coerce").notna().sum()) - - -def _zero_with_positive_count(dataframe: pd.DataFrame, *, zero_column: str, positive_column: str) -> int: - if zero_column not in dataframe.columns or positive_column not in dataframe.columns: - return 0 - zero_values = pd.to_numeric(dataframe[zero_column], errors="coerce") - positive_values = pd.to_numeric(dataframe[positive_column], errors="coerce") - return int(((zero_values == 0) & (positive_values > 0)).sum()) - - -def _sum_prefixed_ints(dataframe: pd.DataFrame, prefix: str) -> dict[str, int]: - totals: dict[str, int] = {} - base_column = prefix.removesuffix(".") - if base_column in dataframe.columns: - for value in dataframe[base_column].dropna().tolist(): - for key, count in _coerce_count_mapping(value).items(): - totals[key] = totals.get(key, 0) + count - for column in sorted(col for col in dataframe.columns if col.startswith(prefix)): - value = _sum_int_or_zero(dataframe, column) - if value: - totals[column.removeprefix(prefix)] = value - return totals - - -def _coerce_count_mapping(value: object) -> dict[str, int]: - payload = value - if isinstance(value, str): - try: - payload = json.loads(value) - except json.JSONDecodeError: - return {} - if not isinstance(payload, dict): - return {} - counts: dict[str, int] = {} - for key, count in payload.items(): - numeric = pd.to_numeric(pd.Series([count]), errors="coerce").dropna() - if not numeric.empty and numeric.iloc[0]: - counts[str(key)] = int(numeric.iloc[0]) - return counts - - def build_group_rows(cases: list[CaseAnalysisRow]) -> list[GroupAnalysisRow]: if not cases: return [] @@ -1404,12 +1057,6 @@ def _float_if_not_nan(value: object) -> float | None: return float(cast(Any, value)) -def _sum_bool_or_zero(dataframe: pd.DataFrame, column: str) -> int: - if column not in dataframe.columns: - return 0 - return int(dataframe[column].fillna(False).astype(bool).sum()) - - def _group_evaluation_metrics(group: pd.DataFrame) -> dict[str, int | float | None]: metrics: dict[str, int | float | None] = {} for rollup in _EVALUATION_ROLLUPS: @@ -1422,54 +1069,6 @@ def _group_evaluation_metrics(group: pd.DataFrame) -> dict[str, int | float | No return metrics -def _sum_int_or_none(dataframe: pd.DataFrame, column: str) -> int | None: - value = _sum_or_none(dataframe, column) - return int(value) if value is not None else None - - -def _request_failure_rate(*, failed: object, total: object) -> float | None: - failed_value = _optional_number(failed) - total_value = _optional_number(total) - if failed_value is None or total_value is None or total_value <= 0: - return None - return failed_value / total_value - - -def _safe_rate(numerator: object, elapsed_sec: object) -> float | None: - numerator_value = _optional_number(numerator) - elapsed_value = _optional_number(elapsed_sec) - if numerator_value is None or elapsed_value is None or elapsed_value <= 0: - return None - return numerator_value / elapsed_value - - -def _safe_ratio(numerator: object, denominator: object) -> float | None: - numerator_value = _optional_number(numerator) - denominator_value = _optional_number(denominator) - if numerator_value is None or denominator_value is None or denominator_value <= 0: - return None - return numerator_value / denominator_value - - -def _sum_optional_numbers(*values: object) -> float | None: - numeric_values = [_optional_number(value) for value in values] - if any(value is None for value in numeric_values): - return None - return sum(cast(float, value) for value in numeric_values) - - -def _f1(precision: float | None, recall: float | None) -> float | None: - if precision is None or recall is None or precision + recall == 0: - return None - return 2 * precision * recall / (precision + recall) - - -def _optional_number(value: object) -> float | None: - if value is None or pd.isna(value): - return None - return float(value) - - def _estimated_validation_chunk_count( artifact_rows: pd.DataFrame, *, diff --git a/tools/measurement/measurement_tools/benchmark_analysis_math.py b/tools/measurement/measurement_tools/benchmark_analysis_math.py new file mode 100644 index 00000000..4701e5b6 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_analysis_math.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared numeric and dataframe operations for benchmark analysis.""" + +from __future__ import annotations + +import json +from typing import cast + +import pandas as pd + +from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero +from measurement_tools.stats import sum_or_none as _sum_or_none + + +def pipeline_stage_rows(dataframe: pd.DataFrame) -> pd.DataFrame: + stages = records_of_type(dataframe, "stage") + if "stage" not in stages.columns: + return stages.iloc[0:0] + return stages[stages["stage"] == "Anonymizer._run_internal"] + +def first_float(frames: list[pd.DataFrame], columns: list[str]) -> float | None: + value = first_value(frames, columns) + return float(value) if value is not None else None + +def non_null_count(dataframe: pd.DataFrame, column: str) -> int: + if column not in dataframe.columns: + return 0 + return int(pd.to_numeric(dataframe[column], errors="coerce").notna().sum()) + +def sum_optional_numbers(*values: object) -> float | None: + numeric_values = [optional_number(value) for value in values] + if any(value is None for value in numeric_values): + return None + return sum(cast(float, value) for value in numeric_values) + +def sum_bool_or_zero(dataframe: pd.DataFrame, column: str) -> int: + if column not in dataframe.columns: + return 0 + return int(dataframe[column].fillna(False).astype(bool).sum()) + +def records_of_type(dataframe: pd.DataFrame, record_type: str) -> pd.DataFrame: + if "record_type" not in dataframe.columns: + return dataframe.iloc[0:0] + return dataframe[dataframe["record_type"] == record_type] + +def first_int(frames: list[pd.DataFrame], columns: list[str]) -> int | None: + value = first_value(frames, columns) + return int(float(value)) if value is not None else None + +def optional_number(value: object) -> float | None: + if value is None or pd.isna(value): + return None + return float(value) + +def safe_rate(numerator: object, elapsed_sec: object) -> float | None: + numerator_value = optional_number(numerator) + elapsed_value = optional_number(elapsed_sec) + if numerator_value is None or elapsed_value is None or elapsed_value <= 0: + return None + return numerator_value / elapsed_value + +def first_value(frames: list[pd.DataFrame], columns: list[str]) -> str | None: + for frame in frames: + for column in columns: + if column not in frame.columns: + continue + values = frame[column].dropna() + if not values.empty: + return str(values.iloc[0]) + return None + +def f1(precision: float | None, recall: float | None) -> float | None: + if precision is None or recall is None or precision + recall == 0: + return None + return 2 * precision * recall / (precision + recall) + +def zero_with_positive_count(dataframe: pd.DataFrame, *, zero_column: str, positive_column: str) -> int: + if zero_column not in dataframe.columns or positive_column not in dataframe.columns: + return 0 + zero_values = pd.to_numeric(dataframe[zero_column], errors="coerce") + positive_values = pd.to_numeric(dataframe[positive_column], errors="coerce") + return int(((zero_values == 0) & (positive_values > 0)).sum()) + +def coalesce_number(*values: float | None) -> float | None: + for value in values: + if value is not None: + return value + return None + +def coerce_count_mapping(value: object) -> dict[str, int]: + payload = value + if isinstance(value, str): + try: + payload = json.loads(value) + except json.JSONDecodeError: + return {} + if not isinstance(payload, dict): + return {} + counts: dict[str, int] = {} + for key, count in payload.items(): + numeric = pd.to_numeric(pd.Series([count]), errors="coerce").dropna() + if not numeric.empty and numeric.iloc[0]: + counts[str(key)] = int(numeric.iloc[0]) + return counts + +def records_of_types(dataframe: pd.DataFrame, record_types: set[str]) -> pd.DataFrame: + if "record_type" not in dataframe.columns: + return dataframe.iloc[0:0] + return dataframe[dataframe["record_type"].isin(record_types)] + +def request_failure_rate(*, failed: object, total: object) -> float | None: + failed_value = optional_number(failed) + total_value = optional_number(total) + if failed_value is None or total_value is None or total_value <= 0: + return None + return failed_value / total_value + +def positive_count(dataframe: pd.DataFrame, column: str) -> int: + if column not in dataframe.columns: + return 0 + values = pd.to_numeric(dataframe[column], errors="coerce").fillna(0) + return int((values > 0).sum()) + +def safe_ratio(numerator: object, denominator: object) -> float | None: + numerator_value = optional_number(numerator) + denominator_value = optional_number(denominator) + if numerator_value is None or denominator_value is None or denominator_value <= 0: + return None + return numerator_value / denominator_value + +def sum_int_or_none(dataframe: pd.DataFrame, column: str) -> int | None: + value = _sum_or_none(dataframe, column) + return int(value) if value is not None else None + +def sum_prefixed_ints(dataframe: pd.DataFrame, prefix: str) -> dict[str, int]: + totals: dict[str, int] = {} + base_column = prefix.removesuffix(".") + if base_column in dataframe.columns: + for value in dataframe[base_column].dropna().tolist(): + for key, count in coerce_count_mapping(value).items(): + totals[key] = totals.get(key, 0) + count + for column in sorted(col for col in dataframe.columns if col.startswith(prefix)): + value = _sum_int_or_zero(dataframe, column) + if value: + totals[column.removeprefix(prefix)] = value + return totals + +def zero_count(dataframe: pd.DataFrame, column: str) -> int: + if column not in dataframe.columns: + return 0 + values = pd.to_numeric(dataframe[column], errors="coerce").dropna() + return int((values == 0).sum()) + +def model_workflow_rows(dataframe: pd.DataFrame) -> pd.DataFrame: + return records_of_types(dataframe, {"ndd_workflow", "model_workflow"}) + + +__all__ = ['coalesce_number', 'coerce_count_mapping', 'f1', 'first_float', 'first_int', 'first_value', 'model_workflow_rows', 'non_null_count', 'optional_number', 'pipeline_stage_rows', 'positive_count', 'records_of_type', 'records_of_types', 'request_failure_rate', 'safe_rate', 'safe_ratio', 'sum_bool_or_zero', 'sum_int_or_none', 'sum_optional_numbers', 'sum_prefixed_ints', 'zero_count', 'zero_with_positive_count'] diff --git a/tools/measurement/measurement_tools/benchmark_analysis_models.py b/tools/measurement/measurement_tools/benchmark_analysis_models.py new file mode 100644 index 00000000..a57c04cd --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_analysis_models.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contracts for benchmark output analysis.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from pydantic import BaseModel, Field, computed_field + + +@dataclass(frozen=True) +class _EvaluationRollup: + prefix: str + valid_column: str + invalid_count_column: str + +_EVALUATION_ROLLUPS = ( + _EvaluationRollup("detection", "detection_valid", "detection_invalid_entity_count"), + _EvaluationRollup("type_fidelity", "type_fidelity_valid", "type_fidelity_invalid_replacement_count"), + _EvaluationRollup( + "relational_consistency", + "relational_consistency_valid", + "relational_consistency_invalid_relation_count", + ), + _EvaluationRollup("attribute_fidelity", "attribute_fidelity_valid", "attribute_fidelity_invalid_entity_count"), +) + +class CaseAnalysisRow(BaseModel): + suite_id: str | None = None + workload_id: str | None = None + workload_category: str | None = None + config_id: str | None = None + experimental_detection_strategy: str | None = None + experimental_replacement_strategy: str | None = None + dd_parser_compat: str | None = None + entity_label_set_id: str | None = None + entity_label_count: int | None = None + gliner_threshold: float | None = None + repetition: int | None = None + case_id: str + run_id: str + case_failed: bool = False + error_stage_count: int = 0 + error_ndd_workflow_count: int = 0 + error_model_workflow_count: int = 0 + pipeline_elapsed_sec: float | None = None + ndd_workflow_count: int = 0 + ndd_elapsed_sec_total: float = 0.0 + observed_total_requests: int = 0 + observed_successful_requests: int = 0 + observed_input_tokens: int = 0 + observed_output_tokens: int = 0 + observed_total_tokens: int = 0 + observed_failed_requests: int = 0 + observed_failed_request_rate: float | None = None + dd_trace_record_count: int = 0 + dd_trace_error_count: int = 0 + dd_trace_sync_client_unavailable_count: int = 0 + observed_bridge_fallback_requests: int | None = None + observed_non_bridge_total_requests: int | None = None + observed_non_bridge_failed_requests: int | None = None + observed_non_bridge_failed_request_rate: float | None = None + record_count: int = 0 + input_text_tokens_total: int | None = None + records_per_pipeline_sec: float | None = None + records_per_ndd_sec: float | None = None + input_text_tokens_per_pipeline_sec: float | None = None + input_text_tokens_per_ndd_sec: float | None = None + topology_endpoint_count: float | None = None + topology_gpu_count: float | None = None + topology_tensor_parallelism: float | None = None + topology_shard_count: float | None = None + input_text_tokens_per_endpoint_sec: float | None = None + input_text_tokens_per_gpu_sec: float | None = None + final_entity_count: float | None = None + empty_detection_count: int = 0 + empty_detection_rate: float | None = None + empty_detection_with_ground_truth_count: int = 0 + empty_detection_with_ground_truth_rate: float | None = None + ground_truth_record_count: int = 0 + ground_truth_entity_count: float | None = None + entity_true_positive_count: float | None = None + entity_false_positive_count: float | None = None + entity_false_negative_count: float | None = None + entity_precision: float | None = None + entity_recall: float | None = None + entity_f1: float | None = None + entity_relaxed_gt_found_count: float | None = None + entity_relaxed_detected_tp_count: float | None = None + entity_relaxed_label_compatible_gt_found_count: float | None = None + entity_relaxed_label_compatible_detected_tp_count: float | None = None + entity_relaxed_precision: float | None = None + entity_relaxed_recall: float | None = None + entity_relaxed_f1: float | None = None + entity_relaxed_label_compatible_precision: float | None = None + entity_relaxed_label_compatible_recall: float | None = None + entity_relaxed_label_compatible_f1: float | None = None + replacement_count: float | None = None + replacement_missing_final_entity_count: float | None = None + replacement_missing_final_entity_label_counts: dict[str, int] = Field(default_factory=dict) + replacement_missing_final_value_count: float | None = None + replacement_synthetic_original_collision_count: float | None = None + replacement_synthetic_original_collision_label_counts: dict[str, int] = Field(default_factory=dict) + replacement_synthetic_original_collision_value_count: float | None = None + original_value_leak_count: float | None = None + original_value_leak_record_count: int = 0 + original_value_leak_label_counts: dict[str, int] = Field(default_factory=dict) + detection_judged_record_count: int = 0 + detection_valid_record_count: int = 0 + detection_valid_rate: float | None = None + detection_invalid_entity_count: int = 0 + type_fidelity_judged_record_count: int = 0 + type_fidelity_valid_record_count: int = 0 + type_fidelity_valid_rate: float | None = None + type_fidelity_invalid_replacement_count: int = 0 + relational_consistency_judged_record_count: int = 0 + relational_consistency_valid_record_count: int = 0 + relational_consistency_valid_rate: float | None = None + relational_consistency_invalid_relation_count: int = 0 + attribute_fidelity_judged_record_count: int = 0 + attribute_fidelity_valid_record_count: int = 0 + attribute_fidelity_valid_rate: float | None = None + attribute_fidelity_invalid_entity_count: int = 0 + validation_max_entities_per_call: int | None = None + detection_artifact_rows: int = 0 + seed_entity_count: float | None = None + seed_validation_candidate_count: float | None = None + estimated_seed_validation_chunk_count: float | None = None + augmented_entity_count: float | None = None + augmented_new_final_value_count: float | None = None + artifact_final_entity_count: float | None = None + artifact_final_detector_entity_count: float | None = None + artifact_final_augmenter_entity_count: float | None = None + artifact_final_entity_signature_count: float | None = None + artifact_final_entity_signature_hashes: list[str] = Field(default_factory=list) + artifact_final_entity_signature_labels: dict[str, str] = Field(default_factory=dict) + artifact_final_entity_signature_details: dict[str, dict[str, Any]] = Field(default_factory=dict) + +class GroupAnalysisRow(BaseModel): + workload_id: str | None = None + workload_category: str | None = None + config_id: str | None = None + experimental_detection_strategy: str | None = None + experimental_replacement_strategy: str | None = None + entity_label_set_id: str | None = None + entity_label_count: int | None = None + gliner_threshold: float | None = None + case_count: int + failed_case_count: int = 0 + failed_case_rate: float | None = None + error_stage_count: int = 0 + error_ndd_workflow_count: int = 0 + error_model_workflow_count: int = 0 + median_pipeline_elapsed_sec: float | None = None + median_ndd_elapsed_sec_total: float | None = None + median_observed_total_requests: float | None = None + median_observed_successful_requests: float | None = None + median_observed_input_tokens: float | None = None + median_observed_output_tokens: float | None = None + median_observed_total_tokens: float | None = None + median_observed_failed_requests: float | None = None + median_observed_failed_request_rate: float | None = None + median_observed_bridge_fallback_requests: float | None = None + median_observed_non_bridge_total_requests: float | None = None + median_observed_non_bridge_failed_requests: float | None = None + median_observed_non_bridge_failed_request_rate: float | None = None + total_record_count: int = 0 + median_record_count: float | None = None + total_input_text_tokens: int | None = None + median_input_text_tokens_total: float | None = None + median_records_per_pipeline_sec: float | None = None + median_records_per_ndd_sec: float | None = None + median_input_text_tokens_per_pipeline_sec: float | None = None + median_input_text_tokens_per_ndd_sec: float | None = None + median_topology_endpoint_count: float | None = None + median_topology_gpu_count: float | None = None + median_topology_tensor_parallelism: float | None = None + median_topology_shard_count: float | None = None + median_input_text_tokens_per_endpoint_sec: float | None = None + median_input_text_tokens_per_gpu_sec: float | None = None + median_final_entity_count: float | None = None + total_empty_detection_count: int = 0 + empty_detection_rate: float | None = None + total_empty_detection_with_ground_truth_count: int = 0 + empty_detection_with_ground_truth_rate: float | None = None + total_ground_truth_record_count: int = 0 + sum_ground_truth_entity_count: float | None = None + sum_entity_true_positive_count: float | None = None + sum_entity_false_positive_count: float | None = None + sum_entity_false_negative_count: float | None = None + micro_entity_precision: float | None = None + micro_entity_recall: float | None = None + micro_entity_f1: float | None = None + sum_entity_relaxed_gt_found_count: float | None = None + sum_entity_relaxed_detected_tp_count: float | None = None + sum_entity_relaxed_label_compatible_gt_found_count: float | None = None + sum_entity_relaxed_label_compatible_detected_tp_count: float | None = None + micro_entity_relaxed_precision: float | None = None + micro_entity_relaxed_recall: float | None = None + micro_entity_relaxed_f1: float | None = None + micro_entity_relaxed_label_compatible_precision: float | None = None + micro_entity_relaxed_label_compatible_recall: float | None = None + micro_entity_relaxed_label_compatible_f1: float | None = None + median_entity_relaxed_f1: float | None = None + median_entity_relaxed_label_compatible_f1: float | None = None + median_replacement_missing_final_entity_count: float | None = None + median_replacement_missing_final_value_count: float | None = None + replacement_missing_final_entity_label_counts: dict[str, int] = Field(default_factory=dict) + median_replacement_synthetic_original_collision_count: float | None = None + median_replacement_synthetic_original_collision_value_count: float | None = None + replacement_synthetic_original_collision_label_counts: dict[str, int] = Field(default_factory=dict) + sum_original_value_leak_count: float | None = None + leaking_case_count: int = 0 + median_original_value_leak_count: float | None = None + sum_detection_judged_record_count: int = 0 + sum_detection_valid_record_count: int = 0 + micro_detection_valid_rate: float | None = None + sum_detection_invalid_entity_count: int = 0 + sum_type_fidelity_judged_record_count: int = 0 + sum_type_fidelity_valid_record_count: int = 0 + micro_type_fidelity_valid_rate: float | None = None + sum_type_fidelity_invalid_replacement_count: int = 0 + sum_relational_consistency_judged_record_count: int = 0 + sum_relational_consistency_valid_record_count: int = 0 + micro_relational_consistency_valid_rate: float | None = None + sum_relational_consistency_invalid_relation_count: int = 0 + sum_attribute_fidelity_judged_record_count: int = 0 + sum_attribute_fidelity_valid_record_count: int = 0 + micro_attribute_fidelity_valid_rate: float | None = None + sum_attribute_fidelity_invalid_entity_count: int = 0 + median_seed_entity_count: float | None = None + median_seed_validation_candidate_count: float | None = None + median_estimated_seed_validation_chunk_count: float | None = None + median_augmented_entity_count: float | None = None + median_augmented_new_final_value_count: float | None = None + median_artifact_final_entity_count: float | None = None + median_artifact_final_detector_entity_count: float | None = None + median_artifact_final_augmenter_entity_count: float | None = None + median_artifact_final_entity_signature_count: float | None = None + +class ModelUsageAnalysisRow(BaseModel): + suite_id: str | None = None + workload_id: str | None = None + config_id: str | None = None + experimental_detection_strategy: str | None = None + experimental_replacement_strategy: str | None = None + dd_parser_compat: str | None = None + repetition: int | None = None + case_id: str + run_id: str + workflow_name: str | None = None + model_alias: str | None = None + model_name: str + model_provider_name: str | None = None + ndd_elapsed_sec: float | None = None + observed_total_requests: int = 0 + observed_successful_requests: int = 0 + observed_failed_requests: int = 0 + observed_input_tokens: int = 0 + observed_output_tokens: int = 0 + observed_total_tokens: int = 0 + observed_reasoning_tokens: int | None = None + observed_failed_request_rate: float | None = None + +class ModelUsageGroupAnalysisRow(BaseModel): + workload_id: str | None = None + config_id: str | None = None + experimental_detection_strategy: str | None = None + experimental_replacement_strategy: str | None = None + dd_parser_compat: str | None = None + workflow_name: str | None = None + model_alias: str | None = None + model_name: str + model_provider_name: str | None = None + case_count: int + workflow_count: int + sum_observed_total_requests: int = 0 + sum_observed_successful_requests: int = 0 + sum_observed_failed_requests: int = 0 + sum_observed_input_tokens: int = 0 + sum_observed_output_tokens: int = 0 + sum_observed_total_tokens: int = 0 + sum_observed_reasoning_tokens: int | None = None + observed_failed_request_rate: float | None = None + median_observed_total_requests: float | None = None + median_observed_failed_requests: float | None = None + median_observed_total_tokens: float | None = None + +class BenchmarkOutputAnalysis(BaseModel): + benchmark_dir: str + detection_artifacts_path: str | None = None + cases: list[CaseAnalysisRow] = Field(default_factory=list) + groups: list[GroupAnalysisRow] = Field(default_factory=list) + model_usage: list[ModelUsageAnalysisRow] = Field(default_factory=list) + model_usage_groups: list[ModelUsageGroupAnalysisRow] = Field(default_factory=list) + + @computed_field + @property + def case_count(self) -> int: + return len(self.cases) + + @computed_field + @property + def group_count(self) -> int: + return len(self.groups) + + @computed_field + @property + def model_usage_count(self) -> int: + return len(self.model_usage) + + @computed_field + @property + def model_usage_group_count(self) -> int: + return len(self.model_usage_groups) + + +__all__ = [ + "BenchmarkOutputAnalysis", + "CaseAnalysisRow", + "GroupAnalysisRow", + "ModelUsageAnalysisRow", + "ModelUsageGroupAnalysisRow", +] From 5200f7c79de81f64b3797020e5e9fbb04ff5a87f Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 22:03:38 +0000 Subject: [PATCH 16/17] refactor(measurement): split model usage and group analysis Signed-off-by: Aaron Gonzales --- tools/measurement/analyze_benchmark_output.py | 411 +----------------- .../benchmark_analysis_math.py | 46 +- .../benchmark_analysis_models.py | 6 + .../benchmark_group_analysis.py | 223 ++++++++++ .../benchmark_model_usage.py | 236 ++++++++++ 5 files changed, 520 insertions(+), 402 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_group_analysis.py create mode 100644 tools/measurement/measurement_tools/benchmark_model_usage.py diff --git a/tools/measurement/analyze_benchmark_output.py b/tools/measurement/analyze_benchmark_output.py index 6735c675..7686855d 100644 --- a/tools/measurement/analyze_benchmark_output.py +++ b/tools/measurement/analyze_benchmark_output.py @@ -17,7 +17,7 @@ import math import sys from pathlib import Path -from typing import Annotated, Any, cast +from typing import Annotated, Any import cyclopts import pandas as pd @@ -60,9 +60,6 @@ from measurement_tools.benchmark_analysis_math import ( safe_ratio as _safe_ratio, ) -from measurement_tools.benchmark_analysis_math import ( - sum_bool_or_zero as _sum_bool_or_zero, -) from measurement_tools.benchmark_analysis_math import ( sum_int_or_none as _sum_int_or_none, ) @@ -99,9 +96,16 @@ from measurement_tools.benchmark_analysis_models import ( _EvaluationRollup as _EvaluationRollup, ) +from measurement_tools.benchmark_group_analysis import ( + build_group_rows as build_group_rows, +) +from measurement_tools.benchmark_model_usage import ( + build_model_usage_group_rows as build_model_usage_group_rows, +) +from measurement_tools.benchmark_model_usage import ( + build_model_usage_rows as build_model_usage_rows, +) from measurement_tools.cli import LogFormat, configure_logging, log_bad_input -from measurement_tools.stats import median_or_none as _median_or_none -from measurement_tools.stats import none_if_nan as _none_if_nan from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero from measurement_tools.stats import sum_or_none as _sum_or_none from measurement_tools.stats import sum_or_zero as _sum_or_zero @@ -586,148 +590,6 @@ def _rows_for_case(dataframe: pd.DataFrame, case_id: str) -> pd.DataFrame: return dataframe[mask] -_MODEL_USAGE_SUFFIXES = { - ".request_usage.total_requests": "observed_total_requests", - ".request_usage.successful_requests": "observed_successful_requests", - ".request_usage.failed_requests": "observed_failed_requests", - ".token_usage.input_tokens": "observed_input_tokens", - ".token_usage.output_tokens": "observed_output_tokens", - ".token_usage.total_tokens": "observed_total_tokens", - ".token_usage.reasoning_tokens": "observed_reasoning_tokens", -} - -_MODEL_USAGE_METADATA_SUFFIXES = { - ".model_alias": "model_alias", - ".model_name": "model_name", - ".model_provider_name": "model_provider_name", -} - - -def build_model_usage_rows(measurements: pd.DataFrame) -> list[ModelUsageAnalysisRow]: - model_rows = _model_workflow_rows(measurements) - if model_rows.empty: - return [] - model_usage_keys = _model_usage_keys(model_rows.columns) - rows: list[ModelUsageAnalysisRow] = [] - for _, measurement in model_rows.iterrows(): - data = measurement.to_dict() - case_id = _string_from_row(data, ["run_tags.case_id", "run_id"]) - run_id = _string_from_row(data, ["run_id", "run_tags.case_id"]) - if case_id is None or run_id is None: - continue - for model_usage_key in model_usage_keys: - usage = _model_usage_metrics(data, model_usage_key) - if not _has_observed_model_usage(usage): - continue - metadata = _model_usage_metadata(data, model_usage_key) - rows.append( - ModelUsageAnalysisRow( - suite_id=_string_from_row(data, ["run_tags.suite_id"]), - workload_id=_string_from_row(data, ["run_tags.workload_id"]), - config_id=_string_from_row(data, ["run_tags.config_id"]), - experimental_detection_strategy=_string_from_row( - data, ["run_tags.experimental_detection_strategy"] - ), - experimental_replacement_strategy=_string_from_row( - data, ["run_tags.experimental_replacement_strategy"] - ), - dd_parser_compat=_string_from_row(data, ["run_tags.dd_parser_compat"]), - repetition=_int_from_row(data, ["run_tags.repetition"]), - case_id=case_id, - run_id=run_id, - workflow_name=_string_from_row(data, ["workflow_name"]), - model_alias=metadata.get("model_alias"), - model_name=metadata.get("model_name") or model_usage_key, - model_provider_name=metadata.get("model_provider_name"), - ndd_elapsed_sec=_float_from_row(data, ["elapsed_sec"]), - **usage, - ) - ) - return rows - - -def _model_usage_keys(columns: pd.Index) -> list[str]: - keys: set[str] = set() - for column in columns: - parsed = _model_usage_column_parts(str(column)) - if parsed is not None: - keys.add(parsed[0]) - return sorted(keys) - - -def _model_usage_column_parts(column: str) -> tuple[str, str] | None: - prefix = "model_usage." - if not column.startswith(prefix): - return None - for suffix, metric in {**_MODEL_USAGE_SUFFIXES, **_MODEL_USAGE_METADATA_SUFFIXES}.items(): - if column.endswith(suffix): - return column[len(prefix) : -len(suffix)], metric - return None - - -def _model_usage_metrics(data: dict[str, Any], model_usage_key: str) -> dict[str, int | float | None]: - values: dict[str, int | float | None] = { - "observed_total_requests": 0, - "observed_successful_requests": 0, - "observed_failed_requests": 0, - "observed_input_tokens": 0, - "observed_output_tokens": 0, - "observed_total_tokens": 0, - "observed_reasoning_tokens": None, - "observed_failed_request_rate": None, - } - for suffix, metric in _MODEL_USAGE_SUFFIXES.items(): - value = data.get(f"model_usage.{model_usage_key}{suffix}") - if value is None or pd.isna(value): - continue - values[metric] = _coerce_int(value) - values["observed_failed_request_rate"] = _request_failure_rate( - failed=values["observed_failed_requests"], - total=values["observed_total_requests"], - ) - return values - - -def _model_usage_metadata(data: dict[str, Any], model_usage_key: str) -> dict[str, str | None]: - values: dict[str, str | None] = { - "model_alias": None, - "model_name": None, - "model_provider_name": None, - } - for suffix, field_name in _MODEL_USAGE_METADATA_SUFFIXES.items(): - value = data.get(f"model_usage.{model_usage_key}{suffix}") - if value is None or pd.isna(value): - continue - values[field_name] = str(value) - return values - - -def _has_observed_model_usage(usage: dict[str, int | float | None]) -> bool: - return any(value not in (None, 0) for value in usage.values()) - - -def _string_from_row(data: dict[str, Any], columns: list[str]) -> str | None: - for column in columns: - value = data.get(column) - if value is not None and not pd.isna(value): - return str(value) - return None - - -def _int_from_row(data: dict[str, Any], columns: list[str]) -> int | None: - value = _string_from_row(data, columns) - return int(float(value)) if value is not None else None - - -def _float_from_row(data: dict[str, Any], columns: list[str]) -> float | None: - value = _string_from_row(data, columns) - return float(value) if value is not None else None - - -def _coerce_int(value: Any) -> int: - return int(float(value)) - - def _artifact_signature_hashes(artifact_rows: pd.DataFrame) -> list[str]: if "final_entity_signature_hashes" not in artifact_rows.columns: return [] @@ -816,259 +678,6 @@ def _signature_count(artifact_rows: pd.DataFrame, *, signature_hashes: list[str] return _sum_or_none(artifact_rows, "final_entity_signature_count") -def build_group_rows(cases: list[CaseAnalysisRow]) -> list[GroupAnalysisRow]: - if not cases: - return [] - table = pd.DataFrame([case.model_dump() for case in cases]) - rows: list[GroupAnalysisRow] = [] - group_columns = [ - "workload_id", - "workload_category", - "config_id", - "experimental_detection_strategy", - "experimental_replacement_strategy", - "entity_label_set_id", - "entity_label_count", - "gliner_threshold", - ] - for keys, group in table.groupby(group_columns, dropna=False): - rows.append(_build_group_row(keys, group)) - return rows - - -def build_model_usage_group_rows(model_usage: list[ModelUsageAnalysisRow]) -> list[ModelUsageGroupAnalysisRow]: - if not model_usage: - return [] - table = pd.DataFrame([row.model_dump() for row in model_usage]) - rows: list[ModelUsageGroupAnalysisRow] = [] - group_columns = [ - "workload_id", - "config_id", - "experimental_detection_strategy", - "experimental_replacement_strategy", - "dd_parser_compat", - "workflow_name", - "model_alias", - "model_name", - "model_provider_name", - ] - for keys, group in table.groupby(group_columns, dropna=False): - rows.append(_build_model_usage_group_row(keys, group)) - return rows - - -def _build_model_usage_group_row(keys: tuple[Any, ...], group: pd.DataFrame) -> ModelUsageGroupAnalysisRow: - ( - workload_id, - config_id, - detection_strategy, - replacement_strategy, - dd_parser_compat, - workflow_name, - model_alias, - model_name, - provider_name, - ) = keys - reasoning_sum = _sum_int_or_none(group, "observed_reasoning_tokens") - total_requests = _sum_int_or_zero(group, "observed_total_requests") - failed_requests = _sum_int_or_zero(group, "observed_failed_requests") - return ModelUsageGroupAnalysisRow( - workload_id=_none_if_nan(workload_id), - config_id=_none_if_nan(config_id), - experimental_detection_strategy=_none_if_nan(detection_strategy), - experimental_replacement_strategy=_none_if_nan(replacement_strategy), - dd_parser_compat=_none_if_nan(dd_parser_compat), - workflow_name=_none_if_nan(workflow_name), - model_alias=_none_if_nan(model_alias), - model_name=str(model_name), - model_provider_name=_none_if_nan(provider_name), - case_count=int(group["case_id"].nunique()), - workflow_count=len(group), - sum_observed_total_requests=total_requests, - sum_observed_successful_requests=_sum_int_or_zero(group, "observed_successful_requests"), - sum_observed_failed_requests=failed_requests, - sum_observed_input_tokens=_sum_int_or_zero(group, "observed_input_tokens"), - sum_observed_output_tokens=_sum_int_or_zero(group, "observed_output_tokens"), - sum_observed_total_tokens=_sum_int_or_zero(group, "observed_total_tokens"), - sum_observed_reasoning_tokens=reasoning_sum, - observed_failed_request_rate=_request_failure_rate(failed=failed_requests, total=total_requests), - median_observed_total_requests=_median_or_none(group, "observed_total_requests"), - median_observed_failed_requests=_median_or_none(group, "observed_failed_requests"), - median_observed_total_tokens=_median_or_none(group, "observed_total_tokens"), - ) - - -def _build_group_row(keys: tuple[Any, ...], group: pd.DataFrame) -> GroupAnalysisRow: - ( - workload_id, - workload_category, - config_id, - detection_strategy, - replacement_strategy, - entity_label_set_id, - entity_label_count, - gliner_threshold, - ) = keys - case_count = int(group["case_id"].nunique()) - failed_case_count = _sum_bool_or_zero(group, "case_failed") - total_record_count = _sum_int_or_zero(group, "record_count") - total_input_text_tokens = _sum_int_or_none(group, "input_text_tokens_total") - total_empty_detection_count = _sum_int_or_zero(group, "empty_detection_count") - total_ground_truth_record_count = _sum_int_or_zero(group, "ground_truth_record_count") - total_empty_detection_with_gt_count = _sum_int_or_zero(group, "empty_detection_with_ground_truth_count") - final_entity_count = _sum_or_none(group, "final_entity_count") - ground_truth_entity_count = _sum_or_none(group, "ground_truth_entity_count") - true_positive = _sum_or_none(group, "entity_true_positive_count") - false_positive = _sum_or_none(group, "entity_false_positive_count") - false_negative = _sum_or_none(group, "entity_false_negative_count") - strict_precision = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_positive)) - strict_recall = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_negative)) - relaxed_gt_found = _sum_or_none(group, "entity_relaxed_gt_found_count") - relaxed_detected_tp = _sum_or_none(group, "entity_relaxed_detected_tp_count") - label_compatible_gt_found = _sum_or_none(group, "entity_relaxed_label_compatible_gt_found_count") - label_compatible_detected_tp = _sum_or_none(group, "entity_relaxed_label_compatible_detected_tp_count") - relaxed_precision = _safe_ratio(relaxed_detected_tp, final_entity_count) - relaxed_recall = _safe_ratio(relaxed_gt_found, ground_truth_entity_count) - label_compatible_precision = _safe_ratio(label_compatible_detected_tp, final_entity_count) - label_compatible_recall = _safe_ratio(label_compatible_gt_found, ground_truth_entity_count) - evaluation_metrics = _group_evaluation_metrics(group) - return GroupAnalysisRow( - workload_id=_none_if_nan(workload_id), - workload_category=_none_if_nan(workload_category), - config_id=_none_if_nan(config_id), - experimental_detection_strategy=_none_if_nan(detection_strategy), - experimental_replacement_strategy=_none_if_nan(replacement_strategy), - entity_label_set_id=_none_if_nan(entity_label_set_id), - entity_label_count=_int_if_not_nan(entity_label_count), - gliner_threshold=_float_if_not_nan(gliner_threshold), - case_count=case_count, - failed_case_count=failed_case_count, - failed_case_rate=_request_failure_rate(failed=failed_case_count, total=case_count), - error_stage_count=_sum_int_or_zero(group, "error_stage_count"), - error_ndd_workflow_count=_sum_int_or_zero(group, "error_ndd_workflow_count"), - error_model_workflow_count=_sum_int_or_zero(group, "error_model_workflow_count"), - median_pipeline_elapsed_sec=_median_or_none(group, "pipeline_elapsed_sec"), - median_ndd_elapsed_sec_total=_median_or_none(group, "ndd_elapsed_sec_total"), - median_observed_total_requests=_median_or_none(group, "observed_total_requests"), - median_observed_successful_requests=_median_or_none(group, "observed_successful_requests"), - median_observed_input_tokens=_median_or_none(group, "observed_input_tokens"), - median_observed_output_tokens=_median_or_none(group, "observed_output_tokens"), - median_observed_total_tokens=_median_or_none(group, "observed_total_tokens"), - median_observed_failed_requests=_median_or_none(group, "observed_failed_requests"), - median_observed_failed_request_rate=_median_or_none(group, "observed_failed_request_rate"), - median_observed_bridge_fallback_requests=_median_or_none(group, "observed_bridge_fallback_requests"), - median_observed_non_bridge_total_requests=_median_or_none(group, "observed_non_bridge_total_requests"), - median_observed_non_bridge_failed_requests=_median_or_none(group, "observed_non_bridge_failed_requests"), - median_observed_non_bridge_failed_request_rate=_median_or_none( - group, - "observed_non_bridge_failed_request_rate", - ), - total_record_count=total_record_count, - median_record_count=_median_or_none(group, "record_count"), - total_input_text_tokens=total_input_text_tokens, - median_input_text_tokens_total=_median_or_none(group, "input_text_tokens_total"), - median_records_per_pipeline_sec=_median_or_none(group, "records_per_pipeline_sec"), - median_records_per_ndd_sec=_median_or_none(group, "records_per_ndd_sec"), - median_input_text_tokens_per_pipeline_sec=_median_or_none(group, "input_text_tokens_per_pipeline_sec"), - median_input_text_tokens_per_ndd_sec=_median_or_none(group, "input_text_tokens_per_ndd_sec"), - median_topology_endpoint_count=_median_or_none(group, "topology_endpoint_count"), - median_topology_gpu_count=_median_or_none(group, "topology_gpu_count"), - median_topology_tensor_parallelism=_median_or_none(group, "topology_tensor_parallelism"), - median_topology_shard_count=_median_or_none(group, "topology_shard_count"), - median_input_text_tokens_per_endpoint_sec=_median_or_none(group, "input_text_tokens_per_endpoint_sec"), - median_input_text_tokens_per_gpu_sec=_median_or_none(group, "input_text_tokens_per_gpu_sec"), - median_final_entity_count=_median_or_none(group, "final_entity_count"), - total_empty_detection_count=total_empty_detection_count, - empty_detection_rate=_safe_ratio(total_empty_detection_count, total_record_count), - total_empty_detection_with_ground_truth_count=total_empty_detection_with_gt_count, - empty_detection_with_ground_truth_rate=_safe_ratio( - total_empty_detection_with_gt_count, - total_ground_truth_record_count, - ), - total_ground_truth_record_count=total_ground_truth_record_count, - sum_ground_truth_entity_count=ground_truth_entity_count, - sum_entity_true_positive_count=true_positive, - sum_entity_false_positive_count=false_positive, - sum_entity_false_negative_count=false_negative, - micro_entity_precision=strict_precision, - micro_entity_recall=strict_recall, - micro_entity_f1=_f1(strict_precision, strict_recall), - sum_entity_relaxed_gt_found_count=relaxed_gt_found, - sum_entity_relaxed_detected_tp_count=relaxed_detected_tp, - sum_entity_relaxed_label_compatible_gt_found_count=label_compatible_gt_found, - sum_entity_relaxed_label_compatible_detected_tp_count=label_compatible_detected_tp, - micro_entity_relaxed_precision=relaxed_precision, - micro_entity_relaxed_recall=relaxed_recall, - micro_entity_relaxed_f1=_f1(relaxed_precision, relaxed_recall), - micro_entity_relaxed_label_compatible_precision=label_compatible_precision, - micro_entity_relaxed_label_compatible_recall=label_compatible_recall, - micro_entity_relaxed_label_compatible_f1=_f1(label_compatible_precision, label_compatible_recall), - median_entity_relaxed_f1=_median_or_none(group, "entity_relaxed_f1"), - median_entity_relaxed_label_compatible_f1=_median_or_none( - group, - "entity_relaxed_label_compatible_f1", - ), - median_replacement_missing_final_entity_count=_median_or_none( - group, - "replacement_missing_final_entity_count", - ), - median_replacement_missing_final_value_count=_median_or_none(group, "replacement_missing_final_value_count"), - replacement_missing_final_entity_label_counts=_sum_prefixed_ints( - group, - "replacement_missing_final_entity_label_counts.", - ), - median_replacement_synthetic_original_collision_count=_median_or_none( - group, - "replacement_synthetic_original_collision_count", - ), - median_replacement_synthetic_original_collision_value_count=_median_or_none( - group, - "replacement_synthetic_original_collision_value_count", - ), - replacement_synthetic_original_collision_label_counts=_sum_prefixed_ints( - group, - "replacement_synthetic_original_collision_label_counts.", - ), - sum_original_value_leak_count=_sum_or_none(group, "original_value_leak_count"), - leaking_case_count=_positive_count(group, "original_value_leak_count"), - median_original_value_leak_count=_median_or_none(group, "original_value_leak_count"), - **evaluation_metrics, - median_seed_entity_count=_median_or_none(group, "seed_entity_count"), - median_seed_validation_candidate_count=_median_or_none(group, "seed_validation_candidate_count"), - median_estimated_seed_validation_chunk_count=_median_or_none(group, "estimated_seed_validation_chunk_count"), - median_augmented_entity_count=_median_or_none(group, "augmented_entity_count"), - median_augmented_new_final_value_count=_median_or_none(group, "augmented_new_final_value_count"), - median_artifact_final_entity_count=_median_or_none(group, "artifact_final_entity_count"), - median_artifact_final_detector_entity_count=_median_or_none(group, "artifact_final_detector_entity_count"), - median_artifact_final_augmenter_entity_count=_median_or_none(group, "artifact_final_augmenter_entity_count"), - median_artifact_final_entity_signature_count=_median_or_none(group, "artifact_final_entity_signature_count"), - ) - - -def _int_if_not_nan(value: object) -> int | None: - if pd.isna(value): - return None - return int(float(cast(Any, value))) - - -def _float_if_not_nan(value: object) -> float | None: - if pd.isna(value): - return None - return float(cast(Any, value)) - - -def _group_evaluation_metrics(group: pd.DataFrame) -> dict[str, int | float | None]: - metrics: dict[str, int | float | None] = {} - for rollup in _EVALUATION_ROLLUPS: - judged_count = _sum_int_or_zero(group, f"{rollup.prefix}_judged_record_count") - valid_count = _sum_int_or_zero(group, f"{rollup.prefix}_valid_record_count") - metrics[f"sum_{rollup.prefix}_judged_record_count"] = judged_count - metrics[f"sum_{rollup.prefix}_valid_record_count"] = valid_count - metrics[f"micro_{rollup.prefix}_valid_rate"] = _safe_ratio(valid_count, judged_count) - metrics[f"sum_{rollup.invalid_count_column}"] = _sum_int_or_zero(group, rollup.invalid_count_column) - return metrics - - def _estimated_validation_chunk_count( artifact_rows: pd.DataFrame, *, diff --git a/tools/measurement/measurement_tools/benchmark_analysis_math.py b/tools/measurement/measurement_tools/benchmark_analysis_math.py index 4701e5b6..0a1452af 100644 --- a/tools/measurement/measurement_tools/benchmark_analysis_math.py +++ b/tools/measurement/measurement_tools/benchmark_analysis_math.py @@ -19,40 +19,48 @@ def pipeline_stage_rows(dataframe: pd.DataFrame) -> pd.DataFrame: return stages.iloc[0:0] return stages[stages["stage"] == "Anonymizer._run_internal"] + def first_float(frames: list[pd.DataFrame], columns: list[str]) -> float | None: value = first_value(frames, columns) return float(value) if value is not None else None + def non_null_count(dataframe: pd.DataFrame, column: str) -> int: if column not in dataframe.columns: return 0 return int(pd.to_numeric(dataframe[column], errors="coerce").notna().sum()) + def sum_optional_numbers(*values: object) -> float | None: numeric_values = [optional_number(value) for value in values] if any(value is None for value in numeric_values): return None return sum(cast(float, value) for value in numeric_values) + def sum_bool_or_zero(dataframe: pd.DataFrame, column: str) -> int: if column not in dataframe.columns: return 0 return int(dataframe[column].fillna(False).astype(bool).sum()) + def records_of_type(dataframe: pd.DataFrame, record_type: str) -> pd.DataFrame: if "record_type" not in dataframe.columns: return dataframe.iloc[0:0] return dataframe[dataframe["record_type"] == record_type] + def first_int(frames: list[pd.DataFrame], columns: list[str]) -> int | None: value = first_value(frames, columns) return int(float(value)) if value is not None else None + def optional_number(value: object) -> float | None: if value is None or pd.isna(value): return None return float(value) + def safe_rate(numerator: object, elapsed_sec: object) -> float | None: numerator_value = optional_number(numerator) elapsed_value = optional_number(elapsed_sec) @@ -60,6 +68,7 @@ def safe_rate(numerator: object, elapsed_sec: object) -> float | None: return None return numerator_value / elapsed_value + def first_value(frames: list[pd.DataFrame], columns: list[str]) -> str | None: for frame in frames: for column in columns: @@ -70,11 +79,13 @@ def first_value(frames: list[pd.DataFrame], columns: list[str]) -> str | None: return str(values.iloc[0]) return None + def f1(precision: float | None, recall: float | None) -> float | None: if precision is None or recall is None or precision + recall == 0: return None return 2 * precision * recall / (precision + recall) + def zero_with_positive_count(dataframe: pd.DataFrame, *, zero_column: str, positive_column: str) -> int: if zero_column not in dataframe.columns or positive_column not in dataframe.columns: return 0 @@ -82,12 +93,14 @@ def zero_with_positive_count(dataframe: pd.DataFrame, *, zero_column: str, posit positive_values = pd.to_numeric(dataframe[positive_column], errors="coerce") return int(((zero_values == 0) & (positive_values > 0)).sum()) + def coalesce_number(*values: float | None) -> float | None: for value in values: if value is not None: return value return None + def coerce_count_mapping(value: object) -> dict[str, int]: payload = value if isinstance(value, str): @@ -104,11 +117,13 @@ def coerce_count_mapping(value: object) -> dict[str, int]: counts[str(key)] = int(numeric.iloc[0]) return counts + def records_of_types(dataframe: pd.DataFrame, record_types: set[str]) -> pd.DataFrame: if "record_type" not in dataframe.columns: return dataframe.iloc[0:0] return dataframe[dataframe["record_type"].isin(record_types)] + def request_failure_rate(*, failed: object, total: object) -> float | None: failed_value = optional_number(failed) total_value = optional_number(total) @@ -116,12 +131,14 @@ def request_failure_rate(*, failed: object, total: object) -> float | None: return None return failed_value / total_value + def positive_count(dataframe: pd.DataFrame, column: str) -> int: if column not in dataframe.columns: return 0 values = pd.to_numeric(dataframe[column], errors="coerce").fillna(0) return int((values > 0).sum()) + def safe_ratio(numerator: object, denominator: object) -> float | None: numerator_value = optional_number(numerator) denominator_value = optional_number(denominator) @@ -129,10 +146,12 @@ def safe_ratio(numerator: object, denominator: object) -> float | None: return None return numerator_value / denominator_value + def sum_int_or_none(dataframe: pd.DataFrame, column: str) -> int | None: value = _sum_or_none(dataframe, column) return int(value) if value is not None else None + def sum_prefixed_ints(dataframe: pd.DataFrame, prefix: str) -> dict[str, int]: totals: dict[str, int] = {} base_column = prefix.removesuffix(".") @@ -146,14 +165,39 @@ def sum_prefixed_ints(dataframe: pd.DataFrame, prefix: str) -> dict[str, int]: totals[column.removeprefix(prefix)] = value return totals + def zero_count(dataframe: pd.DataFrame, column: str) -> int: if column not in dataframe.columns: return 0 values = pd.to_numeric(dataframe[column], errors="coerce").dropna() return int((values == 0).sum()) + def model_workflow_rows(dataframe: pd.DataFrame) -> pd.DataFrame: return records_of_types(dataframe, {"ndd_workflow", "model_workflow"}) -__all__ = ['coalesce_number', 'coerce_count_mapping', 'f1', 'first_float', 'first_int', 'first_value', 'model_workflow_rows', 'non_null_count', 'optional_number', 'pipeline_stage_rows', 'positive_count', 'records_of_type', 'records_of_types', 'request_failure_rate', 'safe_rate', 'safe_ratio', 'sum_bool_or_zero', 'sum_int_or_none', 'sum_optional_numbers', 'sum_prefixed_ints', 'zero_count', 'zero_with_positive_count'] +__all__ = [ + "coalesce_number", + "coerce_count_mapping", + "f1", + "first_float", + "first_int", + "first_value", + "model_workflow_rows", + "non_null_count", + "optional_number", + "pipeline_stage_rows", + "positive_count", + "records_of_type", + "records_of_types", + "request_failure_rate", + "safe_rate", + "safe_ratio", + "sum_bool_or_zero", + "sum_int_or_none", + "sum_optional_numbers", + "sum_prefixed_ints", + "zero_count", + "zero_with_positive_count", +] diff --git a/tools/measurement/measurement_tools/benchmark_analysis_models.py b/tools/measurement/measurement_tools/benchmark_analysis_models.py index a57c04cd..177a99ed 100644 --- a/tools/measurement/measurement_tools/benchmark_analysis_models.py +++ b/tools/measurement/measurement_tools/benchmark_analysis_models.py @@ -16,6 +16,7 @@ class _EvaluationRollup: valid_column: str invalid_count_column: str + _EVALUATION_ROLLUPS = ( _EvaluationRollup("detection", "detection_valid", "detection_invalid_entity_count"), _EvaluationRollup("type_fidelity", "type_fidelity_valid", "type_fidelity_invalid_replacement_count"), @@ -27,6 +28,7 @@ class _EvaluationRollup: _EvaluationRollup("attribute_fidelity", "attribute_fidelity_valid", "attribute_fidelity_invalid_entity_count"), ) + class CaseAnalysisRow(BaseModel): suite_id: str | None = None workload_id: str | None = None @@ -138,6 +140,7 @@ class CaseAnalysisRow(BaseModel): artifact_final_entity_signature_labels: dict[str, str] = Field(default_factory=dict) artifact_final_entity_signature_details: dict[str, dict[str, Any]] = Field(default_factory=dict) + class GroupAnalysisRow(BaseModel): workload_id: str | None = None workload_category: str | None = None @@ -240,6 +243,7 @@ class GroupAnalysisRow(BaseModel): median_artifact_final_augmenter_entity_count: float | None = None median_artifact_final_entity_signature_count: float | None = None + class ModelUsageAnalysisRow(BaseModel): suite_id: str | None = None workload_id: str | None = None @@ -264,6 +268,7 @@ class ModelUsageAnalysisRow(BaseModel): observed_reasoning_tokens: int | None = None observed_failed_request_rate: float | None = None + class ModelUsageGroupAnalysisRow(BaseModel): workload_id: str | None = None config_id: str | None = None @@ -288,6 +293,7 @@ class ModelUsageGroupAnalysisRow(BaseModel): median_observed_failed_requests: float | None = None median_observed_total_tokens: float | None = None + class BenchmarkOutputAnalysis(BaseModel): benchmark_dir: str detection_artifacts_path: str | None = None diff --git a/tools/measurement/measurement_tools/benchmark_group_analysis.py b/tools/measurement/measurement_tools/benchmark_group_analysis.py new file mode 100644 index 00000000..7b715495 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_group_analysis.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Grouped benchmark case analysis and evaluation rollups.""" + +from __future__ import annotations + +from typing import Any, cast + +import pandas as pd + +from measurement_tools.benchmark_analysis_math import ( + f1 as _f1, + positive_count as _positive_count, + request_failure_rate as _request_failure_rate, + safe_ratio as _safe_ratio, + sum_bool_or_zero as _sum_bool_or_zero, + sum_int_or_none as _sum_int_or_none, + sum_optional_numbers as _sum_optional_numbers, + sum_prefixed_ints as _sum_prefixed_ints, +) +from measurement_tools.benchmark_analysis_models import _EVALUATION_ROLLUPS, CaseAnalysisRow, GroupAnalysisRow +from measurement_tools.stats import median_or_none as _median_or_none +from measurement_tools.stats import none_if_nan as _none_if_nan +from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero +from measurement_tools.stats import sum_or_none as _sum_or_none + +def build_group_rows(cases: list[CaseAnalysisRow]) -> list[GroupAnalysisRow]: + if not cases: + return [] + table = pd.DataFrame([case.model_dump() for case in cases]) + rows: list[GroupAnalysisRow] = [] + group_columns = [ + "workload_id", + "workload_category", + "config_id", + "experimental_detection_strategy", + "experimental_replacement_strategy", + "entity_label_set_id", + "entity_label_count", + "gliner_threshold", + ] + for keys, group in table.groupby(group_columns, dropna=False): + rows.append(build_group_row(keys, group)) + return rows + + +def build_group_row(keys: tuple[Any, ...], group: pd.DataFrame) -> GroupAnalysisRow: + ( + workload_id, + workload_category, + config_id, + detection_strategy, + replacement_strategy, + entity_label_set_id, + entity_label_count, + gliner_threshold, + ) = keys + case_count = int(group["case_id"].nunique()) + failed_case_count = _sum_bool_or_zero(group, "case_failed") + total_record_count = _sum_int_or_zero(group, "record_count") + total_input_text_tokens = _sum_int_or_none(group, "input_text_tokens_total") + total_empty_detection_count = _sum_int_or_zero(group, "empty_detection_count") + total_ground_truth_record_count = _sum_int_or_zero(group, "ground_truth_record_count") + total_empty_detection_with_gt_count = _sum_int_or_zero(group, "empty_detection_with_ground_truth_count") + final_entity_count = _sum_or_none(group, "final_entity_count") + ground_truth_entity_count = _sum_or_none(group, "ground_truth_entity_count") + true_positive = _sum_or_none(group, "entity_true_positive_count") + false_positive = _sum_or_none(group, "entity_false_positive_count") + false_negative = _sum_or_none(group, "entity_false_negative_count") + strict_precision = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_positive)) + strict_recall = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_negative)) + relaxed_gt_found = _sum_or_none(group, "entity_relaxed_gt_found_count") + relaxed_detected_tp = _sum_or_none(group, "entity_relaxed_detected_tp_count") + label_compatible_gt_found = _sum_or_none(group, "entity_relaxed_label_compatible_gt_found_count") + label_compatible_detected_tp = _sum_or_none(group, "entity_relaxed_label_compatible_detected_tp_count") + relaxed_precision = _safe_ratio(relaxed_detected_tp, final_entity_count) + relaxed_recall = _safe_ratio(relaxed_gt_found, ground_truth_entity_count) + label_compatible_precision = _safe_ratio(label_compatible_detected_tp, final_entity_count) + label_compatible_recall = _safe_ratio(label_compatible_gt_found, ground_truth_entity_count) + evaluation_metrics = group_evaluation_metrics(group) + return GroupAnalysisRow( + workload_id=_none_if_nan(workload_id), + workload_category=_none_if_nan(workload_category), + config_id=_none_if_nan(config_id), + experimental_detection_strategy=_none_if_nan(detection_strategy), + experimental_replacement_strategy=_none_if_nan(replacement_strategy), + entity_label_set_id=_none_if_nan(entity_label_set_id), + entity_label_count=int_if_not_nan(entity_label_count), + gliner_threshold=float_if_not_nan(gliner_threshold), + case_count=case_count, + failed_case_count=failed_case_count, + failed_case_rate=_request_failure_rate(failed=failed_case_count, total=case_count), + error_stage_count=_sum_int_or_zero(group, "error_stage_count"), + error_ndd_workflow_count=_sum_int_or_zero(group, "error_ndd_workflow_count"), + error_model_workflow_count=_sum_int_or_zero(group, "error_model_workflow_count"), + median_pipeline_elapsed_sec=_median_or_none(group, "pipeline_elapsed_sec"), + median_ndd_elapsed_sec_total=_median_or_none(group, "ndd_elapsed_sec_total"), + median_observed_total_requests=_median_or_none(group, "observed_total_requests"), + median_observed_successful_requests=_median_or_none(group, "observed_successful_requests"), + median_observed_input_tokens=_median_or_none(group, "observed_input_tokens"), + median_observed_output_tokens=_median_or_none(group, "observed_output_tokens"), + median_observed_total_tokens=_median_or_none(group, "observed_total_tokens"), + median_observed_failed_requests=_median_or_none(group, "observed_failed_requests"), + median_observed_failed_request_rate=_median_or_none(group, "observed_failed_request_rate"), + median_observed_bridge_fallback_requests=_median_or_none(group, "observed_bridge_fallback_requests"), + median_observed_non_bridge_total_requests=_median_or_none(group, "observed_non_bridge_total_requests"), + median_observed_non_bridge_failed_requests=_median_or_none(group, "observed_non_bridge_failed_requests"), + median_observed_non_bridge_failed_request_rate=_median_or_none( + group, + "observed_non_bridge_failed_request_rate", + ), + total_record_count=total_record_count, + median_record_count=_median_or_none(group, "record_count"), + total_input_text_tokens=total_input_text_tokens, + median_input_text_tokens_total=_median_or_none(group, "input_text_tokens_total"), + median_records_per_pipeline_sec=_median_or_none(group, "records_per_pipeline_sec"), + median_records_per_ndd_sec=_median_or_none(group, "records_per_ndd_sec"), + median_input_text_tokens_per_pipeline_sec=_median_or_none(group, "input_text_tokens_per_pipeline_sec"), + median_input_text_tokens_per_ndd_sec=_median_or_none(group, "input_text_tokens_per_ndd_sec"), + median_topology_endpoint_count=_median_or_none(group, "topology_endpoint_count"), + median_topology_gpu_count=_median_or_none(group, "topology_gpu_count"), + median_topology_tensor_parallelism=_median_or_none(group, "topology_tensor_parallelism"), + median_topology_shard_count=_median_or_none(group, "topology_shard_count"), + median_input_text_tokens_per_endpoint_sec=_median_or_none(group, "input_text_tokens_per_endpoint_sec"), + median_input_text_tokens_per_gpu_sec=_median_or_none(group, "input_text_tokens_per_gpu_sec"), + median_final_entity_count=_median_or_none(group, "final_entity_count"), + total_empty_detection_count=total_empty_detection_count, + empty_detection_rate=_safe_ratio(total_empty_detection_count, total_record_count), + total_empty_detection_with_ground_truth_count=total_empty_detection_with_gt_count, + empty_detection_with_ground_truth_rate=_safe_ratio( + total_empty_detection_with_gt_count, + total_ground_truth_record_count, + ), + total_ground_truth_record_count=total_ground_truth_record_count, + sum_ground_truth_entity_count=ground_truth_entity_count, + sum_entity_true_positive_count=true_positive, + sum_entity_false_positive_count=false_positive, + sum_entity_false_negative_count=false_negative, + micro_entity_precision=strict_precision, + micro_entity_recall=strict_recall, + micro_entity_f1=_f1(strict_precision, strict_recall), + sum_entity_relaxed_gt_found_count=relaxed_gt_found, + sum_entity_relaxed_detected_tp_count=relaxed_detected_tp, + sum_entity_relaxed_label_compatible_gt_found_count=label_compatible_gt_found, + sum_entity_relaxed_label_compatible_detected_tp_count=label_compatible_detected_tp, + micro_entity_relaxed_precision=relaxed_precision, + micro_entity_relaxed_recall=relaxed_recall, + micro_entity_relaxed_f1=_f1(relaxed_precision, relaxed_recall), + micro_entity_relaxed_label_compatible_precision=label_compatible_precision, + micro_entity_relaxed_label_compatible_recall=label_compatible_recall, + micro_entity_relaxed_label_compatible_f1=_f1(label_compatible_precision, label_compatible_recall), + median_entity_relaxed_f1=_median_or_none(group, "entity_relaxed_f1"), + median_entity_relaxed_label_compatible_f1=_median_or_none( + group, + "entity_relaxed_label_compatible_f1", + ), + median_replacement_missing_final_entity_count=_median_or_none( + group, + "replacement_missing_final_entity_count", + ), + median_replacement_missing_final_value_count=_median_or_none(group, "replacement_missing_final_value_count"), + replacement_missing_final_entity_label_counts=_sum_prefixed_ints( + group, + "replacement_missing_final_entity_label_counts.", + ), + median_replacement_synthetic_original_collision_count=_median_or_none( + group, + "replacement_synthetic_original_collision_count", + ), + median_replacement_synthetic_original_collision_value_count=_median_or_none( + group, + "replacement_synthetic_original_collision_value_count", + ), + replacement_synthetic_original_collision_label_counts=_sum_prefixed_ints( + group, + "replacement_synthetic_original_collision_label_counts.", + ), + sum_original_value_leak_count=_sum_or_none(group, "original_value_leak_count"), + leaking_case_count=_positive_count(group, "original_value_leak_count"), + median_original_value_leak_count=_median_or_none(group, "original_value_leak_count"), + **evaluation_metrics, + median_seed_entity_count=_median_or_none(group, "seed_entity_count"), + median_seed_validation_candidate_count=_median_or_none(group, "seed_validation_candidate_count"), + median_estimated_seed_validation_chunk_count=_median_or_none(group, "estimated_seed_validation_chunk_count"), + median_augmented_entity_count=_median_or_none(group, "augmented_entity_count"), + median_augmented_new_final_value_count=_median_or_none(group, "augmented_new_final_value_count"), + median_artifact_final_entity_count=_median_or_none(group, "artifact_final_entity_count"), + median_artifact_final_detector_entity_count=_median_or_none(group, "artifact_final_detector_entity_count"), + median_artifact_final_augmenter_entity_count=_median_or_none(group, "artifact_final_augmenter_entity_count"), + median_artifact_final_entity_signature_count=_median_or_none(group, "artifact_final_entity_signature_count"), + ) + + +def int_if_not_nan(value: object) -> int | None: + if pd.isna(value): + return None + return int(float(cast(Any, value))) + + +def float_if_not_nan(value: object) -> float | None: + if pd.isna(value): + return None + return float(cast(Any, value)) + + +def group_evaluation_metrics(group: pd.DataFrame) -> dict[str, int | float | None]: + metrics: dict[str, int | float | None] = {} + for rollup in _EVALUATION_ROLLUPS: + judged_count = _sum_int_or_zero(group, f"{rollup.prefix}_judged_record_count") + valid_count = _sum_int_or_zero(group, f"{rollup.prefix}_valid_record_count") + metrics[f"sum_{rollup.prefix}_judged_record_count"] = judged_count + metrics[f"sum_{rollup.prefix}_valid_record_count"] = valid_count + metrics[f"micro_{rollup.prefix}_valid_rate"] = _safe_ratio(valid_count, judged_count) + metrics[f"sum_{rollup.invalid_count_column}"] = _sum_int_or_zero(group, rollup.invalid_count_column) + return metrics + +__all__ = [ + "build_group_row", + "build_group_rows", + "float_if_not_nan", + "group_evaluation_metrics", + "int_if_not_nan", +] diff --git a/tools/measurement/measurement_tools/benchmark_model_usage.py b/tools/measurement/measurement_tools/benchmark_model_usage.py new file mode 100644 index 00000000..5a5de951 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_model_usage.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-model request and token usage analysis for benchmark output.""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd + +from measurement_tools.benchmark_analysis_math import model_workflow_rows as _model_workflow_rows +from measurement_tools.benchmark_analysis_math import request_failure_rate as _request_failure_rate +from measurement_tools.benchmark_analysis_math import sum_int_or_none as _sum_int_or_none +from measurement_tools.benchmark_analysis_models import ModelUsageAnalysisRow, ModelUsageGroupAnalysisRow +from measurement_tools.stats import median_or_none as _median_or_none +from measurement_tools.stats import none_if_nan as _none_if_nan +from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero + +_MODEL_USAGE_SUFFIXES = { + ".request_usage.total_requests": "observed_total_requests", + ".request_usage.successful_requests": "observed_successful_requests", + ".request_usage.failed_requests": "observed_failed_requests", + ".token_usage.input_tokens": "observed_input_tokens", + ".token_usage.output_tokens": "observed_output_tokens", + ".token_usage.total_tokens": "observed_total_tokens", + ".token_usage.reasoning_tokens": "observed_reasoning_tokens", +} + + +_MODEL_USAGE_METADATA_SUFFIXES = { + ".model_alias": "model_alias", + ".model_name": "model_name", + ".model_provider_name": "model_provider_name", +} + + +def build_model_usage_rows(measurements: pd.DataFrame) -> list[ModelUsageAnalysisRow]: + model_rows = _model_workflow_rows(measurements) + if model_rows.empty: + return [] + usage_keys = model_usage_keys(model_rows.columns) + rows: list[ModelUsageAnalysisRow] = [] + for _, measurement in model_rows.iterrows(): + data = measurement.to_dict() + case_id = string_from_row(data, ["run_tags.case_id", "run_id"]) + run_id = string_from_row(data, ["run_id", "run_tags.case_id"]) + if case_id is None or run_id is None: + continue + for model_usage_key in usage_keys: + usage = model_usage_metrics(data, model_usage_key) + if not has_observed_model_usage(usage): + continue + metadata = model_usage_metadata(data, model_usage_key) + rows.append( + ModelUsageAnalysisRow( + suite_id=string_from_row(data, ["run_tags.suite_id"]), + workload_id=string_from_row(data, ["run_tags.workload_id"]), + config_id=string_from_row(data, ["run_tags.config_id"]), + experimental_detection_strategy=string_from_row( + data, ["run_tags.experimental_detection_strategy"] + ), + experimental_replacement_strategy=string_from_row( + data, ["run_tags.experimental_replacement_strategy"] + ), + dd_parser_compat=string_from_row(data, ["run_tags.dd_parser_compat"]), + repetition=int_from_row(data, ["run_tags.repetition"]), + case_id=case_id, + run_id=run_id, + workflow_name=string_from_row(data, ["workflow_name"]), + model_alias=metadata.get("model_alias"), + model_name=metadata.get("model_name") or model_usage_key, + model_provider_name=metadata.get("model_provider_name"), + ndd_elapsed_sec=float_from_row(data, ["elapsed_sec"]), + **usage, + ) + ) + return rows + + +def model_usage_keys(columns: pd.Index) -> list[str]: + keys: set[str] = set() + for column in columns: + parsed = model_usage_column_parts(str(column)) + if parsed is not None: + keys.add(parsed[0]) + return sorted(keys) + + +def model_usage_column_parts(column: str) -> tuple[str, str] | None: + prefix = "model_usage." + if not column.startswith(prefix): + return None + for suffix, metric in {**_MODEL_USAGE_SUFFIXES, **_MODEL_USAGE_METADATA_SUFFIXES}.items(): + if column.endswith(suffix): + return column[len(prefix) : -len(suffix)], metric + return None + + +def model_usage_metrics(data: dict[str, Any], model_usage_key: str) -> dict[str, int | float | None]: + values: dict[str, int | float | None] = { + "observed_total_requests": 0, + "observed_successful_requests": 0, + "observed_failed_requests": 0, + "observed_input_tokens": 0, + "observed_output_tokens": 0, + "observed_total_tokens": 0, + "observed_reasoning_tokens": None, + "observed_failed_request_rate": None, + } + for suffix, metric in _MODEL_USAGE_SUFFIXES.items(): + value = data.get(f"model_usage.{model_usage_key}{suffix}") + if value is None or pd.isna(value): + continue + values[metric] = coerce_int(value) + values["observed_failed_request_rate"] = _request_failure_rate( + failed=values["observed_failed_requests"], + total=values["observed_total_requests"], + ) + return values + + +def model_usage_metadata(data: dict[str, Any], model_usage_key: str) -> dict[str, str | None]: + values: dict[str, str | None] = { + "model_alias": None, + "model_name": None, + "model_provider_name": None, + } + for suffix, field_name in _MODEL_USAGE_METADATA_SUFFIXES.items(): + value = data.get(f"model_usage.{model_usage_key}{suffix}") + if value is None or pd.isna(value): + continue + values[field_name] = str(value) + return values + + +def has_observed_model_usage(usage: dict[str, int | float | None]) -> bool: + return any(value not in (None, 0) for value in usage.values()) + + +def string_from_row(data: dict[str, Any], columns: list[str]) -> str | None: + for column in columns: + value = data.get(column) + if value is not None and not pd.isna(value): + return str(value) + return None + + +def int_from_row(data: dict[str, Any], columns: list[str]) -> int | None: + value = string_from_row(data, columns) + return int(float(value)) if value is not None else None + + +def float_from_row(data: dict[str, Any], columns: list[str]) -> float | None: + value = string_from_row(data, columns) + return float(value) if value is not None else None + + +def coerce_int(value: Any) -> int: + return int(float(value)) + + +def build_model_usage_group_rows(model_usage: list[ModelUsageAnalysisRow]) -> list[ModelUsageGroupAnalysisRow]: + if not model_usage: + return [] + table = pd.DataFrame([row.model_dump() for row in model_usage]) + rows: list[ModelUsageGroupAnalysisRow] = [] + group_columns = [ + "workload_id", + "config_id", + "experimental_detection_strategy", + "experimental_replacement_strategy", + "dd_parser_compat", + "workflow_name", + "model_alias", + "model_name", + "model_provider_name", + ] + for keys, group in table.groupby(group_columns, dropna=False): + rows.append(build_model_usage_group_row(keys, group)) + return rows + + +def build_model_usage_group_row(keys: tuple[Any, ...], group: pd.DataFrame) -> ModelUsageGroupAnalysisRow: + ( + workload_id, + config_id, + detection_strategy, + replacement_strategy, + dd_parser_compat, + workflow_name, + model_alias, + model_name, + provider_name, + ) = keys + reasoning_sum = _sum_int_or_none(group, "observed_reasoning_tokens") + total_requests = _sum_int_or_zero(group, "observed_total_requests") + failed_requests = _sum_int_or_zero(group, "observed_failed_requests") + return ModelUsageGroupAnalysisRow( + workload_id=_none_if_nan(workload_id), + config_id=_none_if_nan(config_id), + experimental_detection_strategy=_none_if_nan(detection_strategy), + experimental_replacement_strategy=_none_if_nan(replacement_strategy), + dd_parser_compat=_none_if_nan(dd_parser_compat), + workflow_name=_none_if_nan(workflow_name), + model_alias=_none_if_nan(model_alias), + model_name=str(model_name), + model_provider_name=_none_if_nan(provider_name), + case_count=int(group["case_id"].nunique()), + workflow_count=len(group), + sum_observed_total_requests=total_requests, + sum_observed_successful_requests=_sum_int_or_zero(group, "observed_successful_requests"), + sum_observed_failed_requests=failed_requests, + sum_observed_input_tokens=_sum_int_or_zero(group, "observed_input_tokens"), + sum_observed_output_tokens=_sum_int_or_zero(group, "observed_output_tokens"), + sum_observed_total_tokens=_sum_int_or_zero(group, "observed_total_tokens"), + sum_observed_reasoning_tokens=reasoning_sum, + observed_failed_request_rate=_request_failure_rate(failed=failed_requests, total=total_requests), + median_observed_total_requests=_median_or_none(group, "observed_total_requests"), + median_observed_failed_requests=_median_or_none(group, "observed_failed_requests"), + median_observed_total_tokens=_median_or_none(group, "observed_total_tokens"), + ) + +__all__ = [ + "build_model_usage_group_row", + "build_model_usage_group_rows", + "build_model_usage_rows", + "coerce_int", + "float_from_row", + "has_observed_model_usage", + "int_from_row", + "model_usage_column_parts", + "model_usage_keys", + "model_usage_metadata", + "model_usage_metrics", + "string_from_row", +] From 48203c94b18d7c9ef376bb3e1f7d4cffc0266746 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 20 Jul 2026 22:13:23 +0000 Subject: [PATCH 17/17] refactor(measurement): split case analysis and analysis I/O Signed-off-by: Aaron Gonzales --- tools/measurement/analyze_benchmark_output.py | 667 +----------------- .../benchmark_analysis_io.py | 137 ++++ .../benchmark_case_analysis.py | 552 +++++++++++++++ .../benchmark_group_analysis.py | 16 + .../benchmark_model_usage.py | 5 +- 5 files changed, 718 insertions(+), 659 deletions(-) create mode 100644 tools/measurement/measurement_tools/benchmark_analysis_io.py create mode 100644 tools/measurement/measurement_tools/benchmark_case_analysis.py diff --git a/tools/measurement/analyze_benchmark_output.py b/tools/measurement/analyze_benchmark_output.py index 7686855d..e478b663 100644 --- a/tools/measurement/analyze_benchmark_output.py +++ b/tools/measurement/analyze_benchmark_output.py @@ -12,71 +12,23 @@ from __future__ import annotations -import json import logging -import math import sys from pathlib import Path -from typing import Annotated, Any +from typing import Annotated import cyclopts -import pandas as pd -from measurement_tools.benchmark_analysis_math import ( - coalesce_number as _coalesce_number, +from measurement_tools.benchmark_analysis_io import ( + analyze_benchmark_output as analyze_benchmark_output, ) -from measurement_tools.benchmark_analysis_math import ( - f1 as _f1, +from measurement_tools.benchmark_analysis_io import ( + read_jsonl_table as read_jsonl_table, ) -from measurement_tools.benchmark_analysis_math import ( - first_float as _first_float, +from measurement_tools.benchmark_analysis_io import ( + read_trace_summary_table as read_trace_summary_table, ) -from measurement_tools.benchmark_analysis_math import ( - first_int as _first_int, -) -from measurement_tools.benchmark_analysis_math import ( - first_value as _first_value, -) -from measurement_tools.benchmark_analysis_math import ( - model_workflow_rows as _model_workflow_rows, -) -from measurement_tools.benchmark_analysis_math import ( - non_null_count as _non_null_count, -) -from measurement_tools.benchmark_analysis_math import ( - pipeline_stage_rows as _pipeline_stage_rows, -) -from measurement_tools.benchmark_analysis_math import ( - positive_count as _positive_count, -) -from measurement_tools.benchmark_analysis_math import ( - records_of_type as _records_of_type, -) -from measurement_tools.benchmark_analysis_math import ( - request_failure_rate as _request_failure_rate, -) -from measurement_tools.benchmark_analysis_math import ( - safe_rate as _safe_rate, -) -from measurement_tools.benchmark_analysis_math import ( - safe_ratio as _safe_ratio, -) -from measurement_tools.benchmark_analysis_math import ( - sum_int_or_none as _sum_int_or_none, -) -from measurement_tools.benchmark_analysis_math import ( - sum_optional_numbers as _sum_optional_numbers, -) -from measurement_tools.benchmark_analysis_math import ( - sum_prefixed_ints as _sum_prefixed_ints, -) -from measurement_tools.benchmark_analysis_math import ( - zero_count as _zero_count, -) -from measurement_tools.benchmark_analysis_math import ( - zero_with_positive_count as _zero_with_positive_count, -) -from measurement_tools.benchmark_analysis_models import ( - _EVALUATION_ROLLUPS, +from measurement_tools.benchmark_analysis_io import ( + write_analysis_tables as write_analysis_tables, ) from measurement_tools.benchmark_analysis_models import ( BenchmarkOutputAnalysis as BenchmarkOutputAnalysis, @@ -106,609 +58,12 @@ build_model_usage_rows as build_model_usage_rows, ) from measurement_tools.cli import LogFormat, configure_logging, log_bad_input -from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero -from measurement_tools.stats import sum_or_none as _sum_or_none -from measurement_tools.stats import sum_or_zero as _sum_or_zero -from measurement_tools.tables import AnalysisExportResult, ExportFormat, ModelTableSpec -from measurement_tools.tables import write_analysis_tables as _write_analysis_table_specs +from measurement_tools.tables import AnalysisExportResult as AnalysisExportResult +from measurement_tools.tables import ExportFormat app = cyclopts.App(help=__doc__) logger = logging.getLogger("measurement.benchmark_output") -_SYNC_CLIENT_UNAVAILABLE_ERROR = "SyncClientUnavailableError" -_SIGNATURE_DETAIL_FIELDS = { - "label", - "source", - "row_index", - "start_position", - "end_position", - "value_length", -} - - -def analyze_benchmark_output( - benchmark_dir: Path, - *, - detection_artifacts: Path | None = None, -) -> BenchmarkOutputAnalysis: - measurements = read_jsonl_table(benchmark_dir / "measurements.jsonl", required=True) - artifacts_path = detection_artifacts or benchmark_dir / "detection-artifacts.jsonl" - artifacts = read_jsonl_table(artifacts_path, required=detection_artifacts is not None) - traces = read_trace_summary_table(benchmark_dir / "traces") - cases = [ - _build_case_row(case_id, measurements, artifacts, traces) - for case_id in _case_ids(measurements, artifacts, traces) - ] - model_usage = build_model_usage_rows(measurements) - return BenchmarkOutputAnalysis( - benchmark_dir=str(benchmark_dir), - detection_artifacts_path=str(artifacts_path) if not artifacts.empty else None, - cases=cases, - groups=build_group_rows(cases), - model_usage=model_usage, - model_usage_groups=build_model_usage_group_rows(model_usage), - ) - - -def read_jsonl_table(path: Path, *, required: bool) -> pd.DataFrame: - if not path.exists(): - if required: - raise ValueError(f"input path does not exist: {path}") - return pd.DataFrame() - if path.is_dir(): - raise ValueError(f"input path is a directory: {path}") - raw = pd.read_json(path, lines=True) - if raw.empty: - return raw - return pd.json_normalize(raw.to_dict("records"), sep=".") - - -def read_trace_summary_table(trace_path: Path) -> pd.DataFrame: - """Read DD trace files into a sanitized table with no prompt/response text.""" - if not trace_path.exists(): - return pd.DataFrame() - if trace_path.is_file(): - paths = [trace_path] - elif trace_path.is_dir(): - paths = sorted(trace_path.rglob("*.jsonl")) - else: - raise ValueError(f"trace path is not a file or directory: {trace_path}") - - rows: list[dict[str, Any]] = [] - for path in paths: - for line in path.read_text(encoding="utf-8").splitlines(): - if not line.strip(): - continue - record = json.loads(line) - if not isinstance(record, dict) or record.get("record_type") != "dd_message_trace": - continue - run_tags = record.get("run_tags") if isinstance(record.get("run_tags"), dict) else {} - rows.append( - { - "record_type": "dd_message_trace", - "run_id": record.get("run_id"), - "run_tags.case_id": run_tags.get("case_id"), - "run_tags.workload_id": run_tags.get("workload_id"), - "run_tags.config_id": run_tags.get("config_id"), - "run_tags.experimental_detection_strategy": run_tags.get("experimental_detection_strategy"), - "run_tags.experimental_replacement_strategy": run_tags.get("experimental_replacement_strategy"), - "run_tags.dd_parser_compat": run_tags.get("dd_parser_compat"), - "run_tags.repetition": run_tags.get("repetition"), - "workflow_name": record.get("workflow_name"), - "model_alias": record.get("model_alias"), - "status": record.get("status"), - "error_type": record.get("error_type"), - "is_async": record.get("is_async"), - } - ) - return pd.DataFrame(rows) - - -def _case_ids(*frames: pd.DataFrame) -> list[str]: - values: set[str] = set() - for dataframe in frames: - for column in ("run_tags.case_id", "case_id", "run_id"): - if column in dataframe.columns: - values.update(str(value) for value in dataframe[column].dropna().tolist()) - return sorted(values) - - -def _build_case_row( - case_id: str, - measurements: pd.DataFrame, - artifacts: pd.DataFrame, - traces: pd.DataFrame, -) -> CaseAnalysisRow: - measurement_rows = _rows_for_case(measurements, case_id) - artifact_rows = _rows_for_case(artifacts, case_id) - trace_rows = _rows_for_case(traces, case_id) - record_rows = _records_of_type(measurement_rows, "record") - evaluation_rows = _records_of_type(measurement_rows, "evaluation_record") - ndd_rows = _records_of_type(measurement_rows, "ndd_workflow") - model_rows = _model_workflow_rows(measurement_rows) - stage_rows = _records_of_type(measurement_rows, "stage") - pipeline_rows = _pipeline_stage_rows(measurement_rows) - validation_max_entities_per_call = _first_int([measurement_rows], ["detect.validation_max_entities_per_call"]) - request_metrics = _case_request_metrics(model_rows) - pipeline_elapsed_sec = _sum_or_none(pipeline_rows, "elapsed_sec") - ndd_elapsed_sec_total = _sum_or_zero(ndd_rows, "elapsed_sec") - record_count = len(record_rows) - input_text_tokens_total = _sum_int_or_none(record_rows, "text_length_tokens") - records_per_pipeline_sec = _safe_rate(record_count, pipeline_elapsed_sec) - records_per_ndd_sec = _safe_rate(record_count, ndd_elapsed_sec_total) - input_text_tokens_per_pipeline_sec = _safe_rate(input_text_tokens_total, pipeline_elapsed_sec) - input_text_tokens_per_ndd_sec = _safe_rate(input_text_tokens_total, ndd_elapsed_sec_total) - final_entity_count = _coalesce_number( - _sum_or_none(record_rows, "final_entity_count"), - _sum_or_none(artifact_rows, "final_entity_count"), - ) - return CaseAnalysisRow( - suite_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_tags.suite_id", "suite_id"]), - workload_id=_first_value( - [measurement_rows, artifact_rows, trace_rows], ["run_tags.workload_id", "workload_id"] - ), - workload_category=_first_value( - [measurement_rows, artifact_rows, trace_rows], - ["run_tags.workload_category", "run_tags.dataset_category", "workload_category", "dataset_category"], - ), - config_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_tags.config_id", "config_id"]), - experimental_detection_strategy=_first_value([measurement_rows], ["run_tags.experimental_detection_strategy"]), - experimental_replacement_strategy=_first_value( - [measurement_rows, trace_rows], - ["run_tags.experimental_replacement_strategy"], - ), - dd_parser_compat=_first_value([measurement_rows], ["run_tags.dd_parser_compat"]), - entity_label_set_id=_first_value( - [measurement_rows], - [ - "run_tags.entity_label_set_id", - "run_tags.entity_label_set", - "run_tags.label_set", - "detect.entity_label_source", - ], - ), - entity_label_count=_first_int([measurement_rows], ["run_tags.entity_label_count", "detect.entity_label_count"]), - gliner_threshold=_first_float([measurement_rows], ["run_tags.gliner_threshold", "detect.gliner_threshold"]), - repetition=_first_int([measurement_rows, artifact_rows, trace_rows], ["run_tags.repetition", "repetition"]), - case_id=case_id, - run_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_id"]) or case_id, - **_case_failure_metrics(stage_rows=stage_rows, ndd_rows=ndd_rows, model_rows=model_rows), - pipeline_elapsed_sec=pipeline_elapsed_sec, - ndd_workflow_count=len(ndd_rows), - ndd_elapsed_sec_total=ndd_elapsed_sec_total, - **request_metrics, - **_case_trace_metrics(trace_rows, request_metrics=request_metrics), - record_count=record_count, - input_text_tokens_total=input_text_tokens_total, - records_per_pipeline_sec=records_per_pipeline_sec, - records_per_ndd_sec=records_per_ndd_sec, - input_text_tokens_per_pipeline_sec=input_text_tokens_per_pipeline_sec, - input_text_tokens_per_ndd_sec=input_text_tokens_per_ndd_sec, - **_case_topology_metrics( - measurement_rows, - input_text_tokens_per_pipeline_sec=input_text_tokens_per_pipeline_sec, - ), - final_entity_count=final_entity_count, - **_case_empty_detection_metrics(record_rows, record_count=record_count), - **_case_ground_truth_metrics(record_rows, final_entity_count=final_entity_count), - replacement_count=_sum_or_none(record_rows, "replacement_count"), - replacement_missing_final_entity_count=_sum_or_none(record_rows, "replacement_missing_final_entity_count"), - replacement_missing_final_entity_label_counts=_sum_prefixed_ints( - record_rows, - "replacement_missing_final_entity_label_counts.", - ), - replacement_missing_final_value_count=_sum_or_none(record_rows, "replacement_missing_final_value_count"), - replacement_synthetic_original_collision_count=_sum_or_none( - record_rows, - "replacement_synthetic_original_collision_count", - ), - replacement_synthetic_original_collision_label_counts=_sum_prefixed_ints( - record_rows, - "replacement_synthetic_original_collision_label_counts.", - ), - replacement_synthetic_original_collision_value_count=_sum_or_none( - record_rows, - "replacement_synthetic_original_collision_value_count", - ), - original_value_leak_count=_sum_or_none(record_rows, "original_value_leak_count"), - original_value_leak_record_count=_positive_count(record_rows, "original_value_leak_count"), - original_value_leak_label_counts=_sum_prefixed_ints(record_rows, "original_value_leak_label_counts."), - **_case_evaluation_metrics(evaluation_rows), - validation_max_entities_per_call=validation_max_entities_per_call, - **_case_artifact_metrics( - artifact_rows, - validation_max_entities_per_call=validation_max_entities_per_call, - ), - ) - - -def _case_request_metrics(model_rows: pd.DataFrame) -> dict[str, int | float | None]: - observed_total_requests = int(_sum_or_zero(model_rows, "observed_total_requests")) - observed_failed_requests = int(_sum_or_zero(model_rows, "observed_failed_requests")) - return { - "observed_total_requests": observed_total_requests, - "observed_successful_requests": int(_sum_or_zero(model_rows, "observed_successful_requests")), - "observed_input_tokens": int(_sum_or_zero(model_rows, "observed_input_tokens")), - "observed_output_tokens": int(_sum_or_zero(model_rows, "observed_output_tokens")), - "observed_total_tokens": int(_sum_or_zero(model_rows, "observed_total_tokens")), - "observed_failed_requests": observed_failed_requests, - "observed_failed_request_rate": _request_failure_rate( - failed=observed_failed_requests, - total=observed_total_requests, - ), - } - - -def _case_trace_metrics( - trace_rows: pd.DataFrame, - *, - request_metrics: dict[str, int | float | None], -) -> dict[str, int | float | None]: - trace_record_count = len(trace_rows) - if trace_record_count == 0: - return { - "dd_trace_record_count": 0, - "dd_trace_error_count": 0, - "dd_trace_sync_client_unavailable_count": 0, - "observed_bridge_fallback_requests": None, - "observed_non_bridge_total_requests": None, - "observed_non_bridge_failed_requests": None, - "observed_non_bridge_failed_request_rate": None, - } - - status = trace_rows["status"].astype(str) if "status" in trace_rows.columns else pd.Series(dtype=str) - error_type = trace_rows["error_type"].astype(str) if "error_type" in trace_rows.columns else pd.Series(dtype=str) - error_count = int((status == "error").sum()) - bridge_fallbacks = int(((status == "error") & (error_type == _SYNC_CLIENT_UNAVAILABLE_ERROR)).sum()) - observed_total = int(request_metrics["observed_total_requests"] or 0) - observed_failed = int(request_metrics["observed_failed_requests"] or 0) - non_bridge_total = max(observed_total - bridge_fallbacks, 0) - non_bridge_failed = max(observed_failed - bridge_fallbacks, 0) - return { - "dd_trace_record_count": trace_record_count, - "dd_trace_error_count": error_count, - "dd_trace_sync_client_unavailable_count": bridge_fallbacks, - "observed_bridge_fallback_requests": bridge_fallbacks, - "observed_non_bridge_total_requests": non_bridge_total, - "observed_non_bridge_failed_requests": non_bridge_failed, - "observed_non_bridge_failed_request_rate": _request_failure_rate( - failed=non_bridge_failed, - total=non_bridge_total, - ), - } - - -def _case_failure_metrics( - *, - stage_rows: pd.DataFrame, - ndd_rows: pd.DataFrame, - model_rows: pd.DataFrame, -) -> dict[str, bool | int]: - error_stage_count = _error_status_count(stage_rows) - error_ndd_workflow_count = _error_status_count(ndd_rows) - error_model_workflow_count = _error_status_count(model_rows) - return { - "case_failed": error_stage_count > 0 or error_ndd_workflow_count > 0 or error_model_workflow_count > 0, - "error_stage_count": error_stage_count, - "error_ndd_workflow_count": error_ndd_workflow_count, - "error_model_workflow_count": error_model_workflow_count, - } - - -def _case_topology_metrics( - measurement_rows: pd.DataFrame, - *, - input_text_tokens_per_pipeline_sec: float | None, -) -> dict[str, float | None]: - endpoint_count = _first_float( - [measurement_rows], - [ - "run_tags.topology_endpoint_count", - "run_tags.endpoint_count", - "run_tags.n_endpoints", - "run_tags.n_llm_endpoints", - ], - ) - gpu_count = _first_float( - [measurement_rows], - [ - "run_tags.topology_gpu_count", - "run_tags.gpu_count", - "run_tags.n_gpus", - "run_tags.n_llm_gpus", - ], - ) - tensor_parallelism = _first_float( - [measurement_rows], - [ - "run_tags.topology_tensor_parallelism", - "run_tags.tensor_parallelism", - "run_tags.gpus_per_endpoint", - "run_tags.tp", - ], - ) - shard_count = _first_float( - [measurement_rows], - ["run_tags.topology_shard_count", "run_tags.shard_count", "run_tags.n_shards"], - ) - return { - "topology_endpoint_count": endpoint_count, - "topology_gpu_count": gpu_count, - "topology_tensor_parallelism": tensor_parallelism, - "topology_shard_count": shard_count, - "input_text_tokens_per_endpoint_sec": _safe_ratio(input_text_tokens_per_pipeline_sec, endpoint_count), - "input_text_tokens_per_gpu_sec": _safe_ratio(input_text_tokens_per_pipeline_sec, gpu_count), - } - - -def _case_empty_detection_metrics(record_rows: pd.DataFrame, *, record_count: int) -> dict[str, int | float | None]: - empty_detection_count = _zero_count(record_rows, "final_entity_count") - ground_truth_record_count = _non_null_count(record_rows, "ground_truth_entity_count") - empty_detection_with_gt_count = _zero_with_positive_count( - record_rows, - zero_column="final_entity_count", - positive_column="ground_truth_entity_count", - ) - return { - "empty_detection_count": empty_detection_count, - "empty_detection_rate": _safe_ratio(empty_detection_count, record_count), - "empty_detection_with_ground_truth_count": empty_detection_with_gt_count, - "empty_detection_with_ground_truth_rate": _safe_ratio( - empty_detection_with_gt_count, - ground_truth_record_count, - ), - "ground_truth_record_count": ground_truth_record_count, - } - - -def _case_ground_truth_metrics( - record_rows: pd.DataFrame, - *, - final_entity_count: float | None, -) -> dict[str, float | None]: - ground_truth_entity_count = _sum_or_none(record_rows, "ground_truth_entity_count") - true_positive = _sum_or_none(record_rows, "entity_true_positive_count") - false_positive = _sum_or_none(record_rows, "entity_false_positive_count") - false_negative = _sum_or_none(record_rows, "entity_false_negative_count") - relaxed_gt_found = _sum_or_none(record_rows, "entity_relaxed_gt_found_count") - relaxed_detected_tp = _sum_or_none(record_rows, "entity_relaxed_detected_tp_count") - label_compatible_gt_found = _sum_or_none(record_rows, "entity_relaxed_label_compatible_gt_found_count") - label_compatible_detected_tp = _sum_or_none( - record_rows, - "entity_relaxed_label_compatible_detected_tp_count", - ) - strict_precision = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_positive)) - strict_recall = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_negative)) - relaxed_precision = _safe_ratio(relaxed_detected_tp, final_entity_count) - relaxed_recall = _safe_ratio(relaxed_gt_found, ground_truth_entity_count) - label_compatible_precision = _safe_ratio(label_compatible_detected_tp, final_entity_count) - label_compatible_recall = _safe_ratio(label_compatible_gt_found, ground_truth_entity_count) - return { - "ground_truth_entity_count": ground_truth_entity_count, - "entity_true_positive_count": true_positive, - "entity_false_positive_count": false_positive, - "entity_false_negative_count": false_negative, - "entity_precision": strict_precision, - "entity_recall": strict_recall, - "entity_f1": _f1(strict_precision, strict_recall), - "entity_relaxed_gt_found_count": relaxed_gt_found, - "entity_relaxed_detected_tp_count": relaxed_detected_tp, - "entity_relaxed_label_compatible_gt_found_count": label_compatible_gt_found, - "entity_relaxed_label_compatible_detected_tp_count": label_compatible_detected_tp, - "entity_relaxed_precision": relaxed_precision, - "entity_relaxed_recall": relaxed_recall, - "entity_relaxed_f1": _f1(relaxed_precision, relaxed_recall), - "entity_relaxed_label_compatible_precision": label_compatible_precision, - "entity_relaxed_label_compatible_recall": label_compatible_recall, - "entity_relaxed_label_compatible_f1": _f1(label_compatible_precision, label_compatible_recall), - } - - -def _error_status_count(rows: pd.DataFrame) -> int: - if "status" not in rows.columns: - return 0 - statuses = rows["status"].astype(str).str.lower() - return int(statuses.isin({"error", "failed"}).sum()) - - -def _case_evaluation_metrics(evaluation_rows: pd.DataFrame) -> dict[str, int | float | None]: - metrics: dict[str, int | float | None] = {} - for rollup in _EVALUATION_ROLLUPS: - judged_count, valid_count = _evaluation_judged_and_valid_counts(evaluation_rows, rollup.valid_column) - metrics[f"{rollup.prefix}_judged_record_count"] = judged_count - metrics[f"{rollup.prefix}_valid_record_count"] = valid_count - metrics[f"{rollup.prefix}_valid_rate"] = _safe_ratio(valid_count, judged_count) - metrics[rollup.invalid_count_column] = _sum_int_or_zero(evaluation_rows, rollup.invalid_count_column) - return metrics - - -def _evaluation_judged_and_valid_counts(evaluation_rows: pd.DataFrame, valid_column: str) -> tuple[int, int]: - if valid_column not in evaluation_rows.columns: - return 0, 0 - verdicts = [_optional_bool(value) for value in evaluation_rows[valid_column].tolist()] - judged_count = sum(verdict is not None for verdict in verdicts) - valid_count = sum(verdict is True for verdict in verdicts) - return judged_count, valid_count - - -def _optional_bool(value: object) -> bool | None: - if value is None or pd.isna(value): - return None - if isinstance(value, bool): - return value - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in {"true", "1", "yes"}: - return True - if normalized in {"false", "0", "no"}: - return False - return None - if isinstance(value, int | float): - return bool(value) - return None - - -def _case_artifact_metrics( - artifact_rows: pd.DataFrame, - *, - validation_max_entities_per_call: int | None, -) -> dict[str, int | float | list[str] | dict[str, str] | None]: - signature_hashes = _artifact_signature_hashes(artifact_rows) - return { - "detection_artifact_rows": len(artifact_rows), - "seed_entity_count": _sum_or_none(artifact_rows, "seed_entity_count"), - "seed_validation_candidate_count": _sum_or_none(artifact_rows, "seed_validation_candidate_count"), - "estimated_seed_validation_chunk_count": _estimated_validation_chunk_count( - artifact_rows, - validation_max_entities_per_call=validation_max_entities_per_call, - ), - "augmented_entity_count": _sum_or_none(artifact_rows, "augmented_entity_count"), - "augmented_new_final_value_count": _sum_or_none(artifact_rows, "augmented_new_final_value_count"), - "artifact_final_entity_count": _sum_or_none(artifact_rows, "final_entity_count"), - "artifact_final_detector_entity_count": _sum_or_none(artifact_rows, "final_source_counts.detector"), - "artifact_final_augmenter_entity_count": _sum_or_none(artifact_rows, "final_source_counts.augmenter"), - "artifact_final_entity_signature_count": _signature_count(artifact_rows, signature_hashes=signature_hashes), - "artifact_final_entity_signature_hashes": signature_hashes, - "artifact_final_entity_signature_labels": _artifact_signature_labels(artifact_rows), - "artifact_final_entity_signature_details": _artifact_signature_details(artifact_rows), - } - - -def _rows_for_case(dataframe: pd.DataFrame, case_id: str) -> pd.DataFrame: - if dataframe.empty: - return dataframe - masks = [ - dataframe[column].astype(str) == case_id - for column in ("run_tags.case_id", "case_id", "run_id") - if column in dataframe.columns - ] - if not masks: - return dataframe.iloc[0:0] - mask = masks[0] - for next_mask in masks[1:]: - mask = mask | next_mask - return dataframe[mask] - - -def _artifact_signature_hashes(artifact_rows: pd.DataFrame) -> list[str]: - if "final_entity_signature_hashes" not in artifact_rows.columns: - return [] - values: set[str] = set() - for raw in artifact_rows["final_entity_signature_hashes"].dropna(): - values.update(_coerce_string_list(raw)) - return sorted(values) - - -def _artifact_signature_labels(artifact_rows: pd.DataFrame) -> dict[str, str]: - labels: dict[str, str] = {} - if "final_entity_signature_labels" in artifact_rows.columns: - for raw in artifact_rows["final_entity_signature_labels"].dropna(): - labels.update(_coerce_string_dict(raw)) - for column in artifact_rows.columns: - prefix = "final_entity_signature_labels." - if not column.startswith(prefix): - continue - signature_hash = column.removeprefix(prefix) - for value in artifact_rows[column].dropna(): - labels[signature_hash] = str(value) - return dict(sorted(labels.items())) - - -def _artifact_signature_details(artifact_rows: pd.DataFrame) -> dict[str, dict[str, Any]]: - details: dict[str, dict[str, Any]] = {} - if "final_entity_signature_details" in artifact_rows.columns: - for raw in artifact_rows["final_entity_signature_details"].dropna(): - details.update(_coerce_detail_map(raw)) - prefix = "final_entity_signature_details." - for column in artifact_rows.columns: - if not column.startswith(prefix): - continue - remainder = column.removeprefix(prefix) - signature_hash, _, field = remainder.partition(".") - if not signature_hash or not field: - continue - if field not in _SIGNATURE_DETAIL_FIELDS: - continue - for value in artifact_rows[column].dropna(): - details.setdefault(signature_hash, {})[field] = _json_scalar(value) - return dict(sorted(details.items())) - - -def _coerce_detail_map(raw: object) -> dict[str, dict[str, Any]]: - if isinstance(raw, str): - try: - raw = json.loads(raw) - except json.JSONDecodeError: - return {} - if not isinstance(raw, dict): - return {} - details: dict[str, dict[str, Any]] = {} - for signature_hash, value in raw.items(): - if isinstance(value, dict): - details[str(signature_hash)] = { - str(key): _json_scalar(item) for key, item in value.items() if str(key) in _SIGNATURE_DETAIL_FIELDS - } - return details - - -def _json_scalar(value: object) -> Any: - if hasattr(value, "item"): - try: - return value.item() - except ValueError: - return value - return value - - -def _coerce_string_list(raw: object) -> list[str]: - if isinstance(raw, list): - return [str(item) for item in raw] - return [] - - -def _coerce_string_dict(raw: object) -> dict[str, str]: - if isinstance(raw, dict): - return {str(key): str(value) for key, value in raw.items()} - return {} - - -def _signature_count(artifact_rows: pd.DataFrame, *, signature_hashes: list[str]) -> float | None: - if signature_hashes: - return float(len(signature_hashes)) - return _sum_or_none(artifact_rows, "final_entity_signature_count") - - -def _estimated_validation_chunk_count( - artifact_rows: pd.DataFrame, - *, - validation_max_entities_per_call: int | None, -) -> float | None: - if validation_max_entities_per_call is None or validation_max_entities_per_call <= 0: - return None - if "seed_validation_candidate_count" not in artifact_rows.columns: - return None - counts = pd.to_numeric(artifact_rows["seed_validation_candidate_count"], errors="coerce").dropna() - if counts.empty: - return None - return float(sum(math.ceil(count / validation_max_entities_per_call) for count in counts if count > 0)) - - -def write_analysis_tables( - result: BenchmarkOutputAnalysis, - output_dir: Path, - export_format: ExportFormat, -) -> AnalysisExportResult: - return _write_analysis_table_specs( - output_dir, - export_format, - [ - ModelTableSpec("case_analysis", result.cases, CaseAnalysisRow), - ModelTableSpec("group_analysis", result.groups, GroupAnalysisRow), - ModelTableSpec("model_analysis", result.model_usage, ModelUsageAnalysisRow), - ModelTableSpec("model_group_analysis", result.model_usage_groups, ModelUsageGroupAnalysisRow), - ], - ) - def render_result(result: BenchmarkOutputAnalysis, *, json_output: bool) -> str: if json_output: diff --git a/tools/measurement/measurement_tools/benchmark_analysis_io.py b/tools/measurement/measurement_tools/benchmark_analysis_io.py new file mode 100644 index 00000000..660242a7 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_analysis_io.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ingestion, assembly, and table export for benchmark output analysis.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pandas as pd + +from measurement_tools.benchmark_analysis_models import ( + BenchmarkOutputAnalysis, + CaseAnalysisRow, + GroupAnalysisRow, + ModelUsageAnalysisRow, + ModelUsageGroupAnalysisRow, +) +from measurement_tools.benchmark_case_analysis import build_case_row +from measurement_tools.benchmark_group_analysis import build_group_rows +from measurement_tools.benchmark_model_usage import build_model_usage_group_rows, build_model_usage_rows +from measurement_tools.tables import AnalysisExportResult, ExportFormat, ModelTableSpec +from measurement_tools.tables import write_analysis_tables as _write_analysis_table_specs + + +def analyze_benchmark_output( + benchmark_dir: Path, + *, + detection_artifacts: Path | None = None, +) -> BenchmarkOutputAnalysis: + measurements = read_jsonl_table(benchmark_dir / "measurements.jsonl", required=True) + artifacts_path = detection_artifacts or benchmark_dir / "detection-artifacts.jsonl" + artifacts = read_jsonl_table(artifacts_path, required=detection_artifacts is not None) + traces = read_trace_summary_table(benchmark_dir / "traces") + cases = [ + build_case_row(case_id, measurements, artifacts, traces) + for case_id in case_ids(measurements, artifacts, traces) + ] + model_usage = build_model_usage_rows(measurements) + return BenchmarkOutputAnalysis( + benchmark_dir=str(benchmark_dir), + detection_artifacts_path=str(artifacts_path) if not artifacts.empty else None, + cases=cases, + groups=build_group_rows(cases), + model_usage=model_usage, + model_usage_groups=build_model_usage_group_rows(model_usage), + ) + + +def read_jsonl_table(path: Path, *, required: bool) -> pd.DataFrame: + if not path.exists(): + if required: + raise ValueError(f"input path does not exist: {path}") + return pd.DataFrame() + if path.is_dir(): + raise ValueError(f"input path is a directory: {path}") + raw = pd.read_json(path, lines=True) + if raw.empty: + return raw + return pd.json_normalize(raw.to_dict("records"), sep=".") + + +def read_trace_summary_table(trace_path: Path) -> pd.DataFrame: + """Read DD trace files into a sanitized table with no prompt/response text.""" + if not trace_path.exists(): + return pd.DataFrame() + if trace_path.is_file(): + paths = [trace_path] + elif trace_path.is_dir(): + paths = sorted(trace_path.rglob("*.jsonl")) + else: + raise ValueError(f"trace path is not a file or directory: {trace_path}") + + rows: list[dict[str, Any]] = [] + for path in paths: + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + record = json.loads(line) + if not isinstance(record, dict) or record.get("record_type") != "dd_message_trace": + continue + run_tags = record.get("run_tags") if isinstance(record.get("run_tags"), dict) else {} + rows.append( + { + "record_type": "dd_message_trace", + "run_id": record.get("run_id"), + "run_tags.case_id": run_tags.get("case_id"), + "run_tags.workload_id": run_tags.get("workload_id"), + "run_tags.config_id": run_tags.get("config_id"), + "run_tags.experimental_detection_strategy": run_tags.get("experimental_detection_strategy"), + "run_tags.experimental_replacement_strategy": run_tags.get("experimental_replacement_strategy"), + "run_tags.dd_parser_compat": run_tags.get("dd_parser_compat"), + "run_tags.repetition": run_tags.get("repetition"), + "workflow_name": record.get("workflow_name"), + "model_alias": record.get("model_alias"), + "status": record.get("status"), + "error_type": record.get("error_type"), + "is_async": record.get("is_async"), + } + ) + return pd.DataFrame(rows) + + +def case_ids(*frames: pd.DataFrame) -> list[str]: + values: set[str] = set() + for dataframe in frames: + for column in ("run_tags.case_id", "case_id", "run_id"): + if column in dataframe.columns: + values.update(str(value) for value in dataframe[column].dropna().tolist()) + return sorted(values) + + +def write_analysis_tables( + result: BenchmarkOutputAnalysis, + output_dir: Path, + export_format: ExportFormat, +) -> AnalysisExportResult: + return _write_analysis_table_specs( + output_dir, + export_format, + [ + ModelTableSpec("case_analysis", result.cases, CaseAnalysisRow), + ModelTableSpec("group_analysis", result.groups, GroupAnalysisRow), + ModelTableSpec("model_analysis", result.model_usage, ModelUsageAnalysisRow), + ModelTableSpec("model_group_analysis", result.model_usage_groups, ModelUsageGroupAnalysisRow), + ], + ) + + +__all__ = [ + "analyze_benchmark_output", + "case_ids", + "read_jsonl_table", + "read_trace_summary_table", + "write_analysis_tables", +] diff --git a/tools/measurement/measurement_tools/benchmark_case_analysis.py b/tools/measurement/measurement_tools/benchmark_case_analysis.py new file mode 100644 index 00000000..83cddd71 --- /dev/null +++ b/tools/measurement/measurement_tools/benchmark_case_analysis.py @@ -0,0 +1,552 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-case metrics and artifact analysis for benchmark output.""" + +from __future__ import annotations + +import json +import math +from typing import Any + +import pandas as pd + +from measurement_tools.benchmark_analysis_math import coalesce_number as _coalesce_number +from measurement_tools.benchmark_analysis_math import f1 as _f1 +from measurement_tools.benchmark_analysis_math import first_float as _first_float +from measurement_tools.benchmark_analysis_math import first_int as _first_int +from measurement_tools.benchmark_analysis_math import first_value as _first_value +from measurement_tools.benchmark_analysis_math import model_workflow_rows as _model_workflow_rows +from measurement_tools.benchmark_analysis_math import non_null_count as _non_null_count +from measurement_tools.benchmark_analysis_math import pipeline_stage_rows as _pipeline_stage_rows +from measurement_tools.benchmark_analysis_math import positive_count as _positive_count +from measurement_tools.benchmark_analysis_math import records_of_type as _records_of_type +from measurement_tools.benchmark_analysis_math import request_failure_rate as _request_failure_rate +from measurement_tools.benchmark_analysis_math import safe_rate as _safe_rate +from measurement_tools.benchmark_analysis_math import safe_ratio as _safe_ratio +from measurement_tools.benchmark_analysis_math import sum_int_or_none as _sum_int_or_none +from measurement_tools.benchmark_analysis_math import sum_optional_numbers as _sum_optional_numbers +from measurement_tools.benchmark_analysis_math import sum_prefixed_ints as _sum_prefixed_ints +from measurement_tools.benchmark_analysis_math import zero_count as _zero_count +from measurement_tools.benchmark_analysis_math import zero_with_positive_count as _zero_with_positive_count +from measurement_tools.benchmark_analysis_models import _EVALUATION_ROLLUPS, CaseAnalysisRow +from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero +from measurement_tools.stats import sum_or_none as _sum_or_none +from measurement_tools.stats import sum_or_zero as _sum_or_zero + +_SYNC_CLIENT_UNAVAILABLE_ERROR = "SyncClientUnavailableError" + + +_SIGNATURE_DETAIL_FIELDS = { + "label", + "source", + "row_index", + "start_position", + "end_position", + "value_length", +} + + +def build_case_row( + case_id: str, + measurements: pd.DataFrame, + artifacts: pd.DataFrame, + traces: pd.DataFrame, +) -> CaseAnalysisRow: + measurement_rows = rows_for_case(measurements, case_id) + artifact_rows = rows_for_case(artifacts, case_id) + trace_rows = rows_for_case(traces, case_id) + record_rows = _records_of_type(measurement_rows, "record") + evaluation_rows = _records_of_type(measurement_rows, "evaluation_record") + ndd_rows = _records_of_type(measurement_rows, "ndd_workflow") + model_rows = _model_workflow_rows(measurement_rows) + stage_rows = _records_of_type(measurement_rows, "stage") + pipeline_rows = _pipeline_stage_rows(measurement_rows) + validation_max_entities_per_call = _first_int([measurement_rows], ["detect.validation_max_entities_per_call"]) + request_metrics = case_request_metrics(model_rows) + pipeline_elapsed_sec = _sum_or_none(pipeline_rows, "elapsed_sec") + ndd_elapsed_sec_total = _sum_or_zero(ndd_rows, "elapsed_sec") + record_count = len(record_rows) + input_text_tokens_total = _sum_int_or_none(record_rows, "text_length_tokens") + records_per_pipeline_sec = _safe_rate(record_count, pipeline_elapsed_sec) + records_per_ndd_sec = _safe_rate(record_count, ndd_elapsed_sec_total) + input_text_tokens_per_pipeline_sec = _safe_rate(input_text_tokens_total, pipeline_elapsed_sec) + input_text_tokens_per_ndd_sec = _safe_rate(input_text_tokens_total, ndd_elapsed_sec_total) + final_entity_count = _coalesce_number( + _sum_or_none(record_rows, "final_entity_count"), + _sum_or_none(artifact_rows, "final_entity_count"), + ) + return CaseAnalysisRow( + suite_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_tags.suite_id", "suite_id"]), + workload_id=_first_value( + [measurement_rows, artifact_rows, trace_rows], ["run_tags.workload_id", "workload_id"] + ), + workload_category=_first_value( + [measurement_rows, artifact_rows, trace_rows], + ["run_tags.workload_category", "run_tags.dataset_category", "workload_category", "dataset_category"], + ), + config_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_tags.config_id", "config_id"]), + experimental_detection_strategy=_first_value([measurement_rows], ["run_tags.experimental_detection_strategy"]), + experimental_replacement_strategy=_first_value( + [measurement_rows, trace_rows], + ["run_tags.experimental_replacement_strategy"], + ), + dd_parser_compat=_first_value([measurement_rows], ["run_tags.dd_parser_compat"]), + entity_label_set_id=_first_value( + [measurement_rows], + [ + "run_tags.entity_label_set_id", + "run_tags.entity_label_set", + "run_tags.label_set", + "detect.entity_label_source", + ], + ), + entity_label_count=_first_int([measurement_rows], ["run_tags.entity_label_count", "detect.entity_label_count"]), + gliner_threshold=_first_float([measurement_rows], ["run_tags.gliner_threshold", "detect.gliner_threshold"]), + repetition=_first_int([measurement_rows, artifact_rows, trace_rows], ["run_tags.repetition", "repetition"]), + case_id=case_id, + run_id=_first_value([measurement_rows, artifact_rows, trace_rows], ["run_id"]) or case_id, + **case_failure_metrics(stage_rows=stage_rows, ndd_rows=ndd_rows, model_rows=model_rows), + pipeline_elapsed_sec=pipeline_elapsed_sec, + ndd_workflow_count=len(ndd_rows), + ndd_elapsed_sec_total=ndd_elapsed_sec_total, + **request_metrics, + **case_trace_metrics(trace_rows, request_metrics=request_metrics), + record_count=record_count, + input_text_tokens_total=input_text_tokens_total, + records_per_pipeline_sec=records_per_pipeline_sec, + records_per_ndd_sec=records_per_ndd_sec, + input_text_tokens_per_pipeline_sec=input_text_tokens_per_pipeline_sec, + input_text_tokens_per_ndd_sec=input_text_tokens_per_ndd_sec, + **case_topology_metrics( + measurement_rows, + input_text_tokens_per_pipeline_sec=input_text_tokens_per_pipeline_sec, + ), + final_entity_count=final_entity_count, + **case_empty_detection_metrics(record_rows, record_count=record_count), + **case_ground_truth_metrics(record_rows, final_entity_count=final_entity_count), + replacement_count=_sum_or_none(record_rows, "replacement_count"), + replacement_missing_final_entity_count=_sum_or_none(record_rows, "replacement_missing_final_entity_count"), + replacement_missing_final_entity_label_counts=_sum_prefixed_ints( + record_rows, + "replacement_missing_final_entity_label_counts.", + ), + replacement_missing_final_value_count=_sum_or_none(record_rows, "replacement_missing_final_value_count"), + replacement_synthetic_original_collision_count=_sum_or_none( + record_rows, + "replacement_synthetic_original_collision_count", + ), + replacement_synthetic_original_collision_label_counts=_sum_prefixed_ints( + record_rows, + "replacement_synthetic_original_collision_label_counts.", + ), + replacement_synthetic_original_collision_value_count=_sum_or_none( + record_rows, + "replacement_synthetic_original_collision_value_count", + ), + original_value_leak_count=_sum_or_none(record_rows, "original_value_leak_count"), + original_value_leak_record_count=_positive_count(record_rows, "original_value_leak_count"), + original_value_leak_label_counts=_sum_prefixed_ints(record_rows, "original_value_leak_label_counts."), + **case_evaluation_metrics(evaluation_rows), + validation_max_entities_per_call=validation_max_entities_per_call, + **case_artifact_metrics( + artifact_rows, + validation_max_entities_per_call=validation_max_entities_per_call, + ), + ) + + +def case_request_metrics(model_rows: pd.DataFrame) -> dict[str, int | float | None]: + observed_total_requests = int(_sum_or_zero(model_rows, "observed_total_requests")) + observed_failed_requests = int(_sum_or_zero(model_rows, "observed_failed_requests")) + return { + "observed_total_requests": observed_total_requests, + "observed_successful_requests": int(_sum_or_zero(model_rows, "observed_successful_requests")), + "observed_input_tokens": int(_sum_or_zero(model_rows, "observed_input_tokens")), + "observed_output_tokens": int(_sum_or_zero(model_rows, "observed_output_tokens")), + "observed_total_tokens": int(_sum_or_zero(model_rows, "observed_total_tokens")), + "observed_failed_requests": observed_failed_requests, + "observed_failed_request_rate": _request_failure_rate( + failed=observed_failed_requests, + total=observed_total_requests, + ), + } + + +def case_trace_metrics( + trace_rows: pd.DataFrame, + *, + request_metrics: dict[str, int | float | None], +) -> dict[str, int | float | None]: + trace_record_count = len(trace_rows) + if trace_record_count == 0: + return { + "dd_trace_record_count": 0, + "dd_trace_error_count": 0, + "dd_trace_sync_client_unavailable_count": 0, + "observed_bridge_fallback_requests": None, + "observed_non_bridge_total_requests": None, + "observed_non_bridge_failed_requests": None, + "observed_non_bridge_failed_request_rate": None, + } + + status = trace_rows["status"].astype(str) if "status" in trace_rows.columns else pd.Series(dtype=str) + error_type = trace_rows["error_type"].astype(str) if "error_type" in trace_rows.columns else pd.Series(dtype=str) + error_count = int((status == "error").sum()) + bridge_fallbacks = int(((status == "error") & (error_type == _SYNC_CLIENT_UNAVAILABLE_ERROR)).sum()) + observed_total = int(request_metrics["observed_total_requests"] or 0) + observed_failed = int(request_metrics["observed_failed_requests"] or 0) + non_bridge_total = max(observed_total - bridge_fallbacks, 0) + non_bridge_failed = max(observed_failed - bridge_fallbacks, 0) + return { + "dd_trace_record_count": trace_record_count, + "dd_trace_error_count": error_count, + "dd_trace_sync_client_unavailable_count": bridge_fallbacks, + "observed_bridge_fallback_requests": bridge_fallbacks, + "observed_non_bridge_total_requests": non_bridge_total, + "observed_non_bridge_failed_requests": non_bridge_failed, + "observed_non_bridge_failed_request_rate": _request_failure_rate( + failed=non_bridge_failed, + total=non_bridge_total, + ), + } + + +def case_failure_metrics( + *, + stage_rows: pd.DataFrame, + ndd_rows: pd.DataFrame, + model_rows: pd.DataFrame, +) -> dict[str, bool | int]: + error_stage_count = error_status_count(stage_rows) + error_ndd_workflow_count = error_status_count(ndd_rows) + error_model_workflow_count = error_status_count(model_rows) + return { + "case_failed": error_stage_count > 0 or error_ndd_workflow_count > 0 or error_model_workflow_count > 0, + "error_stage_count": error_stage_count, + "error_ndd_workflow_count": error_ndd_workflow_count, + "error_model_workflow_count": error_model_workflow_count, + } + + +def case_topology_metrics( + measurement_rows: pd.DataFrame, + *, + input_text_tokens_per_pipeline_sec: float | None, +) -> dict[str, float | None]: + endpoint_count = _first_float( + [measurement_rows], + [ + "run_tags.topology_endpoint_count", + "run_tags.endpoint_count", + "run_tags.n_endpoints", + "run_tags.n_llm_endpoints", + ], + ) + gpu_count = _first_float( + [measurement_rows], + [ + "run_tags.topology_gpu_count", + "run_tags.gpu_count", + "run_tags.n_gpus", + "run_tags.n_llm_gpus", + ], + ) + tensor_parallelism = _first_float( + [measurement_rows], + [ + "run_tags.topology_tensor_parallelism", + "run_tags.tensor_parallelism", + "run_tags.gpus_per_endpoint", + "run_tags.tp", + ], + ) + shard_count = _first_float( + [measurement_rows], + ["run_tags.topology_shard_count", "run_tags.shard_count", "run_tags.n_shards"], + ) + return { + "topology_endpoint_count": endpoint_count, + "topology_gpu_count": gpu_count, + "topology_tensor_parallelism": tensor_parallelism, + "topology_shard_count": shard_count, + "input_text_tokens_per_endpoint_sec": _safe_ratio(input_text_tokens_per_pipeline_sec, endpoint_count), + "input_text_tokens_per_gpu_sec": _safe_ratio(input_text_tokens_per_pipeline_sec, gpu_count), + } + + +def case_empty_detection_metrics(record_rows: pd.DataFrame, *, record_count: int) -> dict[str, int | float | None]: + empty_detection_count = _zero_count(record_rows, "final_entity_count") + ground_truth_record_count = _non_null_count(record_rows, "ground_truth_entity_count") + empty_detection_with_gt_count = _zero_with_positive_count( + record_rows, + zero_column="final_entity_count", + positive_column="ground_truth_entity_count", + ) + return { + "empty_detection_count": empty_detection_count, + "empty_detection_rate": _safe_ratio(empty_detection_count, record_count), + "empty_detection_with_ground_truth_count": empty_detection_with_gt_count, + "empty_detection_with_ground_truth_rate": _safe_ratio( + empty_detection_with_gt_count, + ground_truth_record_count, + ), + "ground_truth_record_count": ground_truth_record_count, + } + + +def case_ground_truth_metrics( + record_rows: pd.DataFrame, + *, + final_entity_count: float | None, +) -> dict[str, float | None]: + ground_truth_entity_count = _sum_or_none(record_rows, "ground_truth_entity_count") + true_positive = _sum_or_none(record_rows, "entity_true_positive_count") + false_positive = _sum_or_none(record_rows, "entity_false_positive_count") + false_negative = _sum_or_none(record_rows, "entity_false_negative_count") + relaxed_gt_found = _sum_or_none(record_rows, "entity_relaxed_gt_found_count") + relaxed_detected_tp = _sum_or_none(record_rows, "entity_relaxed_detected_tp_count") + label_compatible_gt_found = _sum_or_none(record_rows, "entity_relaxed_label_compatible_gt_found_count") + label_compatible_detected_tp = _sum_or_none( + record_rows, + "entity_relaxed_label_compatible_detected_tp_count", + ) + strict_precision = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_positive)) + strict_recall = _safe_ratio(true_positive, _sum_optional_numbers(true_positive, false_negative)) + relaxed_precision = _safe_ratio(relaxed_detected_tp, final_entity_count) + relaxed_recall = _safe_ratio(relaxed_gt_found, ground_truth_entity_count) + label_compatible_precision = _safe_ratio(label_compatible_detected_tp, final_entity_count) + label_compatible_recall = _safe_ratio(label_compatible_gt_found, ground_truth_entity_count) + return { + "ground_truth_entity_count": ground_truth_entity_count, + "entity_true_positive_count": true_positive, + "entity_false_positive_count": false_positive, + "entity_false_negative_count": false_negative, + "entity_precision": strict_precision, + "entity_recall": strict_recall, + "entity_f1": _f1(strict_precision, strict_recall), + "entity_relaxed_gt_found_count": relaxed_gt_found, + "entity_relaxed_detected_tp_count": relaxed_detected_tp, + "entity_relaxed_label_compatible_gt_found_count": label_compatible_gt_found, + "entity_relaxed_label_compatible_detected_tp_count": label_compatible_detected_tp, + "entity_relaxed_precision": relaxed_precision, + "entity_relaxed_recall": relaxed_recall, + "entity_relaxed_f1": _f1(relaxed_precision, relaxed_recall), + "entity_relaxed_label_compatible_precision": label_compatible_precision, + "entity_relaxed_label_compatible_recall": label_compatible_recall, + "entity_relaxed_label_compatible_f1": _f1(label_compatible_precision, label_compatible_recall), + } + + +def error_status_count(rows: pd.DataFrame) -> int: + if "status" not in rows.columns: + return 0 + statuses = rows["status"].astype(str).str.lower() + return int(statuses.isin({"error", "failed"}).sum()) + + +def case_evaluation_metrics(evaluation_rows: pd.DataFrame) -> dict[str, int | float | None]: + metrics: dict[str, int | float | None] = {} + for rollup in _EVALUATION_ROLLUPS: + judged_count, valid_count = evaluation_judged_and_valid_counts(evaluation_rows, rollup.valid_column) + metrics[f"{rollup.prefix}_judged_record_count"] = judged_count + metrics[f"{rollup.prefix}_valid_record_count"] = valid_count + metrics[f"{rollup.prefix}_valid_rate"] = _safe_ratio(valid_count, judged_count) + metrics[rollup.invalid_count_column] = _sum_int_or_zero(evaluation_rows, rollup.invalid_count_column) + return metrics + + +def evaluation_judged_and_valid_counts(evaluation_rows: pd.DataFrame, valid_column: str) -> tuple[int, int]: + if valid_column not in evaluation_rows.columns: + return 0, 0 + verdicts = [optional_bool(value) for value in evaluation_rows[valid_column].tolist()] + judged_count = sum(verdict is not None for verdict in verdicts) + valid_count = sum(verdict is True for verdict in verdicts) + return judged_count, valid_count + + +def optional_bool(value: object) -> bool | None: + if value is None or pd.isna(value): + return None + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + return None + if isinstance(value, int | float): + return bool(value) + return None + + +def case_artifact_metrics( + artifact_rows: pd.DataFrame, + *, + validation_max_entities_per_call: int | None, +) -> dict[str, int | float | list[str] | dict[str, str] | None]: + signature_hashes = artifact_signature_hashes(artifact_rows) + return { + "detection_artifact_rows": len(artifact_rows), + "seed_entity_count": _sum_or_none(artifact_rows, "seed_entity_count"), + "seed_validation_candidate_count": _sum_or_none(artifact_rows, "seed_validation_candidate_count"), + "estimated_seed_validation_chunk_count": estimated_validation_chunk_count( + artifact_rows, + validation_max_entities_per_call=validation_max_entities_per_call, + ), + "augmented_entity_count": _sum_or_none(artifact_rows, "augmented_entity_count"), + "augmented_new_final_value_count": _sum_or_none(artifact_rows, "augmented_new_final_value_count"), + "artifact_final_entity_count": _sum_or_none(artifact_rows, "final_entity_count"), + "artifact_final_detector_entity_count": _sum_or_none(artifact_rows, "final_source_counts.detector"), + "artifact_final_augmenter_entity_count": _sum_or_none(artifact_rows, "final_source_counts.augmenter"), + "artifact_final_entity_signature_count": signature_count(artifact_rows, signature_hashes=signature_hashes), + "artifact_final_entity_signature_hashes": signature_hashes, + "artifact_final_entity_signature_labels": artifact_signature_labels(artifact_rows), + "artifact_final_entity_signature_details": artifact_signature_details(artifact_rows), + } + + +def rows_for_case(dataframe: pd.DataFrame, case_id: str) -> pd.DataFrame: + if dataframe.empty: + return dataframe + masks = [ + dataframe[column].astype(str) == case_id + for column in ("run_tags.case_id", "case_id", "run_id") + if column in dataframe.columns + ] + if not masks: + return dataframe.iloc[0:0] + mask = masks[0] + for next_mask in masks[1:]: + mask = mask | next_mask + return dataframe[mask] + + +def artifact_signature_hashes(artifact_rows: pd.DataFrame) -> list[str]: + if "final_entity_signature_hashes" not in artifact_rows.columns: + return [] + values: set[str] = set() + for raw in artifact_rows["final_entity_signature_hashes"].dropna(): + values.update(coerce_string_list(raw)) + return sorted(values) + + +def artifact_signature_labels(artifact_rows: pd.DataFrame) -> dict[str, str]: + labels: dict[str, str] = {} + if "final_entity_signature_labels" in artifact_rows.columns: + for raw in artifact_rows["final_entity_signature_labels"].dropna(): + labels.update(coerce_string_dict(raw)) + for column in artifact_rows.columns: + prefix = "final_entity_signature_labels." + if not column.startswith(prefix): + continue + signature_hash = column.removeprefix(prefix) + for value in artifact_rows[column].dropna(): + labels[signature_hash] = str(value) + return dict(sorted(labels.items())) + + +def artifact_signature_details(artifact_rows: pd.DataFrame) -> dict[str, dict[str, Any]]: + details: dict[str, dict[str, Any]] = {} + if "final_entity_signature_details" in artifact_rows.columns: + for raw in artifact_rows["final_entity_signature_details"].dropna(): + details.update(coerce_detail_map(raw)) + prefix = "final_entity_signature_details." + for column in artifact_rows.columns: + if not column.startswith(prefix): + continue + remainder = column.removeprefix(prefix) + signature_hash, _, field = remainder.partition(".") + if not signature_hash or not field: + continue + if field not in _SIGNATURE_DETAIL_FIELDS: + continue + for value in artifact_rows[column].dropna(): + details.setdefault(signature_hash, {})[field] = json_scalar(value) + return dict(sorted(details.items())) + + +def coerce_detail_map(raw: object) -> dict[str, dict[str, Any]]: + if isinstance(raw, str): + try: + raw = json.loads(raw) + except json.JSONDecodeError: + return {} + if not isinstance(raw, dict): + return {} + details: dict[str, dict[str, Any]] = {} + for signature_hash, value in raw.items(): + if isinstance(value, dict): + details[str(signature_hash)] = { + str(key): json_scalar(item) for key, item in value.items() if str(key) in _SIGNATURE_DETAIL_FIELDS + } + return details + + +def json_scalar(value: object) -> Any: + if hasattr(value, "item"): + try: + return value.item() + except ValueError: + return value + return value + + +def coerce_string_list(raw: object) -> list[str]: + if isinstance(raw, list): + return [str(item) for item in raw] + return [] + + +def coerce_string_dict(raw: object) -> dict[str, str]: + if isinstance(raw, dict): + return {str(key): str(value) for key, value in raw.items()} + return {} + + +def signature_count(artifact_rows: pd.DataFrame, *, signature_hashes: list[str]) -> float | None: + if signature_hashes: + return float(len(signature_hashes)) + return _sum_or_none(artifact_rows, "final_entity_signature_count") + + +def estimated_validation_chunk_count( + artifact_rows: pd.DataFrame, + *, + validation_max_entities_per_call: int | None, +) -> float | None: + if validation_max_entities_per_call is None or validation_max_entities_per_call <= 0: + return None + if "seed_validation_candidate_count" not in artifact_rows.columns: + return None + counts = pd.to_numeric(artifact_rows["seed_validation_candidate_count"], errors="coerce").dropna() + if counts.empty: + return None + return float(sum(math.ceil(count / validation_max_entities_per_call) for count in counts if count > 0)) + + +__all__ = [ + "artifact_signature_details", + "artifact_signature_hashes", + "artifact_signature_labels", + "build_case_row", + "case_artifact_metrics", + "case_empty_detection_metrics", + "case_evaluation_metrics", + "case_failure_metrics", + "case_ground_truth_metrics", + "case_request_metrics", + "case_topology_metrics", + "case_trace_metrics", + "coerce_detail_map", + "coerce_string_dict", + "coerce_string_list", + "error_status_count", + "estimated_validation_chunk_count", + "evaluation_judged_and_valid_counts", + "json_scalar", + "optional_bool", + "rows_for_case", + "signature_count", +] diff --git a/tools/measurement/measurement_tools/benchmark_group_analysis.py b/tools/measurement/measurement_tools/benchmark_group_analysis.py index 7b715495..1ec46323 100644 --- a/tools/measurement/measurement_tools/benchmark_group_analysis.py +++ b/tools/measurement/measurement_tools/benchmark_group_analysis.py @@ -10,12 +10,26 @@ from measurement_tools.benchmark_analysis_math import ( f1 as _f1, +) +from measurement_tools.benchmark_analysis_math import ( positive_count as _positive_count, +) +from measurement_tools.benchmark_analysis_math import ( request_failure_rate as _request_failure_rate, +) +from measurement_tools.benchmark_analysis_math import ( safe_ratio as _safe_ratio, +) +from measurement_tools.benchmark_analysis_math import ( sum_bool_or_zero as _sum_bool_or_zero, +) +from measurement_tools.benchmark_analysis_math import ( sum_int_or_none as _sum_int_or_none, +) +from measurement_tools.benchmark_analysis_math import ( sum_optional_numbers as _sum_optional_numbers, +) +from measurement_tools.benchmark_analysis_math import ( sum_prefixed_ints as _sum_prefixed_ints, ) from measurement_tools.benchmark_analysis_models import _EVALUATION_ROLLUPS, CaseAnalysisRow, GroupAnalysisRow @@ -24,6 +38,7 @@ from measurement_tools.stats import sum_int_or_zero as _sum_int_or_zero from measurement_tools.stats import sum_or_none as _sum_or_none + def build_group_rows(cases: list[CaseAnalysisRow]) -> list[GroupAnalysisRow]: if not cases: return [] @@ -214,6 +229,7 @@ def group_evaluation_metrics(group: pd.DataFrame) -> dict[str, int | float | Non metrics[f"sum_{rollup.invalid_count_column}"] = _sum_int_or_zero(group, rollup.invalid_count_column) return metrics + __all__ = [ "build_group_row", "build_group_rows", diff --git a/tools/measurement/measurement_tools/benchmark_model_usage.py b/tools/measurement/measurement_tools/benchmark_model_usage.py index 5a5de951..a9aabe16 100644 --- a/tools/measurement/measurement_tools/benchmark_model_usage.py +++ b/tools/measurement/measurement_tools/benchmark_model_usage.py @@ -56,9 +56,7 @@ def build_model_usage_rows(measurements: pd.DataFrame) -> list[ModelUsageAnalysi suite_id=string_from_row(data, ["run_tags.suite_id"]), workload_id=string_from_row(data, ["run_tags.workload_id"]), config_id=string_from_row(data, ["run_tags.config_id"]), - experimental_detection_strategy=string_from_row( - data, ["run_tags.experimental_detection_strategy"] - ), + experimental_detection_strategy=string_from_row(data, ["run_tags.experimental_detection_strategy"]), experimental_replacement_strategy=string_from_row( data, ["run_tags.experimental_replacement_strategy"] ), @@ -220,6 +218,7 @@ def build_model_usage_group_row(keys: tuple[Any, ...], group: pd.DataFrame) -> M median_observed_total_tokens=_median_or_none(group, "observed_total_tokens"), ) + __all__ = [ "build_model_usage_group_row", "build_model_usage_group_rows",