deltacat 0.1.8__py3-none-any.whl → 0.1.11__py3-none-any.whl

deltacat/autoscaler/events/compaction/utils.py

@@ -1,216 +0,0 @@
-import argparse
-import base64
-import gzip
-import json
-import logging
-import pathlib
-from io import BytesIO
-
-import math
-import yaml
-
-from typing import Dict, Any, List, Tuple, Optional, Set, Union, TextIO
-
-from deltacat.autoscaler.events.compaction.cluster import ClusterSizeSuggester
-from deltacat.autoscaler.events.compaction.collections.partition_key_value import PartitionKeyValues, PartitionKeyValue
-from deltacat.autoscaler.events.compaction.input import CompactionInput
-from deltacat.autoscaler.events.session_manager import SessionManager, SESSION_ID_KEY
-from deltacat.compute.compactor.utils import round_completion_file as rcf
-from deltacat.compute.stats.models.delta_stats import DeltaStats
-from deltacat.storage import interface as dcs
-from deltacat import ContentType, logs, SortKey
-from deltacat.compute.compactor import RoundCompletionInfo, compaction_session, PrimaryKeyIndexMeta, \
-    PrimaryKeyIndexLocator
-from deltacat.storage import PartitionLocator
-
-_PRIMARY_KEY_INDEX_ALGORITHM_VERSION: str = "1.0"
-
-
-logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
-
-
-def read_latest_round_completion_file(source_partition_locator,
-                                      compacted_partition_locator,
-                                      compaction_artifact_s3_bucket,
-                                      primary_keys,
-                                      sort_keys: List[SortKey] = None):
-    if sort_keys is None:
-        sort_keys = []
-    # get the root path of a compatible primary key index for this round
-    compatible_primary_key_index_meta = PrimaryKeyIndexMeta.of(
-        compacted_partition_locator,
-        primary_keys,
-        sort_keys,
-        _PRIMARY_KEY_INDEX_ALGORITHM_VERSION,
-    )
-    compatible_primary_key_index_locator = PrimaryKeyIndexLocator.of(
-        compatible_primary_key_index_meta)
-    compatible_primary_key_index_root_path = \
-        compatible_primary_key_index_locator.primary_key_index_root_path
-
-    # read the results from any previously completed compaction round that used
-    # a compatible primary key index
-    round_completion_info = rcf.read_round_completion_file(
-        compaction_artifact_s3_bucket,
-        source_partition_locator,
-        compatible_primary_key_index_root_path,
-    )
-    return round_completion_info
-
-
-def build_partition_locator(partition_data: Dict[str, Any]):
-    return PartitionLocator.at(
-        partition_data["owner"],
-        partition_data["name"],
-        partition_data.get("tableVersion"),
-        partition_data["streamUUID"],
-        None,  # storage type
-        partition_data["partition_values"],  # partition values
-        None  # partition ID
-    )
-
-
-def compact(compaction_input: CompactionInput,
-            hash_bucket_count: Optional[int] = None,
-            deltacat_storage=dcs,
-            **kwargs):
-    compaction_session.compact_partition(
-        compaction_input.source_partition_locator,
-        compaction_input.compacted_partition_locator,
-        set(compaction_input.primary_keys),
-        compaction_input.compaction_artifact_s3_bucket,
-        compaction_input.last_stream_position_to_compact,
-        schema_on_read=compaction_input.schema_on_read,
-        input_deltas_stats=compaction_input.input_deltas_stats,
-        hash_bucket_count=hash_bucket_count if hash_bucket_count else compaction_input.hash_bucket_count,
-        deltacat_storage=deltacat_storage,
-        **kwargs
-    )
-
-
-def calc_new_hash_bucket_count(cluster_memory_bytes: int,
-                               max_memory_per_vcpu: int,
-                               vcpu_per_node: int):
-    new_hash_bucket_count = max(
-        math.ceil(cluster_memory_bytes / max_memory_per_vcpu),
-        min(vcpu_per_node, 256)  # Do not exceed 256 CPUs as a safety measure
-    )
-
-    return new_hash_bucket_count
-
-
-def get_round_completion_file(
-        source_partition_locator: PartitionLocator,
-        compacted_partition_locator: PartitionLocator,
-        primary_keys: Set[str],
-        compaction_artifact_s3_bucket: str
-):
-    return read_latest_round_completion_file(source_partition_locator,
-                                             compacted_partition_locator,
-                                             compaction_artifact_s3_bucket,
-                                             sorted(primary_keys))
-
-
-def calc_compaction_cluster_memory_bytes(compaction_input: CompactionInput,
-                                         new_uncompacted_deltas_bytes: int = 0) -> int:
-    round_completion_file = get_round_completion_file(compaction_input.source_partition_locator,
-                                                      compaction_input.compacted_partition_locator,
-                                                      compaction_input.primary_keys,
-                                                      compaction_input.compaction_artifact_s3_bucket)
-    if round_completion_file is None:
-        # if no previous compaction rounds exist, use the incoming delta size as a place to start for calculations
-        est_incoming_delta_size = new_uncompacted_deltas_bytes * 1.3
-        logger.warning(f"No previous round completion file found for {compaction_input}."
-                       f"Using estimates: {est_incoming_delta_size}")
-        return int(est_incoming_delta_size)
-
-    old_num_records = round_completion_file.compacted_pyarrow_write_result.records
-    sort_keys_bit_width = round_completion_file.sort_keys_bit_width
-    pk_index_row_size = 32 + math.ceil(sort_keys_bit_width / 8)
-    cluster_memory_bytes = max(
-        old_num_records * pk_index_row_size * 1.3,  # object store memory (hash bucketing)
-        round_completion_file.compacted_pyarrow_write_result.pyarrow_bytes + new_uncompacted_deltas_bytes  # dedupe
-    )
-    return int(cluster_memory_bytes)
-
-
-def get_compaction_size_inputs(config: Dict[str, Any],
-                               partition_key_values: PartitionKeyValues,
-                               cluster_memory_bytes: int,
-                               stats_metadata: Dict[int, DeltaStats] = None,
-                               parent_session_id: str = None,
-                               session_id: str = None) -> Tuple[int, TextIO]:
-    suggester = ClusterSizeSuggester(cluster_memory_bytes=cluster_memory_bytes)
-    new_hash_bucket_count = calc_new_hash_bucket_count(cluster_memory_bytes,
-                                                       suggester.get_max_memory_per_vcpu(),
-                                                       suggester.get_num_vcpu_per_node())
-    cluster_cpus = max(
-        new_hash_bucket_count,
-        suggester.get_suggested_vcpu_count()
-    )
-    cluster_nodes = int(math.ceil(cluster_cpus / suggester.get_num_vcpu_per_node()))
-    yaml_file = generate_compaction_session_yaml(config,
-                                                 partition_key_values,
-                                                 worker_node_count=cluster_nodes,
-                                                 instance_type=suggester.instance_type,
-                                                 stats_metadata=stats_metadata,
-                                                 parent_session_id=parent_session_id,
-                                                 session_id=session_id)
-    return new_hash_bucket_count, yaml_file
-
-
-def generate_compaction_session_yaml(config: Dict[str, Any],
-                                     partition_key_values: PartitionKeyValues,
-                                     head_node_count: int = 0,
-                                     worker_node_count: int = 0,
-                                     stats_metadata: Dict[int, DeltaStats] = None,
-                                     instance_type: str = None,
-                                     parent_session_id: str = None,
-                                     session_id: str = None) -> TextIO:
-    # TODO: Remove this workaround when custom AMIs are built with baked-in build files (i.e. wheels, jars)
-    new_config = {**config}
-    for local_path, _ in new_config["file_mounts"].items():
-        new_config["file_mounts"][local_path] = local_path
-    pkv_id = partition_key_values.id
-    new_filename = f"compact.{pkv_id}.yaml"
-    new_config["cluster_name"] = f"compaction-session-{pkv_id}"
-    # Allow child clusters to re-use the same SSH key provided from the parent cluster
-    new_config["auth"]["ssh_private_key"] = f"~/ray_bootstrap_key.pem"
-    new_config["file_mounts"] = {
-        **config["file_mounts"],
-        f"~/{new_filename}": f"~/{new_filename}"
-    }
-    new_config["provider"]["use_internal_ips"] = True
-    new_config["max_workers"] = worker_node_count
-    # TODO: Determine optimal object store memory / worker heap memory allocation ratios?
-    new_config["available_node_types"]["ray.worker.default"]["min_workers"] = \
-        new_config["available_node_types"]["ray.worker.default"]["max_workers"] = worker_node_count
-    new_config["available_node_types"]["ray.worker.default"]["node_config"]["InstanceType"] = instance_type
-    new_config["available_node_types"]["ray.head.default"]["node_config"]["InstanceType"] = instance_type
-
-    # TODO: Formalize supported parameter key/values after initial shadow compaction
-    new_events = {
-        **config["events"],
-        "parameters": {
-            **config["events"]["parameters"],
-            SESSION_ID_KEY: session_id,
-        },
-        "metadata": {
-            "partitionKeyValues": compress(partition_key_values).decode('utf-8')
-        }
-    }
-    if stats_metadata:
-        new_events["metadata"]["statsMetadata"] = compress(stats_metadata).decode('utf-8')
-    new_config["events"] = new_events
-
-    with open(new_filename, "w") as yaml_file:
-        yaml.dump(new_config, yaml_file, default_flow_style=False)
-        return yaml_file
-
-def compress(serializable_obj: Union[Dict, Tuple, List]) -> bytes:
-    json_dict = json.dumps(serializable_obj)
-    out = BytesIO()
-    with gzip.open(out, "wt", encoding="utf-8") as zipfile:
-        zipfile.write(json_dict)
-
-    return base64.b64encode(out.getvalue())

