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

llm_utils/lm/llm.py

@@ -16,6 +16,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from speedy_utils.common.utils_io import jdumps
+from speedy_utils import clean_traceback
 
 from .base_prompt_builder import BasePromptBuilder
 from .mixins import (
@@ -159,6 +160,7 @@ class LLM(
         messages.append({'role': 'user', 'content': user_content})
         return cast(Messages, messages)
 
+    @clean_traceback
     def text_completion(
         self, input_data: str | BaseModel | list[dict], **runtime_kwargs
     ) -> list[dict[str, Any]]:
@@ -214,6 +216,7 @@ class LLM(
             results.append(result_dict)
         return results
 
+    @clean_traceback
     def pydantic_parse(
         self,
         input_data: str | BaseModel | list[dict],

speedy_utils/__init__.py

@@ -58,6 +58,9 @@ from .common.utils_print import (
     fprint,
 )
 
+# Error handling utilities
+from .common.utils_error import clean_traceback, handle_exceptions_with_clean_traceback
+
 # Multi-worker processing
 from .multi_worker.process import multi_process
 from .multi_worker.thread import kill_all_thread, multi_thread
@@ -156,6 +159,9 @@ __all__ = [
     'print_table',
     'setup_logger',
     'log',
+    # Error handling utilities
+    'clean_traceback',
+    'handle_exceptions_with_clean_traceback',
     # Multi-worker processing
     'multi_process',
     'multi_thread',

speedy_utils/common/utils_error.py

@@ -0,0 +1,261 @@
+"""
+Error handling utilities for clean, user-focused tracebacks.
+"""
+
+import inspect
+import linecache
+import sys
+import traceback
+from typing import Any, Callable, TypeVar
+
+try:
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.text import Text
+except ImportError:
+    Console = None
+    Panel = None
+    Text = None
+
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+class CleanTracebackError(Exception):
+    """Exception wrapper that provides clean, user-focused tracebacks."""
+
+    def __init__(
+        self,
+        original_exception: Exception,
+        user_traceback: list[traceback.FrameSummary],
+        caller_frame: traceback.FrameSummary | None = None,
+        func_name: str | None = None,
+    ) -> None:
+        self.original_exception = original_exception
+        self.user_traceback = user_traceback
+        self.caller_frame = caller_frame
+        self.func_name = func_name
+
+        # Create a focused error message
+        tb_str = ''.join(traceback.format_list(user_traceback))
+        func_part = f' in {func_name}' if func_name else ''
+        msg = (
+            f'Error{func_part}:\n'
+            f'\nUser code traceback:\n{tb_str}'
+            f'{type(original_exception).__name__}: {original_exception}'
+        )
+        super().__init__(msg)
+
+    def __str__(self) -> str:
+        return super().__str__()
+
+    def format_rich(self) -> None:
+        """Format and print error with rich panels and code context."""
+        if Console is None or Panel is None or Text is None:
+            print(str(self), file=sys.stderr)
+            return
+
+        console = Console(stderr=True, force_terminal=True)
+
+        # Build traceback display with code context
+        tb_parts: list[str] = []
+
+        # Show caller frame first if available
+        if self.caller_frame and self.caller_frame.lineno is not None:
+            tb_parts.append(
+                f'[cyan]{self.caller_frame.filename}[/cyan]:[yellow]{self.caller_frame.lineno}[/yellow] '
+                f'in [green]{self.caller_frame.name}[/green]'
+            )
+            tb_parts.append('')
+            context = _get_code_context_rich(self.caller_frame.filename, self.caller_frame.lineno, 3)
+            tb_parts.extend(context)
+            tb_parts.append('')
+
+        # Show user code frames with context
+        for frame in self.user_traceback:
+            if frame.lineno is not None:
+                func_name = f' {self.func_name}' if self.func_name else ''
+                tb_parts.append(
+                    f'[cyan]{frame.filename}[/cyan]:[yellow]{frame.lineno}[/yellow] '
+                    f'in [green]{frame.name}{func_name}[/green]'
+                )
+                tb_parts.append('')
+                context = _get_code_context_rich(frame.filename, frame.lineno, 3)
+                tb_parts.extend(context)
+                tb_parts.append('')
+
+        # Print with rich Panel
+        console.print()
+        console.print(
+            Panel(
+                '\n'.join(tb_parts),
+                title='[bold red]Traceback (most recent call last)[/bold red]',
+                border_style='red',
+                expand=False,
+            )
+        )
+        console.print(
+            f'[bold red]{type(self.original_exception).__name__}[/bold red]: '
+            f'{self.original_exception}'
+        )
+        console.print()
+
+
+def _get_code_context(filename: str, lineno: int, context_lines: int = 3) -> list[str]:
+    """Get code context around a line with line numbers and highlighting."""
+    lines: list[str] = []
+    start = max(1, lineno - context_lines)
+    end = lineno + context_lines
+
+    for i in range(start, end + 1):
+        line = linecache.getline(filename, i)
+        if not line:
+            continue
+        line = line.rstrip()
+        marker = '❱' if i == lineno else ' '
+        lines.append(f'  {i:4d} {marker} {line}')
+
+    return lines
+
+
+def _get_code_context_rich(filename: str, lineno: int, context_lines: int = 3) -> list[str]:
+    """Get code context with rich formatting (colors)."""
+    lines: list[str] = []
+    start = max(1, lineno - context_lines)
+    end = lineno + context_lines
+
+    for i in range(start, end + 1):
+        line = linecache.getline(filename, i)
+        if not line:
+            continue
+        line = line.rstrip()
+        num_str = f'{i:4d}'
+
+        if i == lineno:
+            # Highlight error line
+            lines.append(f'[dim]{num_str}[/dim] [red]❱[/red] {line}')
+        else:
+            # Normal context line
+            lines.append(f'[dim]{num_str} │[/dim] {line}')
+
+    return lines
+
+
+def _filter_traceback_frames(tb_list: list[traceback.FrameSummary]) -> list[traceback.FrameSummary]:
+    """Filter traceback frames to show only user code."""
+    user_frames = []
+    skip_patterns = [
+        'site-packages/',
+        'dist-packages/',
+        'python3.',
+        'lib/python',
+        'concurrent/futures/',
+        'threading.py',
+        'multiprocessing/',
+        'urllib/',
+        'httpx/',
+        'httpcore/',
+        'openai/',
+        'requests/',
+        'aiohttp/',
+        'urllib3/',
+        'speedy_utils/common/',
+        'speedy_utils/multi_worker/',
+        'llm_utils/lm/',
+        'llm_utils/chat_format/',
+        'llm_utils/vector_cache/',
+    ]
+
+    skip_functions = [
+        'wrapper',  # Our decorator wrapper
+        '__call__',  # Internal calls
+        '__inner_call__',  # Internal calls
+        '_worker',  # Multi-thread worker
+        '_run_batch',  # Batch runner
+    ]
+
+    for frame in tb_list:
+        # Skip frames matching skip patterns
+        if any(pattern in frame.filename for pattern in skip_patterns):
+            continue
+        # Skip specific function names
+        if frame.name in skip_functions:
+            continue
+        # Skip frames that are too deep in the call stack (internal implementation)
+        if 'speedy_utils' in frame.filename and any(
+            name in frame.name for name in ['__inner_call__', '_worker', '_run_batch']
+        ):
+            continue
+        user_frames.append(frame)
+
+    # If no user frames found, keep frames that are likely user code (not deep library internals)
+    if not user_frames:
+        for frame in reversed(tb_list):
+            # Keep if it's not in site-packages and not in standard library
+            if ('site-packages/' not in frame.filename and 
+                'dist-packages/' not in frame.filename and
+                not frame.filename.startswith('/usr/') and
+                not frame.filename.startswith('/opt/') and
+                'python3.' not in frame.filename and
+                frame.name not in skip_functions):
+                user_frames.append(frame)
+                if len(user_frames) >= 5:  # Keep up to 5 frames
+                    break
+        user_frames.reverse()
+
+    return user_frames
+
+
+def clean_traceback(func: F) -> F:
+    """Decorator to wrap function calls with clean traceback handling."""
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as exc:
+            # Get the current traceback
+            exc_tb = sys.exc_info()[2]
+
+            if exc_tb is not None:
+                tb_list = traceback.extract_tb(exc_tb)
+
+                # Filter to keep only user code frames
+                user_frames = _filter_traceback_frames(tb_list)
+
+                # Get caller frame - walk up the stack to find the original caller
+                caller_context = None
+                frame = inspect.currentframe()
+                while frame:
+                    frame = frame.f_back
+                    if frame and frame.f_code.co_name not in ['wrapper', '__call__', '__inner_call__']:
+                        caller_info = inspect.getframeinfo(frame)
+                        if not any(skip in caller_info.filename for skip in [
+                            'speedy_utils/common/', 'speedy_utils/multi_worker/', 
+                            'llm_utils/lm/', 'site-packages/', 'dist-packages/'
+                        ]):
+                            caller_context = traceback.FrameSummary(
+                                caller_info.filename,
+                                caller_info.lineno,
+                                caller_info.function,
+                            )
+                            break
+
+                # If we have user frames, create and format our custom exception
+                if user_frames:
+                    func_name = getattr(func, '__name__', repr(func))
+                    clean_error = CleanTracebackError(
+                        exc,
+                        user_frames,
+                        caller_context,
+                        func_name,
+                    )
+                    clean_error.format_rich()
+                    sys.exit(1)  # Exit after formatting
+
+            # Fallback: re-raise original if we couldn't extract frames
+            raise
+
+    return wrapper  # type: ignore
+
+
+def handle_exceptions_with_clean_traceback(func: F) -> F:
+    """Alias for clean_traceback decorator."""
+    return clean_traceback(func)

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,19 @@ 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.text import Text
+    from rich.syntax import Syntax
+except ImportError:
+    Console = None
+    Group = None
+    Panel = None
+    Text = None
+    Syntax = None
 
 
 taskset_path = shutil.which('taskset')
@@ -60,18 +73,34 @@ 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):
@@ -96,7 +125,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.44.dist-info → speedy_utils-1.1.46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.44
+Version: 1.1.46
 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
@@ -61,6 +61,7 @@ Description-Content-Type: text/markdown
 ## 🚀 Recent Updates (January 27, 2026)
 
 **Enhanced Error Handling in Parallel Processing:**
+
 - Rich-formatted error tracebacks with code context and syntax highlighting
 - Three error handling modes: 'raise', 'ignore', and 'log'
 - Filtered tracebacks focusing on user code (hiding infrastructure)
@@ -261,6 +262,7 @@ print(results)  # [0, 1, 4, 9, 16, None, 36, 49, 64, 81, None, 121]
 #### Rich Error Tracebacks
 
 When errors occur, you'll see beautifully formatted tracebacks with:
+
 - **Code context**: Lines of code around the error location
 - **Caller information**: Shows where the parallel function was invoked
 - **Filtered frames**: Focuses on user code, hiding infrastructure details
@@ -282,7 +284,7 @@ When using `error_handler='log'`, errors are automatically saved to timestamped
 
 Progress bars now show real-time error and success counts:
 
-```
+```bash
 Multi-thread [8/10] [00:02<00:00, 3.45it/s, success=8, errors=2, pending=0]
 ```
 

{speedy_utils-1.1.44.dist-info → speedy_utils-1.1.46.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=4nTBNvIZfsKUZo5SdllLQOp3w-HohR3B1hcfUnYDl7A,20014
+llm_utils/lm/llm.py,sha256=i6L5aKF6NhzmaFPBA2pCm8TkQmS1nCgORMqP5QyfJ28,20097
 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=1ubAYR6P0cEZJLDt7KQLLxl6ylh-T7WE7HPP94-rVLI,3045
+speedy_utils/__init__.py,sha256=_kSjS816Kv5UZPd4EM_juB68tXM_sHUYt6OFB-RhE6U,3261
 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
@@ -39,6 +39,7 @@ speedy_utils/common/notebook_utils.py,sha256=6mxXZcTHwYob3nobAzbSZnDyXRahFaaSko1
 speedy_utils/common/patcher.py,sha256=Rku-N4DJNue8BCLUx7y3ad_3t_WU2HleHKlbR0vhaRc,2319
 speedy_utils/common/report_manager.py,sha256=78KQ0gUzvbzT6EjHYZ5zgtV41cPRvdX8hnF2oSWA4qA,3849
 speedy_utils/common/utils_cache.py,sha256=1UAqOSb4nFVlhuQRfTEXCN-8Wf6yntXyMA6yp61-98I,26277
+speedy_utils/common/utils_error.py,sha256=KQx2JTZsvsX2DsKRIoVR-4rc-6-l3OzEz9UtnHt8HJg,9108
 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
@@ -49,13 +50,14 @@ speedy_utils/multi_worker/process.py,sha256=U-pjHoWZ3xOeplMl2nSxVeiJE0F9V-eswpSd
 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/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=bf_pdCm5_NaDUV0J652Cg_YFIwva_5AhoVKTGOowENo,6485
 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.44.dist-info/METADATA,sha256=Y5wQ_VbeiPTcqUkOCEbEvOKZ-wnwKCF2nvGhSnU3AJs,13067
-speedy_utils-1.1.44.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-speedy_utils-1.1.44.dist-info/entry_points.txt,sha256=rwn89AYfBUh9SRJtFbpp-u2JIKiqmZ2sczvqyO6s9cI,289
-speedy_utils-1.1.44.dist-info/RECORD,,
+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.44.dist-info → speedy_utils-1.1.46.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