speedy-utils 1.1.17__py3-none-any.whl → 1.1.19__py3-none-any.whl

speedy_utils/multi_worker/process.py

@@ -1,203 +1,191 @@
-import multiprocessing
-import os
-import traceback
-from collections.abc import Callable, Iterable, Iterator
-from concurrent.futures import ProcessPoolExecutor, as_completed
-from itertools import islice
-from typing import Any, TypeVar, Union, cast
-
-T = TypeVar("T")
-
-if hasattr(multiprocessing, "set_start_method"):
-    try:
-        multiprocessing.set_start_method("spawn", force=True)
-    except RuntimeError:
-        pass
-
+# ray_multi_process.py
+import time, os, pickle, uuid, datetime, multiprocessing
+from pathlib import Path
+from typing import Any, Callable
+from tqdm import tqdm
+import psutil
+import threading
+ray: Any
 try:
-    from tqdm import tqdm
-except ImportError:  # pragma: no cover
-    tqdm = None  # type: ignore[assignment]
-
-
-# ──── internal helpers ────────────────────────────────────────────────────
-
-
-def _group_iter(src: Iterable[Any], size: int) -> Iterable[list[Any]]:
-    "Yield *size*-sized chunks from *src*."
-    it = iter(src)
-    while chunk := list(islice(it, size)):
-        yield chunk
-
-
-def _short_tb() -> str:
-    tb = "".join(traceback.format_exc())
-    return "\n".join(ln for ln in tb.splitlines() if "multi_process" not in ln)
-
-
-def _safe_call(func: Callable, obj, fixed):
-    try:
-        return func(obj, **fixed)
-    except Exception as exc:
-        func_name = getattr(func, "__name__", str(func))
-        raise RuntimeError(
-            f"{func_name}({obj!r}) failed: {exc}\n{_short_tb()}"
-        ) from exc
-
+    import ray as ray  # type: ignore
+    _HAS_RAY = True
+except Exception:  # pragma: no cover
+    ray = None  # type: ignore
+    _HAS_RAY = False
+from fastcore.parallel import parallel
+
+
+# ─── cache helpers ──────────────────────────────────────────
+
+def _build_cache_dir(func: Callable, items: list[Any]) -> Path:
+    """Build cache dir with function name + timestamp."""
+    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
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+def wrap_dump(func: Callable, cache_dir: Path | None):
+    """Wrap a function so results are dumped to .pkl when cache_dir is set."""
+    if cache_dir is None:
+        return func
+
+    def wrapped(x, *args, **kwargs):
+        res = func(x, *args, **kwargs)
+        p = cache_dir / f"{uuid.uuid4().hex}.pkl"
+        with open(p, "wb") as fh:
+            pickle.dump(res, fh)
+        return str(p)
+    return wrapped
+
+# ─── 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."""
+    global RAY_WORKER
+    if not ray.is_initialized() or RAY_WORKER != workers:
+        if ray.is_initialized() and pbar:
+            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
+        if pbar:
+            pbar.set_postfix_str(f"ray.init {workers} took {took:.2f}s")
+        RAY_WORKER = workers
+
+# ─── main API ───────────────────────────────────────────────
+from typing import Literal
 
