stats: add output_file option to output the stats to a file

Author: AlexTereshenkov

src/python/pants/goal/stats_aggregator.py

@@ -7,6 +7,7 @@
 import logging
 from collections import Counter
 from dataclasses import dataclass
+from typing import Optional
 
 from pants.engine.internals.scheduler import Workunit
 from pants.engine.rules import collect_rules, rule
@@ -17,7 +18,7 @@
     WorkunitsCallbackFactoryRequest,
 )
 from pants.engine.unions import UnionRule
-from pants.option.option_types import BoolOption
+from pants.option.option_types import BoolOption, StrOption
 from pants.option.subsystem import Subsystem
 from pants.util.collections import deep_getsizeof
 from pants.util.strutil import softwrap
@@ -55,13 +56,21 @@ class StatsAggregatorSubsystem(Subsystem):
         ),
         advanced=True,
     )
+    output_file = StrOption(
+        default=None,
+        metavar="<path>",
+        help="Output the stats to this file. If unspecified, outputs to stdout.",
+    )
 
 
 class StatsAggregatorCallback(WorkunitsCallback):
-    def __init__(self, *, log: bool, memory: bool, has_histogram_module: bool) -> None:
+    def __init__(
+        self, *, log: bool, memory: bool, output_file: Optional[str], has_histogram_module: bool
+    ) -> None:
         super().__init__()
         self.log = log
         self.memory = memory
+        self.output_file = output_file
         self.has_histogram_module = has_histogram_module
 
     @property
@@ -80,6 +89,8 @@ def __call__(
         if not finished:
             return
 
+        output_contents = ""
+
         if self.log:
             # Capture global counters.
             counters = Counter(context.get_metrics())
@@ -93,7 +104,7 @@ def __call__(
             counter_lines = "\n".join(
                 f"  {name}: {count}" for name, count in sorted(counters.items())
             )
-            logger.info(f"Counters:\n{counter_lines}")
+            output_contents += f"Counters:\n{counter_lines}"
 
         if self.memory:
             ids: set[int] = set()
@@ -115,18 +126,20 @@ def __call__(
             memory_lines = "\n".join(
                 f"  {size}\t\t{count}\t\t{name}" for size, count, name in sorted(entries)
             )
-            logger.info(f"Memory summary (total size in bytes, count, name):\n{memory_lines}")
+            output_contents += (
+                f"\nMemory summary (total size in bytes, count, name):\n{memory_lines}"
+            )
 
         if not (self.log and self.has_histogram_module):
             return
         from hdrh.histogram import HdrHistogram  # pants: no-infer-dep
 
         histograms = context.get_observation_histograms()["histograms"]
         if not histograms:
-            logger.info("No observation histogram were recorded.")
+            output_contents += "\nNo observation histogram were recorded."
             return
 
-        logger.info("Observation histogram summaries:")
+        output_contents += "\nObservation histogram summaries:"
         for name, encoded_histogram in histograms.items():
             # Note: The Python library for HDR Histogram will only decode compressed histograms
             # that are further encoded with base64. See
@@ -138,8 +151,8 @@ def __call__(
                     [25, 50, 75, 90, 95, 99]
                 ).items()
             )
-            logger.info(
-                f"Summary of `{name}` observation histogram:\n"
+            output_contents += (
+                f"\nSummary of `{name}` observation histogram:\n"
                 f"  min: {histogram.get_min_value()}\n"
                 f"  max: {histogram.get_max_value()}\n"
                 f"  mean: {histogram.get_mean_value():.3f}\n"
@@ -149,6 +162,12 @@ def __call__(
                 f"{percentile_to_vals}"
             )
 
+        if self.output_file:
+            with open(self.output_file, "w") as fh:
+                fh.write(output_contents)
+        else:
+            logger.info(output_contents)
+
 
 @dataclass(frozen=True)
 class StatsAggregatorCallbackFactoryRequest:
@@ -178,6 +197,7 @@ def construct_callback(
             StatsAggregatorCallback(
                 log=subsystem.log,
                 memory=subsystem.memory_summary,
+                output_file=subsystem.output_file,
                 has_histogram_module=has_histogram_module,
             )
             if subsystem.log or subsystem.memory_summary