deltacat 0.1.10.dev0__py3-none-any.whl → 0.1.12__py3-none-any.whl

deltacat/compute/metastats/stats.py

@@ -1,30 +1,25 @@
-import ray
 import logging
-from typing import Dict, Set, Tuple, List, Optional, Any
 from collections import defaultdict
+from typing import Dict, List, Optional
 
-from deltacat.compute.stats.models.delta_stats import DeltaStats
-from deltacat.compute.stats.models.stats_result import StatsResult
+import ray
 from ray.types import ObjectRef
 
 from deltacat import logs
-from deltacat.compute.stats.models.delta_stats_cache_result import DeltaStatsCacheResult
-
-from deltacat.utils.ray_utils.concurrency import invoke_parallel, \
-    round_robin_options_provider
-from deltacat.compute.metastats.utils.io import collect_stats_by_columns, cache_inflation_rate_data_for_delta_stats_ready, cache_partition_stats_to_s3
-
-from deltacat.storage import PartitionLocator, DeltaLocator, Delta
-from deltacat.storage import interface as unimplemented_deltacat_storage
-
-from deltacat.aws.clients import client_cache
 from deltacat.aws import s3u as s3_utils
-
-
-from deltacat.compute.stats.models.manifest_entry_stats import ManifestEntryStats
-from deltacat.compute.stats.models.delta_column_stats import DeltaColumnStats
-
+from deltacat.aws.clients import client_cache
 from deltacat.compute.compactor import DeltaAnnotated
+from deltacat.compute.metastats.utils.io import (
+    cache_inflation_rate_data_for_delta_stats_ready,
+    cache_partition_stats_to_s3,
+    collect_stats_by_columns,
+)
+from deltacat.compute.stats.models.delta_column_stats import DeltaColumnStats
+from deltacat.compute.stats.models.delta_stats import DeltaStats
+from deltacat.compute.stats.models.manifest_entry_stats import ManifestEntryStats
+from deltacat.compute.stats.models.stats_result import StatsResult
+from deltacat.storage import DeltaLocator, PartitionLocator
+from deltacat.storage import interface as unimplemented_deltacat_storage
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -32,51 +27,67 @@ logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 DEFAULT_CPUS_STATS_CLUSTER_INSTANCE = 32
 
 
-def start_stats_collection(batched_delta_stats_compute_list: List[DeltaAnnotated],
-                           columns: List[str],
-                           stat_results_s3_bucket: Optional[str]=None,
-                           metastats_results_s3_bucket: Optional[str]=None,
-                           deltacat_storage=unimplemented_deltacat_storage) -> Dict[str, List[DeltaStats]]:
+def start_stats_collection(
+    batched_delta_stats_compute_list: List[DeltaAnnotated],
+    columns: List[str],
+    stat_results_s3_bucket: Optional[str] = None,
+    metastats_results_s3_bucket: Optional[str] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+) -> Dict[str, List[DeltaStats]]:
     """Collects statistics on deltas, given a set of delta stream position ranges.
-        Example:
-            >>> collect(locator, set((1, 5), (4, 8), (13, 16)))
-            {
-                1: DeltaStats(),  # DeltaStats for stream positions 1 - 8
-                13: DeltaStats()  # DeltaStats for stream positions 13 - 16
-            }
-        Args:
-            source_partition_locator: Reference to the partition locator tied to the given delta stream positions
-            delta_stream_position_range_set: A set of intervals with an int type representing finite,
-                closed bounded values, and a None type representing unbounded infinity.
-            columns: Columns can be optionally included to collect stats on specific columns.
-                By default, all columns will be calculated.
-            stat_results_s3_bucket: Used as a cache file storage for computed delta stats
-            metastats_results_s3_bucket: Used as cache file storage for inflation rate meta stats
-            deltacat_storage: Client implementation of the DeltaCAT storage interface
-        Returns:
-            A mapping of stream positions to their corresponding delta stats.
+    Example:
+        >>> collect(locator, set((1, 5), (4, 8), (13, 16)))
+        {
+            1: DeltaStats(),  # DeltaStats for stream positions 1 - 8
+            13: DeltaStats()  # DeltaStats for stream positions 13 - 16
+        }
+    Args:
+        source_partition_locator: Reference to the partition locator tied to the given delta stream positions
+        delta_stream_position_range_set: A set of intervals with an int type representing finite,
+            closed bounded values, and a None type representing unbounded infinity.
+        columns: Columns can be optionally included to collect stats on specific columns.
+            By default, all columns will be calculated.
+        stat_results_s3_bucket: Used as a cache file storage for computed delta stats
+        metastats_results_s3_bucket: Used as cache file storage for inflation rate meta stats
+        deltacat_storage: Client implementation of the DeltaCAT storage interface
+    Returns:
+        A mapping of stream positions to their corresponding delta stats.
     """
     # TODO: Add CompactionEventDispatcher for stats collection started event
     delta_stats_compute_pending: List[ObjectRef[Dict[str, List[StatsResult, int]]]] = []
 
     for batched_deltas in batched_delta_stats_compute_list:
-        splitted_annotated_deltas = DeltaAnnotated.split(batched_deltas, DEFAULT_CPUS_STATS_CLUSTER_INSTANCE)
+        splitted_annotated_deltas = DeltaAnnotated.split(
+            batched_deltas, DEFAULT_CPUS_STATS_CLUSTER_INSTANCE
+        )
         for splitted_annotated_delta in splitted_annotated_deltas:
-            delta_stats_compute_pending.append(collect_stats_by_columns.remote(splitted_annotated_delta, columns, deltacat_storage))
+            delta_stats_compute_pending.append(
+                collect_stats_by_columns.remote(
+                    splitted_annotated_delta, columns, deltacat_storage
+                )
+            )
 
     column_stats_map = _process_stats(delta_stats_compute_pending)
 
     if not batched_delta_stats_compute_list:
         logger.info("No new delta need stats collection")
     else:
-        delta_stream_range_stats, partition_canonical_string = resolve_annotated_delta_stats_to_original_deltas_stats(column_stats_map, columns, batched_delta_stats_compute_list[0])
+        (
+            delta_stream_range_stats,
+            partition_canonical_string,
+        ) = resolve_annotated_delta_stats_to_original_deltas_stats(
+            column_stats_map, columns, batched_delta_stats_compute_list[0]
+        )
 
-        _cache_stats_res_to_s3(stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string)
+        _cache_stats_res_to_s3(
+            stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string
+        )
 
         base_path = s3_utils.parse_s3_url(metastats_results_s3_bucket).url
         inflation_rate_stats_s3_url = f"{base_path}/inflation-rates.json"
-        cache_inflation_rate_data_for_delta_stats_ready(delta_stream_range_stats, inflation_rate_stats_s3_url,
-                                                         deltacat_storage)
+        cache_inflation_rate_data_for_delta_stats_ready(
+            delta_stream_range_stats, inflation_rate_stats_s3_url, deltacat_storage
+        )
         # TODO: Add CompactionEventDispatcher for stats collection completed event
         return delta_stream_range_stats
 
@@ -87,13 +98,19 @@ def _get_account_id() -> str:
     return account_id
 
 
-def _process_stats(delta_stats_compute_pending: List[ObjectRef[DeltaStats]]) -> List[DeltaStats]:
-    delta_stats_processed_list: List[DeltaStats] = _resolve_pending_stats(delta_stats_compute_pending)
+def _process_stats(
+    delta_stats_compute_pending: List[ObjectRef[DeltaStats]],
+) -> List[DeltaStats]:
+    delta_stats_processed_list: List[DeltaStats] = _resolve_pending_stats(
+        delta_stats_compute_pending
+    )
 
     return delta_stats_processed_list
 
 
-def _resolve_pending_stats(delta_stats_pending_list: List[ObjectRef[DeltaStats]]) -> List[DeltaStats]:
+def _resolve_pending_stats(
+    delta_stats_pending_list: List[ObjectRef[DeltaStats]],
+) -> List[DeltaStats]:
     delta_stats_processed_list: List[DeltaStats] = []
 
     while delta_stats_pending_list:
