speedy-utils 1.1.46__py3-none-any.whl → 1.1.47__py3-none-any.whl

speedy_utils/multi_worker/progress.py

@@ -4,12 +4,15 @@ 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.
+
+Users can simply call report_progress(n) from within their worker functions,
+and speedy_utils will automatically route it to the centralized tracker.
 """
 import time
 import threading
 from typing import Optional, Callable
 
-__all__ = ['ProgressActor', 'create_progress_tracker', 'get_ray_progress_actor']
+__all__ = ['ProgressActor', 'create_progress_tracker', 'get_ray_progress_actor', 'report_progress', 'set_progress_context']
 
 
 def get_ray_progress_actor():
@@ -119,6 +122,73 @@ class ProgressPoller:
         if self._thread:
             self._thread.join(timeout=2.0)
     
+    def _poll_loop(self):
+        """Background loop that polls progress and updates tqdm bar."""
+        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.refresh()
+            except Exception:
+                pass
+            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
+
+
+# Thread-local storage for progress context
+_progress_context = threading.local()
+
+
+def set_progress_context(progress_actor):
+    """
+    Set the progress actor in thread-local storage.
+    
+    This is called automatically by speedy_utils when setting up worker context.
+    Users don't need to call this directly.
+    
+    Args:
+        progress_actor: Ray actor handle for progress tracking
+    """
+    _progress_context.actor = progress_actor
+
+
+def report_progress(n: int = 1):
+    """
+    Report progress to the centralized tracker.
+    
+    This is a simple helper that workers can call to report items processed.
+    Works automatically when called from within multi_process with Ray backend.
+    Safe to call even when no progress tracking is active (no-op).
+    
+    Args:
+        n: Number of items processed
+        
+    Example:
+        def process_data(items):
+            for i, item in enumerate(items):
+                result = expensive_operation(item)
+                
+                # Report progress every 100 items
+                if i % 100 == 0:
+                    report_progress(100)
+            
+            return results
+    """
+    actor = getattr(_progress_context, 'actor', None)
+    if actor is not None:
+        try:
+            import ray
+            ray.get(actor.update.remote(n))
+        except Exception:
+            # Silently ignore progress reporting errors
+            pass
     def _poll_loop(self):
         """Poll the progress actor and update tqdm."""
         while not self._stop_event.is_set():

speedy_utils/multi_worker/thread.py

@@ -26,6 +26,9 @@ DEFAULT_WORKERS = (os.cpu_count() or 4) * 2
 T = TypeVar('T')
 R = TypeVar('R')
 
+if TYPE_CHECKING:
+    from typing import Literal, overload
+
 SPEEDY_RUNNING_THREADS: list[threading.Thread] = []  # cooperative shutdown tracking
 _SPEEDY_THREADS_LOCK = threading.Lock()
 
@@ -342,6 +345,48 @@ def _cancel_futures(inflight: set[Future[Any]]) -> None:
 # ────────────────────────────────────────────────────────────
 # main API
 # ────────────────────────────────────────────────────────────
+if TYPE_CHECKING:
+    @overload
+    def multi_thread(
+        func: Callable[[T], R],
+        inputs: Iterable[T],
+        *,
+        workers: int | None = ...,
+        batch: int = ...,
+        ordered: bool = ...,
+        progress: bool = ...,
+        progress_update: int = ...,
+        prefetch_factor: int = ...,
+        timeout: float | None = ...,
+        stop_on_error: bool | None = ...,
+        error_handler: Literal['raise'] = ...,
+        max_error_files: int = ...,
+        n_proc: int = ...,
+        store_output_pkl_file: str | None = ...,
+        **fixed_kwargs: Any,
+    ) -> list[R]: ...
+
+    @overload
+    def multi_thread(
+        func: Callable[[T], R],
+        inputs: Iterable[T],
+        *,
+        workers: int | None = ...,
+        batch: int = ...,
+        ordered: bool = ...,
+        progress: bool = ...,
+        progress_update: int = ...,
+        prefetch_factor: int = ...,
+        timeout: float | None = ...,
+        stop_on_error: bool | None = ...,
+        error_handler: Literal['ignore', 'log'] = ...,
+        max_error_files: int = ...,
+        n_proc: int = ...,
+        store_output_pkl_file: str | None = ...,
+        **fixed_kwargs: Any,
+    ) -> list[R | None]: ...
+
+
 def multi_thread(
     func: Callable[[T], R],
     inputs: Iterable[T],

speedy_utils/scripts/mpython.py

@@ -9,11 +9,12 @@ import shutil
 import subprocess
 import sys
 
+
 try:
     from rich.console import Console, Group
     from rich.panel import Panel
-    from rich.text import Text
     from rich.syntax import Syntax
+    from rich.text import Text
 except ImportError:
     Console = None
     Group = None
@@ -81,7 +82,7 @@ def assert_script(python_path):
         )
         if Console and Panel and Text and Syntax and Group:
             console = Console(stderr=True, force_terminal=True)
-            syntax = Syntax(helper_code, "python", theme="monokai", line_numbers=False)
+            syntax = Syntax(helper_code, 'python', theme='monokai', line_numbers=False)
             console.print()
             console.print(
                 Panel(
@@ -93,29 +94,35 @@ def assert_script(python_path):
                 )
             )
             console.print()
-            console.print("```python")
+            console.print('```python')
             console.print(syntax)
-            console.print("```")
-            console.print("-"*80)
+            console.print('```')
+            console.print('-' * 80)
         else:
             # Fallback to plain text
-            print(f'Warning: MP_ID and MP_TOTAL not found in {python_path}, please add them.', file=sys.stderr)
+            print(
+                f'Warning: MP_ID and MP_TOTAL not found in {python_path}, please add them.',
+                file=sys.stderr,
+            )
             print(f'Example:\n{helper_code}', file=sys.stderr)
 
 
 def run_in_tmux(commands_to_run, tmux_name, num_windows):
     with open('/tmp/start_multirun_tmux.sh', 'w') as script_file:
         script_file.write('#!/bin/bash\n\n')
-        script_file.write(f'tmux new-session -d -s {tmux_name}\n')
+        # Use -n to name the first window explicitly (avoids base-index issues)
+        script_file.write(f'tmux new-session -d -s {tmux_name} -n w0\n')
         for i, cmd in enumerate(itertools.cycle(commands_to_run)):
             if i >= num_windows:
                 break
-            window_name = f'{tmux_name}:{i}'
+            window_name = f'w{i}'
+            target = f'{tmux_name}:{window_name}'
             if i == 0:
-                script_file.write(f"tmux send-keys -t {window_name} '{cmd}' C-m\n")
+                script_file.write(f"tmux send-keys -t {target} '{cmd}' C-m\n")
             else:
-                script_file.write(f'tmux new-window -t {tmux_name}\n')
-                script_file.write(f"tmux send-keys -t {window_name} '{cmd}' C-m\n")
+                # Create window with explicit name, then send keys to it
+                script_file.write(f'tmux new-window -t {tmux_name} -n {window_name}\n')
+                script_file.write(f"tmux send-keys -t {target} '{cmd}' C-m\n")
 
         # Make the script executable
         script_file.write('chmod +x /tmp/start_multirun_tmux.sh\n')
@@ -135,7 +142,7 @@ def main():
     parser = argparse.ArgumentParser(
         description='Process fold arguments',
         epilog=f'Helper code for multi-process sharding:\n{helper_code}',
-        formatter_class=argparse.RawDescriptionHelpFormatter
+        formatter_class=argparse.RawDescriptionHelpFormatter,
     )
     parser.add_argument(
         '--total_fold', '-t', default=16, type=int, help='total number of folds'

{speedy_utils-1.1.46.dist-info → speedy_utils-1.1.47.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.46
+Version: 1.1.47
 Summary: Fast and easy-to-use package for data science
 Project-URL: Homepage, https://github.com/anhvth/speedy
 Project-URL: Repository, https://github.com/anhvth/speedy

{speedy_utils-1.1.46.dist-info → speedy_utils-1.1.47.dist-info}/RECORD

@@ -7,7 +7,7 @@ llm_utils/chat_format/transform.py,sha256=PJ2g9KT1GSbWuAs7giEbTpTAffpU9QsIXyRlbf
 llm_utils/chat_format/utils.py,sha256=M2EctZ6NeHXqFYufh26Y3CpSphN0bdZm5xoNaEJj5vg,1251
 llm_utils/lm/__init__.py,sha256=4jYMy3wPH3tg-tHFyWEWOqrnmX4Tu32VZCdzRGMGQsI,778
 llm_utils/lm/base_prompt_builder.py,sha256=_TzYMsWr-SsbA_JNXptUVN56lV5RfgWWTrFi-E8LMy4,12337
-llm_utils/lm/llm.py,sha256=i6L5aKF6NhzmaFPBA2pCm8TkQmS1nCgORMqP5QyfJ28,20097
+llm_utils/lm/llm.py,sha256=KEtrHq5D8ZkeD4iTc_zPgOtgVXpsCRA-A3fmWgYCz0w,21378
 llm_utils/lm/llm_signature.py,sha256=vV8uZgLLd6ZKqWbq0OPywWvXAfl7hrJQnbtBF-VnZRU,1244
 llm_utils/lm/lm_base.py,sha256=Bk3q34KrcCK_bC4Ryxbc3KqkiPL39zuVZaBQ1i6wJqs,9437
 llm_utils/lm/mixins.py,sha256=Nz7CwJFBOvbZNbODUlJC04Pcbac3zWnT8vy7sZG_MVI,24906
@@ -30,7 +30,7 @@ llm_utils/vector_cache/core.py,sha256=VXuYJy1AX22NHKvIXRriETip5RrmQcNp73-g-ZT774
 llm_utils/vector_cache/types.py,sha256=CpMZanJSTeBVxQSqjBq6pBVWp7u2-JRcgY9t5jhykdQ,438
 llm_utils/vector_cache/utils.py,sha256=OsiRFydv8i8HiJtPL9hh40aUv8I5pYfg2zvmtDi4DME,1446
 speedy_utils/__imports.py,sha256=V0YzkDK4-QkK_IDXY1be6C6_STuNhXAKIp4_dM0coQs,7800
-speedy_utils/__init__.py,sha256=_kSjS816Kv5UZPd4EM_juB68tXM_sHUYt6OFB-RhE6U,3261
+speedy_utils/__init__.py,sha256=4RLsC0s9sAkGWcwC85eQQ1HT2Ea_flhbBQUCJuzJAmk,3429
 speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 speedy_utils/common/clock.py,sha256=raLtMGIgzrRej5kUt7hOUm2ZZw2THVPo-q8dMvdZOxw,7354
 speedy_utils/common/function_decorator.py,sha256=GKXqRs_hHFFmhyhql0Br0o52WzekUnpNlm99NfaVwgY,2025
@@ -43,21 +43,25 @@ speedy_utils/common/utils_error.py,sha256=KQx2JTZsvsX2DsKRIoVR-4rc-6-l3OzEz9UtnH
 speedy_utils/common/utils_io.py,sha256=94m_EZ2eIs3w2m0rx-QQWsREPpVJctpweYHco3byczQ,15876
 speedy_utils/common/utils_misc.py,sha256=ZRJCS7OJxybpVm1sasoeCYRW2TaaGCXj4DySYlQeVR8,2227
 speedy_utils/common/utils_print.py,sha256=AGDB7mgJnO00QkJBH6kJb46738q3GzMUZPwtQ248vQw,4763
-speedy_utils/multi_worker/__init__.py,sha256=urcuxzaAJp-Rl3SIwHNre3x2vyHxLR7YGiDdm-Q8GQs,361
+speedy_utils/multi_worker/__init__.py,sha256=JzPNJtYgBgSlAG6yBU0HhFcEwwCrH3iWeo-W7HZ9Pxs,503
+speedy_utils/multi_worker/_multi_process.py,sha256=0ceoXZZAdfSVVqZAgtCLw7Oq_ZLh6y0i6OkPgx6cicc,13343
+speedy_utils/multi_worker/_multi_process_ray.py,sha256=Qri-jbfHMLxUyiV5I2-KehTj1XM2YxagXWMesFkyzBg,9868
+speedy_utils/multi_worker/common.py,sha256=e5FVKtPGrIw4s4-sA_SIzaIFaTeVK598jXBC5U8YcLU,27014
 speedy_utils/multi_worker/dataset_ray.py,sha256=U_l_4Y7CVpaHiApsXQSdNvals8NK87LHPS_XHiJF3qs,10044
+speedy_utils/multi_worker/dataset_sharding.py,sha256=dWGCQiElb8SvLHfh97xT-QA2csmf2W3s8S4cFjZrO9Y,7196
 speedy_utils/multi_worker/parallel_gpu_pool.py,sha256=A7llZcQbRVZqwCqNRku7TpqGCdSoIzpdcTaupgqT5nI,6108
-speedy_utils/multi_worker/process.py,sha256=U-pjHoWZ3xOeplMl2nSxVeiJE0F9V-eswpSdK-8c3dU,45446
-speedy_utils/multi_worker/progress.py,sha256=Ozeca-t-j1224n_dWwZkWzva9DC16SCLgScKeGtXLaQ,4717
-speedy_utils/multi_worker/thread.py,sha256=E7o_iUCIKmgk1tFt7mZAFT7c5q229wVzWj-trmVsxVA,27254
+speedy_utils/multi_worker/process.py,sha256=1P-7J93jt51lMfpHLD9F8DDKAdeMtyYIFVOKEiN5nok,3927
+speedy_utils/multi_worker/progress.py,sha256=V-BN4GSmwOwL_8Y0J7Ya-byJUQXvuYcdqkmQqGHDaVw,7004
+speedy_utils/multi_worker/thread.py,sha256=iVjg4SI3GjzA5DhasT0w29U61fKnWXoLqCgoqQPzyBA,28528
 speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 speedy_utils/scripts/kill_mpython.py,sha256=xYdJIs1KQ4zlyj-E4XN6j1T7dOhzVo65uKzjXmIvUDU,1847
-speedy_utils/scripts/mpython.py,sha256=bf_pdCm5_NaDUV0J652Cg_YFIwva_5AhoVKTGOowENo,6485
+speedy_utils/scripts/mpython.py,sha256=b_Zbytrht2oYfV45P2nl7u5gNsYp4qQ868Pz8QReSD4,6741
 speedy_utils/scripts/openapi_client_codegen.py,sha256=GModmmhkvGnxljK4KczyixKDrk-VEcLaW5I0XT6tzWo,9657
 vision_utils/README.md,sha256=AIDZZj8jo_QNrEjFyHwd00iOO431s-js-M2dLtVTn3I,5740
 vision_utils/__init__.py,sha256=hF54sT6FAxby8kDVhOvruy4yot8O-Ateey5n96O1pQM,284
 vision_utils/io_utils.py,sha256=pI0Va6miesBysJcllK6NXCay8HpGZsaMWwlsKB2DMgA,26510
 vision_utils/plot.py,sha256=HkNj3osA3moPuupP1VguXfPPOW614dZO5tvC-EFKpKM,12028
-speedy_utils-1.1.46.dist-info/METADATA,sha256=nIgSDMrAbHssvJznZtozJZbZVdH4K8dALKRdExipi1g,13073
-speedy_utils-1.1.46.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-speedy_utils-1.1.46.dist-info/entry_points.txt,sha256=QY_2Vn6IcPCaqlY74pDRyZ6UTvPilaNPT7Gxijj7XI8,343
-speedy_utils-1.1.46.dist-info/RECORD,,
+speedy_utils-1.1.47.dist-info/METADATA,sha256=Ol3PRn5VGCiozaLocWrOLiYdO6nTHoEmBYhWODVfIi4,13073
+speedy_utils-1.1.47.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+speedy_utils-1.1.47.dist-info/entry_points.txt,sha256=QY_2Vn6IcPCaqlY74pDRyZ6UTvPilaNPT7Gxijj7XI8,343
+speedy_utils-1.1.47.dist-info/RECORD,,