enhancement/tracking

src/damast/cli/data_visualize.py

@@ -0,0 +1,96 @@
+"""
+Visualize a pipeline - either its structure, or what a run of it did to the data.
+
+Which of the two is drawn follows from the input file: a saved pipeline (``*.damast.ppl``)
+has a structure to render, a run's statistics report (``*.stats.yaml``, written by
+`damast.core.dataprocessing.DataProcessingPipeline.save_stats`) has per-step numbers to plot.
+"""
+
+from argparse import ArgumentParser
+from pathlib import Path
+
+from damast.cli.base import BaseParser
+from damast.core.dataprocessing import (
+    DAMAST_PIPELINE_SUFFIX,
+    DAMAST_STATS_SUFFIX,
+    DataProcessingPipeline,
+)
+from damast.viz.mermaid_export import MermaidExporter
+from damast.viz.stats_export import StatsExporter
+from damast.viz.svg_export import SvgExporter
+
+#: Default file extension per output format
+_EXTENSIONS = {"svg": ".svg", "html": ".html", "mermaid": ".mmd"}
+
+#: Formats each kind of input can be drawn as, the first being its default
+_FORMATS = {
+    "stats": ["svg"],
+    "pipeline": ["svg", "html", "mermaid"],
+}
+
+
+def resolve_kind(filename: Path) -> str:
+    """
+    Decide what the given file holds, from its name.
+
+    :param filename: The file to visualize
+    :return: ``"stats"`` or ``"pipeline"``
+    :raise ValueError: If the suffix matches neither
+    """
+    name = filename.name
+    if name.endswith(DAMAST_STATS_SUFFIX):
+        return "stats"
+    if name.endswith(DAMAST_PIPELINE_SUFFIX):
+        return "pipeline"
+
+    raise ValueError(f"Cannot visualize '{filename}': expected a pipeline ('{DAMAST_PIPELINE_SUFFIX}')"
+                     f" or a run's statistics ('{DAMAST_STATS_SUFFIX}')")
+
+
+class DataVisualizeParser(BaseParser):
+    """Argparser for the 'viz' subcommand."""
+
+    def __init__(self, parser: ArgumentParser):
+        super().__init__(parser=parser)
+
+        parser.description = "damast viz - visualize a pipeline or a run of it"
+
+        parser.add_argument("-f", "--filename", type=str, required=True,
+                            help=f"A saved pipeline ('{DAMAST_PIPELINE_SUFFIX}') to draw the structure"
+                                 f" of, or a run's statistics ('{DAMAST_STATS_SUFFIX}') to plot")
+        parser.add_argument("-o", "--output-file", type=str, default=None,
+                            help="Where to write the result, by default next to the input file")
+        parser.add_argument("--format", type=str, default=None,
+                            choices=sorted({f for formats in _FORMATS.values() for f in formats}),
+                            help="Output format - 'svg' for statistics; 'svg', 'html' or 'mermaid'"
+                                 " for a pipeline (default: svg)")
+
+    def execute(self, args):
+        super().execute(args)
+
+        filename = Path(args.filename)
+        if not filename.exists():
+            raise FileNotFoundError(f"File '{filename}' does not exist")
+
+        kind = resolve_kind(filename)
+        permitted = _FORMATS[kind]
+        output_format = args.format or permitted[0]
+        if output_format not in permitted:
+            raise ValueError(f"Cannot write a {kind} visualization as '{output_format}' -"
+                             f" available: {', '.join(permitted)}")
+
+        output_file = Path(args.output_file) if args.output_file else \
+            filename.parent / f"{filename.name.split('.')[0]}{_EXTENSIONS[output_format]}"
+
+        if kind == "stats":
+            written = StatsExporter.from_report(filename).export_svg(output_file)
+        else:
+            pipeline = DataProcessingPipeline.load(filename)
+            if output_format == "svg":
+                written = SvgExporter(pipeline).export_svg(output_file)
+            elif output_format == "html":
+                written = MermaidExporter(pipeline).export_html(output_file)
+            else:
+                written = MermaidExporter(pipeline).export_mermaid(output_file)
+
+        print(f"Written: {written}")