-def _worker_process(
-    func: Callable, item_batch: Any, fixed_kwargs: dict, batch_size: int
-):
-    """Worker function executed in each process."""
-    if batch_size > 1:
-        results = []
-        for itm in item_batch:
-            try:
-                results.append(_safe_call(func, itm, fixed_kwargs))
-            except Exception:
-                results.append(None)
-        return results
-    return _safe_call(func, item_batch, fixed_kwargs)
-
-
-# ──── public API ──────────────────────────────────────────────────────────
 def multi_process(
     func: Callable[[Any], Any],
-    inputs: Iterable[Any],
+    items: list[Any] | None = None,
     *,
-    workers: Union[int, None] = None,
-    batch: int = 1,
-    ordered: bool = True,
-    progress: bool = False,
-    inflight: Union[int, None] = None,
-    timeout: Union[float, None] = None,
-    stop_on_error: bool = True,
-    process_update_interval=10,
-    for_loop: bool = False,
-    **fixed_kwargs,
+    inputs: list[Any] | None = None,
+    workers: int | None = None,
+    lazy_output: bool = False,
+    progress: bool = True,
+    # backend: str = "ray",   # "seq", "ray", or "fastcore"
+    backend: Literal["seq", "ray", "mp", "threadpool", "safe"] | None = None,
+    # Additional optional knobs (accepted for compatibility)
+    batch: int | None = None,
+    ordered: bool | None = None,
+    process_update_interval: int | None = None,
+    stop_on_error: bool | None = None,
+    **func_kwargs: Any,
 ) -> list[Any]:
     """
-    Simple multi‑processing parallel map that returns a *list*.
-
-    Parameters
-    ----------
-    func          – target callable executed in separate processes, must be of the form f(obj, ...).
-    inputs        – iterable with the objects.
-    workers       – process pool size (defaults to :pyfunc:`os.cpu_count()`).
-    batch         – package *batch* inputs into one call to reduce IPC overhead.
-    ordered       – keep original order; if ``False`` results stream as finished.
-    progress      – show a tqdm bar (requires *tqdm*).
-    inflight      – max logical items concurrently submitted
-                    *(default: ``workers × 4``)*.
-    timeout       – overall timeout for the mapping (seconds).
-    stop_on_error – raise immediately on first exception (default) or
-                    substitute failing result with ``None``.
-    **fixed_kwargs – static keyword args forwarded to every ``func()`` call.
+    Multi-process map with selectable backend.
+
+    backend:
+        - "seq": run sequentially
+        - "ray": run in parallel with Ray
+        - "mp": run in parallel with multiprocessing (uses threadpool to avoid fork warnings)
+        - "threadpool": run in parallel with thread pool
+        - "safe": run in parallel with thread pool (explicitly safe for tests)
+
+    If lazy_output=True, every result is saved to .pkl and
+    the returned list contains file paths.
     """
-    if for_loop:
-        ret = []
-        for arg in inputs:
-            ret.append(func(arg, **fixed_kwargs))
-        return ret
+
+    # default backend selection
+    if backend is None:
+        backend = "ray" if _HAS_RAY else "mp"
+
+    # unify items
+    if items is None and inputs is not None:
+        items = list(inputs)
+    if items is None:
+        raise ValueError("'items' or 'inputs' must be provided")
 
     if workers is None:
         workers = os.cpu_count() or 1
