import logging import os import threading import time import typing from typing import Dict, List, Optional, Tuple from ray.data._internal.actor_autoscaler import ( create_actor_autoscaler, ) from ray.data._internal.cluster_autoscaler import create_cluster_autoscaler from ray.data._internal.execution import create_ranker from ray.data._internal.execution.backpressure_policy import ( BackpressurePolicy, get_backpressure_policies, ) from ray.data._internal.execution.block_ref_counter import BlockRefCounter from ray.data._internal.execution.dataset_state import DatasetState from ray.data._internal.execution.execution_callback import ExecutionCallback from ray.data._internal.execution.interfaces import ( Executor, OutputIterator, PhysicalOperator, RefBundle, ) from ray.data._internal.execution.metadata_fetcher import make_metadata_fetcher from ray.data._internal.execution.operators.base_physical_operator import ( InternalQueueOperatorMixin, ) from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer from ray.data._internal.execution.resource_manager import ( ResourceManager, ) from ray.data._internal.execution.streaming_executor_state import ( OpState, OutputBackpressureGuard, Topology, build_streaming_topology, format_op_state_summary, process_completed_tasks, select_operator_to_run, update_operator_states, ) from ray.data._internal.logging import ( get_log_directory, register_dataset_logger, unregister_dataset_logger, ) from ray.data._internal.metadata_exporter import ( Topology as TopologyMetadata, sanitize_for_struct, ) from ray.data._internal.operator_schema_exporter import ( OperatorSchema, get_operator_schema_exporter, ) from ray.data._internal.progress import get_progress_manager from ray.data._internal.stats import DatasetStats, Timer, _StatsManager from ray.data.context import OK_PREFIX, WARN_PREFIX, DataContext from ray.util.debug import log_once from ray.util.metrics import Gauge if typing.TYPE_CHECKING: from ray.data._internal.issue_detection.issue_detector_manager import ( IssueDetectorManager, ) from ray.data._internal.progress.base_progress import BaseExecutionProgressManager from ray.data.block import Schema logger = logging.getLogger(__name__) # Interval for logging execution progress updates and operator metrics. DEBUG_LOG_INTERVAL_SECONDS = 5 # Maximum string/sequence length for DataContext logging. Set high to avoid truncation # while still protecting against pathological cases. DATA_CONTEXT_LOG_TRUNCATE_LENGTH = 10000 # Visible for testing. _num_shutdown = 0 # Extra environment variables to log that don't start with RAY_DATA. _EXTRA_ENV_VARS_TO_LOG = ( # We historically recommended users configure this value. If a Ray Data job uses # more object store memory than expected, it's worth checking how this environment # variable has been configured. "RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION", ) def _log_ray_data_env_vars() -> None: env_vars = { k: v for k, v in os.environ.items() if k.startswith("RAY_DATA") or k in _EXTRA_ENV_VARS_TO_LOG } if env_vars: formatted = ", ".join(f"{k}={v}" for k, v in sorted(env_vars.items())) logger.debug(f"RAY_DATA environment variables: {formatted}") else: logger.debug("No RAY_DATA environment variables set.") class StreamingExecutor(Executor, threading.Thread): """A streaming Dataset executor. This implementation executes Dataset DAGs in a fully streamed way. It runs by setting up the operator topology, and then routing blocks through operators in a way that maximizes throughput under resource constraints. """ UPDATE_METRICS_INTERVAL_S: float = 5.0 def __init__( self, data_context: DataContext, dataset_id: str = "unknown_dataset", ): self._data_context = data_context self._ranker = create_ranker() self._start_time: Optional[float] = None self._initial_stats: Optional[DatasetStats] = None self._final_stats: Optional[DatasetStats] = None self._progress_manager: Optional["BaseExecutionProgressManager"] = None self._callbacks: List["ExecutionCallback"] = [] # The executor can be shutdown while still running. self._shutdown_lock = threading.RLock() self._execution_started = False self._shutdown = False # Internal execution state shared across thread boundaries. We run the control # loop on a separate thread so that it doesn't become stalled between # generator `yield`s. self._topology: Optional[Topology] = None self._output_node: Optional[Tuple[PhysicalOperator, OpState]] = None self._backpressure_policies: List[BackpressurePolicy] = [] self._op_schema: Dict[PhysicalOperator, Schema] = {} self._dataset_id = dataset_id # Set by IssueDetectionExecutionCallback when issue detection is registered; # otherwise remains None. Access via the issue_detector_manager property. self._issue_detector_manager: Optional["IssueDetectorManager"] = None # Stores if an operator is completed, # used for marking when an op has just completed. self._has_op_completed: Optional[Dict[PhysicalOperator, bool]] = None self._max_errored_blocks = self._data_context.max_errored_blocks self._num_errored_blocks = 0 self._last_debug_log_time = 0 self._data_context.set_dataset_logger_id( register_dataset_logger(self._dataset_id) ) # This stores the last time we updated the metrics. # This allows us to update metrics on some interval, # by comparing it with the current timestamp. self._metrics_last_updated: float = 0.0 self._sched_loop_duration_s = Gauge( "data_sched_loop_duration_s", description="Duration of the scheduling loop in seconds", tag_keys=("dataset",), ) # Resolves pulled (block_ref, meta_ref) pairs into emitted RefBundles. # The threaded fetcher (default) fetches metadata on a background thread # so the scheduling loop never blocks on ``ray.get(meta_refs)``; the # inline fetcher reproduces the synchronous, master-identical path. self._metadata_fetcher = make_metadata_fetcher() Executor.__init__(self, self._data_context.execution_options) thread_name = f"StreamingExecutor-{self._dataset_id}" threading.Thread.__init__(self, daemon=True, name=thread_name) @property def issue_detector_manager(self) -> Optional["IssueDetectorManager"]: """The issue detector manager, or None if issue detection isn't registered.""" return self._issue_detector_manager def execute( self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None, callbacks: Optional[List] = None, ) -> OutputIterator: """Executes the DAG using a streaming execution strategy. We take an event-loop approach to scheduling. We block on the next scheduling event using `ray.wait`, updating operator state and dispatching new tasks. """ if callbacks is not None: self._callbacks = callbacks else: self._callbacks = [] self._initial_stats = initial_stats self._start_time = time.perf_counter() if logger.isEnabledFor(logging.DEBUG): _log_ray_data_env_vars() if not isinstance(dag, InputDataBuffer): if self._data_context.print_on_execution_start: message = f"Starting execution of Dataset {self._dataset_id}." log_path = get_log_directory() if log_path is not None: message += f" Full logs are in {log_path}" logger.info(message) logger.info( f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}" ) # Log the full DataContext for traceability if logger.isEnabledFor(logging.DEBUG) and log_once( f"ray_data_log_context_{self._dataset_id}" ): logger.debug( f"Data Context for dataset {self._dataset_id}:\n%s", sanitize_for_struct( self._data_context, truncate_length=DATA_CONTEXT_LOG_TRUNCATE_LENGTH, ), ) # Setup the streaming DAG topology and start the runner thread. self._block_ref_counter = BlockRefCounter() self._topology = build_streaming_topology( dag, self._options, self._block_ref_counter ) self._resource_manager = ResourceManager( self._topology, self._options, lambda: self._cluster_autoscaler.get_total_resources(), self._data_context, self._block_ref_counter, ) # Constructed once per executor (not per scheduling iteration) so the # guard's idle-detection state accumulates across scheduling iterations. self._output_backpressure_guard = OutputBackpressureGuard( self._topology, self._resource_manager ) # Setup progress manager self._progress_manager = get_progress_manager( self._data_context, self._dataset_id, self._topology, self._options.verbose_progress, ) self._progress_manager.start() self._backpressure_policies = get_backpressure_policies( self._data_context, self._topology, self._resource_manager ) self._cluster_autoscaler = create_cluster_autoscaler( self._topology, self._resource_manager, self._data_context, execution_id=self._dataset_id, ) self._actor_autoscaler = create_actor_autoscaler( self._topology, self._resource_manager, config=self._data_context.autoscaling_config, ) self._has_op_completed = dict.fromkeys(self._topology, False) self._output_node = dag, self._topology[dag] op_to_id = { op: self._get_operator_id(op, i) for i, op in enumerate(self._topology) } _StatsManager.register_dataset_to_stats_actor( self._dataset_id, self._get_operator_tags(), TopologyMetadata.create_topology_metadata(dag, op_to_id), self._data_context, ) for callback in self._callbacks: callback.before_execution_starts(self) self.start() self._execution_started = True return _ClosingIterator(self) def __del__(self): # NOTE: Upon garbage-collection we're allowing running tasks # to be terminated asynchronously (ie avoid unnecessary # synchronization on their completion) self.shutdown(force=False) def shutdown(self, force: bool, exception: Optional[Exception] = None): global _num_shutdown with self._shutdown_lock: if not self._execution_started or self._shutdown: return start = time.perf_counter() status_detail = ( f"failed with {exception}" if exception else "completed successfully" ) logger.debug( f"Shutting down executor for dataset {self._dataset_id} " f"({status_detail})" ) _num_shutdown += 1 self._shutdown = True # Give the scheduling loop some time to finish processing. self.join(timeout=2.0) # Stop the metadata fetcher (after the loop thread that feeds it has # been joined). No-op for the inline fetcher. self._metadata_fetcher.stop() self._update_stats_metrics( state=DatasetState.FINISHED.name if exception is None else DatasetState.FAILED.name, force_update=True, ) # Freeze the stats and save it. self._final_stats = self._generate_stats() stats_summary_string = self._final_stats.to_summary().to_string( include_parent=False ) # Reset the scheduling loop duration gauge + resource manager budgets/usages. self._resource_manager.update_usages() self.update_metrics(0) if self._data_context.enable_auto_log_stats: logger.info(stats_summary_string) # Close the progress manager with a finishing message. if exception is None: desc = ( f"{OK_PREFIX} Dataset {self._dataset_id} execution finished in " f"{self._final_stats.time_total_s:.2f} seconds" ) else: desc = f"{WARN_PREFIX} Dataset {self._dataset_id} execution failed" self._progress_manager.close_with_finishing_description( desc, exception is None ) logger.info(desc) timer = Timer() for op in self._topology.keys(): op.shutdown(timer, force=force) self._clear_topology_queues_post_shutdown(force, exception) # Queues have been drained; any remaining Ray Core callbacks that fire # after this point should be no-ops. self._block_ref_counter.clear() min_ = round(timer.min(), 3) max_ = round(timer.max(), 3) total = round(timer.get(), 3) logger.debug( f"Shut down operator hierarchy for dataset {self._dataset_id}" f" (min/max/total={min_}/{max_}/{total}s)" ) if exception is None: for callback in self._callbacks: callback.after_execution_succeeds(self) else: for callback in self._callbacks: callback.after_execution_fails(self, exception) self._cluster_autoscaler.on_executor_shutdown() dur = time.perf_counter() - start logger.debug( f"Shut down executor for dataset {self._dataset_id} " f"(took {round(dur, 3)}s)" ) # Unregister should be called after all operators are shut down to # capture as many logs as possible. self._data_context.set_dataset_logger_id( unregister_dataset_logger(self._dataset_id) ) def _clear_topology_queues_post_shutdown( self, force: bool, exception: Optional[Exception] = None ) -> None: """Drain topology queues after operator shutdown (releases block refs).""" for op, state in self._topology.items(): if isinstance(op, InternalQueueOperatorMixin): op.clear_internal_input_queue() op.clear_internal_output_queue() # Input queues alias upstream output queues; clears the DAG except the sink. for inqueue in state.input_queues: inqueue.clear() output_op, _ = self._output_node # Clear sink output unless cooperative multi-split success (splits may still read). is_live_multi_split_sink = ( output_op.num_output_splits() > 1 and not force and exception is None ) if not is_live_multi_split_sink: self._topology[output_op].output_queue.clear() def run(self): """Run the control loop in a helper thread. Results are returned via the output node's outqueue. """ exc: Optional[Exception] = None self._metadata_fetcher.start() try: # Run scheduling loop until complete. while True: # Use `perf_counter` rather than `process_time` to ensure we include # time spent on IO, like RPCs to Ray Core. t_start = time.perf_counter() continue_sched = self._scheduling_loop_step(self._topology) sched_loop_duration = time.perf_counter() - t_start self.update_metrics(sched_loop_duration) if self._initial_stats: self._initial_stats.streaming_exec_schedule_s.add( sched_loop_duration ) for callback in self._callbacks: callback.on_execution_step(self) if not continue_sched or self._shutdown: break except Exception as e: # Propagate it to the result iterator. exc = e finally: # Mark state of outputting operator as finished _, state = self._output_node state.mark_finished(exc) def update_metrics(self, sched_loop_duration: int): self._sched_loop_duration_s.set( sched_loop_duration, tags={"dataset": self._dataset_id} ) def get_stats(self): """Return the stats object for the streaming execution. The stats object will be updated as streaming execution progresses. """ if self._final_stats: return self._final_stats else: return self._generate_stats() def set_external_consumer_bytes(self, num_bytes: int) -> None: """Set the bytes buffered by external consumers.""" if self._resource_manager is not None: self._resource_manager.set_external_consumer_bytes(num_bytes) def _generate_stats(self) -> DatasetStats: """Create a new stats object reflecting execution status so far.""" stats = self._initial_stats or DatasetStats(metadata={}, parent=None) for op in self._topology: if isinstance(op, InputDataBuffer): continue builder = stats.child_builder(op.name, override_start_time=self._start_time) stats = builder.build_multioperator(op.get_stats()) stats.extra_metrics = op.metrics.as_dict(skip_internal_metrics=True) # Always assign a ``Timer`` so downstream consumers can call # ``.get()`` / ``.avg()`` / ``.max()`` / ``.percentile()`` # unconditionally. When ``_initial_stats`` is absent we hand # back an empty Timer; zero-sample semantics yield 0 across all # four. stats.streaming_exec_schedule_s = ( self._initial_stats.streaming_exec_schedule_s if self._initial_stats else Timer() ) return stats def _scheduling_loop_step(self, topology: Topology) -> bool: """Run one step of the scheduling loop. This runs a few general phases: 1. Waiting for the next task completion using `ray.wait()`. 2. Pulling completed refs into operator outqueues. 3. Selecting and dispatching new inputs to operators. Args: topology: The :class:`Topology` of operators being executed. Returns: True if we should continue running the scheduling loop. """ self._resource_manager.update_usages() # Note: calling process_completed_tasks() is expensive since it incurs # ray.wait() overhead, so make sure to allow multiple dispatch per call for # greater parallelism. num_errored_blocks = process_completed_tasks( topology, self._backpressure_policies, self._max_errored_blocks, output_backpressure_guard=self._output_backpressure_guard, metadata_fetcher=self._metadata_fetcher, ) if self._max_errored_blocks > 0: self._max_errored_blocks -= num_errored_blocks self._num_errored_blocks += num_errored_blocks self._resource_manager.update_usages() # Dispatch as many operators as we can for completed tasks. self._report_current_usage() i = 0 while True: op = select_operator_to_run( topology, self._resource_manager, self._backpressure_policies, # If consumer is idling (there's nothing for it to consume) # enforce liveness, ie that at least a single task gets scheduled ensure_liveness=self._consumer_idling(), ranker=self._ranker, ) if op is None: break topology[op].dispatch_next_task() self._resource_manager.update_usages() i += 1 if i % self._progress_manager.TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS == 0: self._refresh_progress_manager(topology) # Trigger autoscaling self._cluster_autoscaler.try_trigger_scaling() self._actor_autoscaler.try_trigger_scaling() update_operator_states(topology) self._refresh_progress_manager(topology) self._update_stats_metrics(state=DatasetState.RUNNING.name) if time.time() - self._last_debug_log_time >= DEBUG_LOG_INTERVAL_SECONDS: _log_op_metrics(topology) _debug_dump_topology(topology, self._resource_manager) self._last_debug_log_time = time.time() for op, state in topology.items(): # Export operator schema if it's updated if state._schema is not None and self._op_schema.get(op) != state._schema: self._op_schema[op] = state._schema self._export_operator_schema(op) # Log metrics of newly completed operators. if not op.has_completed(): op.refresh_state() elif not self._has_op_completed[op]: log_str = ( f"Operator {op} completed. " f"Operator Metrics:\n{op._metrics.as_dict(skip_internal_metrics=True)}" ) logger.debug(log_str) self._has_op_completed[op] = True self._validate_operator_queues_empty(op, state) # Keep going until all operators run to completion. return not all(op.has_completed() for op in topology) def _refresh_progress_manager(self, topology: Topology): # Update the progress manager to reflect scheduling decisions. if self._progress_manager: for op_state in topology.values(): if not isinstance(op_state.op, InputDataBuffer): self._progress_manager.update_operator_progress( op_state, self._resource_manager ) self._progress_manager.refresh() def _consumer_idling(self) -> bool: """Returns whether the user thread is blocked on topology execution.""" _, state = self._output_node return len(state.output_queue) == 0 def _export_operator_schema(self, op: PhysicalOperator) -> None: schema = self._op_schema.get(op) operator_schema_exporter = get_operator_schema_exporter() if ( operator_schema_exporter is not None and hasattr(schema, "names") and hasattr(schema, "types") ): names = [str(n) for n in schema.names] types = [str(t) for t in schema.types] operator_schema = OperatorSchema( operator_uuid=op.id, schema_fields=dict(zip(names, types)), ) operator_schema_exporter.export_operator_schema(operator_schema) def _validate_operator_queues_empty( self, op: PhysicalOperator, state: OpState ) -> None: """Validate that all queues are empty when an operator completes. Args: op: The completed operator to validate. state: The operator's execution state. """ error_msg = "Expected {} Queue for {} to be empty, but found {} bundles" if isinstance(op, InternalQueueOperatorMixin): # 1) Check Internal Input Queue is empty assert op.internal_input_queue_num_blocks() == 0, error_msg.format( "Internal Input", op.name, op.internal_input_queue_num_blocks() ) # 2) Check Internal Output Queue is empty assert op.internal_output_queue_num_blocks() == 0, error_msg.format( "Internal Output", op.name, op.internal_output_queue_num_blocks(), ) # 3) Check that External Input Queue is empty for input_q in state.input_queues: assert len(input_q) == 0, error_msg.format( "External Input", op.name, len(input_q) ) def _report_current_usage(self) -> None: # running_usage is the amount of resources that have been requested but # not necessarily available # TODO(sofian) https://github.com/ray-project/ray/issues/47520 # We need to split the reported resources into running, pending-scheduling, # pending-node-assignment. running_usage = self._resource_manager.get_global_running_usage() pending_usage = self._resource_manager.get_global_pending_usage() limits = self._resource_manager.get_global_limits() resources_status = ( f"Active & requested resources: " f"{running_usage.cpu:.4g}/{limits.cpu:.4g} CPU, " ) if running_usage.memory > 0: resources_status += ( f"{running_usage.memory_str()}/{limits.memory_str()} memory, " ) if running_usage.gpu > 0: resources_status += f"{running_usage.gpu:.4g}/{limits.gpu:.4g} GPU, " resources_status += ( f"{running_usage.object_store_memory_str()}/" f"{limits.object_store_memory_str()} object store" ) # Only include pending section when there are pending resources. pending_parts = [] if pending_usage.cpu: pending_parts.append(f"{pending_usage.cpu:.4g} CPU") if pending_usage.memory: pending_parts.append(f"{pending_usage.memory_str()} memory") if pending_usage.gpu: pending_parts.append(f"{pending_usage.gpu:.4g} GPU") if pending_parts: resources_status += f" (pending: {', '.join(pending_parts)})" self._progress_manager.update_total_resource_status(resources_status) def _get_operator_id(self, op: PhysicalOperator, topology_index: int) -> str: return f"{op.name}_{topology_index}" def _get_operator_tags(self): """Returns a list of operator tags.""" return [ f"{self._get_operator_id(op, i)}" for i, op in enumerate(self._topology) ] def _get_state_dict(self, state): last_op, last_state = list(self._topology.items())[-1] return { "state": state, "progress": last_state.num_completed_tasks, "total": last_op.num_outputs_total(), "total_rows": last_op.num_output_rows_total(), "end_time": time.time() if state in (DatasetState.FINISHED.name, DatasetState.FAILED.name) else None, "operators": { f"{self._get_operator_id(op, i)}": { "name": op.name, "progress": op_state.num_completed_tasks, "total": op.num_outputs_total(), "total_rows": op.num_output_rows_total(), "queued_blocks": op_state.total_enqueued_input_blocks(), "state": DatasetState.FINISHED.name if op.has_execution_finished() else state, } for i, (op, op_state) in enumerate(self._topology.items()) }, } def _update_stats_metrics(self, state: str, force_update: bool = False): now = time.time() if ( force_update or (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S ): _StatsManager.update_execution_metrics( self._dataset_id, [op.metrics for op in self._topology], self._get_operator_tags(), self._get_state_dict(state=state), ) self._metrics_last_updated = now def _debug_dump_topology(topology: Topology, resource_manager: ResourceManager) -> None: """Log current execution state for the topology for debugging. Args: topology: The topology to debug. resource_manager: The resource manager for this topology. """ logger.debug("Execution Progress:") for i, (op, state) in enumerate(topology.items()): summary_str = format_op_state_summary(state, resource_manager, verbose=True) logger.debug( f"{i}: {op.name} - {summary_str}, " f"Blocks Outputted: {state.num_completed_tasks}/{op.num_outputs_total()}" ) def _log_op_metrics(topology: Topology) -> None: """Logs the metrics of each operator. Args: topology: The topology to debug. """ log_str = "Operator Metrics:\n" for op in topology: metrics_dict = op.metrics.as_dict(skip_internal_metrics=True) log_str += f"{op.name}: {metrics_dict}\n" logger.debug(log_str) class _ClosingIterator(OutputIterator): """Iterator automatically shutting down executor upon exhausting the iterable sequence. NOTE: If this iterator isn't fully exhausted, executor still have to be closed manually by the caller! """ def __init__(self, executor: StreamingExecutor): self._executor = executor def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle: try: op, state = self._executor._output_node bundle = state.get_output_blocking(output_split_idx) # Update progress-bars if self._executor._progress_manager: self._executor._progress_manager.update_total_progress( bundle.num_rows() or 0, op.num_output_rows_total() ) return bundle # Have to be BaseException to catch ``KeyboardInterrupt`` # # NOTE: This also handles ``StopIteration`` except BaseException as e: # Asynchronously shutdown the executor (ie avoid unnecessary # synchronization on tasks termination) self._executor.shutdown( force=False, exception=e if not isinstance(e, StopIteration) else None ) raise def __del__(self): # NOTE: Upon garbage-collection we're allowing running tasks # to be terminated asynchronously (ie avoid unnecessary # synchronization on their completion) self._executor.shutdown(force=False)