speedy-utils 1.1.40__py3-none-any.whl → 1.1.42__py3-none-any.whl

speedy_utils/multi_worker/process.py

@@ -1,16 +1,23 @@
+import warnings
+import os
+# Suppress Ray FutureWarnings before any imports
+warnings.filterwarnings("ignore", category=FutureWarning, module="ray.*")
+warnings.filterwarnings("ignore", message=".*pynvml.*deprecated.*", category=FutureWarning)
+
+# Set environment variables before Ray is imported anywhere
+os.environ["RAY_ACCEL_ENV_VAR_OVERRI" \
+"DE_ON_ZERO"] = "0"
+os.environ["RAY_DEDUP_LOGS"] = "1"
+os.environ["RAY_LOG_TO_STDERR"] = "0"
+
 from ..__imports import *
+import tempfile
+from .progress import create_progress_tracker, ProgressPoller, get_ray_progress_actor
 
 
 SPEEDY_RUNNING_PROCESSES: list[psutil.Process] = []
 _SPEEDY_PROCESSES_LOCK = threading.Lock()
 
-
-# /mnt/data/anhvth8/venvs/Megatron-Bridge-Host/lib/python3.12/site-packages/ray/_private/worker.py:2046: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0
-# turn off future warning and verbose task logs
-os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0"
-os.environ["RAY_DEDUP_LOGS"] = "0"
-os.environ["RAY_LOG_TO_STDERR"] = "0"
-
 def _prune_dead_processes() -> None:
     """Remove dead processes from tracking list."""
     with _SPEEDY_PROCESSES_LOCK:
@@ -121,35 +128,213 @@ def wrap_dump(func: Callable, cache_dir: Path | None, dump_in_thread: bool = Tru
     return wrapped
 
 
+_LOG_GATE_CACHE: dict[str, bool] = {}
+
+
+def _should_allow_worker_logs(mode: Literal['all', 'zero', 'first'], gate_path: Path | None) -> bool:
+    """Determine if current worker should emit logs for the given mode."""
+    if mode == 'all':
+        return True
+    if mode == 'zero':
+        return False
+    if mode == 'first':
+        if gate_path is None:
+            return True
+        key = str(gate_path)
+        cached = _LOG_GATE_CACHE.get(key)
+        if cached is not None:
+            return cached
+        gate_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            fd = os.open(key, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+        except FileExistsError:
+            allowed = False
+        else:
+            os.close(fd)
+            allowed = True
+        _LOG_GATE_CACHE[key] = allowed
+        return allowed
+    raise ValueError(f'Unsupported log mode: {mode!r}')
+
+
+def _cleanup_log_gate(gate_path: Path | None):
+    if gate_path is None:
+        return
+    try:
+        gate_path.unlink(missing_ok=True)
+    except OSError:
+        pass
+
+
+@contextlib.contextmanager
+def _patch_fastcore_progress_bar(*, leave: bool = True):
+    """Temporarily force fastcore.progress_bar to keep the bar on screen."""
+    try:
+        import fastcore.parallel as _fp
+    except ImportError:
+        yield False
+        return
+
+    orig = getattr(_fp, 'progress_bar', None)
+    if orig is None:
+        yield False
+        return
+
+    def _wrapped(*args, **kwargs):
+        kwargs.setdefault('leave', leave)
+        return orig(*args, **kwargs)
+
+    _fp.progress_bar = _wrapped
+    try:
+        yield True
+    finally:
+        _fp.progress_bar = orig
+
+
+class _PrefixedWriter:
+    """Stream wrapper that prefixes each line with worker id."""
+
+    def __init__(self, stream, prefix: str):
+        self._stream = stream
+        self._prefix = prefix
+        self._at_line_start = True
+
+    def write(self, s):
+        if not s:
+            return 0
+        total = 0
+        for chunk in s.splitlines(True):
+            if self._at_line_start:
+                self._stream.write(self._prefix)
+                total += len(self._prefix)
+            self._stream.write(chunk)
+            total += len(chunk)
+            self._at_line_start = chunk.endswith('\n')
+        return total
+
+    def flush(self):
+        self._stream.flush()
+
+
+def _call_with_log_control(
+    func: Callable,
+    x: Any,
+    func_kwargs: dict[str, Any],
+    log_mode: Literal['all', 'zero', 'first'],
+    gate_path: Path | None,
+):
+    """Call a function, silencing stdout/stderr based on log mode."""
+    allow_logs = _should_allow_worker_logs(log_mode, gate_path)
+    if allow_logs:
+        prefix = f"[worker-{os.getpid()}] "
+        # Route worker logs to stderr to reduce clobbering tqdm/progress output on stdout
+        out = _PrefixedWriter(sys.stderr, prefix)
+        err = out
+        with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
+            return func(x, **func_kwargs)
+    with open(os.devnull, 'w') as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
+        return func(x, **func_kwargs)
+
+
 # ─── ray management ─────────────────────────────────────────
 
 RAY_WORKER = None
 
 
-def ensure_ray(workers: int, pbar: tqdm | None = None):
-    """Initialize or reinitialize Ray with a given worker count, log to bar postfix."""
+def ensure_ray(workers: int | None, pbar: tqdm | None = None, ray_metrics_port: int | None = None):
+    """
+    Initialize or reinitialize Ray safely for both local and cluster environments.
+
+    1. Tries to connect to an existing cluster (address='auto') first.
+    2. If no cluster is found, starts a local Ray instance with 'workers' CPUs.
+    """
     import ray as _ray_module
     import logging
 
     global RAY_WORKER
-    # shutdown when worker count changes or if Ray not initialized
-    if not _ray_module.is_initialized() or workers != RAY_WORKER:
-        if _ray_module.is_initialized() and pbar:
-            pbar.set_postfix_str(f'Restarting Ray {workers} workers')
+    requested_workers = workers
+    if workers is None:
+        workers = os.cpu_count() or 1
+
+    if ray_metrics_port is not None:
+        os.environ['RAY_metrics_export_port'] = str(ray_metrics_port)
+
+    allow_restart = os.environ.get("RESTART_RAY", "0").lower() in ("1", "true")
+    is_cluster_env = "RAY_ADDRESS" in os.environ or os.environ.get("RAY_CLUSTER") == "1"
+
+    # 1. Handle existing session
+    if _ray_module.is_initialized():
+        if not allow_restart:
+            if pbar:
+                pbar.set_postfix_str("Using existing Ray session")
+            return
+
+        # Avoid shutting down shared cluster sessions.
+        if is_cluster_env:
+            if pbar:
+                pbar.set_postfix_str("Cluster active: skipping restart to protect connection")
+            return
+
+        # Local restart: only if worker count changed
+        if workers != RAY_WORKER:
+            if pbar:
+                pbar.set_postfix_str(f'Restarting local Ray with {workers} workers')
             _ray_module.shutdown()
-        t0 = time.time()
+        else:
+            return
+
+    # 2. Initialization logic
+    t0 = time.time()
+    
+    # Try to connect to existing cluster FIRST (address="auto")
+    try:
+        if pbar:
+            pbar.set_postfix_str("Searching for Ray cluster...")
+        
+        # MUST NOT pass num_cpus/num_gpus here to avoid ValueError on existing clusters
+        _ray_module.init(
+            address="auto", 
+            ignore_reinit_error=True, 
+            logging_level=logging.ERROR,
+            log_to_driver=False
+        )
+        
+        if pbar:
+            resources = _ray_module.cluster_resources()
+            cpus = resources.get("CPU", 0)
+            pbar.set_postfix_str(f"Connected to Ray Cluster ({int(cpus)} CPUs)")
+            
+    except Exception:
+        # 3. Fallback: Start a local Ray session
+        if pbar:
+            pbar.set_postfix_str(f"No cluster found. Starting local Ray ({workers} CPUs)...")
+            
         _ray_module.init(
             num_cpus=workers,
             ignore_reinit_error=True,
             logging_level=logging.ERROR,
             log_to_driver=False,
         )
-        took = time.time() - t0
-        _track_ray_processes()  # Track Ray worker processes
+        
         if pbar:
-            pbar.set_postfix_str(f'ray.init {workers} took {took:.2f}s')
-        RAY_WORKER = workers
+            took = time.time() - t0
+            pbar.set_postfix_str(f'ray.init local {workers} took {took:.2f}s')
 
+    _track_ray_processes()
+
+    if requested_workers is None:
+        try:
+            resources = _ray_module.cluster_resources()
+            total_cpus = int(resources.get("CPU", 0))
+            if total_cpus > 0:
+                workers = total_cpus
+        except Exception:
+            pass
+
+    RAY_WORKER = workers
+
+
+# TODO: make smarter backend selection, when shared_kwargs is used, and backend != 'ray', do not raise error but change to ray and warning user about this
 def multi_process(
     func: Callable[[Any], Any],
     items: Iterable[Any] | None = None,
@@ -163,6 +348,10 @@ def multi_process(
     desc: str | None = None,
     shared_kwargs: list[str] | None = None,
     dump_in_thread: bool = True,
+    ray_metrics_port: int | None = None,
+    log_worker: Literal['zero', 'first', 'all'] = 'first',
+    total_items: int | None = None,
+    poll_interval: float = 0.3,
     **func_kwargs: Any,
 ) -> list[Any]:
     """
@@ -185,6 +374,23 @@ def multi_process(
         - Whether to dump results to disk in a separate thread (default: True)
         - If False, dumping is done synchronously, which may block but ensures data is saved before returning
 
+    ray_metrics_port:
+        - Optional port for Ray metrics export (Ray backend only)
+        - Set to 0 to disable Ray metrics
+        - If None, uses Ray's default behavior
+
+    log_worker:
+        - Control worker stdout/stderr noise
+        - 'first': only first worker emits logs, others are silenced (default)
+        - 'all': allow worker prints (may overlap tqdm)
+        - 'zero': silence worker stdout/stderr to keep progress bar clean
+
+    total_items:
+        - Optional item-level total for progress tracking (Ray backend only)
+
+    poll_interval:
+        - Poll interval in seconds for progress actor updates (Ray backend only)
+
     If lazy_output=True, every result is saved to .pkl and
     the returned list contains file paths.
     """
@@ -219,11 +425,13 @@ def multi_process(
                     f"Valid parameters: {valid_params}"
                 )
 
-        # Only allow shared_kwargs with Ray backend
-        if backend != 'ray':
-            raise ValueError(
-                f"shared_kwargs only supported with 'ray' backend, got '{backend}'"
-            )
+    # Prefer Ray backend when shared kwargs are requested
+    if backend != 'ray':
+        warnings.warn(
+            "shared_kwargs only supported with 'ray' backend, switching backend to 'ray'",
+            UserWarning,
+        )
+        backend = 'ray'
 
     # unify items
     # unify items and coerce to concrete list so we can use len() and
@@ -235,35 +443,60 @@ def multi_process(
     if items is None:
         raise ValueError("'items' or 'inputs' must be provided")
 
-    if workers is None:
+    if workers is None and backend != 'ray':
         workers = os.cpu_count() or 1
 
     # build cache dir + wrap func
     cache_dir = _build_cache_dir(func, items) if lazy_output else None
     f_wrapped = wrap_dump(func, cache_dir, dump_in_thread)
 
+    log_gate_path: Path | None = None
+    if log_worker == 'first':
+        log_gate_path = Path(tempfile.gettempdir()) / f'speedy_utils_log_gate_{os.getpid()}_{uuid.uuid4().hex}.gate'
+    elif log_worker not in ('zero', 'all'):
+        raise ValueError(f'Unsupported log_worker: {log_worker!r}')
+
     total = len(items)
     if desc:
         desc = desc.strip() + f'[{backend}]'
     else:
         desc = f'Multi-process [{backend}]'
-    with tqdm(
-        total=total, desc=desc , disable=not progress
-    ) as pbar:
-        # ---- sequential backend ----
-        if backend == 'seq':
+
+    # ---- sequential backend ----
+    if backend == 'seq':
+        results: list[Any] = []
+        with tqdm(total=total, desc=desc, disable=not progress, file=sys.stdout) as pbar:
             for x in items:
-                results.append(f_wrapped(x, **func_kwargs))
+                results.append(
+                    _call_with_log_control(
+                        f_wrapped,
+                        x,
+                        func_kwargs,
+                        log_worker,
+                        log_gate_path,
+                    )
+                )
                 pbar.update(1)
-            return results
+        _cleanup_log_gate(log_gate_path)
+        return results
 
-        # ---- ray backend ----
-        if backend == 'ray':
-            import ray as _ray_module
+    # ---- ray backend ----
+    if backend == 'ray':
+        import ray as _ray_module
 
-            ensure_ray(workers, pbar)
+        results = []
+        gate_path_str = str(log_gate_path) if log_gate_path else None
+        with tqdm(total=total, desc=desc, disable=not progress, file=sys.stdout) as pbar:
+            ensure_ray(workers, pbar, ray_metrics_port)
             shared_refs = {}
             regular_kwargs = {}
+            
+            # Create progress actor for item-level tracking if total_items specified
+            progress_actor = None
+            progress_poller = None
+            if total_items is not None:
+                progress_actor = create_progress_tracker(total_items, desc or "Items")
+                shared_refs['progress_actor'] = progress_actor  # Pass actor handle directly (not via put)
 
             if shared_kwargs:
                 for kw in shared_kwargs:
@@ -283,60 +516,125 @@ def multi_process(
             def _task(x, shared_refs_dict, regular_kwargs_dict):
                 # Dereference shared objects (zero-copy for numpy arrays)
                 import ray as _ray_in_task
-                dereferenced = {k: _ray_in_task.get(v) for k, v in shared_refs_dict.items()}
+                gate = Path(gate_path_str) if gate_path_str else None
+                dereferenced = {}
+                for k, v in shared_refs_dict.items():
+                    if k == 'progress_actor':
+                        # Pass actor handle directly (don't dereference)
+                        dereferenced[k] = v
+                    else:
+                        dereferenced[k] = _ray_in_task.get(v)
                 # Merge with regular kwargs
                 all_kwargs = {**dereferenced, **regular_kwargs_dict}
-                return f_wrapped(x, **all_kwargs)
+                return _call_with_log_control(
+                    f_wrapped,
+                    x,
+                    all_kwargs,
+                    log_worker,
+                    gate,
+                )
 
             refs = [
                 _task.remote(x, shared_refs, regular_kwargs) for x in items
             ]
 
-            results = []
             t_start = time.time()
+            
+            # Start progress poller if using item-level progress
+            if progress_actor is not None:
+                # Update pbar total to show items instead of tasks
+                pbar.total = total_items
+                pbar.refresh()
+                progress_poller = ProgressPoller(progress_actor, pbar, poll_interval)
+                progress_poller.start()
+            
             for r in refs:
                 results.append(_ray_module.get(r))
-                pbar.update(1)
+                if progress_actor is None:
+                    # Only update task-level progress if not using item-level
+                    pbar.update(1)
+            
+            # Stop progress poller
+            if progress_poller is not None:
+                progress_poller.stop()
+                
             t_end = time.time()
-            print(f"Ray processing took {t_end - t_start:.2f}s for {total} items")
-            return results
-
-        # ---- fastcore backend ----
-        if backend == 'mp':
-            results = parallel(
-                f_wrapped, items, n_workers=workers, progress=progress, threadpool=False
+            item_desc = f"{total_items:,} items" if total_items else f"{total} tasks"
+            print(f"Ray processing took {t_end - t_start:.2f}s for {item_desc}")
+        _cleanup_log_gate(log_gate_path)
+        return results
+
+    # ---- fastcore backend ----
+    if backend == 'mp':
+        worker_func = functools.partial(
+            _call_with_log_control,
+            f_wrapped,
+            func_kwargs=func_kwargs,
+            log_mode=log_worker,
+            gate_path=log_gate_path,
+        )
+        with _patch_fastcore_progress_bar(leave=True):
+            results = list(parallel(
+                worker_func, items, n_workers=workers, progress=progress, threadpool=False
+            ))
+        _track_multiprocessing_processes()  # Track multiprocessing workers
+        _prune_dead_processes()  # Clean up dead processes
+        _cleanup_log_gate(log_gate_path)
+        return results
+
+    if backend == 'threadpool':
+        worker_func = functools.partial(
+            _call_with_log_control,
+            f_wrapped,
+            func_kwargs=func_kwargs,
+            log_mode=log_worker,
+            gate_path=log_gate_path,
+        )
+        with _patch_fastcore_progress_bar(leave=True):
+            results = list(parallel(
+                worker_func, items, n_workers=workers, progress=progress, threadpool=True
+            ))
+        _cleanup_log_gate(log_gate_path)
+        return results
+
+    if backend == 'safe':
+        # Completely safe backend for tests - no multiprocessing, no external progress bars
+        import concurrent.futures
+
+        # Import thread tracking from thread module
+        try:
+            from .thread import _prune_dead_threads, _track_executor_threads
+
+            worker_func = functools.partial(
+                _call_with_log_control,
+                f_wrapped,
+                func_kwargs=func_kwargs,
+                log_mode=log_worker,
+                gate_path=log_gate_path,
             )
-            _track_multiprocessing_processes()  # Track multiprocessing workers
-            _prune_dead_processes()  # Clean up dead processes
-            return list(results)
-        if backend == 'threadpool':
-            results = parallel(
-                f_wrapped, items, n_workers=workers, progress=progress, threadpool=True
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=workers
+            ) as executor:
+                _track_executor_threads(executor)  # Track threads
+                results = list(executor.map(worker_func, items))
+            _prune_dead_threads()  # Clean up dead threads
+        except ImportError:
+            # Fallback if thread module not available
+            worker_func = functools.partial(
+                _call_with_log_control,
+                f_wrapped,
+                func_kwargs=func_kwargs,
+                log_mode=log_worker,
+                gate_path=log_gate_path,
             )
-            return list(results)
-        if backend == 'safe':
-            # Completely safe backend for tests - no multiprocessing, no external progress bars
-            import concurrent.futures
-
-            # Import thread tracking from thread module
-            try:
-                from .thread import _prune_dead_threads, _track_executor_threads
-
-                with concurrent.futures.ThreadPoolExecutor(
-                    max_workers=workers
-                ) as executor:
-                    _track_executor_threads(executor)  # Track threads
-                    results = list(executor.map(f_wrapped, items))
-                _prune_dead_threads()  # Clean up dead threads
-            except ImportError:
-                # Fallback if thread module not available
-                with concurrent.futures.ThreadPoolExecutor(
-                    max_workers=workers
-                ) as executor:
-                    results = list(executor.map(f_wrapped, items))
-            return results
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=workers
+            ) as executor:
+                results = list(executor.map(worker_func, items))
+        _cleanup_log_gate(log_gate_path)
+        return results
 
-        raise ValueError(f'Unsupported backend: {backend!r}')
+    raise ValueError(f'Unsupported backend: {backend!r}')
 
 
 def cleanup_phantom_workers():
@@ -396,4 +694,6 @@ __all__ = [
     'SPEEDY_RUNNING_PROCESSES',
     'multi_process',
     'cleanup_phantom_workers',
+    'create_progress_tracker',
+    'get_ray_progress_actor',
 ]

speedy_utils/multi_worker/progress.py

@@ -0,0 +1,140 @@
+"""
+Real-time progress tracking for distributed Ray tasks.
+
+This module provides a ProgressActor that allows workers to report item-level
+progress in real-time, giving users visibility into actual items processed
+rather than just task completion.
+"""
+import time
+import threading
+from typing import Optional, Callable
+
+__all__ = ['ProgressActor', 'create_progress_tracker', 'get_ray_progress_actor']
+
+
+def get_ray_progress_actor():
+    """Get the Ray-decorated ProgressActor class (lazy import to avoid Ray at module load)."""
+    import ray
+    
+    @ray.remote
+    class ProgressActor:
+        """
+        A Ray actor for tracking real-time progress across distributed workers.
+        
+        Workers call `update(n)` to report items processed, and the main process
+        can poll `get_progress()` to update a tqdm bar in real-time.
+        """
+        def __init__(self, total: int, desc: str = "Items"):
+            self.total = total
+            self.processed = 0
+            self.desc = desc
+            self.start_time = time.time()
+            self._lock = threading.Lock()
+        
+        def update(self, n: int = 1) -> int:
+            """Increment processed count by n. Returns new total."""
+            with self._lock:
+                self.processed += n
+                return self.processed
+        
+        def get_progress(self) -> dict:
+            """Get current progress stats."""
+            with self._lock:
+                elapsed = time.time() - self.start_time
+                rate = self.processed / elapsed if elapsed > 0 else 0
+                return {
+                    "processed": self.processed,
+                    "total": self.total,
+                    "elapsed": elapsed,
+                    "rate": rate,
+                    "desc": self.desc,
+                }
+        
+        def set_total(self, total: int):
+            """Update total (useful if exact count unknown at start)."""
+            with self._lock:
+                self.total = total
+        
+        def reset(self):
+            """Reset progress counter."""
+            with self._lock:
+                self.processed = 0
+                self.start_time = time.time()
+    
+    return ProgressActor
+
+
+def create_progress_tracker(total: int, desc: str = "Items"):
+    """
+    Create a progress tracker actor for use with Ray distributed tasks.
+    
+    Args:
+        total: Total number of items to process
+        desc: Description for the progress bar
+        
+    Returns:
+        A Ray actor handle that workers can use to report progress
+        
+    Example:
+        progress_actor = create_progress_tracker(1000000, "Processing items")
+        
+        @ray.remote
+        def worker(items, progress_actor):
+            for item in items:
+                process(item)
+                ray.get(progress_actor.update.remote(1))
+        
+        # In main process, poll progress:
+        while not done:
+            stats = ray.get(progress_actor.get_progress.remote())
+            pbar.n = stats["processed"]
+            pbar.refresh()
+    """
+    import ray
+    ProgressActor = get_ray_progress_actor()
+    return ProgressActor.remote(total, desc)
+
+
+class ProgressPoller:
+    """
+    Background thread that polls a Ray progress actor and updates a tqdm bar.
+    """
+    def __init__(self, progress_actor, pbar, poll_interval: float = 0.5):
+        import ray
+        self._ray = ray
+        self.progress_actor = progress_actor
+        self.pbar = pbar
+        self.poll_interval = poll_interval
+        self._stop_event = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+    
+    def start(self):
+        """Start the polling thread."""
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+    
+    def stop(self):
+        """Stop the polling thread."""
+        self._stop_event.set()
+        if self._thread:
+            self._thread.join(timeout=2.0)
+    
+    def _poll_loop(self):
+        """Poll the progress actor and update tqdm."""
+        while not self._stop_event.is_set():
+            try:
+                stats = self._ray.get(self.progress_actor.get_progress.remote())
+                self.pbar.n = stats["processed"]
+                self.pbar.set_postfix_str(f'{stats["rate"]:.1f} items/s')
+                self.pbar.refresh()
+            except Exception:
+                pass  # Ignore errors during polling
+            self._stop_event.wait(self.poll_interval)
+        
+        # Final update
+        try:
+            stats = self._ray.get(self.progress_actor.get_progress.remote())
+            self.pbar.n = stats["processed"]
+            self.pbar.refresh()
+        except Exception:
+            pass