deltacat 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

deltacat/compute/compactor_v2/utils/task_options.py

@@ -1,7 +1,10 @@
 import botocore
 import logging
-from typing import Dict, Optional, List, Tuple
+from typing import Dict, Optional, List, Tuple, Any
 from deltacat import logs
+from deltacat.compute.compactor_v2.model.merge_file_group import (
+    LocalMergeFileGroupsProvider,
+)
 from deltacat.types.media import ContentEncoding, ContentType
 from deltacat.types.partial_download import PartialParquetParameters
 from deltacat.storage import (
@@ -15,7 +18,6 @@ from deltacat.compute.compactor_v2.utils.primary_key_index import (
     hash_group_index_to_hash_bucket_indices,
 )
 from deltacat.compute.compactor_v2.constants import (
-    TOTAL_MEMORY_BUFFER_PERCENTAGE,
     PARQUET_TO_PYARROW_INFLATION,
 )
 
@@ -133,8 +135,10 @@ def hash_bucket_resource_options_provider(
     item: DeltaAnnotated,
     previous_inflation: float,
     average_record_size_bytes: float,
+    total_memory_buffer_percentage: int,
     primary_keys: List[str] = None,
     ray_custom_resources: Optional[Dict] = None,
+    memory_logs_enabled: Optional[bool] = None,
     **kwargs,
 ) -> Dict:
     debug_memory_params = {"hash_bucket_task_index": index}
@@ -189,10 +193,11 @@ def hash_bucket_resource_options_provider(
     debug_memory_params["average_record_size_bytes"] = average_record_size_bytes
 
     # Consider buffer
-    total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+    total_memory = total_memory * (1 + total_memory_buffer_percentage / 100.0)
     debug_memory_params["total_memory_with_buffer"] = total_memory
-    logger.debug(
-        f"[Hash bucket task {index}]: Params used for calculating hash bucketing memory: {debug_memory_params}"
+    logger.debug_conditional(
+        f"[Hash bucket task {index}]: Params used for calculating hash bucketing memory: {debug_memory_params}",
+        memory_logs_enabled,
     )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)
@@ -204,12 +209,14 @@ def merge_resource_options_provider(
     num_hash_groups: int,
     hash_group_size_bytes: Dict[int, int],
     hash_group_num_rows: Dict[int, int],
+    total_memory_buffer_percentage: int,
     round_completion_info: Optional[RoundCompletionInfo] = None,
     compacted_delta_manifest: Optional[Manifest] = None,
     ray_custom_resources: Optional[Dict] = None,
     primary_keys: Optional[List[str]] = None,
     deltacat_storage=unimplemented_deltacat_storage,
     deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
     **kwargs,
 ) -> Dict:
     debug_memory_params = {"merge_task_index": index}
@@ -224,6 +231,84 @@ def merge_resource_options_provider(
     pk_size_bytes = data_size
     incremental_index_array_size = num_rows * 4
 
+    return get_merge_task_options(
+        index,
+        hb_group_idx,
+        data_size,
+        pk_size_bytes,
+        num_rows,
+        num_hash_groups,
+        total_memory_buffer_percentage,
+        incremental_index_array_size,
+        debug_memory_params,
+        ray_custom_resources,
+        round_completion_info=round_completion_info,
+        compacted_delta_manifest=compacted_delta_manifest,
+        primary_keys=primary_keys,
+        deltacat_storage=deltacat_storage,
+        deltacat_storage_kwargs=deltacat_storage_kwargs,
+        memory_logs_enabled=memory_logs_enabled,
+    )
+
+
+def local_merge_resource_options_provider(
+    estimated_da_size: float,
+    estimated_num_rows: int,
+    total_memory_buffer_percentage: int,
+    round_completion_info: Optional[RoundCompletionInfo] = None,
+    compacted_delta_manifest: Optional[Manifest] = None,
+    ray_custom_resources: Optional[Dict] = None,
+    primary_keys: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
+    **kwargs,
+) -> Dict:
+    index = hb_group_idx = LocalMergeFileGroupsProvider.LOCAL_HASH_BUCKET_INDEX
+    debug_memory_params = {"merge_task_index": index}
+
+    # upper bound for pk size of incremental
+    pk_size_bytes = estimated_da_size
+    incremental_index_array_size = estimated_num_rows * 4
+
+    return get_merge_task_options(
+        index=index,
+        hb_group_idx=hb_group_idx,
+        data_size=estimated_da_size,
+        pk_size_bytes=pk_size_bytes,
+        num_rows=estimated_num_rows,
+        num_hash_groups=1,
+        incremental_index_array_size=incremental_index_array_size,
+        total_memory_buffer_percentage=total_memory_buffer_percentage,
+        debug_memory_params=debug_memory_params,
+        ray_custom_resources=ray_custom_resources,
+        round_completion_info=round_completion_info,
+        compacted_delta_manifest=compacted_delta_manifest,
+        primary_keys=primary_keys,
+        deltacat_storage=deltacat_storage,
+        deltacat_storage_kwargs=deltacat_storage_kwargs,
+        memory_logs_enabled=memory_logs_enabled,
+    )
+
+
+def get_merge_task_options(
+    index: int,
+    hb_group_idx: int,
+    data_size: float,
+    pk_size_bytes: float,
+    num_rows: int,
+    num_hash_groups: int,
+    total_memory_buffer_percentage: int,
+    incremental_index_array_size: int,
+    debug_memory_params: Dict[str, Any],
+    ray_custom_resources: Optional[Dict],
+    round_completion_info: Optional[RoundCompletionInfo] = None,
+    compacted_delta_manifest: Optional[Manifest] = None,
+    primary_keys: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
+) -> Dict[str, Any]:
     if (
         round_completion_info
         and compacted_delta_manifest
@@ -296,10 +381,11 @@ def merge_resource_options_provider(
     debug_memory_params["incremental_index_array_size"] = incremental_index_array_size
     debug_memory_params["total_memory"] = total_memory
 
-    total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+    total_memory = total_memory * (1 + total_memory_buffer_percentage / 100.0)
     debug_memory_params["total_memory_with_buffer"] = total_memory
-    logger.debug(
-        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}"
+    logger.debug_conditional(
+        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}",
+        memory_logs_enabled,
     )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)

deltacat/io/memcached_object_store.py

@@ -181,15 +181,35 @@ class MemcachedObjectStore(IObjectStore):
         for chunk_index in range(chunk_count):
             ref = self._create_ref(uid, ip, chunk_index)
             chunk = client.get(ref)
+            if chunk is None:
+                raise ValueError(
+                    f"Expected uid: {uid}, chunk index: {chunk_index} from client ip: {ip}"
+                    f" to be non-empty."
+                )
             serialized.extend(chunk)
 
         return cloudpickle.loads(serialized)
 
+    def clear(self) -> bool:
+        flushed = all(
+            [
+                self._get_client_by_ip(ip).flush_all(noreply=False)
+                for ip in self.storage_node_ips
+            ]
+        )
+        self.client_cache.clear()
+
+        if flushed:
+            logger.info("Successfully cleared cache contents.")
+
+        return flushed
+
     def close(self) -> None:
         for client in self.client_cache.values():
             client.close()
 
         self.client_cache.clear()
+        logger.info("Successfully closed object store clients.")
 
     def _create_ref(self, uid, ip, chunk_index) -> str:
         return f"{uid}{self.SEPARATOR}{ip}{self.SEPARATOR}{chunk_index}"

deltacat/logs.py

@@ -2,7 +2,7 @@ import logging
 import os
 import pathlib
 from logging import FileHandler, Handler, Logger, LoggerAdapter, handlers
-from typing import Union
+from typing import Any, Dict, Optional, Union
 
 import ray
 from ray.runtime_context import RuntimeContext
@@ -26,7 +26,32 @@ DEFAULT_MAX_BYTES_PER_LOG = 2 ^ 20 * 256  # 256 MiB
 DEFAULT_BACKUP_COUNT = 0
 
 
-class RayRuntimeContextLoggerAdapter(logging.LoggerAdapter):
+class DeltaCATLoggerAdapter(logging.LoggerAdapter):
+    """
+    Logger Adapter class with additional functionality
+    """
+
+    def __init__(self, logger: Logger, extra: Optional[Dict[str, Any]] = {}):
+        super().__init__(logger, extra)
+
+    def debug_conditional(self, msg, do_print: bool, *args, **kwargs):
+        if do_print:
+            self.debug(msg, *args, **kwargs)
+
+    def info_conditional(self, msg, do_print: bool, *args, **kwargs):
+        if do_print:
+            self.info(msg, *args, **kwargs)
+
+    def warning_conditional(self, msg, do_print: bool, *args, **kwargs):
+        if do_print:
+            self.warning(msg, *args, **kwargs)
+
+    def error_conditional(self, msg, do_print: bool, *args, **kwargs):
+        if do_print:
+            self.error(msg, *args, **kwargs)
+
+
+class RayRuntimeContextLoggerAdapter(DeltaCATLoggerAdapter):
     """
     Logger Adapter for injecting Ray Runtime Context into logging messages.
     """
@@ -147,6 +172,8 @@ def _configure_logger(
         ray_runtime_ctx = ray.get_runtime_context()
         if ray_runtime_ctx.worker.connected:
             logger = RayRuntimeContextLoggerAdapter(logger, ray_runtime_ctx)
+    else:
+        logger = DeltaCATLoggerAdapter(logger)
 
     return logger
 

deltacat/tests/compute/test_compact_partition_params.py

@@ -72,6 +72,7 @@ class TestCompactPartitionParams(unittest.TestCase):
                 "partitionValues": [],
                 "partitionId": "79612ea39ac5493eae925abe60767d42",
             },
+            "memory_logs_enabled": True,
             "metrics_config": MetricsConfig("us-east-1", MetricsTarget.CLOUDWATCH_EMF),
         }
 
@@ -135,6 +136,10 @@ class TestCompactPartitionParams(unittest.TestCase):
             json.loads(serialized_params)["destination_partition_locator"]
             == params.destination_partition_locator
         )
+        assert (
+            json.loads(serialized_params)["memory_logs_enabled"]
+            == params.memory_logs_enabled
+        )
         assert (
             json.loads(serialized_params)["metrics_config"]["metrics_target"]
             == params.metrics_config.metrics_target

deltacat/tests/io/test_memcached_object_store.py

@@ -25,6 +25,10 @@ class MockPyMemcacheClient:
     def get(self, key, *args, **kwargs):
         return self.store.get(key)
 
+    def flush_all(self, *args, **kwargs):
+        for key, value in self.store.items():
+            self.store[key] = None
+
 
 class TestMemcachedObjectStore(unittest.TestCase):
 
@@ -192,3 +196,18 @@ class TestMemcachedObjectStore(unittest.TestCase):
         # assert
         result = self.object_store.get(ref)
         self.assertEqual(result, self.TEST_VALUE_LARGE)
+
+    @mock.patch("deltacat.io.memcached_object_store.Client")
+    @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
+    def test_clear_sanity(self, mock_retrying_client, mock_client):
+        # setup
+        mock_client.return_value = MockPyMemcacheClient()
+        mock_retrying_client.return_value = mock_client.return_value
+
+        # action
+        ref = self.object_store.put(self.TEST_VALUE_LARGE)
+        self.object_store.clear()
+
+        # assert
+        with self.assertRaises(ValueError):
+            self.object_store.get(ref)