speedy-utils 1.1.26__py3-none-any.whl → 1.1.28__py3-none-any.whl

speedy_utils/multi_worker/process.py

@@ -1,29 +1,9 @@
 # ray_multi_process.py
-import datetime
-import os
-import pickle
-import threading
-import time
-import uuid
-from pathlib import Path
-from typing import Any, Callable, Iterable
-
-import psutil
-from fastcore.parallel import parallel
-from tqdm import tqdm
-
-ray: Any
-try:
-    import ray as ray  # type: ignore
-    _HAS_RAY = True
-except Exception:  # pragma: no cover
-    ray = None  # type: ignore
-    _HAS_RAY = False
-
-
-# ─── global tracking ──────────────────────────────────────────
-
-# Global tracking for processes and threads
+
+
+from ..__imports import *
+
+
 SPEEDY_RUNNING_PROCESSES: list[psutil.Process] = []
 _SPEEDY_PROCESSES_LOCK = threading.Lock()
 
@@ -31,7 +11,9 @@ _SPEEDY_PROCESSES_LOCK = threading.Lock()
 def _prune_dead_processes() -> None:
     """Remove dead processes from tracking list."""
     with _SPEEDY_PROCESSES_LOCK:
-        SPEEDY_RUNNING_PROCESSES[:] = [p for p in SPEEDY_RUNNING_PROCESSES if p.is_running()]
+        SPEEDY_RUNNING_PROCESSES[:] = [
+            p for p in SPEEDY_RUNNING_PROCESSES if p.is_running()
+        ]
 
 
 def _track_processes(processes: list[psutil.Process]) -> None:
@@ -51,8 +33,7 @@ def _track_processes(processes: list[psutil.Process]) -> None:
 
 def _track_ray_processes() -> None:
     """Track Ray worker processes when Ray is initialized."""
-    if not _HAS_RAY or not ray.is_initialized():
-        return
+
     try:
         # Get Ray worker processes
         current_pid = os.getpid()
@@ -80,7 +61,9 @@ def _track_multiprocessing_processes() -> None:
         for child in parent.children(recursive=False):  # Only direct children
             try:
                 # Basic heuristic: if it's a recent child process, it might be a worker
-                if time.time() - child.create_time() < 5:  # Created within last 5 seconds
+                if (
+                    time.time() - child.create_time() < 5
+                ):  # Created within last 5 seconds
                     new_processes.append(child)
             except (psutil.NoSuchProcess, psutil.AccessDenied):
                 continue
@@ -95,11 +78,11 @@ def _track_multiprocessing_processes() -> None:
 
 def _build_cache_dir(func: Callable, items: list[Any]) -> Path:
     """Build cache dir with function name + timestamp."""
-    func_name = getattr(func, "__name__", "func")
+    func_name = getattr(func, '__name__', 'func')
     now = datetime.datetime.now()
-    stamp = now.strftime("%m%d_%Hh%Mm%Ss")
-    run_id = f"{func_name}_{stamp}_{uuid.uuid4().hex[:6]}"
-    path = Path(".cache") / run_id
+    stamp = now.strftime('%m%d_%Hh%Mm%Ss')
+    run_id = f'{func_name}_{stamp}_{uuid.uuid4().hex[:6]}'
+    path = Path('.cache') / run_id
     path.mkdir(parents=True, exist_ok=True)
     return path
 
@@ -111,8 +94,8 @@ def wrap_dump(func: Callable, cache_dir: Path | None):
 
     def wrapped(x, *args, **kwargs):
         res = func(x, *args, **kwargs)
-        p = cache_dir / f"{uuid.uuid4().hex}.pkl"
-        with open(p, "wb") as fh:
+        p = cache_dir / f'{uuid.uuid4().hex}.pkl'
+        with open(p, 'wb') as fh:
             pickle.dump(res, fh)
         return str(p)
 
@@ -127,23 +110,19 @@ 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."""
     global RAY_WORKER
-    if not ray.is_initialized() or RAY_WORKER != workers:
+    if not ray.is_initialized() or workers != RAY_WORKER:
         if ray.is_initialized() and pbar:
-            pbar.set_postfix_str(f"Restarting Ray {workers} workers")
+            pbar.set_postfix_str(f'Restarting Ray {workers} workers')
             ray.shutdown()
         t0 = time.time()
         ray.init(num_cpus=workers, ignore_reinit_error=True)
         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")
+            pbar.set_postfix_str(f'ray.init {workers} took {took:.2f}s')
         RAY_WORKER = workers
 
 
-# ─── main API ───────────────────────────────────────────────
-from typing import Literal
-
-
 def multi_process(
     func: Callable[[Any], Any],
     items: Iterable[Any] | None = None,
@@ -153,7 +132,8 @@ def multi_process(
     lazy_output: bool = False,
     progress: bool = True,
     # backend: str = "ray",   # "seq", "ray", or "fastcore"
-    backend: Literal["seq", "ray", "mp", "threadpool", "safe"] = "mp",
+    backend: Literal['seq', 'ray', 'mp', 'threadpool', 'safe'] = 'mp',
+    desc: str | None = None,
     **func_kwargs: Any,
 ) -> list[Any]:
     """
