fix(ingest/pipeline): Fix BatchPartitionExecutor Shutdown Race Condition (#14750)

Author: treff7es

metadata-ingestion/src/datahub/ingestion/api/sink.py

@@ -1,7 +1,8 @@
 import datetime
+import logging
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Generic, Optional, Type, TypeVar, cast
+from typing import Any, Callable, Generic, List, Optional, Type, TypeVar, cast
 
 from typing_extensions import Self
 
@@ -12,6 +13,8 @@
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.type_annotations import get_class_from_annotation
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class SinkReport(Report):
@@ -89,6 +92,7 @@ class Sink(Generic[SinkConfig, SinkReportType], Closeable, metaclass=ABCMeta):
     ctx: PipelineContext
     config: SinkConfig
     report: SinkReportType
+    _pre_shutdown_callbacks: List[Callable[[], None]]
 
     @classmethod
     def get_config_class(cls) -> Type[SinkConfig]:
@@ -106,6 +110,7 @@ def __init__(self, ctx: PipelineContext, config: SinkConfig):
         self.ctx = ctx
         self.config = config
         self.report = self.get_report_class()()
+        self._pre_shutdown_callbacks = []
 
         self.__post_init__()
 
@@ -144,8 +149,28 @@ def write_record_async(
     def get_report(self) -> SinkReportType:
         return self.report
 
+    def register_pre_shutdown_callback(self, callback: Callable[[], None]) -> None:
+        """Register a callback to be executed before the sink shuts down.
+
+        This is useful for components that need to send final reports or cleanup
+        operations before the sink's resources are released.
+        """
+        self._pre_shutdown_callbacks.append(callback)
+
     def close(self) -> None:
-        pass
+        """Close the sink and clean up resources.
+
+        This method executes any registered pre-shutdown callbacks before
+        performing the actual shutdown. Subclasses should override this method
+        to provide sink-specific cleanup logic while calling super().close()
+        to ensure callbacks are executed.
+        """
+        # Execute pre-shutdown callbacks before shutdown
+        for callback in self._pre_shutdown_callbacks:
+            try:
+                callback()
+            except Exception as e:
+                logger.warning(f"Pre-shutdown callback failed: {e}", exc_info=True)
 
     def configured(self) -> str:
         """Override this method to output a human-readable and scrubbed version of the configured sink"""

metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -268,4 +268,4 @@ def on_completion(
             entity_urn=self.execution_request_input_urn,
             aspect_value=execution_result_aspect,
         )
-        self.sink.close()
+        # Note: sink.close() is handled by the pipeline's context manager

metadata-ingestion/src/datahub/ingestion/run/pipeline.py

@@ -265,6 +265,11 @@ def __init__(
                 with _add_init_error_context("configure transformers"):
                     self._configure_transforms()
 
+                # Register completion callback with sink to handle final reporting
+                self.sink.register_pre_shutdown_callback(
+                    self._notify_reporters_on_ingestion_completion
+                )
+
             # If all of the initialization succeeds, we can preserve the exit stack until the pipeline run.
             # We need to use an exit stack so that if we have an exception during initialization,
             # things that were already initialized are still cleaned up.
@@ -344,8 +349,8 @@ def _notify_reporters_on_ingestion_start(self) -> None:
         for reporter in self.reporters:
             try:
                 reporter.on_start(ctx=self.ctx)
-            except Exception as e:
-                logger.warning("Reporting failed on start", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on start", exc_info=True)
 
     def _warn_old_cli_version(self) -> None:
         """
@@ -407,8 +412,8 @@ def _notify_reporters_on_ingestion_completion(self) -> None:
                     report=self._get_structured_report(),
                     ctx=self.ctx,
                 )
-            except Exception as e:
-                logger.warning("Reporting failed on completion", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on completion", exc_info=True)
 
     @classmethod
     def create(
@@ -521,10 +526,10 @@ def run(self) -> None:
 
                     except (RuntimeError, SystemExit):
                         raise
-                    except Exception as e:
+                    except Exception:
                         logger.error(
                             "Failed to process some records. Continuing.",
-                            exc_info=e,
+                            exc_info=True,
                         )
                         # TODO: Transformer errors should be reported more loudly / as part of the pipeline report.
 
@@ -553,19 +558,16 @@ def run(self) -> None:
 
                 self.process_commits()
                 self.final_status = PipelineStatus.COMPLETED
-            except (SystemExit, KeyboardInterrupt) as e:
+            except (SystemExit, KeyboardInterrupt):
                 self.final_status = PipelineStatus.CANCELLED
-                logger.error("Caught error", exc_info=e)
+                logger.error("Caught error", exc_info=True)
                 raise
             except Exception as exc:
                 self.final_status = PipelineStatus.ERROR
                 self._handle_uncaught_pipeline_exception(exc)
             finally:
                 clear_global_warnings()
 
-        # This can't be in the finally part because this should happen after context manager exists
-        self._notify_reporters_on_ingestion_completion()
-
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:
         """
         Transforms the given sequence of records by passing the records through the transformers

metadata-ingestion/src/datahub/ingestion/sink/datahub_kafka.py

@@ -74,4 +74,5 @@ def write_record_async(
             callback(err, f"Failed to write record: {err}")
 
     def close(self) -> None:
+        super().close()
         self.emitter.flush()

metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py

@@ -349,6 +349,10 @@ def emit_async(
         )
 
     def close(self):
+        # Execute pre-shutdown callbacks first (handled by parent class)
+        super().close()
+
+        # Then perform sink-specific shutdown
         with self.report.main_thread_blocking_timer:
             self.executor.shutdown()
 

metadata-ingestion/src/datahub/ingestion/sink/file.py

@@ -79,6 +79,7 @@ def write_record_async(
             write_callback.on_success(record_envelope, {})
 
     def close(self):
+        super().close()
         self.file.write("\n]")
         self.file.close()