Updated LocalLab v0.3.3

Author: UtkarshTheDev

locallab/__init__.py

@@ -2,7 +2,7 @@
 LocalLab: Run LLMs locally with a friendly API similar to OpenAI
 """
 
-__version__ = "0.3.2" 
+__version__ = "0.3.3" 
 
 from typing import Dict, Any, Optional
 import logging

locallab/core/app.py

@@ -105,7 +105,7 @@ def init(backend, **kwargs):
 @app.on_event("startup")
 async def startup_event():
     """Initialization tasks when the server starts"""
-    logger.info("Starting LocalLab server...")
+    logger.debug("Initializing LocalLab server...")
 
     # Initialize cache if available
     if FASTAPI_CACHE_AVAILABLE:
@@ -118,12 +118,12 @@ async def startup_event():
     model_to_load = os.environ.get("HUGGINGFACE_MODEL", DEFAULT_MODEL)
 
     # Log model configuration
-    logger.info(f"Model configuration:")
-    logger.info(f" - Model to load: {model_to_load}")
-    logger.info(f" - Quantization: {'Enabled - ' + os.environ.get('LOCALLAB_QUANTIZATION_TYPE', QUANTIZATION_TYPE) if os.environ.get('LOCALLAB_ENABLE_QUANTIZATION', '').lower() == 'true' else 'Disabled'}")
-    logger.info(f" - Attention slicing: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_ATTENTION_SLICING', '').lower() == 'true' else 'Disabled'}")
-    logger.info(f" - Flash attention: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_FLASH_ATTENTION', '').lower() == 'true' else 'Disabled'}")
-    logger.info(f" - Better transformer: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_BETTERTRANSFORMER', '').lower() == 'true' else 'Disabled'}")
+    logger.debug("Model configuration:")
+    logger.debug(" - Model to load: %s", model_to_load)
+    logger.debug(f" - Quantization: {'Enabled - ' + os.environ.get('LOCALLAB_QUANTIZATION_TYPE', QUANTIZATION_TYPE) if os.environ.get('LOCALLAB_ENABLE_QUANTIZATION', '').lower() == 'true' else 'Disabled'}")
+    logger.debug(f" - Attention slicing: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_ATTENTION_SLICING', '').lower() == 'true' else 'Disabled'}")
+    logger.debug(f" - Flash attention: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_FLASH_ATTENTION', '').lower() == 'true' else 'Disabled'}")
+    logger.debug(f" - Better transformer: {'Enabled' if os.environ.get('LOCALLAB_ENABLE_BETTERTRANSFORMER', '').lower() == 'true' else 'Disabled'}")
 
     # Start loading the model in background if specified
     if model_to_load:

locallab/logger/logger.py

@@ -8,6 +8,9 @@
 from typing import Optional, Dict, Any
 from . import get_logger
 
+# Set up logging configuration
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
 # Server start time for uptime calculation
 SERVER_START_TIME = time.time()
 

locallab/server.py

@@ -29,7 +29,12 @@
 from .logger.logger import set_server_status, log_request
 from .utils.system import get_gpu_memory
 from .config import (
-    MIN_FREE_MEMORY
+    DEFAULT_MODEL,
+    ENABLE_QUANTIZATION, 
+    QUANTIZATION_TYPE,
+    ENABLE_ATTENTION_SLICING,
+    ENABLE_BETTERTRANSFORMER, 
+    ENABLE_FLASH_ATTENTION
 )
 
 # Import torch - handle import error gracefully
@@ -42,7 +47,6 @@
 # Get the logger instance
 logger = get_logger("locallab.server")
 
-
 def check_environment() -> List[Tuple[str, str, bool]]:
     """
     Check the environment for potential issues
@@ -73,7 +77,7 @@ def check_environment() -> List[Tuple[str, str, bool]]:
         if not os.environ.get("NGROK_AUTH_TOKEN"):
             issues.append((
                 "Running in Google Colab without NGROK_AUTH_TOKEN set",
-                "Set os.environ['NGROK_AUTH_TOKEN'] = 'your_token' for public URL access. Get your token from https://dashboard.ngrok.com/get-started/your-authtoken",
+                "Set os.environ['NGROK_AUTH_TOKEN'] = 'your_token' for public URL access.",
                 True
             ))
 
