chore: import upstream snapshot with attribution

python/ray/data/_internal/progress/__init__.py

@@ -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
+            )