From 56e94aace12a9a76b0967222092c9f98334288b4 Mon Sep 17 00:00:00 2001 From: Ben Sowell Date: Wed, 19 Jun 2024 12:01:09 -0700 Subject: [PATCH] Disable Ray progress bar by default. Currently the Ray progress bar doesn't update correctly due to this issue: https://github.com/ray-project/ray/issues/44983. Until that is fixed, we disable the progress bar by default. You can explicitly enable by passing enable_progress_bars=True in the sycamore.init function. --- lib/sycamore/sycamore/context.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/lib/sycamore/sycamore/context.py b/lib/sycamore/sycamore/context.py index 0093bd0f8..28499ca12 100644 --- a/lib/sycamore/sycamore/context.py +++ b/lib/sycamore/sycamore/context.py @@ -3,12 +3,13 @@ from typing import Any, Optional import ray +import ray.data from sycamore.rules import Rule class Context: - def __init__(self, ray_args: Optional[dict[str, Any]] = None): + def __init__(self, ray_args: Optional[dict[str, Any]] = None, enable_progress_bars=False): if ray_args is None: ray_args = {} @@ -17,6 +18,9 @@ def __init__(self, ray_args: Optional[dict[str, Any]] = None): ray.init(**ray_args) + # Disable progress bars by default while https://github.com/ray-project/ray/issues/44983 is open. + ray.data.DataContext.get_current().enable_progress_bars = enable_progress_bars + self.extension_rules: list[Rule] = [] self._internal_lock = threading.Lock() @@ -44,7 +48,7 @@ def deregister_rule(self, rule: Rule) -> None: _global_context: Optional[Context] = None -def init(ray_args: Optional[dict[str, Any]] = None) -> Context: +def init(ray_args: Optional[dict[str, Any]] = None, enable_progress_bars=False) -> Context: global _global_context with _context_lock: if _global_context is None: @@ -57,6 +61,6 @@ def init(ray_args: Optional[dict[str, Any]] = None) -> Context: sycamore_logger.setup_logger() - _global_context = Context(ray_args) + _global_context = Context(ray_args, enable_progress_bars=enable_progress_bars) return _global_context