src/damast/cli/main.py

@@ -13,6 +13,7 @@
 from damast.cli.data_converter import DataConvertParser
 from damast.cli.data_inspect import DataInspectParser
 from damast.cli.data_processing import DataProcessingParser
+from damast.cli.data_visualize import DataVisualizeParser
 from damast.cli.data_watch import DataWatchParser
 from damast.cli.experiment import ExperimentParser
 from damast.cli.plugins import PluginsParser
@@ -89,6 +90,10 @@ def run():
                                          help="Process data by running a predefined pipeline",
                                          parser_klass=DataProcessingParser)
 
+    main_parser.attach_subcommand_parser(subcommand="viz",
+                                         help="Visualize a pipeline, or what a run of it did to the data",
+                                         parser_klass=DataVisualizeParser)
+
     main_parser.attach_subcommand_parser(subcommand="plugins",
                                          help="List transformer plugins registered by installed packages",
                                          parser_klass=PluginsParser)

src/damast/core/dataprocessing.py

@@ -9,7 +9,6 @@
 import tempfile
 import traceback as tc
 from collections import OrderedDict
-from datetime import datetime, timezone
 from logging import Logger, getLogger
 from pathlib import Path
 from typing import Any
@@ -26,6 +25,7 @@
 from .dataframe import AnnotatedDataFrame
 from .metadata import DataSpecification, MetaData
 from .pydantic_export import PydanticExporter