@@ -104,29 +121,39 @@ def _resolve_pending_stats(delta_stats_pending_list: List[ObjectRef[DeltaStats]]
     return delta_stats_processed_list
 
 
-def _cache_stats_res_to_s3(stat_results_s3_bucket,
-                           delta_stream_range_stats,
-                           partition_canonical_string):
+def _cache_stats_res_to_s3(
+    stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string
+):
     if stat_results_s3_bucket:
         # Cache the stats into the file store
-        cache_partition_stats_to_s3(stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string)
+        cache_partition_stats_to_s3(
+            stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string
+        )
 
 
-def resolve_annotated_delta_stats_to_original_deltas_stats(column_stats_map, column_names, delta_annotated) -> \
-Dict[int, DeltaStats]:
+def resolve_annotated_delta_stats_to_original_deltas_stats(
+    column_stats_map, column_names, delta_annotated
+) -> Dict[int, DeltaStats]:
 
-    partition_values = delta_annotated["deltaLocator"]["partitionLocator"]["partitionValues"]
+    partition_values = delta_annotated["deltaLocator"]["partitionLocator"][
+        "partitionValues"
+    ]
     partition_id = delta_annotated["deltaLocator"]["partitionLocator"]["partitionId"]
-    stream_locator = delta_annotated["deltaLocator"]["partitionLocator"]["streamLocator"]
-    partition_locator = PartitionLocator.of(stream_locator, partition_values, partition_id)
+    stream_locator = delta_annotated["deltaLocator"]["partitionLocator"][
+        "streamLocator"
+    ]
+    partition_locator = PartitionLocator.of(
+        stream_locator, partition_values, partition_id
+    )
 
     # Dict[stream_position: List[StatsResult]]
     manifest_column_stats_list = defaultdict(lambda: [])
     for i in range(len(column_stats_map)):
         for column_name in column_names:
             for j in range(len(column_stats_map[i][column_name])):
-                manifest_column_stats_list[column_stats_map[i][column_name][j][1]].append(
-                    [column_stats_map[i][column_name][j][0], column_name])
+                manifest_column_stats_list[
+                    column_stats_map[i][column_name][j][1]
+                ].append([column_stats_map[i][column_name][j][0], column_name])
 
     stats_res: Dict[int, List[DeltaStats]] = {}
     for key, value in manifest_column_stats_list.items():
@@ -139,11 +166,15 @@ Dict[int, DeltaStats]:
         delta_ds_column_stats: List[DeltaColumnStats] = []
         for column_name, column_manifest_stats_list in manifest_stats_list.items():
 
-            column_manifest_stats = ManifestEntryStats.of(column_manifest_stats_list, delta_locator)
-            dataset_column_stats = DeltaColumnStats.of(column_name, column_manifest_stats)
+            column_manifest_stats = ManifestEntryStats.of(
+                column_manifest_stats_list, delta_locator
+            )
+            dataset_column_stats = DeltaColumnStats.of(
+                column_name, column_manifest_stats
+            )
             delta_ds_column_stats.append(dataset_column_stats)
 
         dataset_stats: DeltaStats = DeltaStats.of(delta_ds_column_stats)
         stats_res[key] = dataset_stats
 
-    return stats_res, partition_locator.canonical_string()
+    return stats_res, partition_locator.canonical_string()

deltacat/compute/metastats/utils/io.py

@@ -1,28 +1,29 @@
-import logging
 import json
+import logging
+from collections import defaultdict
+from typing import Any, Dict, List, Optional
+
 import pyarrow
 import ray
 
-from deltacat import LocalTable, TableType
-from deltacat.storage import Delta
-from deltacat.compute.compactor import DeltaAnnotated
+from deltacat import LocalTable, TableType, logs
 from deltacat.aws import s3u as s3_utils
-from deltacat.utils.common import sha1_hexdigest
-from deltacat.storage import interface as unimplemented_deltacat_storage
+from deltacat.compute.compactor import DeltaAnnotated
 from deltacat.compute.metastats.model.partition_stats_dict import PartitionStats
-from deltacat.compute.stats.models.delta_stats_cache_result import DeltaStatsCacheResult
 from deltacat.compute.stats.models.delta_column_stats import DeltaColumnStats
 from deltacat.compute.stats.models.delta_stats import DeltaStats, DeltaStatsCacheMiss
+from deltacat.compute.stats.models.delta_stats_cache_result import DeltaStatsCacheResult
 from deltacat.compute.stats.models.stats_result import StatsResult
+from deltacat.storage import Delta
+from deltacat.storage import interface as unimplemented_deltacat_storage
+from deltacat.utils.common import sha1_hexdigest
 
-from typing import Dict, List, Optional, Any
-from collections import defaultdict
-from deltacat import logs
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
-def cache_inflation_rate_data_for_delta_stats_ready(delta_stats_processed_list, inflation_rate_stats_s3_url,
-                                                     deltacat_storage):
+def cache_inflation_rate_data_for_delta_stats_ready(
+    delta_stats_processed_list, inflation_rate_stats_s3_url, deltacat_storage
+):
     meta_stats_processed_list: Dict[int, int] = {}
 
     for key, value in delta_stats_processed_list.items():
@@ -39,17 +40,23 @@ def cache_inflation_rate_data_for_delta_stats_ready(delta_stats_processed_list,
     for key, value in delta_stats_processed_list.items():
         delta_stats_pyarrow_bytes_sum = 0
         delta_stats_row_count = 0
-        for column_stats in delta_stats_processed_list[key].column_stats[0].manifest_stats.stats:
+        for column_stats in (
+            delta_stats_processed_list[key].column_stats[0].manifest_stats.stats
+        ):
             delta_stats_row_count += column_stats.get("rowCount")
         for stats in delta_stats_processed_list[key].get("column_stats"):
 
             delta_stats_pyarrow_bytes_sum += stats.get("stats").get("pyarrowTableBytes")
-        cache_inflation_rate_res[key] = [meta_stats_processed_list[key], delta_stats_row_count,
-                                         delta_stats_pyarrow_bytes_sum]
+        cache_inflation_rate_res[key] = [
+            meta_stats_processed_list[key],
+            delta_stats_row_count,
+            delta_stats_pyarrow_bytes_sum,
+        ]
 
     if inflation_rate_stats_s3_url:
         logger.warning(
-            f"reading previous inflation rate stats from: {inflation_rate_stats_s3_url}")
+            f"reading previous inflation rate stats from: {inflation_rate_stats_s3_url}"
+        )
 
         result = s3_utils.download(inflation_rate_stats_s3_url, fail_if_not_found=False)
 
@@ -57,38 +64,57 @@ def cache_inflation_rate_data_for_delta_stats_ready(delta_stats_processed_list,
         if result:
             json_str = result["Body"].read().decode("utf-8")
             prev_inflation_rate_stats_read = json.loads(json_str)
-            prev_inflation_rate_stats = prev_inflation_rate_stats_read if prev_inflation_rate_stats_read else dict()
-            logger.debug(f"read stats completion info: {prev_inflation_rate_stats_read}")
+            prev_inflation_rate_stats = (
+                prev_inflation_rate_stats_read
+                if prev_inflation_rate_stats_read
+                else dict()
+            )
+            logger.debug(
+                f"read stats completion info: {prev_inflation_rate_stats_read}"
+            )
         logger.debug(
-            f"writing inflation rate info to S3: {inflation_rate_stats_s3_url}")
+            f"writing inflation rate info to S3: {inflation_rate_stats_s3_url}"
+        )
         prev_inflation_rate_stats.update(cache_inflation_rate_res)
-        logger.debug(f"writing current inflation rate info to S3: {prev_inflation_rate_stats}")
+        logger.debug(
+            f"writing current inflation rate info to S3: {prev_inflation_rate_stats}"
+        )
         s3_utils.upload(
-            inflation_rate_stats_s3_url,
-            json.dumps(prev_inflation_rate_stats)
+            inflation_rate_stats_s3_url, json.dumps(prev_inflation_rate_stats)
         )
     else:
-        logger.warning(f"No valid s3 url received to cache inflation rate stats, got {inflation_rate_stats_s3_url}")
+        logger.warning(
+            f"No valid s3 url received to cache inflation rate stats, got {inflation_rate_stats_s3_url}"
+        )
 
 
-def read_cached_partition_stats(partition_canonical_string: str, stat_results_s3_bucket: str):
-    partition_stats_url = get_partition_stats_s3_url(partition_canonical_string, stat_results_s3_bucket)
-    logger.info(
-        f"reading partition stats completion file from: {partition_stats_url}")
+def read_cached_partition_stats(
+    partition_canonical_string: str, stat_results_s3_bucket: str
+):
+    partition_stats_url = get_partition_stats_s3_url(
+        partition_canonical_string, stat_results_s3_bucket
+    )
+    logger.info(f"reading partition stats completion file from: {partition_stats_url}")
 
     result = s3_utils.download(partition_stats_url, fail_if_not_found=False)
     delta_stats_cache_res_map: Dict[int, List[DeltaStatsCacheResult]] = {}
     if result:
         json_str = result["Body"].read().decode("utf-8")
         partition_stats_str = json.loads(json_str)
-        delta_stats_cache_res_map = get_delta_stats_from_partition_stats(partition_stats_str)
+        delta_stats_cache_res_map = get_delta_stats_from_partition_stats(
+            partition_stats_str
+        )
 
     return delta_stats_cache_res_map
 
 
-def get_partition_stats_s3_url(partition_canonical_string: str, stat_results_s3_bucket: str):
+def get_partition_stats_s3_url(
+    partition_canonical_string: str, stat_results_s3_bucket: str
+):
     stats_partition_canonical_string = f"{partition_canonical_string}"
-    stats_partition_hexdigest = sha1_hexdigest(stats_partition_canonical_string.encode("utf-8"))
+    stats_partition_hexdigest = sha1_hexdigest(
+        stats_partition_canonical_string.encode("utf-8")
+    )
     base_path = s3_utils.parse_s3_url(stat_results_s3_bucket).url
 
     return f"{base_path}/{stats_partition_hexdigest}.json"
@@ -109,34 +135,43 @@ def get_delta_stats_from_partition_stats(partition_stats_str: str):
                 missed_columns.append(cs.column)
 
         delta_locator = delta_stats.column_stats[0].manifest_stats.delta_locator
-        found_stats: Optional[DeltaStats] = DeltaStats.of(found_columns_stats) if found_columns_stats else None
-        missed_stats: Optional[DeltaStatsCacheMiss] = DeltaStatsCacheMiss(missed_columns, delta_locator) \
-            if missed_columns else None
+        found_stats: Optional[DeltaStats] = (
+            DeltaStats.of(found_columns_stats) if found_columns_stats else None
+        )
+        missed_stats: Optional[DeltaStatsCacheMiss] = (
+            DeltaStatsCacheMiss(missed_columns, delta_locator)
+            if missed_columns
+            else None
+        )
         delta_stats_cache_res = DeltaStatsCacheResult.of(found_stats, missed_stats)
         found_columns_stats_map[int(stream_position)] = delta_stats_cache_res
     return found_columns_stats_map
 
 
-def cache_partition_stats_to_s3(stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string):
-    partition_stats = PartitionStats.of(delta_stream_range_stats, partition_canonical_string)
-    logger.info(
-        f"writing partition stats completion for {partition_canonical_string}")
+def cache_partition_stats_to_s3(
+    stat_results_s3_bucket, delta_stream_range_stats, partition_canonical_string
+):
+    partition_stats = PartitionStats.of(
+        delta_stream_range_stats, partition_canonical_string
+    )
+    logger.info(f"writing partition stats completion for {partition_canonical_string}")
     partition_stats_completion_file_s3_url = get_partition_stats_s3_url(
-        partition_canonical_string,
-        stat_results_s3_bucket
+        partition_canonical_string, stat_results_s3_bucket
     )
     s3_utils.upload(
-        partition_stats_completion_file_s3_url,
-        str(json.dumps(partition_stats))
+        partition_stats_completion_file_s3_url, str(json.dumps(partition_stats))
     )
     logger.debug(
-        f"stats completion file written to: {partition_stats_completion_file_s3_url}")
+        f"stats completion file written to: {partition_stats_completion_file_s3_url}"
+    )
 
 
 @ray.remote
