rapidfireai 0.0.1__py3-none-any.whl → 0.9.9__py3-none-any.whl

rapidfireai/cli.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python3
+"""
+Command-line interface for RapidFire AI
+"""
+
+import os
+import sys
+import subprocess
+import argparse
+import platform
+import shutil
+import re
+from pathlib import Path
+from .version import __version__
+
+
+def get_script_path():
+    """Get the path to the start.sh script."""
+    # Get the directory where this package is installed
+    package_dir = Path(__file__).parent
+    script_path = package_dir / "start.sh"
+    
+    if not script_path.exists():
+        # Fallback: try to find it relative to the current working directory
+        script_path = Path.cwd() / "rapidfireai" / "start.sh"
+        if not script_path.exists():
+            raise FileNotFoundError(f"Could not find start.sh script at {script_path}")
+    
+    return script_path
+
+
+def run_script(args):
+    """Run the start.sh script with the given arguments."""
+    script_path = get_script_path()
+    
+    # Make sure the script is executable
+    if not os.access(script_path, os.X_OK):
+        os.chmod(script_path, 0o755)
+    
+    # Run the script with the provided arguments
+    try:
+        result = subprocess.run([str(script_path)] + args, check=True)
+        return result.returncode
+    except subprocess.CalledProcessError as e:
+        print(f"Error running start.sh: {e}", file=sys.stderr)
+        return e.returncode
+    except FileNotFoundError:
+        print(f"Error: start.sh script not found at {script_path}", file=sys.stderr)
+        return 1
+
+
+def get_python_info():
+    """Get comprehensive Python information."""
+    info = {}
+    
+    # Python version and implementation
+    info['version'] = sys.version
+    info['implementation'] = platform.python_implementation()
+    info['executable'] = sys.executable
+    
+    # Environment information
+    info['conda_env'] = os.environ.get('CONDA_DEFAULT_ENV', 'none')
+    info['venv'] = 'yes' if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix) else 'no'
+    
+    return info
+
+
+def get_pip_packages():
+    """Get list of installed pip packages."""
+    try:
+        result = subprocess.run([sys.executable, '-m', 'pip', 'list'], 
+                              capture_output=True, text=True, check=True)
+        return result.stdout
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return "Failed to get pip packages"
+
+
+def get_gpu_info():
+    """Get comprehensive GPU and CUDA information."""
+    info = {}
+    
+    # Check for nvidia-smi
+    nvidia_smi_path = shutil.which('nvidia-smi')
+    info['nvidia_smi'] = 'found' if nvidia_smi_path else 'not found'
+    
+    if nvidia_smi_path:
+        try:
+            # Get driver and CUDA runtime version from the full nvidia-smi output
+            result = subprocess.run(['nvidia-smi'], 
+                                  capture_output=True, text=True, check=True)
+            if result.stdout.strip():
+                lines = result.stdout.strip().split('\n')
+                # Look for the header line that contains CUDA version
+                for line in lines:
+                    if 'CUDA Version:' in line:
+                        # Extract CUDA version from line like "NVIDIA-SMI 535.183.06 Driver Version: 535.183.06 CUDA Version: 12.2"
+                        cuda_version = line.split('CUDA Version:')[1].split()[0]
+                        info['cuda_runtime'] = cuda_version
+                        # Also extract driver version from the same line
+                        if 'Driver Version:' in line:
+                            driver_version = line.split('Driver Version:')[1].split('CUDA Version:')[0].strip()
+                            info['driver_version'] = driver_version
+                        break
+                else:
+                    info['driver_version'] = 'unknown'
+                    info['cuda_runtime'] = 'unknown'
+        except (subprocess.CalledProcessError, ValueError):
+            info['driver_version'] = 'unknown'
+            info['cuda_runtime'] = 'unknown'
+        
+        # Get GPU count, models, and VRAM
+        try:
+            result = subprocess.run(['nvidia-smi', '--query-gpu=count,name,memory.total', '--format=csv,noheader,nounits'], 
+                                  capture_output=True, text=True, check=True)
+            if result.stdout.strip():
+                lines = result.stdout.strip().split('\n')
+                if lines:
+                    count, name, memory = lines[0].split(', ')
+                    info['gpu_count'] = int(count)
+                    info['gpu_model'] = name.strip()
+                    # Convert memory from MiB to GB
+                    memory_mib = int(memory.split()[0])
+                    memory_gb = memory_mib / 1024
+                    info['gpu_memory_gb'] = f"{memory_gb:.1f}"
+                    
+                    # Get detailed info for multiple GPUs if present
+                    if info['gpu_count'] > 1:
+                        info['gpu_details'] = []
+                        for line in lines:
+                            count, name, memory = line.split(', ')
+                            memory_mib = int(memory.split()[0])
+                            memory_gb = memory_mib / 1024
+                            info['gpu_details'].append({
+                                'name': name.strip(),
+                                'memory_gb': f"{memory_gb:.1f}"
+                            })
+        except (subprocess.CalledProcessError, ValueError):
+            info['gpu_count'] = 0
+            info['gpu_model'] = 'unknown'
+            info['gpu_memory_gb'] = 'unknown'
+    else:
+        info['driver_version'] = 'N/A'
+        info['cuda_runtime'] = 'N/A'
+        info['gpu_count'] = 0
+        info['gpu_model'] = 'N/A'
+        info['gpu_memory_gb'] = 'N/A'
+    
+    # Check for nvcc (CUDA compiler)
+    nvcc_path = shutil.which('nvcc')
+    info['nvcc'] = 'found' if nvcc_path else 'not found'
+    
+    if nvcc_path:
+        try:
+            result = subprocess.run(['nvcc', '--version'], 
+                                  capture_output=True, text=True, check=True)
+            # Extract version from output like "Cuda compilation tools, release 11.8, V11.8.89"
+            version_line = result.stdout.split('\n')[0]
+            if 'release' in version_line:
+                version = version_line.split('release')[1].split(',')[0].strip()
+                info['nvcc_version'] = version
+            else:
+                info['nvcc_version'] = 'unknown'
+        except subprocess.CalledProcessError:
+            info['nvcc_version'] = 'unknown'
+    else:
+        info['nvcc_version'] = 'N/A'
+    
+    # Check CUDA installation paths
+    cuda_paths = [
+        '/usr/local/cuda',
+        '/opt/cuda',
+        '/usr/cuda',
+        os.path.expanduser('~/cuda')
+    ]
+    
+    cuda_installed = False
+    for path in cuda_paths:
+        if os.path.exists(path):
+            cuda_installed = True
+            break
+    
+    info['cuda_installation'] = 'present' if cuda_installed else 'not present'
+    
+    # Check if CUDA is on PATH
+    cuda_on_path = any('cuda' in p.lower() for p in os.environ.get('PATH', '').split(os.pathsep))
+    info['cuda_on_path'] = 'yes' if cuda_on_path else 'no'
+    
+    return info
+
+
+def run_doctor():
+    """Run the doctor command to diagnose system issues."""
+    print("🔍 RapidFire AI System Diagnostics")
+    print("=" * 50)
+    
+    # Python Information
+    print("\n🐍 Python Environment:")
+    print("-" * 30)
+    python_info = get_python_info()
+    print(f"Version: {python_info['version'].split()[0]}")
+    print(f"Implementation: {python_info['implementation']}")
+    print(f"Executable: {python_info['executable']}")
+    print(f"Conda Environment: {python_info['conda_env']}")
+    print(f"Virtual Environment: {python_info['venv']}")
+    
+    # Pip Packages
+    print("\n📦 Installed Packages:")
+    print("-" * 30)
+    pip_output = get_pip_packages()
+    if pip_output != "Failed to get pip packages":
+        # Show only relevant packages
+        relevant_packages = ['rapidfireai', 'mlflow', 'torch', 'transformers', 'flask', 'gunicorn', 'peft', 'trl', 'bitsandbytes', 'nltk', 'evaluate', 'rouge-score', 'sentencepiece']
+        lines = pip_output.split('\n')
+        for line in lines:
+            if any(pkg.lower() in line.lower() for pkg in relevant_packages):
+                print(line)
+        print("... (showing only relevant packages)")
+    else:
+        print(pip_output)
+    
+    # GPU Information
+    print("\n🚀 GPU & CUDA Information:")
+    print("-" * 30)
+    gpu_info = get_gpu_info()
+    print(f"nvidia-smi: {gpu_info['nvidia_smi']}")
+    
+    if gpu_info['nvidia_smi'] == 'found':
+        print(f"Driver Version: {gpu_info['driver_version']}")
+        print(f"CUDA Runtime: {gpu_info['cuda_runtime']}")
+        print(f"GPU Count: {gpu_info['gpu_count']}")
+        
+        if gpu_info['gpu_count'] > 0:
+            if 'gpu_details' in gpu_info:
+                print("GPU Details:")
+                for i, gpu in enumerate(gpu_info['gpu_details']):
+                    print(f"  GPU {i}: {gpu['name']} ({gpu['memory_gb']} GB)")
+            else:
+                print(f"GPU Model: {gpu_info['gpu_model']}")
+                print(f"Total VRAM: {gpu_info['gpu_memory_gb']} GB")
+    
+    print(f"nvcc: {gpu_info['nvcc']}")
+    if gpu_info['nvcc'] == 'found':
+        print(f"nvcc Version: {gpu_info['nvcc_version']}")
+    
+    print(f"CUDA Installation: {gpu_info['cuda_installation']}")
+    print(f"CUDA on PATH: {gpu_info['cuda_on_path']}")
+    
+    # System Information
+    print("\n💻 System Information:")
+    print("-" * 30)
+    print(f"Platform: {platform.platform()}")
+    print(f"Architecture: {platform.machine()}")
+    print(f"Processor: {platform.processor()}")
+    
+    # Environment Variables
+    print("\n🔧 Environment Variables:")
+    print("-" * 30)
+    relevant_vars = ['CUDA_HOME', 'CUDA_PATH', 'LD_LIBRARY_PATH', 'PATH']
+    for var in relevant_vars:
+        value = os.environ.get(var, 'not set')
+        if value != 'not set' and len(value) > 100:
+            value = value[:100] + "..."
+        print(f"{var}: {value}")
+    
+    print("\n✅ Diagnostics complete!")
+    return 0
+
+def get_cuda_version():
+    """Detect CUDA version from nvcc or nvidia-smi"""
+    try:
+        result = subprocess.run(['nvcc', '--version'], 
+                              capture_output=True, text=True, check=True)
+        match = re.search(r'release (\d+)\.(\d+)', result.stdout)
+        if match:
+            return int(match.group(1))
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        try:
+            result = subprocess.run(['nvidia-smi'], 
+                                  capture_output=True, text=True, check=True)
+            match = re.search(r'CUDA Version: (\d+)\.(\d+)', result.stdout)
+            if match:
+                return int(match.group(1))
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            pass
+    return None
+
+def get_compute_capability():
+    """Get compute capability from nvidia-smi"""
+    try:
+        result = subprocess.run(['nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader,nounits'], 
+                              capture_output=True, text=True, check=True)
+        match = re.search(r'(\d+)\.(\d+)', result.stdout)
+        if match:
+            return int(match.group(1))
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return None
+
+def install_packages():
+    """Install packages for the RapidFire AI project."""
+    packages = []
+    # Generate CUDA requirements file
+    cuda_major = get_cuda_version()
+    compute_capability = get_compute_capability()
+    print(f"CUDA major version: {cuda_major}")
+    print(f"Compute capability: {compute_capability}")
+    if cuda_major == 12:
+        print(f"\n🎯 Detected CUDA {cuda_major}.x")
+        packages.append({"package": "vllm==0.10.1.1", "extra_args": ["--torch-backend=cu126"]})
+    elif cuda_major == 11:
+        print(f"\n🎯 Detected CUDA {cuda_major}.x")
+        packages.append({"package": "vllm==0.10.1.1", "extra_args": ["--torch-backend=cu118"]})
+    else:
+        print("\n⚠️  CUDA version not detected or unsupported.")
+    if compute_capability == 7:
+        print(f"\n🎯 Detected CUDA Compute Capability {compute_capability}.x")
+        packages.append({"package": "flash-attn==1.0.9", "extra_args": ["--no-build-isolation"]})
+    elif compute_capability == 8:
+        print(f"\n🎯 Detected CUDA Compute Capability {compute_capability}.x")
+        packages.append({"package": "flash-attn==2.8.3", "extra_args": ["--no-build-isolation"]})
+    else:
+        print("\n⚠️  CUDA Compute Capability not detected or unsupported.")
+
+    for package_info in packages:
+        try:
+            package = package_info["package"]
+            cmd = [sys.executable, "-m", "uv", "pip", "install", package] + package_info["extra_args"]
+            print(f"   Installing {package}...")
+            subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
+            print(f"✅ Successfully installed {package}")
+        except subprocess.CalledProcessError as e:
+            print(f"❌ Failed to install {package}")
+            print(f"   Error: {e}")
+            print(f"   You may need to install {package} manually")
+    return 0
+
+def run_init():
+    """Run the init command to initialize the project."""
+    print("🔧 Initializing RapidFire AI project...")
+    print("-" * 30)
+    print("Initializing project...")
+    install_packages()
+    return 0
+
+def main():
+    """Main entry point for the rapidfire-start command."""
+    parser = argparse.ArgumentParser(
+        description="RapidFire AI - Start/stop/manage services",
+        prog="rapidfire"
+    )
+    
+    parser.add_argument(
+        "command",
+        nargs="?",
+        default="start",
+        choices=["start", "stop", "status", "restart", "setup", "doctor", "init"],
+        help="Command to execute (default: start)"
+    )
+    
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"RapidFire AI {__version__}"
+    )
+    
+    args = parser.parse_args()
+    
+    # Handle doctor command separately
+    if args.command == "doctor":
+        return run_doctor()
+
+    # Handle init command separately
+    if args.command == "init":
+        return run_init()
+    
+    # Run the script with the specified command
+    return run_script([args.command])
+
+
+if __name__ == "__main__":
+    sys.exit(main()) 

