Restructure deltacat metrics to allow for decorators

Author: yankevn

deltacat/compute/compactor/compaction_session.py

@@ -53,6 +53,7 @@
 from deltacat.compute.compactor.model.compactor_version import CompactorVersion
 from deltacat.compute.compactor.utils.sort_key import validate_sort_keys
 from deltacat.utils.resources import get_current_node_peak_memory_usage_in_bytes
+from deltacat.utils.metrics import MetricsConfigSingleton
 
 
 if importlib.util.find_spec("memray"):
@@ -139,6 +140,10 @@ def compact_partition(
     if s3_client_kwargs is None:
         s3_client_kwargs = {}
 
+    if metrics_config:
+        # Initialize MetricsConfigSingleton
+        MetricsConfigSingleton.instance(metrics_config)
+
     # memray official documentation link:
     # https://bloomberg.github.io/memray/getting_started.html
     with memray.Tracker(
@@ -166,7 +171,6 @@ def compact_partition(
             rebase_source_partition_locator,
             rebase_source_partition_high_watermark,
             enable_profiler,
-            metrics_config,
             list_deltas_kwargs,
             read_kwargs_provider,
             s3_table_writer_kwargs,
@@ -217,7 +221,6 @@ def _execute_compaction_round(
     rebase_source_partition_locator: Optional[PartitionLocator],
     rebase_source_partition_high_watermark: Optional[int],
     enable_profiler: Optional[bool],
-    metrics_config: Optional[MetricsConfig],
     list_deltas_kwargs: Optional[Dict[str, Any]],
     read_kwargs_provider: Optional[ReadKwargsProvider],
     s3_table_writer_kwargs: Optional[Dict[str, Any]],
@@ -438,7 +441,6 @@ def _execute_compaction_round(
         num_buckets=hash_bucket_count,
         num_groups=max_parallelism,
         enable_profiler=enable_profiler,
-        metrics_config=metrics_config,
         read_kwargs_provider=read_kwargs_provider,
         object_store=object_store,
         deltacat_storage=deltacat_storage,
@@ -453,7 +455,8 @@ def _execute_compaction_round(
     hb_end = time.monotonic()
     hb_results_retrieved_at = time.time()
 
-    telemetry_time_hb = compaction_audit.save_step_stats(
+    cluster_util_after_task_latency = 0
+    cluster_util_after_task_latency += compaction_audit.save_step_stats(
         CompactionSessionAuditInfo.HASH_BUCKET_STEP_NAME,
         hb_results,
         hb_results_retrieved_at,
@@ -528,7 +531,6 @@ def _execute_compaction_round(
         sort_keys=sort_keys,
         num_materialize_buckets=num_materialize_buckets,
         enable_profiler=enable_profiler,
-        metrics_config=metrics_config,
         object_store=object_store,
     )
 
@@ -546,7 +548,7 @@ def _execute_compaction_round(
     total_dd_record_count = sum([ddr.deduped_record_count for ddr in dd_results])
     logger.info(f"Deduped {total_dd_record_count} records...")
 
-    telemetry_time_dd = compaction_audit.save_step_stats(
+    cluster_util_after_task_latency += compaction_audit.save_step_stats(
         CompactionSessionAuditInfo.DEDUPE_STEP_NAME,
         dd_results,
         dedupe_results_retrieved_at,
@@ -605,7 +607,6 @@ def _execute_compaction_round(
         max_records_per_output_file=records_per_compacted_file,
         compacted_file_content_type=compacted_file_content_type,
         enable_profiler=enable_profiler,
-        metrics_config=metrics_config,
         read_kwargs_provider=read_kwargs_provider,
         s3_table_writer_kwargs=s3_table_writer_kwargs,
         object_store=object_store,
@@ -623,7 +624,7 @@ def _execute_compaction_round(
     materialize_end = time.monotonic()
     materialize_results_retrieved_at = time.time()
 
-    telemetry_time_materialize = compaction_audit.save_step_stats(
+    cluster_util_after_task_latency += compaction_audit.save_step_stats(
         CompactionSessionAuditInfo.MATERIALIZE_STEP_NAME,
         mat_results,
         materialize_results_retrieved_at,
@@ -685,8 +686,15 @@ def _execute_compaction_round(
         session_peak_memory
     )
 
+    metrics_telemetry_time = 0
+    try:
+        metrics_telemetry_time = MetricsConfigSingleton.instance().total_telemetry_time
+    except Exception as e:
+        logger.warn(
+            f"Skipping calculating metrics telemetry time due to exception: {e}"
+        )
     compaction_audit.save_round_completion_stats(
-        mat_results, telemetry_time_hb + telemetry_time_dd + telemetry_time_materialize
+        mat_results, cluster_util_after_task_latency + metrics_telemetry_time
     )
 
     s3_utils.upload(

deltacat/compute/compactor/model/compact_partition_params.py

@@ -23,7 +23,6 @@
 )
 from deltacat.constants import PYARROW_INFLATION_MULTIPLIER
 from deltacat.compute.compactor.utils.sort_key import validate_sort_keys
-from deltacat.utils.metrics import MetricsConfig
 
 
 class CompactPartitionParams(dict):
@@ -90,8 +89,6 @@ def of(params: Optional[Dict]) -> CompactPartitionParams:
         result.drop_duplicates = params.get("drop_duplicates", DROP_DUPLICATES)
         result.ray_custom_resources = params.get("ray_custom_resources")
 
-        result.metrics_config = params.get("metrics_config")
-
         if not importlib.util.find_spec("memray"):
             result.enable_profiler = False
 
@@ -354,14 +351,6 @@ def sort_keys(self) -> Optional[List[SortKey]]:
     def sort_keys(self, keys: List[SortKey]) -> None:
         self["sort_keys"] = keys
 
-    @property
-    def metrics_config(self) -> Optional[MetricsConfig]:
-        return self.get("metrics_config")
-
-    @metrics_config.setter
-    def metrics_config(self, config: MetricsConfig) -> None:
-        self["metrics_config"] = config
-
     @staticmethod
     def json_handler_for_compact_partition_params(obj):
         """

deltacat/compute/compactor/model/compaction_session_audit_info.py

@@ -832,7 +832,6 @@ def save_step_stats(
             f"{step_name}PostObjectStoreMemoryUsedBytes"
         ] = cluster_utilization_after_task.used_object_store_memory_bytes
 
-        telemetry_time = 0
         if task_results:
             last_task_completed_at = max(
                 result.task_completed_at for result in task_results
@@ -846,13 +845,9 @@ def save_step_stats(
                 result.peak_memory_usage_bytes for result in task_results
             )
 
-            telemetry_time = sum(
-                result.telemetry_time_in_seconds for result in task_results
-            )
-
             self[f"{step_name}TaskPeakMemoryUsedBytes"] = peak_task_memory.item()
 
-        return cluster_util_after_task_latency + telemetry_time
+        return cluster_util_after_task_latency
 
     def save_round_completion_stats(
         self, mat_results: List[MaterializeResult], total_telemetry_time: float

deltacat/compute/compactor/model/dedupe_result.py

@@ -7,5 +7,4 @@ class DedupeResult(NamedTuple):
     mat_bucket_idx_to_obj_id: Dict[int, Tuple]
     deduped_record_count: np.int64
     peak_memory_usage_bytes: np.double
-    telemetry_time_in_seconds: np.double
     task_completed_at: np.double

deltacat/compute/compactor/model/hash_bucket_result.py

@@ -7,5 +7,4 @@ class HashBucketResult(NamedTuple):
     hash_bucket_group_to_obj_id: np.ndarray
     hb_record_count: np.int64
     peak_memory_usage_bytes: np.double
-    telemetry_time_in_seconds: np.double
     task_completed_at: np.double

deltacat/compute/compactor/model/materialize_result.py

@@ -16,7 +16,6 @@ def of(
         pyarrow_write_result: PyArrowWriteResult,
         referenced_pyarrow_write_result: Optional[PyArrowWriteResult] = None,
         peak_memory_usage_bytes: Optional[np.double] = None,
-        telemetry_time_in_seconds: Optional[np.double] = None,
         task_completed_at: Optional[np.double] = None,
     ) -> MaterializeResult:
         materialize_result = MaterializeResult()
@@ -25,7 +24,6 @@ def of(
         materialize_result["paWriteResult"] = pyarrow_write_result
         materialize_result["referencedPaWriteResult"] = referenced_pyarrow_write_result
         materialize_result["peakMemoryUsageBytes"] = peak_memory_usage_bytes
-        materialize_result["telemetryTimeInSeconds"] = telemetry_time_in_seconds
         materialize_result["taskCompletedAt"] = task_completed_at
         return materialize_result
 
@@ -44,10 +42,6 @@ def task_index(self) -> int:
     def peak_memory_usage_bytes(self) -> Optional[np.double]:
         return self["peakMemoryUsageBytes"]
 
-    @property
-    def telemetry_time_in_seconds(self) -> Optional[np.double]:
-        return self["telemetryTimeInSeconds"]
-
     @property
     def pyarrow_write_result(self) -> PyArrowWriteResult:
         val: Dict[str, Any] = self.get("paWriteResult")

deltacat/compute/compactor/repartition_session.py

@@ -30,7 +30,7 @@
     PartitionLocator,
     interface as unimplemented_deltacat_storage,
 )
-from deltacat.utils.metrics import MetricsConfig
+from deltacat.utils.metrics import MetricsConfig, MetricsConfigSingleton
 from deltacat.compute.compactor.utils.sort_key import validate_sort_keys
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
@@ -59,6 +59,8 @@ def repartition(
     deltacat_storage=unimplemented_deltacat_storage,
     **kwargs,
 ) -> Optional[str]:
+    # Initialize MetricsConfigSingleton
+    MetricsConfigSingleton.instance(metrics_config)
 
     node_resource_keys = None
     if pg_config:  # use resource in each placement group
@@ -130,7 +132,6 @@ def repartition(
         max_records_per_output_file=records_per_repartitioned_file,
         destination_partition=partition,
         enable_profiler=enable_profiler,
-        metrics_config=metrics_config,
         read_kwargs_provider=read_kwargs_provider,
         s3_table_writer_kwargs=s3_table_writer_kwargs,
         repartitioned_file_content_type=repartitioned_file_content_type,

deltacat/compute/compactor/steps/dedupe.py

@@ -23,7 +23,7 @@
     get_current_ray_worker_id,
 )
 from deltacat.utils.performance import timed_invocation
-from deltacat.utils.metrics import emit_timer_metrics, MetricsConfig
+from deltacat.utils.metrics import emit_timer_metrics
 from deltacat.io.object_store import IObjectStore
 from deltacat.utils.resources import get_current_node_peak_memory_usage_in_bytes
 
@@ -100,15 +100,18 @@ def delta_file_locator_to_mat_bucket_index(
     return int.from_bytes(digest, "big") % materialize_bucket_count
 
 
-def _timed_dedupe(
+@ray.remote
+@emit_timer_metrics(metrics_name="dedupe")
+def dedupe(
     object_ids: List[Any],
     sort_keys: List[SortKey],
     num_materialize_buckets: int,
     dedupe_task_index: int,
     enable_profiler: bool,
     object_store: Optional[IObjectStore],
     **kwargs,
-):
+) -> DedupeResult:
+    logger.info(f"[Dedupe task {dedupe_task_index}] Starting dedupe task...")
     task_id = get_current_ray_task_id()
     worker_id = get_current_ray_worker_id()
     with memray.Tracker(
@@ -229,53 +232,10 @@ def _timed_dedupe(
         )
 
         peak_memory_usage_bytes = get_current_node_peak_memory_usage_in_bytes()
+        logger.info(f"[Dedupe task index {dedupe_task_index}] Finished dedupe task...")
         return DedupeResult(
             mat_bucket_to_dd_idx_obj_id,
             np.int64(total_deduped_records),
             np.double(peak_memory_usage_bytes),
-            np.double(0.0),
             np.double(time.time()),
         )
-
-
-@ray.remote
-def dedupe(
-    object_ids: List[Any],
-    sort_keys: List[SortKey],
-    num_materialize_buckets: int,
-    dedupe_task_index: int,
-    enable_profiler: bool,
-    metrics_config: MetricsConfig,
-    object_store: Optional[IObjectStore],
-    **kwargs,
-) -> DedupeResult:
-    logger.info(f"[Dedupe task {dedupe_task_index}] Starting dedupe task...")
-    dedupe_result, duration = timed_invocation(
-        func=_timed_dedupe,
-        object_ids=object_ids,
-        sort_keys=sort_keys,
-        num_materialize_buckets=num_materialize_buckets,
-        dedupe_task_index=dedupe_task_index,
-        enable_profiler=enable_profiler,
-        object_store=object_store,
-        **kwargs,
-    )
-
-    emit_metrics_time = 0.0
-    if metrics_config:
-        emit_result, latency = timed_invocation(
-            func=emit_timer_metrics,
-            metrics_name="dedupe",
-            value=duration,
-            metrics_config=metrics_config,
-        )
-        emit_metrics_time = latency
-
-    logger.info(f"[Dedupe task index {dedupe_task_index}] Finished dedupe task...")
-    return DedupeResult(
-        dedupe_result[0],
-        dedupe_result[1],
-        dedupe_result[2],
-        np.double(emit_metrics_time),
-        dedupe_result[4],
-    )