speedy-utils 1.1.21__py3-none-any.whl → 1.1.23__py3-none-any.whl

llm_utils/lm/utils.py

@@ -1,123 +1,345 @@
-import fcntl
 import os
-import tempfile
+import signal
 import time
-from typing import List, Dict
-import numpy as np
+from typing import Any, List, Optional, cast
+
 from loguru import logger
 
 
-def _atomic_save(array: np.ndarray, filename: str):
-    tmp_dir = os.path.dirname(filename) or "."
-    with tempfile.NamedTemporaryFile(dir=tmp_dir, delete=False) as tmp:
-        np.save(tmp, array)
-        temp_name = tmp.name
-    os.replace(temp_name, filename)
+# Additional imports for VLLM utilities
+import re
+import subprocess
+import requests
+from openai import OpenAI
+
+try:
+    import psutil
+    HAS_PSUTIL = True
+except ImportError:
+    HAS_PSUTIL = False
+    psutil = cast(Any, None)
+    logger.warning("psutil not available. Some VLLM process management features may be limited.")
+
+# Global tracking of VLLM processes
+_VLLM_PROCESSES: List[subprocess.Popen] = []
+
+
+def _extract_port_from_vllm_cmd(vllm_cmd: str) -> int:
+    """Extract port from VLLM command string."""
+    port_match = re.search(r'--port\s+(\d+)', vllm_cmd)
+    if port_match:
+        return int(port_match.group(1))
+    return 8000
+
+
+def _parse_env_vars_from_cmd(cmd: str) -> tuple[dict[str, str], str]:
+    """Parse environment variables from command string.
+    
+    Args:
+        cmd: Command string that may contain environment variables like 'VAR=value command...'
+        
+    Returns:
+        Tuple of (env_dict, cleaned_cmd) where env_dict contains parsed env vars
+        and cleaned_cmd is the command without the env vars.
+    """
+    import shlex
+    
+    # Split the command while preserving quoted strings
+    parts = shlex.split(cmd)
+    
+    env_vars = {}
+    cmd_parts = []
+    
+    for part in parts:
+        if '=' in part and not part.startswith('-'):
+            # Check if this looks like an environment variable
+            # Should be KEY=VALUE format, not contain spaces (unless quoted), and KEY should be uppercase
+            key_value = part.split('=', 1)
+            if len(key_value) == 2:
+                key, value = key_value
+                if key.isupper() and key.replace('_', '').isalnum():
+                    env_vars[key] = value
+                    continue
+        
+        # Not an env var, add to command parts
+        cmd_parts.append(part)
+    
+    # Reconstruct the cleaned command
+    cleaned_cmd = ' '.join(cmd_parts)
+    
+    return env_vars, cleaned_cmd
+
 