-    if inflight is None:
-        inflight = workers * 4
-    if batch < 1:
-        raise ValueError("batch must be ≥ 1")
-
-    try:
-        n_inputs = len(inputs)  # type: ignore[arg-type]
-    except Exception:
-        n_inputs = None
-
-    src_iter: Iterator[Any] = iter(inputs)
-    if batch > 1:
-        src_iter = cast(Iterator[Any], _group_iter(src_iter, batch))
-
-    logical_total = n_inputs
-    bar = None
-    last_bar = 0
-    if progress and tqdm is not None and logical_total is not None:
-        bar = tqdm(
-            total=logical_total,
-            ncols=80,
-            colour="green",
-            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"
-            " [{elapsed}<{remaining}, {rate_fmt}{postfix}]",
-        )
-
-    if ordered and logical_total is not None:
-        results: list[Any] = [None] * logical_total
-    else:
-        results = []
-
-    completed = 0
-    next_idx = 0
-
-    with ProcessPoolExecutor(max_workers=workers) as pool:
-        futures = set()
-
-        for _ in range(min(inflight, workers)):
-            try:
-                arg = next(src_iter)
-            except StopIteration:
-                break
-            fut = pool.submit(_worker_process, func, arg, fixed_kwargs, batch)
-            fut.idx = next_idx  # type: ignore[attr-defined]
-            futures.add(fut)
-            next_idx += len(arg) if batch > 1 else 1
-
-        while futures:
-            for fut in as_completed(futures, timeout=timeout):
-                futures.remove(fut)
-                idx = fut.idx  # type: ignore[attr-defined]
-                try:
-                    res = fut.result()
-                except Exception:
-                    if stop_on_error:
-                        raise
-                    num_items = batch if batch > 1 else 1
-                    res = [None] * num_items if batch > 1 else None
-
-                out_items = res if batch > 1 else [res]
-                if out_items is None:
-                    out_items = []
-
-                if ordered and logical_total is not None:
-                    if isinstance(out_items, list) and len(out_items) > 0:
-                        for i, item in enumerate(out_items):
-                            if idx + i < len(results):
-                                results[idx + i] = item
-                else:
-                    if isinstance(out_items, list):
-                        results.extend(out_items)
-
-                completed += len(out_items)
-
-                if bar and completed - last_bar >= process_update_interval:
-                    bar.update(completed - last_bar)
-                    last_bar = completed
-
-                try:
-                    while next_idx - completed < inflight:
-                        arg = next(src_iter)
-                        fut2 = pool.submit(
-                            _worker_process, func, arg, fixed_kwargs, batch
-                        )
-                        fut2.idx = next_idx  # type: ignore[attr-defined]
-                        futures.add(fut2)
-                        next_idx += len(arg) if batch > 1 else 1
-                except StopIteration:
-                    pass
-                break
-
-    if bar:
-        bar.update(completed - last_bar)
-        bar.close()
-
-    return results
-
-
-__all__ = ["multi_process"]
+
+    # build cache dir + wrap func
+    cache_dir = _build_cache_dir(func, items) if lazy_output else None
+    f_wrapped = wrap_dump(func, cache_dir)
+
+    total = len(items)
+    with tqdm(total=total, desc=f"multi_process [{backend}]", disable=not progress) as pbar:
+
+        # ---- sequential backend ----
+        if backend == "seq":
+            pbar.set_postfix_str("backend=seq")
+            results = []
+            for x in items:
+                results.append(f_wrapped(x, **func_kwargs))
+                pbar.update(1)
+            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")
+            ensure_ray(workers, pbar)
+
+            @ray.remote
+            def _task(x):
+                return f_wrapped(x, **func_kwargs)
+
+            refs = [_task.remote(x) for x in items]
+
+            results = []
+            for r in refs:
+                results.append(ray.get(r))
+                pbar.update(1)
+            return results
+
+        # ---- fastcore backend ----
+        if backend == "mp":
+            # Use threadpool instead of multiprocessing to avoid fork warnings
+            # in multi-threaded environments like pytest
+            results = parallel(f_wrapped, items, n_workers=workers, progress=progress, threadpool=True)
+            return list(results)
+        if backend == "threadpool":
+            results = parallel(f_wrapped, items, n_workers=workers, progress=progress, threadpool=True)
+            return list(results)
+        if backend == "safe":
+            # Completely safe backend for tests - no multiprocessing, no external progress bars
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                results = list(executor.map(f_wrapped, items))
+        raise ValueError(f"Unsupported backend: {backend!r}")
+
+
+
+def cleanup_phantom_workers():
+    """
+    Kill all child processes (phantom workers) without killing the Jupyter kernel itself.
+    Also lists non-daemon threads that remain.
+    """
+    parent = psutil.Process(os.getpid())
+    
+    # Kill only children, never the current process
+    for child in parent.children(recursive=True):
+        try:
+            print(f"🔪 Killing child process {child.pid} ({child.name()})")
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+    
+    # Report stray threads (can't hard-kill them in Python)
+    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("✅ Cleaned up child processes (kernel untouched).")
+
+# Usage: run this anytime after cancelling a cell
+

speedy_utils/multi_worker/thread.py