+from .tracking import PipelineElementTracker, RowCountTracker, TimingTracker
 
 __all__ = [
     "DataProcessingPipeline",
@@ -38,6 +38,9 @@
 """Suffix of :class:`DataProcessingPipeline` files created in :func:`DataProcessingPipeline.save`
 and used by :func:`DataProcessingPipeline.load`"""
 
+DAMAST_STATS_SUFFIX: str = ".stats.yaml"
+"""Suffix of the per-run report created by :func:`DataProcessingPipeline.save_stats`"""
+
 VAEX_STATE_SUFFIX: str = ".vaex-state.json"
 """Suffix of :class:`DataProcessingPipeline` files created in :func:`DataProcessingPipeline.save_state`
 and used by :func:`DataProcessingPipeline.load_state`"""
@@ -53,6 +56,12 @@ class DataProcessingPipeline(PipelineElement):
     :param inplace_transformation: If true, the input :class:`damast.core.dataframe.AnnotatedDataFrame` is not
         copied when calling :func:`transform`. Else input data-frame is untouched
     :param name_mappings: Name mappings that should be applied to individual transformations.
+    :param trackers: Observers of what each step does to the data, see
+        :class:`damast.core.tracking.PipelineElementTracker`. What they report ends up in
+        :func:`processing_stats`. Defaults to a :class:`damast.core.tracking.TimingTracker` and a
+        :class:`damast.core.tracking.RowCountTracker`; pass a list to add to or replace those, e.g.
+        a :class:`damast.core.tracking.KeyTracker` to also count distinct entities. An empty list
+        collects nothing.
 
     :raises ValueError: If any of the transformer names are `None`
     :raises AttributeError: If the transformer is missing the :func:`transform` function
@@ -80,6 +89,11 @@ class DataProcessingPipeline(PipelineElement):
     _name_mappings: dict[str, dict[str, str]]
     _processing_stats: dict[str, dict[str, Any]]
 
+    #: Observers of what each step does to the data
+    _trackers: list[PipelineElementTracker]
+    #: uuids of the steps between their start and their end, to keep a tracker from seeing one twice
+    _steps_in_flight: set[Any]
+
     _meta: dict[str, str]
 
     def __init__(self, *,
@@ -90,6 +104,7 @@ def __init__(self, *,
                  inplace_transformation: bool = False,
                  name_mappings: dict[str, dict[str, str]] = { DAMAST_DEFAULT_DATASOURCE: {}},
                  meta: dict[str, str] | None = None,
+                 trackers: list[PipelineElementTracker] | None = None,
                  ):
         super().__init__()
 
@@ -103,6 +118,9 @@ def __init__(self, *,
         self._name_mappings = name_mappings
         self._processing_stats = {}
 
+        self._trackers = list(trackers) if trackers is not None else [TimingTracker(), RowCountTracker()]
+        self._steps_in_flight = set()
+
         self.processing_graph = ProcessingGraph()
         if processing_graph:
             if isinstance(processing_graph, ProcessingGraph):
@@ -159,15 +177,30 @@ def base_dir(self, value):
     @property
     def processing_stats(self) -> dict[str, dict[str, Any]]:
         """
-        Per-step timing and row-count statistics collected by the last :func:`transform` run.
+        What the :attr:`trackers` reported about each step of the last :func:`transform` run.
 
-        Keyed by step name, e.g. ``{"my_step": {"processing_time_in_s": 0.1,
-        "input_dataframe_length": {...}, "output_dataframe_length": 123, ...}}``. Populated
-        incrementally by :func:`on_transform_start`/:func:`on_transform_end` as steps run, so
-        it also reflects partial progress if a pipeline fails mid-run.
+        Keyed by step name, holding the merged result of every tracker, e.g.
+        ``{"my_step": {"processing_time_in_s": 0.1, "input_dataframe_length": {...},
+        "output_dataframe_length": 123, "rows_removed": 7, ...}}``. Filled in as steps run, so it
+        also reflects partial progress if a pipeline fails mid-run.
         """
         return self._processing_stats
 
+    @property
+    def trackers(self) -> list[PipelineElementTracker]:
+        """
+        The observers of what each step does to the data, see
+        :class:`damast.core.tracking.PipelineElementTracker`.
+
+        Settable, so that a pipeline from :func:`load` - which is built from the saved file alone,
+        and whose trackers are therefore the defaults - can be given different ones.
+        """
+        return self._trackers
+
+    @trackers.setter
+    def trackers(self, value: list[PipelineElementTracker]):
+        self._trackers = list(value)
+
     @property
     def output_specs(self):
         if not self.is_ready:
@@ -512,6 +545,26 @@ def save(self, dir: str | Path) -> Path:
             yaml.dump(dict(self), f)
         return filename
 
+    def save_stats(self, dir: str | Path | None = None) -> Path:
+        """
+        Save the statistics of the last :func:`transform` run - see :func:`processing_stats`.
+
+        Written automatically at the end of :func:`transform`, so that a run leaves an audit trail
+        of what each step did even without an experiment tracker.
+
+        :param dir: Directory to write to, :attr:`base_dir` if not given
+        :return: The path of the written report
+        """
+        base_dir = Path(dir) if dir is not None else self.base_dir
+        base_dir.mkdir(parents=True, exist_ok=True)
+        filename = base_dir / f"{self.name}{DAMAST_STATS_SUFFIX}"
+
+        with open(filename, "w") as f:
+            # sort_keys=False: the steps are a sequence of events, so alphabetising them would
+            # lose the execution order the statistics only make sense in
+            yaml.dump({"name": self.name, "steps": self._processing_stats}, f, sort_keys=False)
+        return filename
+
     def __iter__(self):
         yield "name", self.name
         yield "description", self.description
@@ -675,6 +728,10 @@ def transform(self,
         :param df: The input dataframe
         :returns: The transformed dataframe
         """
+        # the statistics describe this run, so anything a previous one left behind has to go
+        self._processing_stats = {}
+        self._steps_in_flight = set()
+
         dataframes = { DAMAST_DEFAULT_DATASOURCE: df }
         for x in self.processing_graph.get_joins():
             if x.name not in kwargs:
@@ -707,6 +764,8 @@ def transform(self,
         adf = pipeline._run(in_dataframes, verbose=verbose)
         assert isinstance(adf, AnnotatedDataFrame)
         adf.validate_metadata(validation_mode=damast.core.ValidationMode.UPDATE_METADATA)
+
+        self.save_stats()
         return adf
 
     def _run(self,
@@ -760,47 +819,49 @@ def on_transform_start(self,
                            step: PipelineElement,
                            dataframes: dict[str, AnnotatedDataFrame]):
         """
-        Default implementation of the on_transform_start callback.
+        Let every tracker see a step's input before it runs - see :attr:`trackers`.
 
-        This output is a message in INFO log level
+        This also outputs a message in INFO log level.
         """
-        if hasattr(step, "transform_start"):
+        # a step is announced once per datasource, so without this a tracker would see a join twice
+        if step.uuid in self._steps_in_flight:
             return
+        self._steps_in_flight.add(step.uuid)
 
         logger.info(f"[transform] start: {step.__class__.__name__} - {step.name_mappings}")
-        start_time = datetime.now(timezone.utc)
-        step.transform_start = start_time
 
         step_name = self.processing_graph[step.uuid].name
-        self._processing_stats[step_name] = {
-            "input_dataframe_length": {x: y.shape[0] for x,y in dataframes.items()},
-            "start_time": start_time
-        }
+        self._processing_stats[step_name] = {}
+        for tracker in self._trackers:
+            tracker.on_step_start(step, dataframes)
 
     def on_transform_end(self,
                          step: PipelineElement,
                          adf: AnnotatedDataFrame):
         """
-        Default implementation of the on_transform_end callback.
+        Collect what every tracker has to report about a step that just ran - see :attr:`trackers`.
 
-        This outputs a message in INFO log level
+        The trackers' results are merged into :func:`processing_stats` under the step's name.
+        This also outputs a message in INFO log level.
         """
-        if not hasattr(step, "transform_start"):
+        if step.uuid not in self._steps_in_flight:
             return
-
-        start = step.transform_start
-        end_time = datetime.now(timezone.utc)
-        delta = (end_time - start).total_seconds()
-        delattr(step, "transform_start")
-        logger.info(f"[transform] end: {step.__class__.__name__} - {step.name_mappings}: "
-                  f"{delta} seconds, {adf.shape[0]} remaining rows)")
+        self._steps_in_flight.discard(step.uuid)
 
         step_name = self.processing_graph[step.uuid].name
-        self._processing_stats[step_name].update({
-            "processing_time_in_s": delta,
-            "output_dataframe_length": adf.shape[0],
-            "end_time": end_time
-        })
+        statistics = self._processing_stats.setdefault(step_name, {})
+        for tracker in self._trackers:
+            statistics.update(tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf}))
+
+        # Report whatever the trackers happened to measure, so that a custom set of them degrades
+        # the message instead of breaking it
+        details = []
+        if "processing_time_in_s" in statistics:
+            details.append(f"{statistics['processing_time_in_s']} seconds")
+        if "output_dataframe_length" in statistics:
+            details.append(f"{statistics['output_dataframe_length']} remaining rows")
+        logger.info(f"[transform] end: {step.__class__.__name__} - {step.name_mappings}"
+                    f"{': ' + ', '.join(details) if details else ''}")
 
     def __repr__(self) -> str:
         """
@@ -820,6 +881,9 @@ def __deepcopy__(self, memo) -> DataProcessingPipeline:
         pipeline._processing_stats = copy.deepcopy(self._processing_stats)
         pipeline.processing_graph = copy.deepcopy(self.processing_graph)
         pipeline._output_specs = [copy.deepcopy(x) for x in self._output_specs]
+        # '_trackers' stays shared on purpose: 'prepare' copies the pipeline, but the transformers
+        # keep calling the original's hooks, so the caller's own tracker instances are the ones
+        # that run and hold state
         return pipeline
 
     def describe(self) -> str: