wehub-resource-sync
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
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
|
|
)
|