-def collect_stats_by_columns(delta_annotated: DeltaAnnotated,
-                              columns_to_compute: Optional[List[str]] = None,
-                              deltacat_storage=unimplemented_deltacat_storage) -> Dict[str, Any]:
+def collect_stats_by_columns(
+    delta_annotated: DeltaAnnotated,
+    columns_to_compute: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+) -> Dict[str, Any]:
     """Materializes one manifest entry at a time to save memory usage and calculate stats from each of its columns.
 
     Args:
@@ -150,15 +185,25 @@ def collect_stats_by_columns(delta_annotated: DeltaAnnotated,
     total_tables_size = 0
 
     # Mapping of column_name -> [stats_file_idx_1, stats_file_idx_2, ... stats_file_idx_n]
-    column_stats_map = defaultdict(lambda: [[None, None]] * len(delta_annotated["manifest"].get("entries")))
+    column_stats_map = defaultdict(
+        lambda: [[None, None]] * len(delta_annotated["manifest"].get("entries"))
+    )
     src_da_entries = delta_annotated["manifest"].get("entries")
     manifest_annotations = delta_annotated["annotations"]
     for file_idx, manifest in enumerate(src_da_entries):
-        entry_pyarrow_table: LocalTable = \
-            deltacat_storage.download_delta_manifest_entry(delta_annotated, file_idx, TableType.PYARROW, columns_to_compute, equivalent_table_types="uncompacted")
-        assert isinstance(entry_pyarrow_table, pyarrow.Table), \
-            f"Stats collection is only supported for PyArrow tables, but received a table of " \
+        entry_pyarrow_table: LocalTable = (
+            deltacat_storage.download_delta_manifest_entry(
+                delta_annotated,
+                file_idx,
+                TableType.PYARROW,
+                columns_to_compute,
+                equivalent_table_types="uncompacted",
+            )
+        )
+        assert isinstance(entry_pyarrow_table, pyarrow.Table), (
+            f"Stats collection is only supported for PyArrow tables, but received a table of "
             f"type '{type(entry_pyarrow_table)}' for manifest entry {file_idx} of delta: {delta_annotated.locator}."
+        )
         total_tables_size += entry_pyarrow_table.nbytes
         if not columns_to_compute:
             columns_to_compute = entry_pyarrow_table.column_names
@@ -166,7 +211,9 @@ def collect_stats_by_columns(delta_annotated: DeltaAnnotated,
         for column_idx, pyarrow_column in enumerate(entry_pyarrow_table.columns):
             column_name = columns_to_compute[column_idx]
             origin_delta_stream_position = manifest_annotations[file_idx][-1]
-            column_stats_map[column_name][file_idx] = [StatsResult.of(len(pyarrow_column), pyarrow_column.nbytes),
-                                                       origin_delta_stream_position]
+            column_stats_map[column_name][file_idx] = [
+                StatsResult.of(len(pyarrow_column), pyarrow_column.nbytes),
+                origin_delta_stream_position,
+            ]
 
-    return column_stats_map
+    return column_stats_map

deltacat/compute/metastats/utils/pyarrow_memory_estimation_function.py

@@ -1,6 +1,9 @@
 from deltacat.constants import PYARROW_INFLATION_MULTIPLIER_ALL_COLUMNS
 
-def estimation_function(content_length, content_type, content_encoding, *args, **kwargs):
+
+def estimation_function(
+    content_length, content_type, content_encoding, *args, **kwargs
+):
     # TODO(zyiqin): update the estimation here to be consistent with number of required worker nodes estimate.
     #  Current implementation is only a rough guess using the PYARROW_INFLATION_MULTIPLIER(content_length to pyarrow butes(all columns).
     #  The full implementation logic should be:
@@ -12,4 +15,4 @@ def estimation_function(content_length, content_type, content_encoding, *args, *
     if content_length:
         return content_length * PYARROW_INFLATION_MULTIPLIER_ALL_COLUMNS
     else:
-        return 0
+        return 0

deltacat/compute/metastats/utils/ray_utils.py

@@ -1,15 +1,18 @@
+import errno
+import logging
 import os
 import subprocess
+from subprocess import run
+from typing import Any
+
 import ray
-import errno
-import logging
+from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
 
 from deltacat import logs
-from tenacity import retry, stop_after_attempt
-from typing import Any
-from deltacat.compute.metastats.utils.constants import WORKER_NODE_OBJECT_STORE_MEMORY_RESERVE_RATIO, MAX_WORKER_MULTIPLIER
-from tenacity import Retrying, stop_after_attempt, wait_fixed, RetryError
-from subprocess import run, PIPE
+from deltacat.compute.metastats.utils.constants import (
+    MAX_WORKER_MULTIPLIER,
+    WORKER_NODE_OBJECT_STORE_MEMORY_RESERVE_RATIO,
+)
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
@@ -24,13 +27,10 @@ def run_cmd_exit_code(cmd: str) -> int:
 
 def run_cmd_with_retry(cmd: str) -> None:
     retrying = Retrying(
-        wait=wait_fixed(2),
-        stop=stop_after_attempt(RAY_DOWN_DEFAULT_RETRY_ATTEMPTS)
+        wait=wait_fixed(2), stop=stop_after_attempt(RAY_DOWN_DEFAULT_RETRY_ATTEMPTS)
     )
     try:
-        retrying(
-           run_cmd_exit_code(cmd)
-        )
+        retrying(run_cmd_exit_code(cmd))
     except RetryError:
         logger.info(f"{cmd} failed after {RAY_DOWN_DEFAULT_RETRY_ATTEMPTS} retries.")
 
@@ -38,8 +38,9 @@ def run_cmd_with_retry(cmd: str) -> None:
 def run_cmd(cmd: str) -> None:
     result = run(cmd, shell=True, capture_output=True)
     exit_code = int(result.returncode)
-    assert exit_code == 0, f"`{cmd}` failed. Exit code: {exit_code} " \
-                           f"Error Trace: {result.stderr}"
+    assert exit_code == 0, (
+        f"`{cmd}` failed. Exit code: {exit_code} " f"Error Trace: {result.stderr}"
+    )
 
 
 def ray_up(cluster_cfg: str) -> None:
@@ -67,7 +68,8 @@ def get_head_node_ip(cluster_cfg: str) -> str:
         shell=True,
         capture_output=True,
         text=True,
-        check=True)
+        check=True,
+    )
     # the head node IP should be the last line printed to stdout
     head_node_ip = proc.stdout.splitlines()[-1]
     logger.info(f"Ray cluster head node IP for '{cluster_cfg}': {head_node_ip}")
@@ -83,14 +85,15 @@ def ray_init(host, port) -> Any:
 
 
 def replace_cluster_cfg_vars(
-        partition_canonical_string: str,
-        trace_id: str,
-        file_path: str,
-        min_workers: int,
-        head_type: str,
-        worker_type: str,
-        head_object_store_memory_pct: int,
-        worker_object_store_memory_pct: int) -> str:
+    partition_canonical_string: str,
+    trace_id: str,
+    file_path: str,
+    min_workers: int,
+    head_type: str,
+    worker_type: str,
+    head_object_store_memory_pct: int,
+    worker_object_store_memory_pct: int,
+) -> str:
 
     head_object_store_memory_pct = head_object_store_memory_pct if not None else 30
     worker_object_store_memory_pct = WORKER_NODE_OBJECT_STORE_MEMORY_RESERVE_RATIO * 100
@@ -98,18 +101,20 @@ def replace_cluster_cfg_vars(
     max_workers = int(min_workers * MAX_WORKER_MULTIPLIER)
     with open(file_path, "r+") as file:
         contents = file.read().replace("{{use-internal-ips}}", "True")
-        contents = contents.replace("{{partition_canonical_string}}", partition_canonical_string)
-        contents = contents.replace("{{trace_id}}", trace_id)
-        contents = contents.replace("{{min-workers}}", str(min_workers))
-        contents = contents.replace("{{max-workers}}", str(max_workers))
-        contents = contents.replace("{{head-instance-type}}", head_type)
-        contents = contents.replace("{{worker-instance-type}}", worker_type)
         contents = contents.replace(
-            "{{head-object-store-memory-pct}}",
-            str(head_object_store_memory_pct))
+            "{{partition_canonical_string}}", partition_canonical_string
+        )
+        contents = contents.replace("'{{trace_id}}'", trace_id)
+        contents = contents.replace("'{{min-workers}}'", str(min_workers))
+        contents = contents.replace("'{{max-workers}}'", str(max_workers))
+        contents = contents.replace("'{{head-instance-type}}'", head_type)
+        contents = contents.replace("'{{worker-instance-type}}'", worker_type)
         contents = contents.replace(
-            "{{worker-object-store-memory-pct}}",
-            str(worker_object_store_memory_pct))
+            "'{{head-object-store-memory-pct}}'", str(head_object_store_memory_pct)
+        )
+        contents = contents.replace(
+            "'{{worker-object-store-memory-pct}}'", str(worker_object_store_memory_pct)
+        )
     partition_id = partition_canonical_string.split("|")[-1]
     out_file_name = f"{trace_id}-{partition_id}.{os.path.basename(file_path)}"
     out_file_dir = os.path.join(os.path.dirname(file_path), "tmp")