@@ -77,7 +77,9 @@
 # ============================================================================= #
 """
 
+import ctypes
 import os
+import threading
 import time
 import traceback
 from collections.abc import Callable, Iterable
@@ -98,6 +100,42 @@ DEFAULT_WORKERS = (os.cpu_count() or 4) * 2
 T = TypeVar("T")
 R = TypeVar("R")
 
+SPEEDY_RUNNING_THREADS: list[threading.Thread] = []
+_SPEEDY_THREADS_LOCK = threading.Lock()
+
+_PY_SET_ASYNC_EXC = ctypes.pythonapi.PyThreadState_SetAsyncExc
+try:
+    _PY_SET_ASYNC_EXC.argtypes = (ctypes.c_ulong, ctypes.py_object)  # type: ignore[attr-defined]
+    _PY_SET_ASYNC_EXC.restype = ctypes.c_int  # type: ignore[attr-defined]
+except AttributeError:  # pragma: no cover - platform specific
+    pass
+
+
+def _prune_dead_threads() -> None:
+    with _SPEEDY_THREADS_LOCK:
+        SPEEDY_RUNNING_THREADS[:] = [t for t in SPEEDY_RUNNING_THREADS if t.is_alive()]
+
+
+def _track_threads(threads: Iterable[threading.Thread]) -> None:
+    if not threads:
+        return
+    with _SPEEDY_THREADS_LOCK:
+        living = [t for t in SPEEDY_RUNNING_THREADS if t.is_alive()]
+        for candidate in threads:
+            if not candidate.is_alive():
+                continue
+            if any(existing is candidate for existing in living):
+                continue
+            living.append(candidate)
+        SPEEDY_RUNNING_THREADS[:] = living
+
+
+def _track_executor_threads(pool: ThreadPoolExecutor) -> None:
+    thread_set = getattr(pool, "_threads", None)
+    if not thread_set:
+        return
+    _track_threads(tuple(thread_set))
+
 
 def _group_iter(src: Iterable[T], size: int) -> Iterable[list[T]]:
     """Yield successive chunks from iterable of specified size."""
@@ -273,11 +311,13 @@ def multi_thread(
                 fut.idx = next_logical_idx  # type: ignore[attr-defined]
                 inflight.add(fut)
                 next_logical_idx += len(arg)
+                _track_executor_threads(pool)
             else:
                 fut = pool.submit(_worker, arg, func, fixed_kwargs)
                 fut.idx = next_logical_idx  # type: ignore[attr-defined]
                 inflight.add(fut)
                 next_logical_idx += 1
+                _track_executor_threads(pool)
 
         try:
             # Process futures as they complete and add new ones to keep the pool busy
@@ -347,11 +387,13 @@ def multi_thread(
                                 fut2.idx = next_logical_idx  # type: ignore[attr-defined]
                                 inflight.add(fut2)
                                 next_logical_idx += len(arg)
+                                _track_executor_threads(pool)
                             else:
                                 fut2 = pool.submit(_worker, arg, func, fixed_kwargs)
                                 fut2.idx = next_logical_idx  # type: ignore[attr-defined]
                                 inflight.add(fut2)
                                 next_logical_idx += 1
+                                _track_executor_threads(pool)
                     except StopIteration:
                         pass
 
@@ -370,6 +412,7 @@ def multi_thread(
                 bar.close()
     if store_output_pkl_file:
         dump_json_or_pickle(results, store_output_pkl_file)
+    _prune_dead_threads()
     return results
 
 
@@ -396,9 +439,58 @@ def multi_thread_standard(
         Results in same order as input items.
     """
     with ThreadPoolExecutor(max_workers=workers) as executor:
-        futures = [executor.submit(fn, item) for item in items]
+        futures = []
+        for item in items:
+            futures.append(executor.submit(fn, item))
+            _track_executor_threads(executor)
         results = [fut.result() for fut in futures]
+    _prune_dead_threads()
     return results
 
 