+def _start_vllm_server(vllm_cmd: str, timeout: int = 120) -> subprocess.Popen:
+    """Start VLLM server and wait for ready."""
+    # Parse environment variables from command
+    env_vars, cleaned_cmd = _parse_env_vars_from_cmd(vllm_cmd)
+    
+    port = _extract_port_from_vllm_cmd(cleaned_cmd)
+    
+    logger.info(f"Starting VLLM server: {cleaned_cmd}")
+    if env_vars:
+        logger.info(f"Environment variables: {env_vars}")
+    logger.info(f"VLLM output logged to: /tmp/vllm_{port}.txt")
+    
+    with open(f'/tmp/vllm_{port}.txt', 'w') as log_file:
+        log_file.write(f"VLLM Server started at {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
+        log_file.write(f"Command: {cleaned_cmd}\n")
+        if env_vars:
+            log_file.write(f"Environment: {env_vars}\n")
+        log_file.write(f"Port: {port}\n")
+        log_file.write("-" * 50 + "\n")
+    
+    # Prepare environment for subprocess
+    env = os.environ.copy()
+    env.update(env_vars)
 
-def _update_port_use(port: int, increment: int) -> None:
-    file_counter: str = f"/tmp/port_use_counter_{port}.npy"
-    file_counter_lock: str = f"/tmp/port_use_counter_{port}.lock"
-    with open(file_counter_lock, "w") as lock_file:
-        fcntl.flock(lock_file, fcntl.LOCK_EX)
+    with open(f'/tmp/vllm_{port}.txt', 'a') as log_file:
+        process = subprocess.Popen(
+            cleaned_cmd.split(),
+            stdout=log_file,
+            stderr=subprocess.STDOUT,
+            text=True,
+            preexec_fn=os.setsid,
+            env=env
+        )
+    
+    _VLLM_PROCESSES.append(process)
+    
+    start_time = time.time()
+    while time.time() - start_time < timeout:
         try:
-            if os.path.exists(file_counter):
-                try:
-                    counter = np.load(file_counter)
-                except Exception as e:
-                    logger.warning(f"Corrupted usage file {file_counter}: {e}")
-                    counter = np.array([0])
-            else:
-                counter: np.ndarray = np.array([0], dtype=np.int64)
-            counter[0] += increment
-            _atomic_save(counter, file_counter)
-        finally:
-            fcntl.flock(lock_file, fcntl.LOCK_UN)
+            response = requests.get(f"http://localhost:{port}/health", timeout=2)
+            if response.status_code == 200:
+                logger.info(f"VLLM server ready on port {port}")
+                return process
+        except requests.RequestException:
+            pass
+        
+        if process.poll() is not None:
+            stdout, stderr = process.communicate()
+            raise RuntimeError(
+                f"VLLM server terminated unexpectedly. "
+                f"Return code: {process.returncode}, "
+                f"stderr: {stderr[:200]}..."
+            )
+        
+        time.sleep(2)
+    
+    process.terminate()
+    try:
+        process.wait(timeout=5)
+    except subprocess.TimeoutExpired:
+        process.kill()
+    
+    if process in _VLLM_PROCESSES:
+        _VLLM_PROCESSES.remove(process)
+    
+    raise RuntimeError(f"VLLM server failed to start within {timeout}s on port {port}")
 
 
-def _pick_least_used_port(ports: List[int]) -> int:
-    global_lock_file = "/tmp/ports.lock"
-    with open(global_lock_file, "w") as lock_file:
-        fcntl.flock(lock_file, fcntl.LOCK_EX)
+def _kill_vllm_on_port(port: int) -> bool:
+    """Kill VLLM server on port."""
+    killed = False
+    logger.info(f"Checking VLLM server on port {port}")
+    
+    processes_to_remove = []
+    for process in _VLLM_PROCESSES:
         try:
-            port_use: Dict[int, int] = {}
-            for port in ports:
-                file_counter = f"/tmp/port_use_counter_{port}.npy"
-                if os.path.exists(file_counter):
+            if process.poll() is None:
+                killed_process = False
+                if HAS_PSUTIL:
                     try:
-                        counter = np.load(file_counter)
-                    except Exception as e:
-                        logger.warning(f"Corrupted usage file {file_counter}: {e}")
-                        counter = np.array([0])
-                else:
-                    counter = np.array([0])
-                port_use[port] = counter[0]
-            if not port_use:
-                if ports:
-                    raise ValueError("Port usage data is empty, cannot pick a port.")
-                else:
-                    raise ValueError("No ports provided to pick from.")
-            lsp = min(port_use, key=lambda k: port_use[k])
-            _update_port_use(lsp, 1)
-        finally:
-            fcntl.flock(lock_file, fcntl.LOCK_UN)
-    return lsp
-
-
-def retry_on_exception(max_retries=10, exceptions=(Exception,), sleep_time=3):
-    def decorator(func):
-        from functools import wraps
-
-        def wrapper(self, *args, **kwargs):
-            retry_count = kwargs.get("retry_count", 0)
-            last_exception = None
-            while retry_count <= max_retries:
+                        proc = psutil.Process(process.pid)
+                        cmdline = ' '.join(proc.cmdline())
+                        if f'--port {port}' in cmdline or f'--port={port}' in cmdline:
+                            logger.info(f"Killing tracked VLLM process {process.pid} on port {port}")
+                            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+                            try:
+                                process.wait(timeout=5)
+                            except subprocess.TimeoutExpired:
+                                os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+                                process.wait()
+                            killed = True
+                            killed_process = True
+                    except (psutil.NoSuchProcess, psutil.AccessDenied):
+                        pass
+                
+                if not HAS_PSUTIL or not killed_process:
+                    logger.info(f"Killing tracked VLLM process {process.pid}")
+                    try:
+                        os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+                        try:
+                            process.wait(timeout=5)
+                        except subprocess.TimeoutExpired:
+                            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+                            process.wait()
+                        killed = True
+                    except (ProcessLookupError, OSError):
+                        pass
+                
+                processes_to_remove.append(process)
+            else:
+                processes_to_remove.append(process)
+        except (ProcessLookupError, OSError):
+            processes_to_remove.append(process)
+    
+    for process in processes_to_remove:
+        if process in _VLLM_PROCESSES:
+            _VLLM_PROCESSES.remove(process)
+    
+    if not killed and HAS_PSUTIL:
+        try:
+            for proc in psutil.process_iter(['pid', 'cmdline']):
                 try:
-                    return func(self, *args, **kwargs)
-                except exceptions as e:
-                    import litellm  # type: ignore
-
-                    if isinstance(
-                        e, (litellm.exceptions.APIError, litellm.exceptions.Timeout)
-                    ):
-                        base_url_info = kwargs.get(
-                            "base_url", getattr(self, "base_url", None)
-                        )
-                        logger.warning(
-                            f"[{base_url_info=}] {type(e).__name__}: {str(e)[:100]}, will sleep for {sleep_time}s and retry"
-                        )
-                        time.sleep(sleep_time)
-                        retry_count += 1
-                        kwargs["retry_count"] = retry_count
-                        last_exception = e
-                        continue
-                    elif hasattr(
-                        litellm.exceptions, "ContextWindowExceededError"
-                    ) and isinstance(e, litellm.exceptions.ContextWindowExceededError):
-                        logger.error(f"Context window exceeded: {e}")
-                        raise
-                    else:
-                        logger.error(f"Generic error during LLM call: {e}")
-                        import traceback
-
-                        traceback.print_exc()
-                        raise
-            logger.error(f"Retry limit exceeded, error: {last_exception}")
-            if last_exception:
-                raise last_exception
-            raise ValueError("Retry limit exceeded with no specific error.")
-
-        return wraps(func)(wrapper)
-
-    return decorator
-
-
-def forward_only(func):
-    from functools import wraps
-
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        kwargs["retry_count"] = 0
-        return func(self, *args, **kwargs)
-
-    return wrapper
+                    cmdline = ' '.join(proc.info['cmdline'] or [])
+                    if ('vllm' in cmdline.lower() and 
+                        (f'--port {port}' in cmdline or f'--port={port}' in cmdline)):
+                        logger.info(f"Killing untracked VLLM process {proc.info['pid']} on port {port}")
+                        proc.terminate()
+                        try:
+                            proc.wait(timeout=5)
+                        except psutil.TimeoutExpired:
+                            proc.kill()
+                        killed = True
+                        break
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    continue
+        except Exception as e:
+            logger.warning(f"Error searching processes on port {port}: {e}")
+    
+    if killed:
+        logger.info(f"Killed VLLM server on port {port}")
+        time.sleep(2)
+    else:
+        logger.info(f"No VLLM server on port {port}")
+    
+    return killed
+
+
+def stop_vllm_process(process: subprocess.Popen, wait_timeout: int = 10) -> None:
+    """Terminate a tracked VLLM process and remove it from tracking."""
+    logger.info(f"Stopping VLLM process {process.pid}")
+    try:
+        if process.poll() is None:
+            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+            try:
+                process.wait(timeout=wait_timeout)
+                logger.info("VLLM process stopped gracefully")
+            except subprocess.TimeoutExpired:
+                logger.warning("VLLM process didn't stop gracefully, forcing kill")
+                os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+                process.wait()
+        else:
+            logger.info("VLLM process already terminated")
+    except (ProcessLookupError, OSError) as exc:
+        logger.warning(f"Process may have already terminated: {exc}")
+    finally:
+        if process in _VLLM_PROCESSES:
+            _VLLM_PROCESSES.remove(process)
+
+
+def kill_all_vllm_processes() -> int:
+    """Kill all tracked VLLM processes."""
+    killed_count = 0
+    for process in list(_VLLM_PROCESSES):
+        if process.poll() is None:
+            logger.info(f"Killing VLLM process with PID {process.pid}")
+            stop_vllm_process(process, wait_timeout=5)
+            killed_count += 1
+        else:
+            _VLLM_PROCESSES.remove(process)
+    logger.info(f"Killed {killed_count} VLLM processes")
+    return killed_count
+
+
+def _is_server_running(port: int) -> bool:
+    """Check if server is running on port."""
+    try:
+        response = requests.get(f"http://localhost:{port}/health", timeout=2)
+        return response.status_code == 200
+    except requests.RequestException:
+        return False
+
+
+def get_base_client(
+    client=None, 
+    cache: bool = True, 
+    api_key="abc",
+    vllm_cmd=None,
+    vllm_process=None
+) -> OpenAI:
+    """Get OpenAI client from various inputs."""
+    from llm_utils import MOpenAI
+
+    open_ai_class = OpenAI if not cache else MOpenAI
+    
+    if client is None:
+        if vllm_cmd is not None:
+            # Parse environment variables from command to get clean command for port extraction
+            _, cleaned_cmd = _parse_env_vars_from_cmd(vllm_cmd)
+            port = _extract_port_from_vllm_cmd(cleaned_cmd)
+            return open_ai_class(base_url=f"http://localhost:{port}/v1", api_key=api_key)
+        else:
+            return open_ai_class()
+    elif isinstance(client, int):
+        return open_ai_class(base_url=f"http://localhost:{client}/v1", api_key=api_key)
+    elif isinstance(client, str):
+        return open_ai_class(base_url=client, api_key=api_key)
+    elif isinstance(client, OpenAI):
+        return client
+    else:
+        raise ValueError("Invalid client type. Must be OpenAI, port (int), base_url (str), or None.")
+
+
+def _is_lora_path(path: str) -> bool:
+    """Check if path is LoRA adapter directory."""
+    if not os.path.isdir(path):
+        return False
+    adapter_config_path = os.path.join(path, 'adapter_config.json')
+    return os.path.isfile(adapter_config_path)
+
+
+def _get_port_from_client(client: OpenAI) -> Optional[int]:
+    """Extract port from OpenAI client base_url."""
+    if hasattr(client, 'base_url') and client.base_url:
+        base_url = str(client.base_url)
+        if 'localhost:' in base_url:
+            try:
+                port_part = base_url.split('localhost:')[1].split('/')[0]
+                return int(port_part)
+            except (IndexError, ValueError):
+                pass
+    return None
+
+
+def _load_lora_adapter(lora_path: str, port: int) -> str:
+    """Load LoRA adapter from path."""
+    lora_name = os.path.basename(lora_path.rstrip('/\\'))
+    if not lora_name:
+        lora_name = os.path.basename(os.path.dirname(lora_path))
+    
+    response = requests.post(
+        f'http://localhost:{port}/v1/load_lora_adapter',
+        headers={'accept': 'application/json', 'Content-Type': 'application/json'},
+        json={"lora_name": lora_name, "lora_path": os.path.abspath(lora_path)}
+    )
+    response.raise_for_status()
+    return lora_name
+
+
+def _unload_lora_adapter(lora_path: str, port: int) -> None:
+    """Unload LoRA adapter."""
+    try:
+        lora_name = os.path.basename(lora_path.rstrip('/\\'))
+        if not lora_name:
+            lora_name = os.path.basename(os.path.dirname(lora_path))
+            
+        response = requests.post(
+            f'http://localhost:{port}/v1/unload_lora_adapter',
+            headers={'accept': 'application/json', 'Content-Type': 'application/json'},
+            json={"lora_name": lora_name, "lora_int_id": 0}
+        )
+        response.raise_for_status()
+    except requests.RequestException as e:
+        logger.warning(f"Error unloading LoRA adapter: {str(e)[:100]}")

speedy_utils/multi_worker/process.py

@@ -1,15 +1,17 @@
 # ray_multi_process.py
-import time, os, pickle, uuid, datetime, multiprocessing
 import datetime
 import os
 import pickle
+import threading
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Callable
-from tqdm import tqdm
+from typing import Any, Callable, Iterable
+
 import psutil
-import threading
+from fastcore.parallel import parallel
+from tqdm import tqdm
+
 ray: Any
 try:
     import ray as ray  # type: ignore
@@ -17,11 +19,75 @@ try:
 except Exception:  # pragma: no cover
     ray = None  # type: ignore
     _HAS_RAY = False
-from typing import Any, Callable, Iterable
 
-import ray
-from fastcore.parallel import parallel
-from tqdm import tqdm
+
+# ─── global tracking ──────────────────────────────────────────
+
+# Global tracking for processes and threads
+SPEEDY_RUNNING_PROCESSES: list[psutil.Process] = []
+_SPEEDY_PROCESSES_LOCK = threading.Lock()
+
+
+def _prune_dead_processes() -> None:
+    """Remove dead processes from tracking list."""
+    with _SPEEDY_PROCESSES_LOCK:
+        SPEEDY_RUNNING_PROCESSES[:] = [p for p in SPEEDY_RUNNING_PROCESSES if p.is_running()]
+
+
+def _track_processes(processes: list[psutil.Process]) -> None:
+    """Add processes to global tracking list."""
+    if not processes:
+        return
+    with _SPEEDY_PROCESSES_LOCK:
+        living = [p for p in SPEEDY_RUNNING_PROCESSES if p.is_running()]
+        for candidate in processes:
+            if not candidate.is_running():
+                continue
+            if any(existing.pid == candidate.pid for existing in living):
+                continue
+            living.append(candidate)
+        SPEEDY_RUNNING_PROCESSES[:] = living
+
+
+def _track_ray_processes() -> None:
+    """Track Ray worker processes when Ray is initialized."""
+    if not _HAS_RAY or not ray.is_initialized():
+        return
+    try:
+        # Get Ray worker processes
+        current_pid = os.getpid()
+        parent = psutil.Process(current_pid)
+        ray_processes = []
+        for child in parent.children(recursive=True):
+            try:
+                if 'ray' in child.name().lower() or 'worker' in child.name().lower():
+                    ray_processes.append(child)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+        _track_processes(ray_processes)
+    except Exception:
+        # Don't fail if process tracking fails
+        pass
+
+
+def _track_multiprocessing_processes() -> None:
+    """Track multiprocessing worker processes."""
+    try:
+        # Find recently created child processes that might be multiprocessing workers
+        current_pid = os.getpid()
+        parent = psutil.Process(current_pid)
+        new_processes = []
+        for child in parent.children(recursive=False):  # Only direct children
+            try:
+                # Basic heuristic: if it's a recent child process, it might be a worker
+                if time.time() - child.create_time() < 5:  # Created within last 5 seconds
+                    new_processes.append(child)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+        _track_processes(new_processes)
+    except Exception:
+        # Don't fail if process tracking fails
+        pass
 
 
 # ─── cache helpers ──────────────────────────────────────────
@@ -68,6 +134,7 @@ def ensure_ray(workers: int, pbar: tqdm | None = None):
         t0 = time.time()
         ray.init(num_cpus=workers, ignore_reinit_error=True)
         took = time.time() - t0
+        _track_ray_processes()  # Track Ray worker processes
         if pbar:
             pbar.set_postfix_str(f"ray.init {workers} took {took:.2f}s")
         RAY_WORKER = workers
@@ -87,11 +154,6 @@ def multi_process(
     progress: bool = True,
     # backend: str = "ray",   # "seq", "ray", or "fastcore"
     backend: Literal["seq", "ray", "mp", "threadpool", "safe"] = "mp",
-    # 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]:
     """
@@ -171,6 +233,8 @@ def multi_process(
             results = parallel(
                 f_wrapped, items, n_workers=workers, progress=progress, threadpool=False
             )
+            _track_multiprocessing_processes()  # Track multiprocessing workers
+            _prune_dead_processes()  # Clean up dead processes
             return list(results)
         if backend == "threadpool":
             results = parallel(
@@ -180,8 +244,17 @@ def multi_process(
         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))
+            # Import thread tracking from thread module
+            try:
+                from .thread import _track_executor_threads, _prune_dead_threads
+                with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                    _track_executor_threads(executor)  # Track threads
+                    results = list(executor.map(f_wrapped, items))
+                _prune_dead_threads()  # Clean up dead threads
+            except ImportError:
+                # Fallback if thread module not available
+                with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
+                    results = list(executor.map(f_wrapped, items))
             return results
         
         raise ValueError(f"Unsupported backend: {backend!r}")
@@ -190,12 +263,24 @@ def multi_process(
 
 def cleanup_phantom_workers():
     """
-    Kill all child processes (phantom workers) without killing the Jupyter kernel itself.
+    Kill all tracked processes and threads (phantom workers) without killing the Jupyter kernel itself.
     Also lists non-daemon threads that remain.
     """
-    parent = psutil.Process(os.getpid())
+    # Clean up tracked processes first
+    _prune_dead_processes()
+    killed_processes = 0
+    with _SPEEDY_PROCESSES_LOCK:
+        for process in SPEEDY_RUNNING_PROCESSES[:]:  # Copy to avoid modification during iteration
+            try:
+                print(f"🔪 Killing tracked process {process.pid} ({process.name()})")
+                process.kill()
+                killed_processes += 1
+            except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
+                print(f"⚠️ Could not kill process {process.pid}: {e}")
+        SPEEDY_RUNNING_PROCESSES.clear()
     
-    # Kill only children, never the current process
+    # Also kill any remaining child processes (fallback)
+    parent = psutil.Process(os.getpid())
     for child in parent.children(recursive=True):
         try:
             print(f"🔪 Killing child process {child.pid} ({child.name()})")
@@ -203,14 +288,29 @@ def cleanup_phantom_workers():
         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).")
+    # Try to clean up threads using thread module functions if available
+    try:
+        from .thread import SPEEDY_RUNNING_THREADS, kill_all_thread, _prune_dead_threads
+        _prune_dead_threads()
+        killed_threads = kill_all_thread()
+        if killed_threads > 0:
+            print(f"🔪 Killed {killed_threads} tracked threads")
+    except ImportError:
+        # Fallback: just report stray threads
+        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).")
+    print(f"✅ Cleaned up {killed_processes} tracked processes and child processes (kernel untouched).")
 
 # Usage: run this anytime after cancelling a cell
 
+
+__all__ = [
+    "SPEEDY_RUNNING_PROCESSES",
+    "multi_process",
+    "cleanup_phantom_workers",
+]
+