@@ -127,7 +131,7 @@ def check_environment() -> List[Tuple[str, str, bool]]:
         total_gb = memory.total / (1024 * 1024 * 1024)
         available_gb = memory.available / (1024 * 1024 * 1024)
 
-        if available_gb < MIN_FREE_MEMORY / 1024:  # Convert MB to GB
+        if available_gb < 2.0:  # Less than 2GB available
             issues.append((
                 f"Low system memory: Only {available_gb:.1f}GB available",
                 "Models may require 2-8GB of system memory. Consider closing other applications",
@@ -163,7 +167,6 @@ def check_environment() -> List[Tuple[str, str, bool]]:
 
     return issues
 
-
 def signal_handler(signum, frame):
     """Handle system signals for graceful shutdown"""
     print(f"\n{Fore.YELLOW}Received signal {signum}, shutting down server...{Style.RESET_ALL}")
@@ -173,7 +176,6 @@ def signal_handler(signum, frame):
 
     # Attempt to run shutdown tasks
     try:
-        # Import here to avoid circular imports
         from .core.app import shutdown_event
 
         loop = asyncio.get_event_loop()
@@ -189,106 +191,69 @@ def delayed_exit():
 
     threading.Thread(target=delayed_exit, daemon=True).start()
 
-
-def start_server(
-    host: str = None,
-    port: int = None,
-    use_ngrok: bool = False,
-    ngrok_auth_token: str = None,
-    **kwargs
-):
-    """Start the FastAPI server with optional ngrok tunnel"""
-    from .config import HOST, PORT
+def start_server(use_ngrok: bool = False, port=8000, ngrok_auth_token: Optional[str] = None):
+    """Start the LocalLab server directly in the main process"""
 
-    host = host or HOST
-    port = port or PORT
+    # Set initial server status
+    set_server_status("initializing")
 
-    public_url = None
-    
-    # Set up ngrok if requested
-    if use_ngrok:
-        try:
-            # Import pyngrok
-            from pyngrok import ngrok, conf
-            import nest_asyncio
-            
-            # Apply nest_asyncio for Colab environment
-            nest_asyncio.apply()
-            
-            # Set ngrok auth token if provided
-            if ngrok_auth_token:
-                ngrok.set_auth_token(ngrok_auth_token)
-            
-            # Start ngrok tunnel
-            public_url = ngrok.connect(port).public_url
-            print(f"Ngrok tunnel established! Public URL: {public_url}")
-            
-            # Store the public URL in environment variables for access across modules
-            os.environ["LOCALLAB_PUBLIC_URL"] = public_url
-            
-        except ImportError:
-            print("Error: pyngrok not installed. Install with 'pip install pyngrok'.")
-            return
-        except Exception as e:
-            print(f"Error setting up ngrok: {str(e)}")
-            return
+    # Display startup banner with INITIALIZING status
+    print_initializing_banner(__version__)
 
-    # Set up uvicorn config with public_url for on_startup function to access
-    from .core.app import app
+    # Check environment for issues
+    issues = check_environment()
+    if issues:
+        print(f"\n{Fore.YELLOW}⚠️ Environment Check Results:{Style.RESET_ALL}")
+        for issue, suggestion, is_critical in issues:
+            prefix = f"{Fore.RED}CRITICAL:" if is_critical else f"{Fore.YELLOW}WARNING:"
+            print(f"{prefix} {issue}{Style.RESET_ALL}")
+            print(f"  {Fore.CYAN}Suggestion: {suggestion}{Style.RESET_ALL}\n")
 
-    # Inject public_url into app state for use in on_startup
-    app.state.public_url = public_url
-    
-    # Start uvicorn server
-    import uvicorn
-    uvicorn.run(
-        "locallab.core.app:app",
-        host=host,
-        port=port,
-        reload=False,
-        log_level="info",
-        **kwargs
-    )
-
-async def on_startup():
-    """Initialize server and display startup banner"""
-    from .core.app import app
-    from .logger import update_server_status, get_logger
-    
-    # Get logger
-    logger = get_logger("locallab.server")
-    
-    # Update server status
-    update_server_status("running")
-    
-    # Determine server URL - check app.state first for ngrok url
-    server_url = getattr(app.state, "public_url", None)
-    
-    if not server_url:
-        # Check environment variable
-        server_url = os.environ.get("LOCALLAB_PUBLIC_URL")
-        
-    # If still no URL, use the local URL
-    if not server_url:
-        from .config import HOST, PORT
-        server_url = f"http://{HOST}:{PORT}"
+    # Check if port is already in use
+    if is_port_in_use(port):
+        logger.warning(f"Port {port} is already in use. Trying to find another port...")
+        for p in range(port+1, port+100):
+            if not is_port_in_use(p):
+                port = p
+                logger.info(f"Using alternative port: {port}")
+                break
+        else:
+            raise RuntimeError(f"Could not find an available port in range {port}-{port+100}")
 
-    # Log server URL
-    logger.info(f"Server running at: {server_url}")
+    # Set up ngrok before starting server if requested
+    public_url = None
+    if use_ngrok:
+        logger.info(f"{Fore.CYAN}Setting up ngrok tunnel to port {port}...{Style.RESET_ALL}")
+        public_url = setup_ngrok(port=port, auth_token=ngrok_auth_token)
+        if public_url:
+            ngrok_section = f"\n{Fore.CYAN}┌────────────────────────── Ngrok Tunnel Details ─────────────────────────────┐{Style.RESET_ALL}\n│\n│  🚀 Ngrok Public URL: {Fore.GREEN}{public_url}{Style.RESET_ALL}\n│\n{Fore.CYAN}└──────────────────────────────────────────────────────────────────────────────┘{Style.RESET_ALL}\n"
+            logger.info(ngrok_section)
+            print(ngrok_section)
+        else:
+            logger.warning(f"{Fore.YELLOW}Failed to set up ngrok tunnel. Server will run locally on port {port}.{Style.RESET_ALL}")
 
-    # Print system instructions and banners
-    from .ui.banners import print_startup_banner, print_system_instructions, print_model_info, print_api_docs
+    # Server info section
+    server_section = f"\n{Fore.CYAN}┌────────────────────────── Server Details ─────────────────────────────┐{Style.RESET_ALL}\n│\n│  🖥️ Local URL: {Fore.GREEN}http://localhost:{port}{Style.RESET_ALL}\n│  ⚙️ Status: {Fore.GREEN}Starting{Style.RESET_ALL}\n│\n{Fore.CYAN}└──────────────────────────────────────────────────────────────────────────────┘{Style.RESET_ALL}\n"
+    print(server_section, flush=True)
 
-    print_startup_banner()
-    print_system_instructions(server_url)
+    # Set up signal handlers for graceful shutdown
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
 
-    # Print model info if a model is already loaded
-    from .core.app import model_manager
-    if model_manager.current_model:
-        print_model_info(model_manager.current_model)
+    # Import app here to avoid circular imports
+    try:
+        from .core.app import app
+    except ImportError as e:
+        logger.error(f"{Fore.RED}Failed to import app: {str(e)}{Style.RESET_ALL}")
+        raise
 
-    # Print API docs with the correct server URL
-    print_api_docs(server_url)
+    # Start the server
+    try:
+        logger.info(f"Starting server on port {port}...")
+        uvicorn.run(app, host="0.0.0.0", port=port)
+    except Exception as e:
+        logger.error(f"Server startup failed: {str(e)}")
+        raise 
 
 def cli():
     """Command line interface entry point for the package"""
@@ -305,4 +270,4 @@ def run(use_ngrok, port, ngrok_auth_token):
     run()
 
 if __name__ == "__main__":
-    cli() 
+    cli()

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.3.2",
+    version="0.3.3",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",