speedy-utils 1.1.45__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/kill_mpython.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+"""Script to kill all tmux sessions matching 'mpython*' pattern."""
+
+import subprocess
+import sys
+
+
+def main():
+    """Kill all tmux sessions with names starting with 'mpython'."""
+    try:
+        # Get list of tmux sessions matching the pattern
+        result = subprocess.run(
+            ["tmux", "ls"],
+            capture_output=True,
+            text=True
+        )
+        if result.returncode != 0:
+            if "no server running" in result.stderr.lower():
+                print("No tmux server running.")
+                return
+            print(f"Error listing tmux sessions: {result.stderr}")
+            sys.exit(result.returncode)
+
+        sessions = []
+        for line in result.stdout.strip().split('\n'):
+            if line.strip():
+                session_name = line.split(':')[0]
+                if session_name.startswith('mpython'):
+                    sessions.append(session_name)
+
+        if not sessions:
+            print("No tmux sessions found matching 'mpython*'")
+            return
+
+        print(f"Found {len(sessions)} tmux session(s) to kill: {', '.join(sessions)}")
+
+        # Kill each session
+        for session in sessions:
+            kill_result = subprocess.run(
+                ["tmux", "kill-session", "-t", session],
+                capture_output=True,
+                text=True
+            )
+            if kill_result.returncode == 0:
+                print(f"Successfully killed tmux session '{session}'")
+            else:
+                print(f"Error killing tmux session '{session}': {kill_result.stderr}")
+
+    except FileNotFoundError:
+        print("Error: tmux command not found. Please ensure tmux is installed.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

speedy_utils/scripts/mpython.py

@@ -7,6 +7,20 @@ import re
 import shlex  # To properly escape command line arguments
 import shutil
 import subprocess
+import sys
+
+
+try:
+    from rich.console import Console, Group
+    from rich.panel import Panel
+    from rich.syntax import Syntax
+    from rich.text import Text
+except ImportError:
+    Console = None
+    Group = None
+    Panel = None
+    Text = None
+    Syntax = None
 
 
 taskset_path = shutil.which('taskset')
@@ -60,33 +74,55 @@ def assert_script(python_path):
     with open(python_path) as f:
         code_str = f.read()
     if 'MP_ID' not in code_str or 'MP_TOTAL' not in code_str:
-        example_code = (
-            'import os; MP_TOTAL = int(os.environ.get("MP_TOTAL"));MP_ID = int(os.environ.get("MP_ID"))\n'
+        helper_code = (
+            'import os\n'
+            'MP_ID = int(os.getenv("MP_ID", "0"))\n'
+            'MP_TOTAL = int(os.getenv("MP_TOTAL", "1"))\n'
             'inputs = list(inputs[MP_ID::MP_TOTAL])'
         )
-        # ANSI escape codes for coloring
-        YELLOW = '\033[93m'
-        RESET = '\033[0m'
-        raise_msg = (
-            f'MP_ID and MP_TOTAL not found in {python_path}, please add them.\n\n'
-            f'Example:\n{YELLOW}{example_code}{RESET}'
-        )
-        raise Exception(raise_msg)
+        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)
+            console.print()
+            console.print(
+                Panel(
+                    f'Your script {python_path} is missing MP_ID and/or MP_TOTAL variables.\n\n'
+                    f'Add the following code to enable multi-process sharding:',
+                    title='[bold yellow]Warning: Missing Multi-Process Variables[/bold yellow]',
+                    border_style='yellow',
+                    expand=False,
+                )
+            )
+            console.print()
+            console.print('```python')
+            console.print(syntax)
+            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'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')
@@ -96,7 +132,18 @@ def run_in_tmux(commands_to_run, tmux_name, num_windows):
 def main():
     # Assert that MP_ID and MP_TOTAL are not already set
 
-    parser = argparse.ArgumentParser(description='Process fold arguments')
+    helper_code = (
+        'import os\n'
+        'MP_ID = int(os.getenv("MP_ID", "0"))\n'
+        'MP_TOTAL = int(os.getenv("MP_TOTAL", "1"))\n'
+        'inputs = list(inputs[MP_ID::MP_TOTAL])'
+    )
+
+    parser = argparse.ArgumentParser(
+        description='Process fold arguments',
+        epilog=f'Helper code for multi-process sharding:\n{helper_code}',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
     parser.add_argument(
         '--total_fold', '-t', default=16, type=int, help='total number of folds'
     )

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.45
+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.45.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,20 +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/mpython.py,sha256=uJ-QPG43cyHUptMP4MnyH7IdSTh4mzVQAejy9o1cQoE,5068
+speedy_utils/scripts/kill_mpython.py,sha256=xYdJIs1KQ4zlyj-E4XN6j1T7dOhzVo65uKzjXmIvUDU,1847
+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.45.dist-info/METADATA,sha256=yXr9vEuXiRpvZ3VibGULqr3X1a832aQZJAb-lJLM6mM,13073
-speedy_utils-1.1.45.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-speedy_utils-1.1.45.dist-info/entry_points.txt,sha256=rwn89AYfBUh9SRJtFbpp-u2JIKiqmZ2sczvqyO6s9cI,289
-speedy_utils-1.1.45.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,,

{speedy_utils-1.1.45.dist-info → speedy_utils-1.1.47.dist-info}/entry_points.txt

@@ -1,5 +1,6 @@
 [console_scripts]
 fast-vllm = llm_utils.scripts.fast_vllm:main
+kill-mpython = speedy_utils.scripts.kill_mpython:main
 mpython = speedy_utils.scripts.mpython:main
 openapi_client_codegen = speedy_utils.scripts.openapi_client_codegen:main
 svllm = llm_utils.scripts.vllm_serve:main