chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,97 @@
|
||||
import logging
|
||||
import sys
|
||||
import typing
|
||||
|
||||
import ray
|
||||
from ray.util.debug import log_once
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ray.data._internal.execution.streaming_executor_state import Topology
|
||||
from ray.data._internal.progress.base_progress import BaseExecutionProgressManager
|
||||
from ray.data.context import DataContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_progress_manager(
|
||||
ctx: "DataContext", dataset_id: str, topology: "Topology", verbose_progress: bool
|
||||
) -> "BaseExecutionProgressManager":
|
||||
"""Obtain the appropriate progress manager for the given DataContext."""
|
||||
show_op_progress = ctx.enable_operator_progress_bars
|
||||
|
||||
if not ctx.enable_progress_bars:
|
||||
from ray.data._internal.progress.base_progress import (
|
||||
NoopExecutionProgressManager,
|
||||
)
|
||||
|
||||
if log_once("ray_data_progress_manager_disabled"):
|
||||
logger.warning(
|
||||
"Progress bars disabled. To enable, set "
|
||||
"`ray.data.DataContext.get_current()."
|
||||
"enable_progress_bars = True`."
|
||||
)
|
||||
return NoopExecutionProgressManager(
|
||||
dataset_id, topology, show_op_progress, verbose_progress
|
||||
)
|
||||
|
||||
if not show_op_progress:
|
||||
if log_once("ray_data_progress_manager_global"):
|
||||
logger.warning(
|
||||
"Progress bars for operators disabled. To enable, "
|
||||
"set `ray.data.DataContext.get_current()."
|
||||
"enable_operator_progress_bars = True`."
|
||||
)
|
||||
|
||||
rich_enabled = ctx.enable_rich_progress_bars
|
||||
use_ray_tqdm = ctx.use_ray_tqdm
|
||||
worker = ray._private.worker
|
||||
in_ray_worker = worker.global_worker.mode == worker.WORKER_MODE
|
||||
|
||||
if not sys.stdout.isatty() and not (use_ray_tqdm and in_ray_worker):
|
||||
from ray.data._internal.progress.logging_progress import (
|
||||
LoggingExecutionProgressManager,
|
||||
)
|
||||
|
||||
if log_once("ray_data_logging_progress_activated"):
|
||||
logger.info(
|
||||
"Progress will be logged because stdout is a non-interactive terminal."
|
||||
)
|
||||
return LoggingExecutionProgressManager(
|
||||
dataset_id, topology, show_op_progress, verbose_progress
|
||||
)
|
||||
|
||||
if not rich_enabled or use_ray_tqdm:
|
||||
from ray.data._internal.progress.tqdm_progress import (
|
||||
TqdmExecutionProgressManager,
|
||||
)
|
||||
|
||||
if log_once("ray_data_rich_progress_disabled"):
|
||||
logger.info(
|
||||
"[dataset]: A new progress UI is available. To enable, "
|
||||
"set `ray.data.DataContext.get_current()."
|
||||
"enable_rich_progress_bars = True` and `ray.data."
|
||||
"DataContext.get_current().use_ray_tqdm = False`."
|
||||
)
|
||||
return TqdmExecutionProgressManager(
|
||||
dataset_id, topology, show_op_progress, verbose_progress
|
||||
)
|
||||
else:
|
||||
try:
|
||||
from ray.data._internal.progress.rich_progress import (
|
||||
RichExecutionProgressManager,
|
||||
)
|
||||
|
||||
return RichExecutionProgressManager(
|
||||
dataset_id, topology, show_op_progress, verbose_progress
|
||||
)
|
||||
except ImportError:
|
||||
from ray.data._internal.progress.base_progress import (
|
||||
NoopExecutionProgressManager,
|
||||
)
|
||||
|
||||
logger.warning(
|
||||
"[dataset]: Run `pip install rich` to enable progress reporting."
|
||||
)
|
||||
return NoopExecutionProgressManager(
|
||||
dataset_id, topology, show_op_progress, verbose_progress
|
||||
)
|
||||
@@ -0,0 +1,258 @@
|
||||
import logging
|
||||
import threading
|
||||
import typing
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
|
||||
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
|
||||
from ray.types import ObjectRef
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Used a signal to cancel execution.
|
||||
_canceled_threads = set()
|
||||
_canceled_threads_lock = threading.Lock()
|
||||
|
||||
|
||||
def _extract_num_rows(result: Any) -> int:
|
||||
"""Extract the number of rows from a result object.
|
||||
|
||||
Args:
|
||||
result: The result object from which to extract the number of rows.
|
||||
|
||||
Returns:
|
||||
The number of rows, defaulting to 1 if it cannot be determined.
|
||||
"""
|
||||
if hasattr(result, "num_rows"):
|
||||
return result.num_rows
|
||||
elif hasattr(result, "__len__"):
|
||||
# For output is DataFrame,i.e. sort_sample
|
||||
return len(result)
|
||||
else:
|
||||
return 1
|
||||
|
||||
|
||||
class BaseProgressBar(ABC):
|
||||
"""Base class to define a progress bar."""
|
||||
|
||||
def block_until_complete(self, remaining: List["ObjectRef"]) -> None:
|
||||
t = threading.current_thread()
|
||||
while remaining:
|
||||
done, remaining = ray.wait(
|
||||
remaining, num_returns=len(remaining), fetch_local=False, timeout=0.1
|
||||
)
|
||||
total_rows_processed = 0
|
||||
for _, result in zip(done, ray.get(done)):
|
||||
num_rows = _extract_num_rows(result)
|
||||
total_rows_processed += num_rows
|
||||
self.update(total_rows_processed)
|
||||
|
||||
with _canceled_threads_lock:
|
||||
if t in _canceled_threads:
|
||||
break
|
||||
|
||||
def fetch_until_complete(self, refs: List["ObjectRef"]) -> List[Any]:
|
||||
ref_to_result = {}
|
||||
remaining = refs
|
||||
t = threading.current_thread()
|
||||
# Triggering fetch_local redundantly for the same object is slower.
|
||||
# We only need to trigger the fetch_local once for each object,
|
||||
# raylet will persist these fetch requests even after ray.wait returns.
|
||||
# See https://github.com/ray-project/ray/issues/30375.
|
||||
fetch_local = True
|
||||
while remaining:
|
||||
done, remaining = ray.wait(
|
||||
remaining,
|
||||
num_returns=len(remaining),
|
||||
fetch_local=fetch_local,
|
||||
timeout=0.1,
|
||||
)
|
||||
if fetch_local:
|
||||
fetch_local = False
|
||||
total_rows_processed = 0
|
||||
for ref, result in zip(done, ray.get(done)):
|
||||
ref_to_result[ref] = result
|
||||
num_rows = _extract_num_rows(result)
|
||||
total_rows_processed += num_rows
|
||||
self.update(total_rows_processed)
|
||||
|
||||
with _canceled_threads_lock:
|
||||
if t in _canceled_threads:
|
||||
break
|
||||
|
||||
return [ref_to_result[ref] for ref in refs]
|
||||
|
||||
@abstractmethod
|
||||
def set_description(self, name: str) -> None:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_description(self) -> str:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
||||
...
|
||||
|
||||
def refresh(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class BaseExecutionProgressManager(ABC):
|
||||
"""Base Data Execution Progress Display Manager"""
|
||||
|
||||
# If the name/description of the progress bar exceeds this length,
|
||||
# it will be truncated.
|
||||
MAX_NAME_LENGTH = 100
|
||||
|
||||
# Total progress refresh rate (update interval in scheduling step)
|
||||
# refer to `streaming_executor.py::StreamingExecutor::_scheduling_loop_step`
|
||||
TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS = 50
|
||||
|
||||
@abstractmethod
|
||||
def __init__(
|
||||
self,
|
||||
dataset_id: str,
|
||||
topology: "Topology",
|
||||
show_op_progress: bool,
|
||||
verbose_progress: bool,
|
||||
):
|
||||
"""Initialize the progress manager, create all necessary progress bars
|
||||
and sub-progress bars for the given topology. Sub-progress bars are
|
||||
created for operators that implement the SubProgressBarMixin.
|
||||
|
||||
Args:
|
||||
dataset_id: id of Dataset
|
||||
topology: operation topology built via `build_streaming_topology`
|
||||
show_op_progress: whether to show individual operator progress
|
||||
(only for non-AllToAll by default).
|
||||
verbose_progress: whether to show individual operator progress for
|
||||
non-AllToAll operators as well.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def start(self) -> None:
|
||||
"""Start the progress manager."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def refresh(self) -> None:
|
||||
"""Refresh displayed progress."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def close_with_finishing_description(self, desc: str, success: bool) -> None:
|
||||
"""Close the progress manager with a finishing message.
|
||||
|
||||
Args:
|
||||
desc: description to display
|
||||
success: whether the dataset execution was successful
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update_total_progress(self, new_rows: int, total_rows: Optional[int]) -> None:
|
||||
"""Update the total progress rows.
|
||||
|
||||
Args:
|
||||
new_rows: new rows processed by the streaming_executor
|
||||
total_rows: total rows to be processed (if known)
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update_total_resource_status(self, resource_status: str) -> None:
|
||||
"""Update the total resource usage statistics.
|
||||
|
||||
Args:
|
||||
resource_status: resource status information string.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update_operator_progress(
|
||||
self, opstate: "OpState", resource_manager: "ResourceManager"
|
||||
) -> None:
|
||||
"""Update individual operator progress.
|
||||
|
||||
Args:
|
||||
opstate: opstate of the operator.
|
||||
resource_manager: the ResourceManager.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class NoopSubProgressBar(BaseProgressBar):
|
||||
"""Sub-Progress Bar for Noop (Disabled) Progress Manager"""
|
||||
|
||||
def __init__(self, name: str, max_name_length: int):
|
||||
self._max_name_length = max_name_length
|
||||
self._desc = truncate_operator_name(name, self._max_name_length)
|
||||
|
||||
def set_description(self, name: str) -> None:
|
||||
self._desc = truncate_operator_name(name, self._max_name_length)
|
||||
|
||||
def get_description(self) -> str:
|
||||
return self._desc
|
||||
|
||||
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
||||
pass
|
||||
|
||||
def refresh(self):
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class NoopExecutionProgressManager(BaseExecutionProgressManager):
|
||||
"""Noop Data Execution Progress Display Manager (Progress Display Disabled)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_id: str,
|
||||
topology: "Topology",
|
||||
show_op_progress: bool,
|
||||
verbose_progress: bool,
|
||||
):
|
||||
for state in topology.values():
|
||||
op = state.op
|
||||
if not isinstance(op, SubProgressBarMixin):
|
||||
continue
|
||||
sub_pg_names = op.get_sub_progress_bar_names()
|
||||
if sub_pg_names is not None:
|
||||
for name in sub_pg_names:
|
||||
pg = NoopSubProgressBar(
|
||||
name=name, max_name_length=self.MAX_NAME_LENGTH
|
||||
)
|
||||
op.set_sub_progress_bar(name, pg)
|
||||
|
||||
def start(self) -> None:
|
||||
pass
|
||||
|
||||
def refresh(self) -> None:
|
||||
pass
|
||||
|
||||
def close_with_finishing_description(self, desc: str, success: bool) -> None:
|
||||
pass
|
||||
|
||||
def update_total_progress(self, new_rows: int, total_rows: Optional[int]) -> None:
|
||||
pass
|
||||
|
||||
def update_total_resource_status(self, resource_status: str) -> None:
|
||||
pass
|
||||
|
||||
def update_operator_progress(
|
||||
self, opstate: "OpState", resource_manager: "ResourceManager"
|
||||
) -> None:
|
||||
pass
|
||||
@@ -0,0 +1,233 @@
|
||||
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}")
|
||||
@@ -0,0 +1,110 @@
|
||||
from typing import Optional
|
||||
|
||||
from ray.data._internal.progress.base_progress import BaseProgressBar
|
||||
from ray.data._internal.progress.utils import truncate_operator_name
|
||||
from ray.experimental import tqdm_ray
|
||||
|
||||
try:
|
||||
import tqdm
|
||||
|
||||
needs_warning = False
|
||||
except ImportError:
|
||||
tqdm = None
|
||||
needs_warning = True
|
||||
|
||||
|
||||
class ProgressBar(BaseProgressBar):
|
||||
"""Thin wrapper around tqdm to handle soft imports.
|
||||
|
||||
If `total` is `None` known (for example, it is unknown
|
||||
because no tasks have finished yet), doesn't display the full
|
||||
progress bar. Still displays basic progress stats from tqdm."""
|
||||
|
||||
# If the name/description of the progress bar exceeds this length,
|
||||
# it will be truncated.
|
||||
MAX_NAME_LENGTH = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
total: Optional[int],
|
||||
unit: str,
|
||||
position: int = 0,
|
||||
enabled: Optional[bool] = None,
|
||||
):
|
||||
from ray.data.context import DataContext
|
||||
|
||||
self._desc = truncate_operator_name(name, self.MAX_NAME_LENGTH)
|
||||
self._progress = 0
|
||||
# Prepend a space to the unit for better formatting.
|
||||
if unit[0] != " ":
|
||||
unit = " " + unit
|
||||
|
||||
if enabled is None:
|
||||
# When enabled is None (not explicitly set by the user),
|
||||
# check DataContext setting
|
||||
enabled = DataContext.get_current().enable_progress_bars
|
||||
|
||||
use_ray_tqdm = DataContext.get_current().use_ray_tqdm
|
||||
|
||||
if not enabled:
|
||||
self._bar = None
|
||||
elif use_ray_tqdm:
|
||||
self._bar = tqdm_ray.tqdm(total=total, unit=unit, position=position)
|
||||
self._bar.set_description(self._desc)
|
||||
elif tqdm:
|
||||
self._bar = tqdm.tqdm(
|
||||
total=total or 0,
|
||||
position=position,
|
||||
dynamic_ncols=True,
|
||||
unit=unit,
|
||||
unit_scale=True,
|
||||
)
|
||||
self._bar.set_description(self._desc)
|
||||
else:
|
||||
global needs_warning
|
||||
if needs_warning:
|
||||
print("[dataset]: Run `pip install tqdm` to enable progress reporting.")
|
||||
needs_warning = False
|
||||
self._bar = None
|
||||
|
||||
def set_description(self, name: str) -> None:
|
||||
name = truncate_operator_name(name, self.MAX_NAME_LENGTH)
|
||||
if self._bar and name != self._desc:
|
||||
self._desc = name
|
||||
self._bar.set_description(self._desc)
|
||||
|
||||
def get_description(self) -> str:
|
||||
return self._desc
|
||||
|
||||
def refresh(self):
|
||||
if self._bar:
|
||||
self._bar.refresh()
|
||||
|
||||
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
||||
if self._bar and (increment != 0 or self._bar.total != total):
|
||||
self._progress += increment
|
||||
if total is not None:
|
||||
self._bar.total = total
|
||||
if self._bar.total is not None and self._progress > self._bar.total:
|
||||
# If the progress goes over 100%, update the total.
|
||||
self._bar.total = self._progress
|
||||
self._bar.update(increment)
|
||||
|
||||
def close(self):
|
||||
if self._bar:
|
||||
if self._bar.total is not None and self._progress != self._bar.total:
|
||||
# If the progress is not complete, update the total.
|
||||
self._bar.total = self._progress
|
||||
self._bar.refresh()
|
||||
self._bar.close()
|
||||
self._bar = None
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def __getstate__(self):
|
||||
return {}
|
||||
|
||||
def __setstate__(self, state):
|
||||
self._bar = None # Progress bar is disabled on remote nodes.
|
||||
@@ -0,0 +1,427 @@
|
||||
import dataclasses
|
||||
import logging
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
import typing
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from rich.console import Console
|
||||
from rich.live import Live
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TaskID,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
from rich.table import Column, Table
|
||||
from rich.text import Text
|
||||
|
||||
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__)
|
||||
|
||||
_TREE_BRANCH = " ├─"
|
||||
_TREE_VERTICAL = "│"
|
||||
_TREE_VERTICAL_SUB_PROGRESS = " │ -"
|
||||
_TREE_VERTICAL_INDENT = f" {_TREE_VERTICAL} "
|
||||
_TOTAL_PROGRESS_TOTAL = 1.0
|
||||
_RESOURCE_REPORT_HEADER = f" {_TREE_VERTICAL} Active/total resources: "
|
||||
|
||||
|
||||
class RichSubProgressBar(BaseProgressBar):
|
||||
"""Thin wrapper to provide identical interface to the ProgressBar.
|
||||
|
||||
Updates RichExecutionProgressManager internally.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
total: Optional[int] = None,
|
||||
progress: Progress = None,
|
||||
tid: TaskID = 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.
|
||||
progress: rich.Progress instance for the corresponding
|
||||
sub-progress bar.
|
||||
tid: rich.TaskId for the corresponding sub-progress bar task.
|
||||
max_name_length: maximum operator name length.
|
||||
"""
|
||||
self._total = total
|
||||
self._completed = 0
|
||||
self._start_time = None
|
||||
self._enabled = True
|
||||
self._progress = progress
|
||||
self._tid = tid
|
||||
self._max_name_length = max_name_length
|
||||
self._desc = truncate_operator_name(name, self._max_name_length)
|
||||
|
||||
def set_description(self, name: str) -> None:
|
||||
self._desc = truncate_operator_name(name, self._max_name_length)
|
||||
if self._enabled:
|
||||
self._progress.update(self._tid, description=self._desc)
|
||||
|
||||
def get_description(self) -> str:
|
||||
return self._desc
|
||||
|
||||
def _update(self, completed: int, total: Optional[int] = None) -> None:
|
||||
assert self._enabled
|
||||
if self._start_time is None:
|
||||
self._start_time = time.time()
|
||||
metrics = _get_progress_metrics(self._start_time, completed, total)
|
||||
self._progress.update(
|
||||
self._tid,
|
||||
completed=metrics.completed,
|
||||
total=metrics.total,
|
||||
rate_str=metrics.rate_str,
|
||||
count_str=metrics.count_str,
|
||||
)
|
||||
|
||||
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
||||
if self._enabled and increment != 0:
|
||||
if total is not None:
|
||||
self._total = total
|
||||
self._completed += increment
|
||||
self._update(self._completed, self._total)
|
||||
|
||||
def complete(self) -> None:
|
||||
if self._enabled:
|
||||
self._update(self._completed, self._completed)
|
||||
|
||||
def __getstate__(self):
|
||||
return {"max_name_length": self._max_name_length}
|
||||
|
||||
def __setstate__(self, state):
|
||||
self._enabled = False # Progress bar is disabled on remote nodes.
|
||||
self._max_name_length = state.get("max_name_length", 100)
|
||||
|
||||
|
||||
class RichExecutionProgressManager(BaseExecutionProgressManager):
|
||||
"""Execution progress display using rich."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_id: str,
|
||||
topology: "Topology",
|
||||
show_op_progress: bool,
|
||||
verbose_progress: bool,
|
||||
):
|
||||
self._dataset_id = dataset_id
|
||||
self._sub_progress_bars: List[BaseProgressBar] = []
|
||||
self._show_op_progress = show_op_progress
|
||||
self._verbose_progress = verbose_progress
|
||||
self._start_time: Optional[float] = None
|
||||
|
||||
# rich
|
||||
self._console = Console(file=sys.stderr)
|
||||
self._total = self._make_progress_bar(" ", "•", 15)
|
||||
self._current_rows = 0
|
||||
self._total_resources = Text(
|
||||
f"{_RESOURCE_REPORT_HEADER}Initializing...", no_wrap=True
|
||||
)
|
||||
|
||||
self._op_display: Dict[
|
||||
"OpState", Tuple[Optional[TaskID], Optional[Progress], Optional[Text]]
|
||||
] = {}
|
||||
|
||||
self._layout_table = Table.grid(padding=(0, 1, 0, 0), expand=True)
|
||||
self._layout_table.add_row(self._total)
|
||||
self._layout_table.add_row(self._total_resources)
|
||||
|
||||
self._setup_operator_progress(topology)
|
||||
|
||||
# empty new line to prevent "packed" feeling
|
||||
self._layout_table.add_row(Text())
|
||||
|
||||
# rich.Live is the auto-refreshing rich component display.
|
||||
# refreshing/closing is all done through rich.Live
|
||||
self._live = Live(
|
||||
self._layout_table,
|
||||
console=self._console,
|
||||
refresh_per_second=2,
|
||||
vertical_overflow="visible",
|
||||
)
|
||||
|
||||
self._total_task_id = self._total.add_task(
|
||||
f"Dataset {self._dataset_id} running:",
|
||||
total=_TOTAL_PROGRESS_TOTAL,
|
||||
rate_str="? rows/s",
|
||||
count_str="0/?",
|
||||
)
|
||||
|
||||
def _setup_operator_progress(self, topology: "Topology"):
|
||||
rows = []
|
||||
for state in topology.values():
|
||||
op = state.op
|
||||
if isinstance(op, InputDataBuffer):
|
||||
continue
|
||||
|
||||
contains_sub_progress_bars = isinstance(op, SubProgressBarMixin)
|
||||
sub_progress_bar_enabled = self._show_op_progress and (
|
||||
contains_sub_progress_bars or self._verbose_progress
|
||||
)
|
||||
|
||||
if sub_progress_bar_enabled:
|
||||
progress = self._make_progress_bar(_TREE_BRANCH, " ", 10)
|
||||
stats = Text(f"{_TREE_VERTICAL_INDENT}Initializing...", no_wrap=True)
|
||||
total = state.op.num_output_rows_total()
|
||||
name = truncate_operator_name(state.op.name, self.MAX_NAME_LENGTH)
|
||||
tid = progress.add_task(
|
||||
name,
|
||||
total=total if total is not None else 1,
|
||||
start=True,
|
||||
rate_str="? rows/s",
|
||||
count_str="0/?",
|
||||
)
|
||||
rows.append(progress)
|
||||
rows.append(stats)
|
||||
self._op_display[state] = (tid, progress, stats)
|
||||
|
||||
if not contains_sub_progress_bars:
|
||||
continue
|
||||
|
||||
sub_progress_bar_names = op.get_sub_progress_bar_names()
|
||||
if sub_progress_bar_names is None:
|
||||
continue
|
||||
|
||||
for name in sub_progress_bar_names:
|
||||
if sub_progress_bar_enabled:
|
||||
progress = self._make_progress_bar(
|
||||
_TREE_VERTICAL_SUB_PROGRESS, "", 10
|
||||
)
|
||||
total = state.op.num_output_rows_total()
|
||||
tid = progress.add_task(
|
||||
name,
|
||||
total=total if total is not None else 1,
|
||||
start=True,
|
||||
rate_str="? rows/s",
|
||||
count_str="0/?",
|
||||
)
|
||||
rows.append(progress)
|
||||
pg = RichSubProgressBar(
|
||||
name=name,
|
||||
total=total,
|
||||
progress=progress,
|
||||
tid=tid,
|
||||
max_name_length=self.MAX_NAME_LENGTH,
|
||||
)
|
||||
else:
|
||||
pg = NoopSubProgressBar(
|
||||
name=name, max_name_length=self.MAX_NAME_LENGTH
|
||||
)
|
||||
op.set_sub_progress_bar(name, pg)
|
||||
self._sub_progress_bars.append(pg)
|
||||
if rows:
|
||||
self._layout_table.add_row(Text(f" {_TREE_VERTICAL}", no_wrap=True))
|
||||
for row in rows:
|
||||
self._layout_table.add_row(row)
|
||||
|
||||
def _make_progress_bar(self, indent_str: str, spinner_finish: str, bar_width: int):
|
||||
return Progress(
|
||||
TextColumn(indent_str, table_column=Column(no_wrap=True)),
|
||||
SpinnerColumn(finished_text=spinner_finish),
|
||||
TextColumn(
|
||||
"{task.description} {task.percentage:>3.0f}%",
|
||||
table_column=Column(no_wrap=True),
|
||||
),
|
||||
BarColumn(bar_width=bar_width),
|
||||
TextColumn("{task.fields[count_str]}", table_column=Column(no_wrap=True)),
|
||||
TextColumn("["),
|
||||
TimeElapsedColumn(),
|
||||
TextColumn(","),
|
||||
TextColumn("{task.fields[rate_str]}", table_column=Column(no_wrap=True)),
|
||||
TextColumn("]"),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
expand=False,
|
||||
)
|
||||
|
||||
# Management
|
||||
def start(self):
|
||||
if not self._live.is_started:
|
||||
self._live.start()
|
||||
|
||||
def refresh(self):
|
||||
if self._live.is_started:
|
||||
self._live.refresh()
|
||||
|
||||
def close_with_finishing_description(self, desc: str, success: bool):
|
||||
if self._live.is_started:
|
||||
kwargs = {}
|
||||
if success:
|
||||
# set everything to completed
|
||||
kwargs["completed"] = 1.0
|
||||
kwargs["total"] = 1.0
|
||||
for pg in self._sub_progress_bars:
|
||||
if isinstance(pg, RichSubProgressBar):
|
||||
pg.complete()
|
||||
if self._start_time is None:
|
||||
self._start_time = time.time()
|
||||
for tid, progress, _ in self._op_display.values():
|
||||
completed = progress.tasks[tid].completed or 0
|
||||
metrics = _get_progress_metrics(
|
||||
self._start_time, completed, completed
|
||||
)
|
||||
_update_with_conditional_rate(progress, tid, metrics)
|
||||
self._total.update(self._total_task_id, description=desc, **kwargs)
|
||||
self.refresh()
|
||||
# need this sleep delay to ensure that changes are rendered to screen
|
||||
# before rich Live module is stopped.
|
||||
time.sleep(0.02)
|
||||
self._live.stop()
|
||||
|
||||
# Total Progress
|
||||
def _can_update_total(self) -> bool:
|
||||
return (
|
||||
self._total_task_id is not None
|
||||
and self._total_task_id in self._total.task_ids
|
||||
)
|
||||
|
||||
def update_total_progress(self, new_rows: int, total_rows: Optional[int]):
|
||||
if not self._can_update_total():
|
||||
return
|
||||
if self._live.is_started:
|
||||
if self._start_time is None:
|
||||
self._start_time = time.time()
|
||||
if new_rows is not None:
|
||||
self._current_rows += new_rows
|
||||
metrics = _get_progress_metrics(
|
||||
self._start_time, self._current_rows, total_rows
|
||||
)
|
||||
_update_with_conditional_rate(self._total, self._total_task_id, metrics)
|
||||
|
||||
def update_total_resource_status(self, resource_status: str):
|
||||
if not self._can_update_total():
|
||||
return
|
||||
if self._live.is_started:
|
||||
self._total_resources.plain = _RESOURCE_REPORT_HEADER + resource_status
|
||||
|
||||
def _can_update_operator(self, op_state: "OpState") -> bool:
|
||||
if op_state not in self._op_display:
|
||||
return False
|
||||
tid, progress, stats = self._op_display[op_state]
|
||||
if tid is None or not progress or not stats or tid not in progress.task_ids:
|
||||
return False
|
||||
return True
|
||||
|
||||
def update_operator_progress(
|
||||
self, op_state: "OpState", resource_manager: "ResourceManager"
|
||||
):
|
||||
if not self._can_update_operator(op_state):
|
||||
return
|
||||
if self._start_time is None:
|
||||
self._start_time = time.time()
|
||||
tid, progress, stats = self._op_display[op_state]
|
||||
|
||||
# progress
|
||||
current_rows = op_state.op.metrics.row_outputs_taken
|
||||
total_rows = op_state.op.num_output_rows_total()
|
||||
metrics = _get_progress_metrics(self._start_time, current_rows, total_rows)
|
||||
_update_with_conditional_rate(progress, tid, metrics)
|
||||
# stats
|
||||
stats_str = format_op_state_summary(op_state, resource_manager)
|
||||
stats.plain = f"{_TREE_VERTICAL_INDENT}{stats_str}"
|
||||
|
||||
|
||||
# utilities
|
||||
def _format_k(val: int) -> str:
|
||||
if val >= 1000:
|
||||
fval = val / 1000.0
|
||||
fval_str = f"{int(fval)}" if fval.is_integer() else f"{fval:.2f}"
|
||||
return fval_str + "k"
|
||||
return str(val)
|
||||
|
||||
|
||||
def _format_row_count(completed: int, total: Optional[int]) -> str:
|
||||
"""Formats row counts with k units."""
|
||||
cstr = _format_k(completed)
|
||||
if total is None or math.isinf(total):
|
||||
tstr = "?k" if cstr.endswith("k") else "?"
|
||||
else:
|
||||
tstr = _format_k(total)
|
||||
return f"{cstr}/{tstr}"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _ProgressMetrics:
|
||||
completed: int
|
||||
total: int
|
||||
rate_str: str
|
||||
count_str: str
|
||||
|
||||
|
||||
def _get_progress_metrics(
|
||||
start_time: float, completed_rows: int, total_rows: Optional[int]
|
||||
) -> _ProgressMetrics:
|
||||
"""
|
||||
Args:
|
||||
start_time: time when progress tracking started
|
||||
completed_rows: cumulative rows outputted
|
||||
total_rows: total rows expected (can be unknown)
|
||||
Returns:
|
||||
_ProgressMetrics instance containing the calculated data.
|
||||
"""
|
||||
# Note, when total is unknown, we default the progress bar to 0.
|
||||
# Will properly have estimates for rate and count strings though.
|
||||
total = 1 if total_rows is None or total_rows < 1 else total_rows
|
||||
completed = 0 if total_rows is None else completed_rows
|
||||
|
||||
if total_rows is None:
|
||||
rate_str = "? row/s"
|
||||
else:
|
||||
elapsed = time.time() - start_time
|
||||
rate_val = completed_rows / elapsed if elapsed > 1 else 0
|
||||
rate_unit = "row/s"
|
||||
if rate_val >= 1000:
|
||||
rate_val /= 1000
|
||||
rate_unit = "k row/s"
|
||||
rate_str = f"{rate_val:.2f} {rate_unit}"
|
||||
count_str = _format_row_count(completed_rows, total_rows)
|
||||
|
||||
return _ProgressMetrics(
|
||||
completed=completed, total=total, rate_str=rate_str, count_str=count_str
|
||||
)
|
||||
|
||||
|
||||
def _update_with_conditional_rate(
|
||||
progress: Progress, tid: TaskID, metrics: _ProgressMetrics
|
||||
):
|
||||
# not doing type checking because rich is imported conditionally.
|
||||
# progress: rich.Progress
|
||||
# tid: rich.TaskId
|
||||
# metrics: _ProgressMetrics
|
||||
task = progress.tasks[tid]
|
||||
kwargs = {
|
||||
"completed": metrics.completed,
|
||||
"total": metrics.total,
|
||||
"count_str": metrics.count_str,
|
||||
}
|
||||
if task.completed != metrics.completed:
|
||||
# update rate string only if there are new rows.
|
||||
# this allows updates to other metric data while
|
||||
# preserving the right rate notation.
|
||||
kwargs["rate_str"] = metrics.rate_str
|
||||
progress.update(tid, **kwargs)
|
||||
@@ -0,0 +1,162 @@
|
||||
import logging
|
||||
import typing
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
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.progress_bar import ProgressBar
|
||||
|
||||
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__)
|
||||
|
||||
|
||||
class TqdmSubProgressBar(ProgressBar):
|
||||
"""Thin wrapper to provide helper interface for TqdmExecutionProgressManager"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
total: Optional[int],
|
||||
unit: str,
|
||||
position: int = 0,
|
||||
enabled: Optional[bool] = None,
|
||||
max_name_length: int = 100,
|
||||
):
|
||||
# patch to make max_name_length configurable from ProgressManager.
|
||||
self.MAX_NAME_LENGTH = max_name_length
|
||||
super().__init__(name, total, unit, position, enabled)
|
||||
|
||||
def update_absolute(self, completed: int, total_rows: Optional[int] = None) -> None:
|
||||
if self._bar:
|
||||
self._progress = completed
|
||||
if total_rows is not None:
|
||||
self._bar.total = total_rows
|
||||
if self._bar.total is not None and self._progress > self._bar.total:
|
||||
# If the progress goes over 100%, update the total.
|
||||
self._bar.total = self._progress
|
||||
self._bar.n = self._progress
|
||||
|
||||
|
||||
class TqdmExecutionProgressManager(BaseExecutionProgressManager):
|
||||
"""Execution progress display using tqdm."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_id: str,
|
||||
topology: "Topology",
|
||||
show_op_progress: bool,
|
||||
verbose_progress: bool,
|
||||
):
|
||||
self._dataset_id = dataset_id
|
||||
|
||||
self._sub_progress_bars: List[BaseProgressBar] = []
|
||||
self._op_display: Dict["OpState", TqdmSubProgressBar] = {}
|
||||
|
||||
num_progress_bars = 0
|
||||
|
||||
self._total = TqdmSubProgressBar(
|
||||
name=f"Running Dataset {self._dataset_id}.",
|
||||
total=None,
|
||||
unit="row",
|
||||
position=num_progress_bars,
|
||||
max_name_length=self.MAX_NAME_LENGTH,
|
||||
enabled=True,
|
||||
)
|
||||
num_progress_bars += 1
|
||||
|
||||
for state in 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
|
||||
)
|
||||
|
||||
# create operator progress bar
|
||||
if sub_progress_bar_enabled:
|
||||
pg = TqdmSubProgressBar(
|
||||
name=f"- {op.name}",
|
||||
total=total,
|
||||
unit="row",
|
||||
position=num_progress_bars,
|
||||
max_name_length=self.MAX_NAME_LENGTH,
|
||||
)
|
||||
num_progress_bars += 1
|
||||
self._op_display[state] = pg
|
||||
self._sub_progress_bars.append(pg)
|
||||
|
||||
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 = TqdmSubProgressBar(
|
||||
name=f" *- {name}",
|
||||
total=total,
|
||||
unit="row",
|
||||
position=num_progress_bars,
|
||||
max_name_length=self.MAX_NAME_LENGTH,
|
||||
enabled=True,
|
||||
)
|
||||
num_progress_bars += 1
|
||||
else:
|
||||
pg = NoopSubProgressBar(
|
||||
name=f" *- {name}",
|
||||
max_name_length=self.MAX_NAME_LENGTH,
|
||||
)
|
||||
op.set_sub_progress_bar(name, pg)
|
||||
self._sub_progress_bars.append(pg)
|
||||
|
||||
# Management
|
||||
def start(self):
|
||||
# tqdm is automatically started
|
||||
pass
|
||||
|
||||
def refresh(self):
|
||||
self._total.refresh()
|
||||
for pg in self._sub_progress_bars:
|
||||
pg.refresh()
|
||||
|
||||
def close_with_finishing_description(self, desc: str, success: bool):
|
||||
del success # unused
|
||||
self._total.set_description(desc)
|
||||
self._total.close()
|
||||
for pg in self._sub_progress_bars:
|
||||
pg.close()
|
||||
|
||||
# Total Progress
|
||||
def update_total_progress(self, new_rows: int, total_rows: Optional[int]):
|
||||
self._total.update(new_rows, total_rows)
|
||||
|
||||
def update_total_resource_status(self, resource_status: str):
|
||||
desc = f"Running Dataset: {self._dataset_id}. {resource_status}"
|
||||
self._total.set_description(desc)
|
||||
|
||||
# Operator Progress
|
||||
def update_operator_progress(
|
||||
self, opstate: "OpState", resource_manager: "ResourceManager"
|
||||
):
|
||||
pg = self._op_display.get(opstate)
|
||||
if pg is not None:
|
||||
pg.update_absolute(
|
||||
opstate.op.metrics.row_outputs_taken, opstate.op.num_output_rows_total()
|
||||
)
|
||||
summary_str = format_op_state_summary(opstate, resource_manager)
|
||||
pg.set_description(f"- {opstate.op.name}: {summary_str}")
|
||||
@@ -0,0 +1,42 @@
|
||||
import logging
|
||||
|
||||
from ray.util.debug import log_once
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def truncate_operator_name(name: str, max_name_length: int) -> str:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
ctx = DataContext.get_current()
|
||||
if not ctx.enable_progress_bar_name_truncation or len(name) <= max_name_length:
|
||||
return name
|
||||
|
||||
op_names = name.split("->")
|
||||
if len(op_names) == 1:
|
||||
return op_names[0]
|
||||
|
||||
# Include as many operators as possible without approximately
|
||||
# exceeding `MAX_NAME_LENGTH`. Always include the first and
|
||||
# last operator names so it is easy to identify the DAG.
|
||||
truncated_op_names = [op_names[0]]
|
||||
for op_name in op_names[1:-1]:
|
||||
if (
|
||||
len("->".join(truncated_op_names))
|
||||
+ len("->")
|
||||
+ len(op_name)
|
||||
+ len("->")
|
||||
+ len(op_names[-1])
|
||||
) > max_name_length:
|
||||
truncated_op_names.append("...")
|
||||
if log_once("ray_data_truncate_operator_name"):
|
||||
logger.warning(
|
||||
f"Truncating long operator name to {max_name_length} "
|
||||
"characters. To disable this behavior, set "
|
||||
"`ray.data.DataContext.get_current()."
|
||||
"DEFAULT_ENABLE_PROGRESS_BAR_NAME_TRUNCATION = False`."
|
||||
)
|
||||
break
|
||||
truncated_op_names.append(op_name)
|
||||
truncated_op_names.append(op_names[-1])
|
||||
return "->".join(truncated_op_names)
|
||||
Reference in New Issue
Block a user