Merge pull request #50 from izzet/feature/console-additional-metrics-output

Add view-scoped additional metrics and console reporting

Author: hariharan-devarajan

python/dftracer/analyzer/analyzer.py

@@ -951,6 +951,10 @@ def _analyze_hlm(
             _hlms=hlms,
             _main_views=main_views,
             _metric_boundaries=metric_boundaries,
+            additional_metrics={
+                view_type: list(metrics.keys())
+                for view_type, metrics in (self.preset.additional_metrics or {}).items()
+            },  # type: ignore
             checkpoint_dir=self.checkpoint_dir,
             flat_views=flat_views,
             layers=self.layers,
@@ -1112,16 +1116,22 @@ def _process_flat_view(
                 time_boundary_layer=self.get_time_boundary_layer(),
             )
         with log_block("set_additional_metrics", view_key=view_key):
-            flat_view = self._set_additional_metrics(flat_view, is_view_process_based=is_view_process_based)
+            flat_view = self._set_additional_metrics(
+                flat_view,
+                view_key=view_key,
+            )
         return flat_view.sort_index(axis=1)
 
     @staticmethod
     def _save_flat_view(view: pd.DataFrame, view_path: str):
         view.to_parquet(f"{view_path}.parquet")
 
-    def _set_additional_metrics(self, view: pd.DataFrame, is_view_process_based: bool, epsilon=1e-9) -> pd.DataFrame:
+    def _set_additional_metrics(self, view: pd.DataFrame, view_key: ViewKey, epsilon=1e-9) -> pd.DataFrame:
+        view_type = view_key[-1]
+        is_view_process_based = self.is_view_process_based(view_key)
         time_metric = "time_sum" if is_view_process_based else "time_max"
-        for metric, eval_condition in self.preset.additional_metrics.items():
+        view_additional_metrics = (self.preset.additional_metrics or {}).get(view_type, {})
+        for metric, eval_condition in view_additional_metrics.items():
             eval_condition = eval_condition.format(
                 epsilon=epsilon,
                 time_interval=self.time_granularity,

python/dftracer/analyzer/config.py

@@ -31,7 +31,7 @@
 
 @dc.dataclass
 class AnalyzerPresetConfig:
-    additional_metrics: Optional[Dict[str, Optional[str]]] = dc.field(default_factory=dict)
+    additional_metrics: Optional[Dict[str, Dict[str, str]]] = dc.field(default_factory=dict)
     async_layers: Optional[List[str]] = dc.field(default_factory=list)
     derived_metrics: Optional[Dict[str, Dict[str, str]]] = dc.field(default_factory=dict)
     layer_defs: Dict[str, Optional[str]] = MISSING

python/dftracer/analyzer/output.py

@@ -3,12 +3,13 @@
 import dask
 import dataclasses as dc
 import inflect
+import numpy as np
 import pandas as pd
 from rich.console import Console
 from rich.table import Table
 from typing import Dict, List, Optional
 
-from .constants import COL_PROC_NAME, HUMANIZED_LAYERS, Layer, MiB
+from .constants import COL_PROC_NAME, HUMANIZED_LAYERS, GiB, Layer, MiB
 from .types import (
     AnalyzerResultType,
     RawStats,
@@ -155,6 +156,9 @@ def handle_result(self, result: AnalyzerResultType):
             summary_table = self._create_summary_table(summary=summary, view_key=view_key)
             layer_breakdown_table = self._create_layer_breakdown_table(summary=summary, view_key=view_key)
             print_objects.append(summary_table)
+            additional_metrics_table = self._create_additional_metrics_table(result=result, view_key=view_key)
+            if additional_metrics_table is not None:
+                print_objects.append(additional_metrics_table)
             print_objects.append(layer_breakdown_table)
         console = Console(record=True)
         console.print(*print_objects)
@@ -232,6 +236,63 @@ def _create_summary_table(self, summary: OutputSummary, view_key: ViewKey) -> Ta
 
         return summary_table
 
+    def _create_additional_metrics_table(self, result: AnalyzerResultType, view_key: ViewKey) -> Optional[Table]:
+        if not result.additional_metrics:
+            return None
+
+        flat_view = result.flat_views[view_key]
+        view_type = view_key[-1]
+        view_additional_metrics = result.additional_metrics.get(view_type, [])
+        if not view_additional_metrics:
+            return None
+        view_name = humanized_view_name(view_key, ' ')
+        additional_table = Table(title=f"{view_name} Additional Metrics", title_style='bold magenta', expand=True)
+        additional_table.add_column(header='Metric', style='bold')
+        additional_table.add_column(header='Unit', style='italic')
+        additional_table.add_column(header='Non-null', justify='right')
+        additional_table.add_column(header='Min', justify='right')
+        additional_table.add_column(header='Mean', justify='right')
+        additional_table.add_column(header='Max', justify='right')
+
+        found_metric = False
+        for metric in view_additional_metrics:
+            if metric not in flat_view.columns:
+                continue
+            metric_series = pd.to_numeric(flat_view[metric], errors='coerce').replace([np.inf, -np.inf], pd.NA)
+            scale, unit = self._additional_metric_scale_and_unit(metric)
+            metric_series = metric_series / scale
+            non_null = int(metric_series.notna().sum())
+            if non_null == 0:
+                additional_table.add_row(metric, unit, "0", "-", "-", "-")
+                found_metric = True
+                continue
+            additional_table.add_row(
+                metric,
+                unit,
+                f"{non_null:,}",
+                f"{float(metric_series.min()):.3f}",
+                f"{float(metric_series.mean()):.3f}",
+                f"{float(metric_series.max()):.3f}",
+            )
+            found_metric = True
+
+        if not found_metric:
+            return None
+        return additional_table
+
+    @staticmethod
+    def _additional_metric_scale_and_unit(metric: str):
+        metric_lower = metric.lower()
+        if metric_lower.endswith('_gbps'):
+            return GiB, 'GB/s'
+        if metric_lower.endswith('_mbps'):
+            return MiB, 'MB/s'
+        if metric_lower.endswith('_gb'):
+            return GiB, 'GB'
+        if metric_lower.endswith('_mb'):
+            return MiB, 'MB'
+        return 1.0, '-'
+
     def _format_val(self, value: float, fmt_int=False) -> str:
         if value is None or value == 0:
             return '-'

python/dftracer/analyzer/types.py

@@ -192,6 +192,7 @@ class OutputType:
 
 @dc.dataclass
 class AnalyzerResultType:
+    additional_metrics: Dict[ViewType, List[str]]
     checkpoint_dir: str
     flat_views: Dict[ViewKey, pd.DataFrame]
     layers: List[Layer]