deltacat/autoscaler/events/compaction/workflow.py

@@ -1,303 +0,0 @@
-import json
-import logging
-import uuid
-from typing import List, Union, Dict, Set, Any, Optional
-
-from deltacat import logs
-from deltacat.autoscaler.events.compaction.input import CompactionInput
-from deltacat.autoscaler.events.compaction.process import CompactionProcess
-from deltacat.autoscaler.events.compaction.utils import calc_compaction_cluster_memory_bytes, get_compaction_size_inputs
-from deltacat.autoscaler.events.event_store import EventStoreClient
-from deltacat.autoscaler.events.exceptions import WorkflowException
-from deltacat.autoscaler.events.session_manager import PARENT_SESSION_ID_KEY, SESSION_ID_KEY
-from deltacat.autoscaler.events.workflow import EventWorkflow, StateTransitionMap
-from deltacat.autoscaler.events.compaction.dispatcher import CompactionEventDispatcher
-from deltacat.autoscaler.events.states import ScriptStartedEvent, ScriptInProgressEvent, \
-    ScriptCompletedEvent, States, ScriptInProgressCustomEvent, RayJobRequestEvent
-from deltacat.storage import PartitionLocator
-
-logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
-
-
-NEW_REQUEST = "NEW"
-STATS_METADATA_COLLECTION_STARTED = "STATS_METADATA_COLLECTION_STARTED"
-STATS_METADATA_COLLECTION_COMPLETED = "STATS_METADATA_COLLECTION_COMPLETED"
-COMPACTION_SESSION_STARTED = "COMPACTION_SESSION_STARTED"
-COMPACTION_SESSION_PARTITIONS_COMPACTED = "COMPACTION_SESSION_PARTITIONS_COMPACTED"
-COMPACTION_SESSION_PARTITIONS_FAILURE = "COMPACTION_SESSION_PARTITIONS_FAILURE"
-COMPACTION_SESSION_PROGRESS_UPDATE = "COMPACTION_SESSION_PROGRESS_UPDATE"
-COMPACTION_SESSION_COMPLETED = "COMPACTION_SESSION_COMPLETED"
-COMPACTION_METADATA_SESSION_WORKFLOW_FAILURE = "COMPACTION_METADATA_SESSION_WORKFLOW_FAILURE"
-
-
-# TODO: Might be worth renaming this later to avoid confusion with Ray Workflows
-class CompactionWorkflow(EventWorkflow):
-
-    def __init__(self,
-                 config: Dict[str, Any],
-                 event_dispatcher: CompactionEventDispatcher,
-                 event_store: EventStoreClient = None,
-                 compaction_inputs: List[CompactionInput] = None):
-        self.config = config
-        self.event_dispatcher = event_dispatcher
-        assert self.event_dispatcher is not None, f"Event dispatcher must be provided to build and transition " \
-                                                  f"to different job states."
-        self.event_store = event_store
-        self._compaction_inputs = compaction_inputs
-
-        # Initialization
-        self._add_state_handlers()
-        self.state_transitions = self._build_state_transitions()
-        self._metastats = {}
-
-    @property
-    def state_transition_map(self) -> StateTransitionMap:
-        return self.state_transitions
-
-    def _add_state_handlers(self):
-        in_progress_custom_events = [
-            ScriptInProgressCustomEvent(COMPACTION_METADATA_SESSION_WORKFLOW_FAILURE, 0),
-            ScriptInProgressCustomEvent(STATS_METADATA_COLLECTION_STARTED, 1),
-            ScriptInProgressCustomEvent(STATS_METADATA_COLLECTION_COMPLETED, 2),
-            ScriptInProgressCustomEvent(COMPACTION_SESSION_STARTED, 3),
-            ScriptInProgressCustomEvent(COMPACTION_SESSION_PARTITIONS_COMPACTED, 4),
-            ScriptInProgressCustomEvent(COMPACTION_SESSION_PROGRESS_UPDATE, 5),
-            ScriptInProgressCustomEvent(COMPACTION_SESSION_COMPLETED, 6),
-            ScriptInProgressCustomEvent(COMPACTION_SESSION_PARTITIONS_FAILURE, 7),
-        ]
-
-        self._event_map = event_map = {event.name: event for event in in_progress_custom_events}
-        self.event_dispatcher.add_event_handlers([val for _, val in event_map.items()])
-
-        # This callback is added for tracking new child jobs to be launched
-        self.event_dispatcher.add_event_handlers([RayJobRequestEvent.new_request_delivered])
-
-    def _build_state_transitions(self) -> StateTransitionMap:
-        """Builds a mapping of event states to state transitioning callbacks, or
-        a dictionary of state transitioning callbacks.
-
-        If an event has state sequences, a dictionary of callbacks is provided
-        with sequences as keys and callback functions as values.
-
-        Returns: a map of event states to callbacks or a dictionary of callbacks
-        """
-        init_sequence = 0
-        in_progress_sequence = {name: event.state_sequence for name, event in self._event_map.items()}
-        return {
-            States.IN_PROGRESS.name: {
-                init_sequence: self.begin_stats_metadata_collection,
-                in_progress_sequence[STATS_METADATA_COLLECTION_COMPLETED]: self.begin_compaction,
-                in_progress_sequence[COMPACTION_SESSION_STARTED]: self.wait_or_mark_compaction_complete,
-                in_progress_sequence[COMPACTION_SESSION_PARTITIONS_COMPACTED]: self.wait_or_mark_compaction_complete,
-                in_progress_sequence[COMPACTION_SESSION_PROGRESS_UPDATE]: self.wait_or_mark_compaction_complete,
-                in_progress_sequence[COMPACTION_SESSION_PARTITIONS_FAILURE]: self.wait_or_mark_compaction_complete,
-                in_progress_sequence[COMPACTION_SESSION_COMPLETED]: self.complete_job
-            },
-        }
-
-    def register_compaction_inputs(self, compaction_inputs: List[CompactionInput]):
-        """Extracts and registers a set of partition IDs that need compaction from the compaction inputs.
-        """
-        self._compaction_inputs = compaction_inputs
-        compaction_source_partition_locators = [task.source_partition_locator for task in self._compaction_inputs]
-        compaction_source_partition_ids = [loc.partition_id for loc in compaction_source_partition_locators]
-        self._partition_ids_to_compact = set(compaction_source_partition_ids)
-
-    def start_workflow(self):
-        """Publish a job state event that indicates that a request to start a job run has been successfully received,
-                but the job run has not yet finished prerequisite initialization steps.
-        """
-        self.event_dispatcher.dispatch_event(ScriptInProgressEvent.in_progress)
-
-    def begin_stats_metadata_collection(self):
-        """Publish a job state event that indicates that stats metadata collection has started.
-        """
-        event = self._event_map[STATS_METADATA_COLLECTION_STARTED]
-        self.event_dispatcher.dispatch_event(event,
-                                             event_data={
-                                                 "eventName": event.name,
-                                                 "stateDetailDescription": "Running stats metadata session",
-                                             })
-        if self.session_manager:
-            self._metastats = self.session_manager.launch_stats_metadata_collection(
-                [compact.source_partition_locator for compact in self._compaction_inputs]
-            )
-            self.stats_metadata_collection_completed()
-
-    def stats_metadata_collection_completed(self):
-        """Publish a job state event that indicates that stats metadata collection is complete.
-        """
-        event = self._event_map[STATS_METADATA_COLLECTION_COMPLETED]
-        self.event_dispatcher.dispatch_event(event,
-                                             event_data={
-                                                 "eventName": event.name,
-                                                 "stateDetailDescription": "Finished collecting stats metadata",
-                                             })
-
-    def begin_compaction(self):
-        """Publish a job state event that indicates that the compaction run has started.
-        """
-        event = self._event_map[COMPACTION_SESSION_STARTED]
-        self.event_dispatcher.dispatch_event(event,
-                                             event_data={
-                                                 "eventName": event.name,
-                                                 "stateDetailDescription": "Running compaction session",
-                                             })
-        if self.session_manager:
-            processes = self.build_compaction_processes()
-            self.session_manager.launch_compaction(processes)
-
-    def build_compaction_processes(self) -> List[CompactionProcess]:
-        processes = []
-        partition_stats_metadata = self._metastats
-        for compaction_input in self._compaction_inputs:
-            stats_metadata = partition_stats_metadata.get(compaction_input.source_partition_locator.partition_id, {})
-            stats_metadata = {stream_pos: delta_stats for stream_pos, delta_stats in stats_metadata.items()
-                              if stream_pos <= compaction_input.last_stream_position_to_compact}
-            total_pyarrow_table_bytes = sum([stats_result.stats.pyarrow_table_bytes
-                                             for stream_pos, stats_result in stats_metadata.items()
-                                             if stats_result.stats is not None])
-            cluster_memory_bytes = calc_compaction_cluster_memory_bytes(compaction_input, total_pyarrow_table_bytes)
-            new_session_id = str(uuid.uuid4())
-            self.event_dispatcher.dispatch_event(RayJobRequestEvent.new_request_delivered,
-                                                 event_data={
-                                                     PARENT_SESSION_ID_KEY: self.session_manager.session_id,
-                                                     SESSION_ID_KEY: new_session_id
-                                                 })
-            new_hash_bucket_count, yaml_file = get_compaction_size_inputs(self.config,
-                                                                          compaction_input.partition_key_values,
-                                                                          cluster_memory_bytes,
-                                                                          stats_metadata=stats_metadata,
-                                                                          parent_session_id=self.session_manager.session_id,
-                                                                          session_id=new_session_id)
-            compaction_process = CompactionProcess(compaction_input.source_partition_locator,
-                                                   yaml_file.name,
-                                                   new_hash_bucket_count,
-                                                   compaction_input.last_stream_position_to_compact,
-                                                   compaction_input.partition_key_values,
-                                                   cluster_memory_bytes=cluster_memory_bytes,
-                                                   input_delta_total_bytes=total_pyarrow_table_bytes)
-
-            # TODO: Increase file descriptor limit on host (up to ~60k)
-            # TODO: Emit metrics for compaction jobs with very high number of partitions
-            processes.append(compaction_process)
-        return processes
-
-    def partitions_compacted(self,
-                             partition_locators: List[PartitionLocator]):
-        """Publish a job state event that indicates that a single partition has finished compaction.
-        A compaction session can have 1...N partitions to compact.
-        """
-        partition_completed_event = self._event_map[COMPACTION_SESSION_PARTITIONS_COMPACTED]
-        self.event_dispatcher.dispatch_event(partition_completed_event,
-                                             event_data={
-                                                 "eventName": partition_completed_event.name,
-                                                 "stateDetailDescription":
-                                                     f"Finished compaction on partitions: "
-                                                     f"{[pl.partition_id for pl in partition_locators]}",
-                                                 "stateDetailMetadata": {
-                                                     **{pl.partition_id: json.dumps(pl.partition_values)
-                                                        for pl in partition_locators}
-                                                 }
-                                             })
-
-    def partitions_compaction_failure(self,
-                                      partition_locators: List[PartitionLocator],
-                                      error_trace: Optional[str] = None ):
-        """Publish a job state event that indicates failure to compact a list of partitions.
-        """
-        failed_event = self._event_map[COMPACTION_SESSION_PARTITIONS_FAILURE]
-        self.event_dispatcher.dispatch_event(failed_event,
-                                             event_data={
-                                                 "eventName": failed_event.name,
-                                                 "errorMessage":
-                                                     f"Failure to compact partitions: "
-                                                     f"{[pl.partition_id for pl in partition_locators]}",
-                                                 "errorStackTrace": error_trace,
-                                                 "stateDetailDescription":
-                                                     f"Failure to compact partitions: "
-                                                     f"{[pl.partition_id for pl in partition_locators]}",
-                                                 "stateDetailMetadata": {
-                                                     **{pl.partition_id: json.dumps(pl.partition_values)
-                                                        for pl in partition_locators}
-                                                 }
-                                             })
-
-    def update_compaction_job_progress(self,
-                                       partition_locator: PartitionLocator,
-                                       session_id: str):
-        """Dispatch a compaction job update event for a given session ID.
-
-        :param partition_locator: Locator for a partition
-        :param session_id: Session ID to dispatch the event for
-        :return:
-        """
-        progress_event = self._event_map[COMPACTION_SESSION_PROGRESS_UPDATE]
-        partition_id = partition_locator.partition_id
-        self.event_dispatcher.dispatch_event(progress_event,
-                                             event_data={
-                                                 PARENT_SESSION_ID_KEY: session_id,
-                                                 SESSION_ID_KEY: session_id,
-                                                 "eventName": progress_event.name,
-                                                 "stateDetailDescription": f"Compaction Update",
-                                                 "stateDetailMetadata": {
-                                                     partition_id: str(partition_locator.partition_values)
-                                                 }
-                                             })
-
-    def wait_or_mark_compaction_complete(self):
-        """Publish a job state event that indicates that the compaction run is complete.
-        """
-        if self.event_store is None or self._partition_ids_to_compact is None:
-            # TODO: Separate this workflow out into multiple workflows, for different applications
-            raise WorkflowException(f"Event store and partition IDs must be defined in a workflow."
-                                    f"Event store: {self.event_store}"
-                                    f"Partition IDs to compact: {self._partition_ids_to_compact}")
-
-        partition_ids_failed = set(self.event_store.get_failed_partition_ids(self.trace_id))
-        if len(partition_ids_failed) > 0:
-            raise WorkflowException(f"Compaction workflow failed due to partition errors: {partition_ids_failed}")
-
-
-        partition_ids_completed = set(self.event_store.get_compacted_partition_ids(self.trace_id))
-        if partition_ids_completed == self._partition_ids_to_compact:
-            logger.info(f"Compaction run complete.")
-            event = self._event_map[COMPACTION_SESSION_COMPLETED]
-            self.event_dispatcher.dispatch_event(event,
-                                                 event_data={
-                                                     "eventName": event.name,
-                                                     "stateDetailDescription": "Finished compaction run",
-                                                 })
-        else:
-            logger.info(f"Compaction is in progress: {len(partition_ids_completed)} "
-                        f"out of {len(self._partition_ids_to_compact)} partitions completed...")
-
-    def complete_job(self):
-        """Publish a job state event that indicates that the job run has completed.
-        """
-        self.event_dispatcher.dispatch_event(ScriptCompletedEvent.completed)
-
-    def workflow_failure(
-            self,
-            error_message: Optional[str] = None,
-            error_trace: Optional[str] = None):
-        """Publish a job state event that indicates failure to compact a list of partitions.
-        """
-        failed_workflow_event = self._event_map[COMPACTION_METADATA_SESSION_WORKFLOW_FAILURE]
-        self.event_dispatcher.dispatch_event(failed_workflow_event,
-                                             event_data={
-                                                 "eventName": failed_workflow_event.name,
-                                                 "errorMessage": error_message,
-                                                 "errorStackTrace": error_trace,
-                                                 "stateDetailDescription":
-                                                     f"Workflow encountered a failure.",
-                                                 "stateDetailStatus": "FAILED",
-                                             })
-
-    @property
-    def session_manager(self):
-        return self.event_dispatcher.session_manager
-
-    @property
-    def trace_id(self):
-        return self.event_dispatcher.events_publisher.event_base_params.get("traceId", "UNKNOWN_TRACE_ID")

