Ensuring error handling works. Delegating ProgressIndicator for ticking of metrics.

Author: angelosantos4

nodestream/cli/operations/run_pipeline.py

@@ -9,6 +9,7 @@
 from ...utils import StringSuggester
 from ..commands.nodestream_command import NodestreamCommand
 from .operation import Operation
+from logging import getLogger
 
 STATS_TABLE_COLS = ["Statistic", "Value"]
 
@@ -113,7 +114,7 @@ def get_progress_indicator(
         self, command: NodestreamCommand, pipeline_name: str
     ) -> "ProgressIndicator":
         if command.has_json_logging_set:
-            return ProgressIndicator(command, pipeline_name)
+            return JsonProgressIndicator(command, pipeline_name)
 
         return SpinnerProgressIndicator(command, pipeline_name)
 
@@ -122,6 +123,7 @@ def create_progress_reporter(
     ) -> PipelineProgressReporter:
         indicator = self.get_progress_indicator(command, pipeline_name)
         return PipelineProgressReporter(
+            logger=indicator.logger,
             reporting_frequency=int(command.option("reporting-frequency")),
             callback=indicator.progress_callback,
             on_start_callback=indicator.on_start,
@@ -134,6 +136,7 @@ class ProgressIndicator:
     def __init__(self, command: NodestreamCommand, pipeline_name: str) -> None:
         self.command = command
         self.pipeline_name = pipeline_name
+        self.logger = getLogger()
 
     def on_start(self):
         pass
@@ -160,14 +163,15 @@ def on_start(self):
         self.progress = self.command.progress_indicator()
         self.progress.start(f"Running pipeline: '{self.pipeline_name}'")
 
-    def progress_callback(self, index, _):
+    def progress_callback(self, index, metrics: Metrics):
         self.progress.set_message(
             f"Currently processing record at index: <info>{index}</info>"
         )
+        metrics.tick()
 
     def on_finish(self, metrics: Metrics):
         self.progress.finish(f"Finished running pipeline: '{self.pipeline_name}'")
-
+        metrics.tick()
         if self.exception:
             raise self.exception
 
@@ -176,3 +180,26 @@ def on_fatal_error(self, exception: Exception):
             "<error>Encountered a fatal error while running pipeline</error>"
         )
         self.exception = exception
+
+
+class JsonProgressIndicator(ProgressIndicator):
+    def __init__(self, command: NodestreamCommand, pipeline_name: str) -> None:
+        super().__init__(command, pipeline_name)
+        self.exception = None
+
+    def on_start(self):
+        self.logger.info("Starting Pipeline")
+
+    def progress_callback(self, index, metrics: Metrics):
+        self.logger.info("Processing Record", extra={"index": index})
+        metrics.tick()
+
+    def on_finish(self, metrics: Metrics):
+        self.logger.info("Pipeline Completed")
+        metrics.tick()
+        if self.exception:
+            raise self.exception
+
+    def on_fatal_error(self, exception: Exception):
+        self.logger.error("Pipeline Failed", exc_info=exception)
+        self.exception = exception

nodestream/metrics.py

@@ -113,6 +113,7 @@ class NodestreamMetricRegistry(MetricRegistry):
     # 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")
     RELATIONSHIPS_UPSERTED = Metric(
         "relationships_upserted", "Number of relationships upserted to the graph"
@@ -241,7 +242,7 @@ def tick(self):
         self.render()
 
     def stop(self):
-        self.render()
+        pass
 
 
 class JsonLogMetricHandler(MetricHandler):
@@ -265,7 +266,7 @@ def render(self):
         )
 
     def stop(self):
-        self.render()
+        pass
 
     def tick(self):
         self.render()

nodestream/pipeline/channel.py

@@ -37,15 +37,32 @@ class Channel:
     prevent one step from overwhelming another step with too many records.
     """
 
-    __slots__ = ("queue", "input_dropped", "metric")
+    __slots__ = ("queue", "input_dropped", "_metric", "input_name", "output_name")
 
-    def __init__(self, size: int, input_name: str, output_name: str) -> None:
+    def __init__(self, size: int) -> None:
         self.queue = Queue(maxsize=size)
         self.input_dropped = False
-        self.metric = Metric(
-            f"buffered_{input_name}_to_{output_name}",
-            f"Records buffered: {input_name} → {output_name}",
-        )
+        self.input_name = "Void"
+        self.output_name = "Void"
+        self._metric = None
+
+    @property
+    def metric(self) -> Metric:
+        """Get the metric for the channel."""
+        if self._metric is None:
+            self._metric = Metric(
+                f"buffered_{self.output_name}_to_{self.input_name}",
+                f"Records buffered: {self.output_name} → {self.input_name}",
+            )
+        return self._metric
+
+    def register_input(self, name: str) -> None:
+        """Register the name of the step that will consume from this channel."""
+        self.input_name = name
+
+    def register_output(self, name: str) -> None:
+        """Register the name of the step that will produce to this channel."""
+        self.output_name = name
 
     async def get(self):
         """Get an object from the channel.
@@ -58,6 +75,7 @@ async def get(self):
             object: The object that was retrieved from the channel.
         """
         object = await self.queue.get()
+        Metrics.get().decrement(self.metric)
         return object
 
     async def put(self, obj) -> bool:
@@ -92,6 +110,10 @@ class StepOutput:
     def __init__(self, channel: Channel) -> None:
         self.channel = channel
 
+    def register(self, name: str) -> None:
+        """Register the name of the step that will produce to this channel."""
+        self.channel.register_output(name)
+
     async def done(self):
         """Mark the output channel as done.
 
