speedy-utils 1.1.45__tar.gz → 1.1.46__tar.gz

{speedy_utils-1.1.45 → speedy_utils-1.1.46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.45
+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

{speedy_utils-1.1.45 → speedy_utils-1.1.46}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "speedy-utils"
-version = "1.1.45"
+version = "1.1.46"
 description = "Fast and easy-to-use package for data science"
 authors = [{ name = "AnhVTH", email = "anhvth.226@gmail.com" }]
 readme = "README.md"
@@ -65,6 +65,7 @@ svllm = "llm_utils.scripts.vllm_serve:main"
 svllm-lb = "llm_utils.scripts.vllm_load_balancer:run_load_balancer"
 openapi_client_codegen = "speedy_utils.scripts.openapi_client_codegen:main"
 fast-vllm = "llm_utils.scripts.fast_vllm:main"
+kill-mpython = "speedy_utils.scripts.kill_mpython:main"
 
 [build-system]
 requires = ["hatchling"]

speedy_utils-1.1.46/src/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-1.1.45 → speedy_utils-1.1.46}/src/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.46/test.py

@@ -0,0 +1,8 @@
+import os                                                                                                                                                                                   
+MP_ID = int(os.getenv("MP_ID", "0"))                                                                                                                                                        
+MP_TOTAL = int(os.getenv("MP_TOTAL", "1"))                                                                                                                                                  
+
+inputs = range(1000)
+def f(x):
+    print(x)
+inputs = list(inputs[MP_ID::MP_TOTAL])