ray-project--ray/python/ray/data/_internal/progress/logging_progress.py

import dataclasses
import logging
import time
import typing
from collections import defaultdict
from typing import Callable, Dict, List, Optional

from ray._common.utils import env_integer
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
from ray.data._internal.execution.streaming_executor_state import (
    format_op_state_summary,
)
from ray.data._internal.progress.base_progress import (
    BaseExecutionProgressManager,
    BaseProgressBar,
    NoopSubProgressBar,
)
from ray.data._internal.progress.utils import truncate_operator_name

if typing.TYPE_CHECKING:
    from ray.data._internal.execution.resource_manager import ResourceManager
    from ray.data._internal.execution.streaming_executor_state import OpState, Topology


logger = logging.getLogger(__name__)


@dataclasses.dataclass
class _LoggingMetrics:
    name: str
    desc: Optional[str]
    completed: int
    total: Optional[int]


class LoggingSubProgressBar(BaseProgressBar):
    """Thin wrapper to provide identical interface to the ProgressBar.

    Internally passes relevant logging metrics to `LoggingExecutionProgressManager`.
    Sub-progress is actually handled by Ray through operators, while operator-level
    and total progress is handled by the `StreamingExecutor`. To ensure log-order,
    this class helps to pass metric data to the progress manager so progress metrics
    are logged centrally.
    """

    def __init__(
        self,
        name: str,
        total: Optional[int] = None,
        max_name_length: int = 100,
    ):
        """Initialize sub-progress bar

        Args:
            name: name of sub-progress bar
            total: total number of output rows. None for unknown.
            max_name_length: maximum operator name length (unused).
        """
        del max_name_length  # unused
        self._total = total
        self._completed = 0
        self._name = name

    def set_description(self, name: str) -> None:
        pass  # unused

    def get_description(self) -> str:
        return ""  # unused

    def update(self, increment: int = 0, total: Optional[int] = None):
        if total is not None:
            self._total = total
        self._completed += increment

    def get_logging_metrics(self) -> _LoggingMetrics:
        return _LoggingMetrics(
            name=f"    - {self._name}",
            desc=None,
            completed=self._completed,
            total=self._total,
        )


class LoggingExecutionProgressManager(BaseExecutionProgressManager):
    """Execution progress display for non-tty situations, preventing
    spamming of progress reporting."""

    # Refer to following issues for more context about this feature:
    # https://github.com/ray-project/ray/issues/60083
    # https://github.com/ray-project/ray/issues/57734

    # This progress manager needs to refresh (log) based on elapsed time
    # not scheduling steps. This elapsed time handling is done within
    # this class.
    TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS = 1

    # Time interval (seconds) in which progress is logged to console again.
    LOG_REPORT_INTERVAL_SEC = env_integer("RAY_DATA_NON_TTY_PROGRESS_LOG_INTERVAL", 10)

    def __init__(
        self,
        dataset_id: str,
        topology: "Topology",
        show_op_progress: bool,
        verbose_progress: bool,
        *,
        _get_time: Callable[[], float] = time.time,
    ):
        self._dataset_id = dataset_id
        self._topology = topology
        self._get_time = _get_time
        self._last_log_time = self._get_time() - self.LOG_REPORT_INTERVAL_SEC

        self._global_progress_metric = _LoggingMetrics(
            name="Total Progress", desc=None, completed=0, total=None
        )
        self._op_progress_metrics: Dict["OpState", _LoggingMetrics] = {}
        self._sub_progress_metrics: Dict[
            "OpState", List[LoggingSubProgressBar]
        ] = defaultdict(list)

        for state in self._topology.values():
            op = state.op
            if isinstance(op, InputDataBuffer):
                continue
            total = op.num_output_rows_total() or 1

            contains_sub_progress_bars = isinstance(op, SubProgressBarMixin)
            sub_progress_bar_enabled = show_op_progress and (
                contains_sub_progress_bars or verbose_progress
            )

            if sub_progress_bar_enabled:
                self._op_progress_metrics[state] = _LoggingMetrics(
                    name=truncate_operator_name(op.name, self.MAX_NAME_LENGTH),
                    desc=None,
                    completed=0,
                    total=total,
                )

            if not contains_sub_progress_bars:
                continue

            sub_pg_names = op.get_sub_progress_bar_names()
            if sub_pg_names is None:
                continue
            for name in sub_pg_names:
                if sub_progress_bar_enabled:
                    pg = LoggingSubProgressBar(
                        name=name, total=total, max_name_length=self.MAX_NAME_LENGTH
                    )
                    self._sub_progress_metrics[state].append(pg)
                else:
                    pg = NoopSubProgressBar(
                        name=name, max_name_length=self.MAX_NAME_LENGTH
                    )
                op.set_sub_progress_bar(name, pg)

    # Management
    def start(self):
        # logging progress manager doesn't need separate start
        pass

    def refresh(self):
        current_time = self._get_time()
        if current_time - self._last_log_time < self.LOG_REPORT_INTERVAL_SEC:
            return
        self._last_log_time = current_time

        # starting delimiter
        firstline = f"======= Running Dataset: {self._dataset_id} ======="
        lastline = "=" * len(firstline)
        logger.info(firstline)

        # log global progress
        _log_global_progress(self._global_progress_metric)

        # log operator-level progress
        if len(self._op_progress_metrics.keys()) > 0:
            logger.info("")

        for opstate in self._topology.values():
            metrics = self._op_progress_metrics.get(opstate)
            if metrics is None:
                continue
            _log_op_or_sub_progress(metrics)
            for pg in self._sub_progress_metrics[opstate]:
                _log_op_or_sub_progress(pg.get_logging_metrics())

        # finish logging
        logger.info(lastline)

    def close_with_finishing_description(self, desc: str, success: bool):
        # We log in StreamingExecutor. No need for duplicate logging.
        pass

    # Total Progress
    def update_total_progress(self, new_rows: int, total_rows: Optional[int]):
        if total_rows is not None:
            self._global_progress_metric.total = total_rows
        self._global_progress_metric.completed += new_rows

    def update_total_resource_status(self, resource_status: str):
        self._global_progress_metric.desc = resource_status

    # Operator Progress
    def update_operator_progress(
        self, opstate: "OpState", resource_manager: "ResourceManager"
    ):
        op_metrics = self._op_progress_metrics.get(opstate)
        if op_metrics is not None:
            op_metrics.completed = opstate.op.metrics.row_outputs_taken
            total = opstate.op.num_output_rows_total()
            if total is not None:
                op_metrics.total = total
            op_metrics.desc = format_op_state_summary(opstate, resource_manager)


def _format_progress(m: _LoggingMetrics) -> str:
    return f"{m.name}: {m.completed}/{m.total or '?'}"


def _log_global_progress(m: _LoggingMetrics):
    logger.info(_format_progress(m))
    if m.desc is not None:
        logger.info(m.desc)


def _log_op_or_sub_progress(m: _LoggingMetrics):
    logger.info(_format_progress(m))
    if m.desc is not None:
        logger.info(f"  {m.desc}")