@@ -172,7 +152,7 @@ def multi_process(
 
     # default backend selection
     if backend is None:
-        backend = "ray" if _HAS_RAY else "mp"
+        backend = 'ray' if _HAS_RAY else 'mp'
 
     # unify items
     # unify items and coerce to concrete list so we can use len() and
@@ -192,12 +172,16 @@ def multi_process(
     f_wrapped = wrap_dump(func, cache_dir)
 
     total = len(items)
+    if desc:
+        desc = desc.strip() + f'[{backend}]'
+    else:
+        desc = f'Multi-process [{backend}]'
     with tqdm(
-        total=total, desc=f"multi_process [{backend}]", disable=not progress
+        total=total, desc=desc , disable=not progress
     ) as pbar:
         # ---- sequential backend ----
-        if backend == "seq":
-            pbar.set_postfix_str("backend=seq")
+        if backend == 'seq':
+            pbar.set_postfix_str('backend=seq')
             results = []
             for x in items:
                 results.append(f_wrapped(x, **func_kwargs))
@@ -205,15 +189,8 @@ def multi_process(
             return results
 
         # ---- ray backend ----
-        if backend == "ray":
-            if not _HAS_RAY:
-                msg = (
-                    "Ray backend requested but 'ray' is not installed. "
-                    "Install extra: pip install 'speedy-utils[ray]' or "
-                    "poetry install -E ray."
-                )
-                raise RuntimeError(msg)
-            pbar.set_postfix_str("backend=ray")
+        if backend == 'ray':
+            pbar.set_postfix_str('backend=ray')
             ensure_ray(workers, pbar)
 
             @ray.remote
@@ -229,36 +206,41 @@ def multi_process(
             return results
 
         # ---- fastcore backend ----
-        if backend == "mp":
+        if backend == 'mp':
             results = parallel(
                 f_wrapped, items, n_workers=workers, progress=progress, threadpool=False
             )
             _track_multiprocessing_processes()  # Track multiprocessing workers
             _prune_dead_processes()  # Clean up dead processes
             return list(results)
-        if backend == "threadpool":
+        if backend == 'threadpool':
             results = parallel(
                 f_wrapped, items, n_workers=workers, progress=progress, threadpool=True
             )
             return list(results)
-        if backend == "safe":
+        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 _track_executor_threads, _prune_dead_threads
-                with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                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:
+                with concurrent.futures.ThreadPoolExecutor(
+                    max_workers=workers
+                ) as executor:
                     results = list(executor.map(f_wrapped, items))
             return results
-        
-        raise ValueError(f"Unsupported backend: {backend!r}")
 
+        raise ValueError(f'Unsupported backend: {backend!r}')
 
 
 def cleanup_phantom_workers():
@@ -270,47 +252,52 @@ def cleanup_phantom_workers():
     _prune_dead_processes()
     killed_processes = 0
     with _SPEEDY_PROCESSES_LOCK:
-        for process in SPEEDY_RUNNING_PROCESSES[:]:  # Copy to avoid modification during iteration
+        for process in SPEEDY_RUNNING_PROCESSES[
+            :
+        ]:  # Copy to avoid modification during iteration
             try:
-                print(f"🔪 Killing tracked process {process.pid} ({process.name()})")
+                print(f'🔪 Killing tracked process {process.pid} ({process.name()})')
                 process.kill()
                 killed_processes += 1
             except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
-                print(f"⚠️ Could not kill process {process.pid}: {e}")
+                print(f'⚠️ Could not kill process {process.pid}: {e}')
         SPEEDY_RUNNING_PROCESSES.clear()
-    
+
     # Also kill any remaining child processes (fallback)
     parent = psutil.Process(os.getpid())
     for child in parent.children(recursive=True):
         try:
-            print(f"🔪 Killing child process {child.pid} ({child.name()})")
+            print(f'🔪 Killing child process {child.pid} ({child.name()})')
             child.kill()
         except psutil.NoSuchProcess:
             pass
-    
+
     # Try to clean up threads using thread module functions if available
     try:
-        from .thread import SPEEDY_RUNNING_THREADS, kill_all_thread, _prune_dead_threads
+        from .thread import SPEEDY_RUNNING_THREADS, _prune_dead_threads, kill_all_thread
+
         _prune_dead_threads()
         killed_threads = kill_all_thread()
         if killed_threads > 0:
-            print(f"🔪 Killed {killed_threads} tracked threads")
+            print(f'🔪 Killed {killed_threads} tracked threads')
     except ImportError:
         # Fallback: just report stray threads
         for t in threading.enumerate():
             if t is threading.current_thread():
                 continue
             if not t.daemon:
-                print(f"⚠️ Thread {t.name} is still running (cannot be force-killed).")
-    
-    print(f"✅ Cleaned up {killed_processes} tracked processes and child processes (kernel untouched).")
+                print(f'⚠️ Thread {t.name} is still running (cannot be force-killed).')
+
+    print(
+        f'✅ Cleaned up {killed_processes} tracked processes and child processes (kernel untouched).'
+    )
+
 
 # Usage: run this anytime after cancelling a cell
 
 
 __all__ = [
-    "SPEEDY_RUNNING_PROCESSES",
-    "multi_process",
-    "cleanup_phantom_workers",
+    'SPEEDY_RUNNING_PROCESSES',
+    'multi_process',
+    'cleanup_phantom_workers',
 ]
-