From 38bcd9c38197eb77b4693980bea0d57ffe23bbfe Mon Sep 17 00:00:00 2001
From: Harshal Sheth <hsheth2@gmail.com>
Date: Mon, 16 Sep 2024 23:11:58 -0700
Subject: [PATCH] feat(ingest): default to ASYNC_BATCH mode in datahub-rest
 sink (#11369)

---
 docs/how/updating-datahub.md                  |   1 +
 metadata-ingestion/sink_docs/datahub.md       |  47 ++++----
 .../src/datahub/emitter/rest_emitter.py       |  36 +++++--
 .../datahub/ingestion/sink/datahub_rest.py    |  12 ++-
 .../datahub/utilities/partition_executor.py   | 101 ++++++++++++++++--
 .../tests/unit/test_pipeline.py               |   4 +-
 .../unit/utilities/test_partition_executor.py |   8 ++
 7 files changed, 168 insertions(+), 41 deletions(-)

diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md
index abb6bcd32a554f..d8a6e4c6bdca06 100644
--- a/docs/how/updating-datahub.md
+++ b/docs/how/updating-datahub.md
@@ -37,6 +37,7 @@ This file documents any backwards-incompatible changes in DataHub and assists pe
   Re-running with stateful ingestion should automatically clear up the entities with old URNS and add entities with new URNs, therefore not duplicating the containers or jobs.
 
 - #11313 - `datahub get` will no longer return a key aspect for entities that don't exist.
+- #11369 - The default datahub-rest sink mode has been changed to `ASYNC_BATCH`. This requires a server with version 0.14.0+.
 
 ### Potential Downtime
 
diff --git a/metadata-ingestion/sink_docs/datahub.md b/metadata-ingestion/sink_docs/datahub.md
index 8ddcf6aff10352..a60356fc3789f7 100644
--- a/metadata-ingestion/sink_docs/datahub.md
+++ b/metadata-ingestion/sink_docs/datahub.md
@@ -29,6 +29,7 @@ sink:
 ```
 
 If you are connecting to a hosted DataHub Cloud instance, your sink will look like
+
 ```yml
 source:
   # source configs
@@ -68,16 +69,17 @@ If you are using [UI based ingestion](../../docs/ui-ingestion.md) then where GMS
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
 | Field                      | Required | Default              | Description                                                                                        |
-|----------------------------|----------|----------------------|----------------------------------------------------------------------------------------------------|
-| `server`                   | ✅        |                      | URL of DataHub GMS endpoint.                                                                       |
+| -------------------------- | -------- | -------------------- | -------------------------------------------------------------------------------------------------- |
+| `server`                   | ✅       |                      | URL of DataHub GMS endpoint.                                                                       |
+| `token`                    |          |                      | Bearer token used for authentication.                                                              |
 | `timeout_sec`              |          | 30                   | Per-HTTP request timeout.                                                                          |
 | `retry_max_times`          |          | 1                    | Maximum times to retry if HTTP request fails. The delay between retries is increased exponentially |
 | `retry_status_codes`       |          | [429, 502, 503, 504] | Retry HTTP request also on these status codes                                                      |
-| `token`                    |          |                      | Bearer token used for authentication.                                                              |
 | `extra_headers`            |          |                      | Extra headers which will be added to the request.                                                  |
-| `max_threads`              |          | `15`                 | Experimental: Max parallelism for REST API calls                                                   |
-| `ca_certificate_path`      |          |                      | Path to server's CA certificate for verification of HTTPS communications                                                    |
-| `client_certificate_path`      |          |                      | Path to client's CA certificate for HTTPS communications                                                    |
+| `max_threads`              |          | `15`                 | Max parallelism for REST API calls                                                                 |
+| `mode`                     |          | `ASYNC_BATCH`        | [Advanced] Mode of operation - `SYNC`, `ASYNC`, or `ASYNC_BATCH`                                   |
+| `ca_certificate_path`      |          |                      | Path to server's CA certificate for verification of HTTPS communications                           |
+| `client_certificate_path`  |          |                      | Path to client's CA certificate for HTTPS communications                                           |
 | `disable_ssl_verification` |          | false                | Disable ssl certificate validation                                                                 |
 
 ## DataHub Kafka
@@ -115,14 +117,14 @@ sink:
 
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
-| Field                                        | Required | Default | Description                                                                                                                                              |
-| -------------------------------------------- | -------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `connection.bootstrap`                       | ✅       |         | Kafka bootstrap URL.                                                                                                                                     |
-| `connection.producer_config.<option>`        |          |         | Passed to https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.SerializingProducer                  |
-| `connection.schema_registry_url`             | ✅       |         | URL of schema registry being used.                                                                                                                       |
-| `connection.schema_registry_config.<option>` |          |         | Passed to https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.schema_registry.SchemaRegistryClient |
-| `topic_routes.MetadataChangeEvent`           |          | MetadataChangeEvent     | Overridden Kafka topic name for the MetadataChangeEvent |
-| `topic_routes.MetadataChangeProposal`        |          | MetadataChangeProposal  | Overridden Kafka topic name for the MetadataChangeProposal |
+| Field                                        | Required | Default                | Description                                                                                                                                              |
+| -------------------------------------------- | -------- | ---------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `connection.bootstrap`                       | ✅       |                        | Kafka bootstrap URL.                                                                                                                                     |
+| `connection.producer_config.<option>`        |          |                        | Passed to https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.SerializingProducer                  |
+| `connection.schema_registry_url`             | ✅       |                        | URL of schema registry being used.                                                                                                                       |
+| `connection.schema_registry_config.<option>` |          |                        | Passed to https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#confluent_kafka.schema_registry.SchemaRegistryClient |
+| `topic_routes.MetadataChangeEvent`           |          | MetadataChangeEvent    | Overridden Kafka topic name for the MetadataChangeEvent                                                                                                  |
+| `topic_routes.MetadataChangeProposal`        |          | MetadataChangeProposal | Overridden Kafka topic name for the MetadataChangeProposal                                                                                               |
 
 The options in the producer config and schema registry config are passed to the Kafka SerializingProducer and SchemaRegistryClient respectively.
 
@@ -178,15 +180,14 @@ DataHub Lite currently doesn't support stateful ingestion, so you'll have to tur
 
 Note that a `.` is used to denote nested fields in the YAML recipe.
 
-| Field                      | Required | Default              | Description                                                                                        |
-|----------------------------|----------|----------------------|----------------------------------------------------------------------------------------------------|
-| `type`                   |       |      duckdb                | Type of DataHub Lite implementation to use |
-| `config`              |          | `{"file": "~/.datahub/lite/datahub.duckdb"}`                   | Config dictionary to pass through to the DataHub Lite implementation. See below for fields accepted by the DuckDB implementation |
+| Field    | Required | Default                                      | Description                                                                                                                      |
+| -------- | -------- | -------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------- |
+| `type`   |          | duckdb                                       | Type of DataHub Lite implementation to use                                                                                       |
+| `config` |          | `{"file": "~/.datahub/lite/datahub.duckdb"}` | Config dictionary to pass through to the DataHub Lite implementation. See below for fields accepted by the DuckDB implementation |
 
 #### DuckDB Config Details
 
-| Field                      | Required | Default              | Description                                                                                        |
-|----------------------------|----------|----------------------|----------------------------------------------------------------------------------------------------|
-| `file`                   |       |   `"~/.datahub/lite/datahub.duckdb"`     | File to use for DuckDB storage |
-| `options`                |          | `{}`                   | Options dictionary to pass through to DuckDB library. See [the official spec](https://duckdb.org/docs/sql/configuration.html) for the options supported by DuckDB. |
-
+| Field     | Required | Default                            | Description                                                                                                                                                        |
+| --------- | -------- | ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `file`    |          | `"~/.datahub/lite/datahub.duckdb"` | File to use for DuckDB storage                                                                                                                                     |
+| `options` |          | `{}`                               | Options dictionary to pass through to DuckDB library. See [the official spec](https://duckdb.org/docs/sql/configuration.html) for the options supported by DuckDB. |
diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
index f4bc3666a18c82..e370ad3562a06b 100644
--- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py
+++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
@@ -45,6 +45,9 @@
     os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
 )
 
+# The limit is 16mb. We will use a max of 15mb to have some space for overhead.
+_MAX_BATCH_INGEST_PAYLOAD_SIZE = 15 * 1024 * 1024
+
 
 class DataHubRestEmitter(Closeable, Emitter):
     _gms_server: str
@@ -269,20 +272,37 @@ def emit_mcps(
         self,
         mcps: List[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
         async_flag: Optional[bool] = None,
-    ) -> None:
+    ) -> int:
         url = f"{self._gms_server}/aspects?action=ingestProposalBatch"
         for mcp in mcps:
             ensure_has_system_metadata(mcp)
 
         mcp_objs = [pre_json_transform(mcp.to_obj()) for mcp in mcps]
-        payload_dict: dict = {"proposals": mcp_objs}
-
-        if async_flag is not None:
-            payload_dict["async"] = "true" if async_flag else "false"
-
-        payload = json.dumps(payload_dict)
 
-        self._emit_generic(url, payload)
+        # As a safety mechanism, we need to make sure we don't exceed the max payload size for GMS.
+        # If we will exceed the limit, we need to break it up into chunks.
+        mcp_obj_chunks: List[List[str]] = []
+        current_chunk_size = _MAX_BATCH_INGEST_PAYLOAD_SIZE
+        for mcp_obj in mcp_objs:
+            mcp_obj_size = len(json.dumps(mcp_obj))
+
+            if mcp_obj_size + current_chunk_size > _MAX_BATCH_INGEST_PAYLOAD_SIZE:
+                mcp_obj_chunks.append([])
+                current_chunk_size = 0
+            mcp_obj_chunks[-1].append(mcp_obj)
+            current_chunk_size += mcp_obj_size
+
+        for mcp_obj_chunk in mcp_obj_chunks:
+            # TODO: We're calling json.dumps on each MCP object twice, once to estimate
+            # the size when chunking, and again for the actual request.
+            payload_dict: dict = {"proposals": mcp_obj_chunk}
+            if async_flag is not None:
+                payload_dict["async"] = True if async_flag else False
+
+            payload = json.dumps(payload_dict)
+            self._emit_generic(url, payload)
+
+        return len(mcp_obj_chunks)
 
     @deprecated
     def emit_usage(self, usageStats: UsageAggregation) -> None:
diff --git a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
index 7090dc5cb6e10b..3d50ef5f254a0f 100644
--- a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
+++ b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
@@ -57,7 +57,7 @@ class RestSinkMode(ConfigEnum):
 
 
 _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
-    RestSinkMode, os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE", RestSinkMode.ASYNC)
+    RestSinkMode, os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE", RestSinkMode.ASYNC_BATCH)
 )
 
 
@@ -78,6 +78,8 @@ class DataHubRestSinkReport(SinkReport):
     gms_version: Optional[str] = None
     pending_requests: int = 0
 
+    async_batches_split: int = 0
+
     main_thread_blocking_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
 
     def compute_stats(self) -> None:
@@ -255,7 +257,13 @@ def _emit_batch_wrapper(
             else:
                 events.append(event)
 
-        self.emitter.emit_mcps(events)
+        chunks = self.emitter.emit_mcps(events)
+        if chunks > 1:
+            self.report.async_batches_split += chunks
+            logger.info(
+                f"In async_batch mode, the payload was split into {chunks} batches. "
+                "If there's many of these issues, consider decreasing `max_per_batch`."
+            )
 
     def write_record_async(
         self,
diff --git a/metadata-ingestion/src/datahub/utilities/partition_executor.py b/metadata-ingestion/src/datahub/utilities/partition_executor.py
index 05e81da47285d0..92413eeb674f54 100644
--- a/metadata-ingestion/src/datahub/utilities/partition_executor.py
+++ b/metadata-ingestion/src/datahub/utilities/partition_executor.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import atexit
 import collections
 import functools
 import logging
 import queue
 import threading
 import time
+import weakref
 from concurrent.futures import Future, ThreadPoolExecutor
 from datetime import datetime, timedelta, timezone
 from threading import BoundedSemaphore
@@ -26,7 +28,6 @@
 
 logger = logging.getLogger(__name__)
 _R = TypeVar("_R")
-_Args = TypeVar("_Args", bound=tuple)
 _PARTITION_EXECUTOR_FLUSH_SLEEP_INTERVAL = 0.05
 _DEFAULT_BATCHER_MIN_PROCESS_INTERVAL = timedelta(seconds=30)
 
@@ -173,6 +174,55 @@ def _now() -> datetime:
     return datetime.now(tz=timezone.utc)
 
 
+_live_executors: weakref.WeakSet[BatchPartitionExecutor] = weakref.WeakSet()
+
+
+def _shutdown_executors() -> None:
+    # Create a copy of the set, to avoid issues with the set being mutated while we're iterating over it.
+    executors = list(_live_executors)
+
+    for executor in executors:
+        executor.shutdown()
+
+
+# The order of Python shutdown hooks is:
+# 1. execute threading._register_atexit handlers (only exists in Python 3.9+)
+# 2. calls `thread.join` on all non-daemon threads.
+# 3. interpreter shutdown: atexit handlers run.
+#
+# There was also discussion of making the threading._register_atexit method public,
+# but that hasn't happened yet. See https://github.com/python/cpython/issues/86128.
+#
+# The BatchPartitionExecutor uses a ThreadPoolExecutor internally.
+# As of Python 3.9, the ThreadPoolExecutor no longer uses daemon threads.
+# See https://bugs.python.org/issue39812. In 3.9+, it uses the newly
+# added threading._register_atexit method to register its own shutdown
+# logic, and uses the standard atexit module on prior versions of Python.
+#
+# Our main requirement is that our shutdown logic must run before
+# the thread pool executor's shutdown logic. We need this because the
+# clearinghouse thread must get a signal to shut down before it is
+# joined. Otherwise it'll stall indefinitely while waiting for work items.
+# For all shutdown hooks, they're executed in reverse order of registration.
+# As long as we use the same shutdown method as the ThreadPoolExecutor and
+# register our shutdown hook after the ThreadPoolExecutor's shutdown hook,
+# we're guaranteed that our shutdown logic runs first.
+#
+# Some other posts had suggested using threading.main_thread().join().
+# See https://stackoverflow.com/a/63075281.
+# I couldn't get this to work reliably, and so opted for this approach.
+# Based on my reading of https://github.com/python/cpython/pull/19149,
+# it should've worked. But not worth the effort to debug.
+#
+# This entire shutdown hook is largely a backstop mechanism to protect against
+# improper usage of the BatchPartitionExecutor. In proper usage that uses
+# a context manager or calls shutdown() explicitly, this will be a no-op.
+if hasattr(threading, "_register_atexit"):
+    threading._register_atexit(_shutdown_executors)
+else:
+    atexit.register(_shutdown_executors)
+
+
 class BatchPartitionExecutor(Closeable):
     def __init__(
         self,
@@ -220,9 +270,14 @@ def __init__(
             maxsize=max_pending
         )
 
+        # If this is true, that means shutdown() has been called.
+        self._shutting_down = False
+
         # If this is true, that means shutdown() has been called and self._pending is empty.
         self._queue_empty_for_shutdown = False
 
+        _live_executors.add(self)
+
     def _clearinghouse_worker(self) -> None:  # noqa: C901
         # This worker will pull items off the queue, and submit them into the executor
         # in batches. Only this worker will submit process commands to the executor thread pool.
@@ -291,11 +346,19 @@ def _build_batch() -> List[_BatchPartitionWorkItem]:
                     blocking = False
 
                 try:
-                    next_item: Optional[_BatchPartitionWorkItem] = self._pending.get(
-                        block=blocking,
-                        timeout=self.min_process_interval.total_seconds(),
-                    )
-                    if next_item is None:
+                    next_item: Optional[_BatchPartitionWorkItem]
+                    if not blocking:
+                        next_item = self._pending.get_nowait()
+                    else:
+                        # Why 3 seconds? It's somewhat arbitrary.
+                        # We don't want it to be too high, since then liveness suffers,
+                        # particularly during a dirty shutdown. If it's too low, then we'll
+                        # waste CPU cycles rechecking the timer, only to call get again.
+                        next_item = self._pending.get(
+                            timeout=3,  # seconds
+                        )
+
+                    if next_item is None:  # None is the shutdown signal
                         self._queue_empty_for_shutdown = True
                         break
 
@@ -349,6 +412,21 @@ def _submit_batch(next_batch: List[_BatchPartitionWorkItem]) -> None:
             # We could wait for the in-flight items to complete,
             # but the executor will take care of waiting for them to complete.
         except Exception as e:
+            if isinstance(e, RuntimeError) and "shutdown" in str(e):
+                # This comes from the thread pool executor. It happens when
+                # (1) shutdown() is not called explicitly, or via the context manager,
+                # (2) submit() was called at least once, and
+                # (3) the program is now terminating (either normally or via an exception like KeyboardInterrupt)
+                # This means that the submit calls will not go through, and instead
+                # will be lost. That's ok though, since it's the caller's fault for not
+                # calling shutdown() properly.
+                logger.error(
+                    f"{self.__class__.__name__}: submit() was called but the executor was not cleaned up properly. "
+                    "The data from the submit() calls will be lost. Use a context manager or call shutdown() explicitly "
+                    "to ensure all submitted work is processed."
+                )
+                return
+
             # This represents a fatal error that makes the entire executor defunct.
             logger.exception(
                 "Threaded executor's clearinghouse worker failed.", exc_info=e
@@ -357,6 +435,11 @@ def _submit_batch(next_batch: List[_BatchPartitionWorkItem]) -> None:
             self._clearinghouse_started = False
 
     def _ensure_clearinghouse_started(self) -> None:
+        if self._shutting_down:
+            raise RuntimeError(
+                f"{self.__class__.__name__} is shutting down; cannot submit new work items."
+            )
+
         # Lazily start the clearinghouse worker.
         if not self._clearinghouse_started:
             self._clearinghouse_started = True
@@ -376,6 +459,8 @@ def submit(
         self._pending.put(_BatchPartitionWorkItem(key, args, done_callback))
 
     def shutdown(self) -> None:
+        self._shutting_down = True
+
         if not self._clearinghouse_started:
             # This is required to make shutdown() idempotent, which is important
             # when it's called explicitly and then also by a context manager.
@@ -398,7 +483,9 @@ def shutdown(self) -> None:
         while self._clearinghouse_started:
             time.sleep(_PARTITION_EXECUTOR_FLUSH_SLEEP_INTERVAL)
 
-        self._executor.shutdown(wait=False)
+        self._executor.shutdown(wait=True)
+        assert not self._clearinghouse_started
+        _live_executors.remove(self)
 
     def close(self) -> None:
         self.shutdown()
diff --git a/metadata-ingestion/tests/unit/test_pipeline.py b/metadata-ingestion/tests/unit/test_pipeline.py
index defa117e370f00..a462f281367973 100644
--- a/metadata-ingestion/tests/unit/test_pipeline.py
+++ b/metadata-ingestion/tests/unit/test_pipeline.py
@@ -142,7 +142,9 @@ def test_configure_with_rest_sink_initializes_graph(
                         "token": "foo",
                     },
                 },
-            }
+            },
+            # We don't want to make actual API calls during this test.
+            report_to=None,
         )
         # assert that the default sink config is for a DatahubRestSink
         assert isinstance(pipeline.config.sink, DynamicTypedConfig)
diff --git a/metadata-ingestion/tests/unit/utilities/test_partition_executor.py b/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
index 81c5b898caf2bb..2901b4bba107ea 100644
--- a/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
+++ b/metadata-ingestion/tests/unit/utilities/test_partition_executor.py
@@ -148,3 +148,11 @@ def process_batch(batch):
     assert len(batches_processed) == 3
     for batch in batches_processed:
         assert len(batch) <= 2, "Batch size exceeded max_per_batch limit"
+
+
+def test_empty_batch_partition_executor():
+    # We want to test that even if no submit() calls are made, cleanup works fine.
+    with BatchPartitionExecutor(
+        max_workers=5, max_pending=20, process_batch=lambda batch: None, max_per_batch=2
+    ) as executor:
+        assert executor is not None