rapidfireai/db/db_interface.py

@@ -0,0 +1,135 @@
+"""Interface for the database."""
+
+import functools
+import os
+import sqlite3
+import time
+from typing import Any, Callable
+
+from rapidfireai.utils.constants import DBConfig
+from rapidfireai.utils.exceptions import DBException
+
+
+class DatabaseInterface:
+    """Interface for the database."""
+
+    def __init__(self):
+        try:
+            if not os.path.exists(DBConfig.DB_PATH):
+                path = os.path.dirname(DBConfig.DB_PATH)
+                os.makedirs(path, exist_ok=True)
+                print(f"Created directory for database at {path}")
+
+            self.conn: sqlite3.Connection = sqlite3.connect(
+                DBConfig.DB_PATH,
+                timeout=DBConfig.CONNECTION_TIMEOUT,
+                check_same_thread=False,
+                isolation_level=None,
+            )
+
+            # Configure database with all PRAGMA settings
+            pragma_sql = f"""
+            PRAGMA cache_size={DBConfig.CACHE_SIZE};
+            PRAGMA mmap_size={DBConfig.MMAP_SIZE};
+            PRAGMA page_size={DBConfig.PAGE_SIZE};
+            PRAGMA busy_timeout={DBConfig.BUSY_TIMEOUT};
+            PRAGMA journal_mode=WAL;
+            PRAGMA synchronous=NORMAL;
+            PRAGMA temp_store=MEMORY;
+            PRAGMA foreign_keys=ON;
+            """
+            _ = self.conn.executescript(pragma_sql)
+
+            self.cursor: sqlite3.Cursor = self.conn.cursor()
+
+        except sqlite3.Error as e:
+            raise DBException(f"Failed to initialize database connection: {e}") from e
+        except Exception as e:
+            raise DBException(f"Unexpected error during database initialization: {e}") from e
+
+    @staticmethod
+    def retry_on_locked(
+        max_retries: int = DBConfig.DEFAULT_MAX_RETRIES,
+        base_delay: float = DBConfig.DEFAULT_BASE_DELAY,
+        max_delay: float = DBConfig.DEFAULT_MAX_DELAY,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator to retry operations when database is locked"""
+
+        def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+            @functools.wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                last_exception = None
+                for attempt in range(max_retries):
+                    try:
+                        return func(*args, **kwargs)
+                    except sqlite3.OperationalError as e:
+                        if "database is locked" in str(e).lower():
+                            last_exception = e
+                            if attempt < max_retries - 1:
+                                # Exponential backoff with jitter
+                                delay = min(base_delay * (2**attempt), max_delay)
+                                delay += time.time() % 0.1  # Add small jitter
+                                time.sleep(delay)
+                                continue
+                        # Re-raise if it's not a "database is locked" error
+                        raise
+                # If we get here, all retries failed
+                if last_exception:
+                    raise last_exception
+                else:
+                    raise RuntimeError("All retries failed but no exception was captured")
+
+            return wrapper
+
+        return decorator
+
+    def close(self) -> None:
+        """Close the database connection properly"""
+        try:
+            if self.conn:
+                self.conn.close()
+        except sqlite3.Error as e:
+            raise DBException(f"Error closing database connection: {e}") from e
+        except Exception as e:
+            raise DBException(f"Unexpected error closing database connection: {e}") from e
+
+    def optimize_periodically(self) -> None:
+        """Run periodic optimization - call this occasionally, not on every query"""
+        try:
+            _ = self.conn.execute("PRAGMA optimize")
+        except sqlite3.Error as e:
+            raise DBException(f"Failed to optimize database: {e}") from e
+        except Exception as e:
+            raise DBException(f"Unexpected error during database optimization: {e}") from e
+
+    def execute(
+        self,
+        query: str,
+        params: dict[str, Any] | tuple[Any, ...] | None = None,
+        fetch: bool = False,
+        commit: bool = False,
+    ) -> list[Any] | tuple[Any] | None:
+        """Execute a query with automatic retry on database locked errors"""
+        # Validate that either fetch or commit is True
+        if not fetch and not commit:
+            raise ValueError("Either fetch or commit must be True")
+
+        try:
+            # Execute the query with parameters if provided
+            if params:
+                result = self.cursor.execute(query, params)
+            else:
+                result = self.cursor.execute(query)
+
+            # Commit the transaction if commit is True
+            if commit:
+                self.conn.commit()
+
+            # Return the result if fetch is True
+            if fetch:
+                return result.fetchall()
+
+        except sqlite3.Error as e:
+            raise DBException(f"Database error executing query '{query[:50]}...': {e}") from e
+        except Exception as e:
+            raise DBException(f"Unexpected error executing query '{query[:50]}...': {e}") from e