deltacat/autoscaler/events/dispatcher.py

@@ -1,95 +0,0 @@
-import logging
-from typing import Dict, Any, Optional, Callable, List
-
-from deltacat.autoscaler.events.session_manager import SessionManager, SESSION_ID_KEY, PARENT_SESSION_ID_KEY
-from deltacat.autoscaler.events.states import event_enum_values
-from ray.autoscaler._private.event_system import RayEvent, EventPublisher
-
-from deltacat import logs
-from deltacat.storage import interface as unimplemented_deltacat_storage
-
-logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
-
-
-class EventDispatcher:
-    def __init__(self,
-                 events_publisher: EventPublisher,
-                 deltacat_storage: unimplemented_deltacat_storage,
-                 session_manager: SessionManager = None):
-        """Constructor for the event dispatcher.
-
-        Intended for usage by Ray parent and child clusters running managed jobs.
-
-        Args:
-            events_publisher: Events manager for publishing events through a cloud provider
-            session_manager: Manager for tracking and launching Ray sessions
-            deltacat_storage: Storage interface for deltacat
-        """
-        self.events_publisher = events_publisher
-        self.deltacat_storage = deltacat_storage
-        self.session_manager = session_manager
-
-        # Setup event callbacks in the constructor
-        self._add_base_event_handlers()
-
-    def dispatch_event(self,
-                       event: RayEvent,
-                       event_data: Optional[Dict[str, Any]] = None):
-        """Generic helper method to dispatch Ray job events
-
-        Args:
-            event: Ray job event to dispatch
-            event_data: Additional metadata for the given event. Optional.
-
-        Returns:
-
-        """
-        if event_data is None:
-            event_data = {}
-
-        event_data["event_name"] = event
-        if self.session_manager:
-            event_data.setdefault(PARENT_SESSION_ID_KEY, self.session_manager.session_id)
-            event_data.setdefault(SESSION_ID_KEY, self.session_manager.session_id)
-
-            logger.info(f"Dispatching event {event.name} "
-                        f"with parent Ray session ID = {event_data[PARENT_SESSION_ID_KEY]} "
-                        f"and current Ray session ID = {event_data[SESSION_ID_KEY]}")
-
-        event_payload = {
-            **self.events_publisher.config["parameters"],
-            **event_data
-        }
-
-        # Trim un-required, space intensive data from payload
-        if "statsMetadata" in event_payload:
-            event_payload.pop("statsMetadata")
-        if "partitionsToCompact" in event_payload:
-            event_payload.pop("partitionsToCompact")
-
-        self._publish_event(event_payload)
-
-    def _add_base_event_handlers(self):
-        """Add callback handlers for base job events
-        """
-        publisher = self.events_publisher
-        if publisher:
-            for event in event_enum_values:
-                logger.info(f"[{publisher.__class__.__name__}]: Adding callback for event {event.name}")
-                publisher.add_callback(event)
-
-    def add_event_handlers(self, custom_events: List[RayEvent]):
-        """Add callback handlers for custom job events
-        """
-        publisher = self.events_publisher
-        if publisher:
-            for event in custom_events:
-                logger.info(f"[{publisher.__class__.__name__}]: Adding callback for event {event.name}")
-                publisher.add_callback(event)
-
-    def _publish_event(self, event_data: Dict[str, Any]):
-        publisher = self.events_publisher
-        if publisher and event_data and event_data.get("event_name"):
-            event: RayEvent = event_data["event_name"]
-            logger.info(f"[{publisher.__class__.__name__}]: Publishing event {event.name}")
-            publisher.publish(event, event_data)