-__all__ = ["multi_thread", "multi_thread_standard"]
+def _async_raise(thread_id: int, exc_type: type[BaseException]) -> bool:
+    if thread_id <= 0:
+        return False
+    if not issubclass(exc_type, BaseException):
+        raise TypeError("exc_type must derive from BaseException")
+    res = _PY_SET_ASYNC_EXC(ctypes.c_ulong(thread_id), ctypes.py_object(exc_type))
+    if res == 0:
+        return False
+    if res > 1:  # pragma: no cover - defensive branch
+        _PY_SET_ASYNC_EXC(ctypes.c_ulong(thread_id), None)
+        raise SystemError("PyThreadState_SetAsyncExc failed")
+    return True
+
+
+def kill_all_thread(exc_type: type[BaseException] = SystemExit, join_timeout: float = 0.1) -> int:
+    """Forcefully stop tracked worker threads. Returns number of threads signalled."""
+    _prune_dead_threads()
+    current = threading.current_thread()
+    with _SPEEDY_THREADS_LOCK:
+        targets = [t for t in SPEEDY_RUNNING_THREADS if t.is_alive()]
+
+    terminated = 0
+    for thread in targets:
+        if thread is current:
+            continue
+        ident = thread.ident
+        if ident is None:
+            continue
+        try:
+            if _async_raise(ident, exc_type):
+                terminated += 1
+                thread.join(timeout=join_timeout)
+            else:
+                logger.warning("Unable to signal thread %s", thread.name)
+        except Exception as exc:  # pragma: no cover - defensive
+            logger.error("Failed to stop thread %s: %s", thread.name, exc)
+    _prune_dead_threads()
+    return terminated
+
+
+__all__ = [
+    "SPEEDY_RUNNING_THREADS",
+    "multi_thread",
+    "multi_thread_standard",
+    "kill_all_thread",
+]

{speedy_utils-1.1.17.dist-info → speedy_utils-1.1.19.dist-info}/METADATA

@@ -1,10 +1,14 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.17
+Version: 1.1.19
 Summary: Fast and easy-to-use package for data science
-Author: AnhVTH
-Author-email: anhvth.226@gmail.com
-Requires-Python: >=3.8
+Project-URL: Homepage, https://github.com/anhvth/speedy
+Project-URL: Repository, https://github.com/anhvth/speedy
+Author-email: AnhVTH <anhvth.226@gmail.com>
+License: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
@@ -12,29 +16,35 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.8
+Requires-Dist: aiohttp>=3.10.11
 Requires-Dist: bump2version
 Requires-Dist: cachetools
 Requires-Dist: debugpy
 Requires-Dist: fastcore
 Requires-Dist: fastprogress
-Requires-Dist: freezegun (>=1.5.1,<2.0.0)
+Requires-Dist: freezegun>=1.5.1
 Requires-Dist: ipdb
 Requires-Dist: ipywidgets
-Requires-Dist: json-repair (>=0.25.0,<0.31.0)
+Requires-Dist: json-repair<0.31.0,>=0.25.0
 Requires-Dist: jupyterlab
 Requires-Dist: loguru
 Requires-Dist: matplotlib
 Requires-Dist: numpy
-Requires-Dist: openai (>=1.106.0,<2.0.0)
-Requires-Dist: packaging (>=23.2,<25)
+Requires-Dist: openai>=1.106.0
+Requires-Dist: packaging<25,>=23.2
 Requires-Dist: pandas
 Requires-Dist: pydantic
+Requires-Dist: pytest>=8.3.5
+Requires-Dist: ray>=2.36.1
 Requires-Dist: requests
 Requires-Dist: scikit-learn
 Requires-Dist: tabulate
 Requires-Dist: tqdm
 Requires-Dist: xxhash
-Project-URL: Homepage, https://github.com/anhvth/speedy
+Provides-Extra: ray
+Requires-Dist: ray>=2.49.1; (python_version >= '3.9') and extra == 'ray'
 Description-Content-Type: text/markdown
 
 # Speedy Utils
@@ -83,6 +93,19 @@ cd speedy-utils
 pip install .
 ```
 
+### Extras
+
+Optional dependencies can be installed via extras. For the `ray` backend
+support (requires Python >= 3.9):
+
+```bash
+# pip
+pip install 'speedy-utils[ray]'
+
+# Poetry (for developing this repo)
+poetry install -E ray
+```
+
 ## Updating from previous versions
 
 To update from previous versions or switch to v1.x, first uninstall any old
@@ -281,9 +304,8 @@ python speedy_utils/common/dataclass_parser.py
 
 Example output:
 
-| Field              | Value                                 |
-|--------------------|---------------------------------------|
-| from_peft          | ./outputs/llm_hn_qw32b/hn_results_r3/ |
+| Field     | Value                                 |
+| --------- | ------------------------------------- |
+| from_peft | ./outputs/llm_hn_qw32b/hn_results_r3/ |
 
 Please ensure your code adheres to the project's coding standards and includes appropriate tests.
-