Merge pull request #427 from nodestream-proj/MetricsPivot

Optionally adding time intervals as a reporter frequency option. Adding accumulator metrics.

Author: angelosantos4

nodestream/cli/commands/run.py

@@ -57,6 +57,12 @@ class Run(NodestreamCommand):
             description="Storage backend to use for checkpointing",
             flag=False,
         ),
+        option(
+            "metrics-interval-in-seconds",
+            description="Time interval to report metrics in seconds",
+            default=None,
+            flag=False,
+        ),
         *PROMETHEUS_OPTIONS,
     ]
 

nodestream/cli/operations/run_pipeline.py

@@ -122,8 +122,14 @@ def create_progress_reporter(
         self, command: NodestreamCommand, pipeline_name: str
     ) -> PipelineProgressReporter:
         indicator = self.get_progress_indicator(command, pipeline_name)
+        metrics_interval_in_seconds = (
+            float(command.option("metrics-interval-in-seconds"))
+            if command.option("metrics-interval-in-seconds")
+            else None
+        )
         return PipelineProgressReporter(
             reporting_frequency=int(command.option("reporting-frequency")),
+            metrics_interval_in_seconds=metrics_interval_in_seconds,
             callback=indicator.progress_callback,
             on_start_callback=indicator.on_start,
             on_finish_callback=indicator.on_finish,

nodestream/metrics.py

@@ -14,6 +14,7 @@
 class Metric:
     name: str
     description: Optional[str] = None
+    accumulate: bool = False
 
     def increment_on(self, handler: "MetricHandler", value: Number = 1):
         """Increment this metric on the given handler."""
@@ -27,6 +28,18 @@ def register(self, handler: "MetricHandler"):
         handler.increment(self, 0)
         return self
 
+    def __str__(self):
+        return self.name
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def __eq__(self, other):
+        return self.name == other.name
+
+    def __ne__(self, other):
+        return self.name != other.name
+
 
 class MetricHandler(ABC):
     """Recieves metrics and handles them in some way.
@@ -74,21 +87,28 @@ def tick(self):
 
 
 # Core metrics
-RECORDS = Metric("records", "Number of records processed")
-NON_FATAL_ERRORS = Metric("non_fatal_errors", "Number of non-fatal errors")
-FATAL_ERRORS = Metric("fatal_errors", "Number of fatal errors")
-NODES_UPSERTED = Metric("nodes_upserted", "Number of nodes upserted to the graph")
+RECORDS = Metric("records", "Number of records processed", accumulate=True)
+NON_FATAL_ERRORS = Metric(
+    "non_fatal_errors", "Number of non-fatal errors", accumulate=True
+)
+FATAL_ERRORS = Metric("fatal_errors", "Number of fatal errors", accumulate=True)
+NODES_UPSERTED = Metric(
+    "nodes_upserted", "Number of nodes upserted to the graph", accumulate=True
+)
 RELATIONSHIPS_UPSERTED = Metric(
-    "relationships_upserted", "Number of relationships upserted to the graph"
+    "relationships_upserted",
+    "Number of relationships upserted to the graph",
+    accumulate=True,
 )
 TIME_TO_LIVE_OPERATIONS = Metric(
-    "time_to_live_operations", "Number of time-to-live operations executed"
+    "time_to_live_operations",
+    "Number of time-to-live operations executed",
+    accumulate=True,
 )
 INGEST_HOOKS_EXECUTED = Metric(
-    "ingest_hooks_executed", "Number of ingest hooks executed to the graph"
-)
-BUFFERED_RECORDS = Metric(
-    "buffered_records", "Number of records buffered between steps"
+    "ingest_hooks_executed",
+    "Number of ingest hooks executed to the graph",
+    accumulate=True,
 )
 STEPS_RUNNING = Metric(
     "steps_running", "Number of steps currently running in the pipeline"
@@ -184,6 +204,8 @@ def tick(self):
 class ConsoleMetricHandler(MetricHandler):
     """A metric handler that prints metrics to the console."""
 
+    # Note: TODO: Tie incremental values to the metrics rather than the handler.
+
     def __init__(self, command: Command):
         self.metrics: dict[Metric, Number] = {}
         self.command = command
@@ -194,8 +216,17 @@ def increment(self, metric: Metric, value: Number):
     def decrement(self, metric: Metric, value: Number):
         self.metrics[metric] = self.metrics.get(metric, 0) - value
 
+    def discharge(self) -> dict[Metric, Number]:
+        metrics = {}
+        for metric, value in self.metrics.items():
+            metrics[metric.name] = value
+            if metric.accumulate:
+                self.metrics[metric] = 0
+        return metrics
+
     def render(self):
-        stats = ((k.name, str(v)) for k, v in self.metrics.items() if v > 0)
+        metrics = self.discharge()
+        stats = ((k, str(v)) for k, v in metrics.items())
         table = self.command.table(STATS_TABLE_COLS, stats)
         table.render()
 
@@ -219,9 +250,20 @@ def increment(self, metric: Metric, value: Number):
     def decrement(self, metric: Metric, value: Number):
         self.metrics[metric] = self.metrics.get(metric, 0) - value
 
+    def discharge(self) -> dict[Metric, Number]:
+        metrics = {}
+        for metric, value in self.metrics.items():
+            metrics[metric.name] = value
+            if metric.accumulate:
+                self.metrics[metric] = 0
+
+        return metrics
+
     def render(self):
+        metrics = self.discharge()
         self.logger.info(
-            "Metrics Report", extra={k.name: v for k, v in self.metrics.items()}
+            "Metrics Report",
+            extra=metrics,
         )
 
     def stop(self):

nodestream/pipeline/progress_reporter.py

@@ -1,7 +1,8 @@
 import os
+import time
 from dataclasses import dataclass, field
 from logging import Logger, getLogger
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 from psutil import Process
 
@@ -32,13 +33,15 @@ class PipelineProgressReporter:
     """A `PipelineProgressReporter` is a utility that can be used to report on the progress of a pipeline."""
 
     reporting_frequency: int = 10000
+    metrics_interval_in_seconds: Optional[float] = None
     logger: Logger = field(default_factory=getLogger)
     callback: Callable[[int, Metrics], None] = field(default=no_op)
     on_start_callback: Callable[[], None] = field(default=no_op)
     on_finish_callback: Callable[[Metrics], None] = field(default=no_op)
     on_fatal_error_callback: Callable[[Exception], None] = field(default=no_op)
     encountered_fatal_error: bool = field(default=False)
     observability_callback: Callable[[Any], None] = field(default=no_op)
+    last_report_time: float = field(default=0)
 
     def on_fatal_error(self, exception: Exception):
         self.encountered_fatal_error = True
@@ -68,7 +71,12 @@ def for_testing(cls, results_list: list) -> "PipelineProgressReporter":
         )
 
     def report(self, index, metrics: Metrics):
-        if index % self.reporting_frequency == 0:
+        if self.metrics_interval_in_seconds is not None:
+            current_time = time.time()
+            if current_time - self.last_report_time >= self.metrics_interval_in_seconds:
+                self.callback(index, metrics)
+                self.last_report_time = current_time
+        elif index % self.reporting_frequency == 0:
             self.callback(index, metrics)
 
     def observe(self, record: Any):

tests/unit/cli/operations/test_run_pipeline.py

@@ -32,7 +32,18 @@ def test_make_run_request(run_pipeline_operation, mocker):
     targets = ["t1", "t2"]
     pipeline_name = "my_pipeline"
     command = mocker.Mock()
-    command.option.side_effect = ["my-storage", annotations, "10001", targets, "10000"]
+    # Updated to handle all the option calls including time-interval-seconds
+    option_responses = {
+        "storage-backend": "my-storage",
+        "annotations": annotations,
+        "step-outbox-size": "10001",
+        "target": targets,
+        "time-interval-seconds": None,  # No time interval provided
+        "reporting-frequency": "10000",
+    }
+    command.option.side_effect = lambda opt: option_responses.get(opt)
+    command.has_json_logging_set = False  # Required for create_progress_reporter
+    command.is_very_verbose = False  # Required for make_run_request
     command.argument.return_value = [pipeline_name]
     pipeline = mocker.patch("nodestream.project.PipelineDefinition")
     pipeline.name = pipeline_name
@@ -43,6 +54,129 @@ def test_make_run_request(run_pipeline_operation, mocker):
     assert_that(result.progress_reporter.reporting_frequency, equal_to(10000))
 
 
+def test_create_progress_reporter_with_time_interval_seconds(
+    run_pipeline_operation, mocker
+):
+    """Test that time_interval_seconds gets properly converted to float and passed to PipelineProgressReporter"""
+    command = mocker.Mock()
+    command.option.side_effect = lambda opt: {
+        "metrics-interval-in-seconds": "30.5",
+        "reporting-frequency": "1000",
+    }.get(opt)
+    command.has_json_logging_set = False
+
+    # Mock PipelineProgressReporter to capture arguments
+    mock_progress_reporter = mocker.patch(
+        "nodestream.cli.operations.run_pipeline.PipelineProgressReporter"
+    )
+
+    _ = run_pipeline_operation.create_progress_reporter(command, "test_pipeline")
+
+    # Verify PipelineProgressReporter was called with correct time_interval_seconds
+    mock_progress_reporter.assert_called_once()
+    call_args = mock_progress_reporter.call_args
+    assert_that(call_args.kwargs["metrics_interval_in_seconds"], equal_to(30.5))
+    assert_that(call_args.kwargs["reporting_frequency"], equal_to(1000))
+
+
+def test_create_progress_reporter_without_time_interval_seconds(
+    run_pipeline_operation, mocker
+):
+    """Test that time_interval_seconds is None when not provided"""
+    command = mocker.Mock()
+    command.option.side_effect = lambda opt: {
+        "metrics-interval-in-seconds": None,
+        "reporting-frequency": "2000",
+    }.get(opt)
+    command.has_json_logging_set = False
+
+    # Mock PipelineProgressReporter to capture arguments
+    mock_progress_reporter = mocker.patch(
+        "nodestream.cli.operations.run_pipeline.PipelineProgressReporter"
+    )
+
+    _ = run_pipeline_operation.create_progress_reporter(command, "test_pipeline")
+
+    # Verify PipelineProgressReporter was called with None for time_interval_seconds
+    mock_progress_reporter.assert_called_once()
+    call_args = mock_progress_reporter.call_args
+    assert_that(call_args.kwargs["metrics_interval_in_seconds"], equal_to(None))
+    assert_that(call_args.kwargs["reporting_frequency"], equal_to(2000))
+
+
+def test_create_progress_reporter_with_json_indicator(run_pipeline_operation, mocker):
+    """Test that create_progress_reporter works correctly with JSON progress indicator"""
+    command = mocker.Mock()
+    command.option.side_effect = lambda opt: {
+        "metrics-interval-in-seconds": "15.0",
+        "reporting-frequency": "500",
+    }.get(opt)
+    command.has_json_logging_set = True
+
+    # Mock PipelineProgressReporter to capture arguments
+    mock_progress_reporter = mocker.patch(
+        "nodestream.cli.operations.run_pipeline.PipelineProgressReporter"
+    )
+
+    _ = run_pipeline_operation.create_progress_reporter(command, "test_pipeline")
+
+    # Verify PipelineProgressReporter was called with correct arguments
+    mock_progress_reporter.assert_called_once()
+    call_args = mock_progress_reporter.call_args
+    assert_that(call_args.kwargs["metrics_interval_in_seconds"], equal_to(15.0))
+    assert_that(call_args.kwargs["reporting_frequency"], equal_to(500))
+
+
+def test_make_run_request_with_time_interval_seconds_integration(
+    run_pipeline_operation, mocker
+):
+    """Integration test to ensure make_run_request properly handles time_interval_seconds through create_progress_reporter"""
+    annotations = ["annotation1"]
+    targets = ["t1"]
+    pipeline_name = "my_pipeline"
+    command = mocker.Mock()
+
+    # Setup command.option to handle all the different option calls made by make_run_request
+    option_responses = {
+        "storage-backend": "my-storage",
+        "annotations": annotations,
+        "step-outbox-size": "10001",
+        "target": targets,
+        "metrics-interval-in-seconds": "45.0",
+        "reporting-frequency": "5000",
+    }
+    command.option.side_effect = lambda opt: option_responses.get(opt)
+    command.has_json_logging_set = False
+    command.is_very_verbose = False
+    command.argument.return_value = [pipeline_name]
+
+    pipeline = mocker.Mock()
+    pipeline.name = pipeline_name
+    pipeline.configuration = PipelineConfiguration()
+
+    # Mock the project's get_object_storage_by_name method
+    run_pipeline_operation.project.get_object_storage_by_name.return_value = None
+    run_pipeline_operation.project.get_target_by_name.return_value = None
+
+    # Mock PipelineProgressReporter to capture its arguments
+    mock_progress_reporter = mocker.patch(
+        "nodestream.cli.operations.run_pipeline.PipelineProgressReporter"
+    )
+
+    result = run_pipeline_operation.make_run_request(command, pipeline)
+
+    # Verify the progress reporter was created with correct time_interval_seconds
+    mock_progress_reporter.assert_called_once()
+    call_args = mock_progress_reporter.call_args
+    assert_that(call_args.kwargs["metrics_interval_in_seconds"], equal_to(45.0))
+    assert_that(call_args.kwargs["reporting_frequency"], equal_to(5000))
+
+    # Verify other parts of the request are still correct
+    assert_that(result.pipeline_name, equal_to(pipeline_name))
+    assert_that(result.initialization_arguments.annotations, equal_to(annotations))
+    assert_that(result.initialization_arguments.step_outbox_size, equal_to(10001))
+
+
 def test_spinner_on_start(mocker):
     spinner = SpinnerProgressIndicator(mocker.Mock(), "pipeline_name")
     spinner.on_start()

tests/unit/pipeline/test_pipeline_progress_reporter.py

@@ -1,6 +1,9 @@
+from unittest.mock import Mock
+
 import pytest
 from hamcrest import assert_that, equal_to
 
+from nodestream.metrics import Metrics
 from nodestream.pipeline import IterableExtractor, Pipeline, PipelineProgressReporter
 
 
@@ -17,3 +20,35 @@ async def test_pipeline_progress_reporter_for_testing(mocker):
     result = PipelineProgressReporter.for_testing([])
     assert_that(result.reporting_frequency, equal_to(1))
     assert_that(result.logger.name, equal_to("test"))
+
+
+def test_pipeline_progress_reporter_with_metrics_interval_in_seconds(mocker):
+    """Test that metrics_interval_in_seconds works correctly"""
+    mock_callback = Mock()
+    reporter = PipelineProgressReporter(
+        metrics_interval_in_seconds=0.1, callback=mock_callback
+    )
+
+    mock_time = mocker.patch("nodestream.pipeline.progress_reporter.time.time")
+    mock_time.side_effect = [0.15, 0.2]  # 150ms, 200ms
+
+    metrics = Metrics()
+    reporter.report(1, metrics)  # Should report (150ms >= 100ms from 0)
+    reporter.report(2, metrics)  # Should not report (200ms - 150ms = 50ms < 100ms)
+
+    assert_that(mock_callback.call_count, equal_to(1))
+
+
+def test_pipeline_progress_reporter_without_metrics_interval_in_seconds_uses_frequency():
+    """Test that None metrics_interval_in_seconds falls back to frequency-based reporting"""
+    mock_callback = Mock()
+    reporter = PipelineProgressReporter(
+        metrics_interval_in_seconds=None, reporting_frequency=3, callback=mock_callback
+    )
+
+    metrics = Metrics()
+    for i in range(1, 7):  # 1,2,3,4,5,6
+        reporter.report(i, metrics)
+
+    # Should report on multiples of 3: indices 3, 6
+    assert_that(mock_callback.call_count, equal_to(2))