ray-project--ray/python/ray/data/_internal/execution/streaming_executor.py

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)