@@ -138,6 +160,10 @@ class StepInput:
     def __init__(self, channel: Channel) -> None:
         self.channel = channel
 
+    def register(self, name: str) -> None:
+        """Register the name of the step that will consume from this channel."""
+        self.channel.register_input(name)
+
     async def get(self) -> Optional[object]:
         """Get an object from the input channel.
 
@@ -161,15 +187,13 @@ def done(self):
         self.channel.input_dropped = True
 
 
-def channel(
-    size: int, input_name: str, output_name: str
-) -> Tuple[StepInput, StepOutput]:
+def channel(size: int) -> Tuple[StepInput, StepOutput]:
     """Create a new input and output channel.
 
     Args:
         size: The size of the channel.
         input_name: The name of the input step.
         output_name: The name of the output step.
     """
-    channel = Channel(size, input_name, output_name)
+    channel = Channel(size)
     return StepInput(channel), StepOutput(channel)

nodestream/pipeline/pipeline.py

@@ -113,9 +113,9 @@ async def run(self):
         self.call_handling_errors(self.reporter.on_start_callback)
 
         index = 0
-        while (obj := await self.input.get()) is not None:
+        while (await self.input.get()) is not None:
             metrics.increment(NodestreamMetricRegistry.RECORDS)
-            self.call_handling_errors(self.reporter.report, index, obj)
+            self.call_handling_errors(self.reporter.report, index, metrics)
             index += 1
 
         self.call_handling_errors(self.reporter.on_finish_callback, metrics)
@@ -161,8 +161,6 @@ async def run(self, reporter: PipelineProgressReporter):
         `PipelineProgressReporter`.
 
         Args:
-            channel_size: The size of the channels used to pass records between
-                steps in the pipeline.
             reporter: The `PipelineProgressReporter` used to report on the
                 progress of the pipeline.
         """
@@ -173,15 +171,9 @@ async def run(self, reporter: PipelineProgressReporter):
         # the steps in the pipeline. The channels have a fixed size to control
         # the flow of records between the steps.
         executors: List[StepExecutor] = []
-        current_input, current_output = channel(
-            self.step_outbox_size,
-            input_name=self.steps[-1].__class__.__name__ + f"_{len(self.steps) - 1}",
-            output_name="Void",
-        )
+        current_input, current_output = channel(self.step_outbox_size)
         pipeline_output = PipelineOutput(current_input, reporter)
 
-        print([step.__class__.__name__ for step in self.steps])
-
         # Create the executors for the steps in the pipeline. The executors
         # will be used to run the steps concurrently. The steps are created in
         # reverse order so that the output of each step is connected to the
@@ -190,20 +182,9 @@ async def run(self, reporter: PipelineProgressReporter):
             index = len(self.steps) - reversed_index - 1
             storage = self.object_store.namespaced(str(index))
             context = StepContext(step.__class__.__name__, index, reporter, storage)
-
-            previous_step = (
-                self.steps[reversed_index - 1] if reversed_index > 0 else None
-            )
-
-            current_input, next_output = channel(
-                self.step_outbox_size,
-                input_name=(
-                    previous_step.__class__.__name__ + f"_{reversed_index - 1}"
-                    if previous_step
-                    else "Void"
-                ),
-                output_name=step.__class__.__name__ + f"_{reversed_index}",
-            )
+            current_input, next_output = channel(self.step_outbox_size)
+            current_input.register(step.__class__.__name__)
+            current_output.register(step.__class__.__name__)
             exec = StepExecutor(step, current_input, current_output, context)
             current_output = next_output
             executors.append(exec)
@@ -218,6 +199,4 @@ async def run(self, reporter: PipelineProgressReporter):
         # concurrently. This will block until all steps are finished.
         running_steps = (create_task(executor.run()) for executor in executors)
 
-        self.logger.info("Starting Pipeline")
         await gather(*running_steps, create_task(pipeline_output.run()))
-        self.logger.info("Pipeline Completed")

nodestream/pipeline/progress_reporter.py

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass, field
 from logging import Logger, getLogger
-from typing import Any, Callable
+from typing import Callable
 
 from psutil import Process
 
@@ -33,7 +33,7 @@ class PipelineProgressReporter:
 
     reporting_frequency: int = 10000
     logger: Logger = field(default_factory=getLogger)
-    callback: Callable[[int, Any], None] = field(default=no_op)
+    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)
@@ -66,11 +66,6 @@ def for_testing(cls, results_list: list) -> "PipelineProgressReporter":
             on_fatal_error_callback=raise_exception,
         )
 
-    def report(self, index, record):
+    def report(self, index, metrics: Metrics):
         if index % 10000 == 0:
-            # self.logger.info(
-            #     "Records Processed",
-            #     extra={"index": index, "max_memory": get_max_mem_mb()},
-            # )
-            # self.callback(index, record)
-            Metrics.get().tick()
+            self.callback(index, metrics)

nodestream/pipeline/step.py

@@ -57,6 +57,7 @@ def report_error(
         )
         if fatal:
             self.reporter.on_fatal_error(exception)
+            Metrics.get().increment(NodestreamMetricRegistry.FATAL_ERRORS)
         else:
             Metrics.get().increment(NodestreamMetricRegistry.NON_FATAL_ERRORS)