Fixed some issues

Author: UtkarshTheDev

locallab/config.py

@@ -1,48 +1,34 @@
 import os
 import json
-from typing import Dict, Any, Optional, List, Type
-from locallab.cli.config import ensure_config_dir
-import torch
-import psutil
-from huggingface_hub import model_info, HfApi
 import logging
+from typing import Dict, Any, Optional, List, Type
 from pathlib import Path
 import click
+from huggingface_hub import model_info, HfApi
 
+# Standard environment variable names
+NGROK_TOKEN_ENV = "NGROK_AUTHTOKEN"  # Official ngrok env var name
+HF_TOKEN_ENV = "HUGGINGFACE_TOKEN"   # Standard HF token env var
 
 def get_env_var(key: str, *, default: Any = None, var_type: Type = str) -> Any:
-    """Get environment variable with type conversion and validation.
-
-    Args:
-        key: Environment variable key
-        default: Default value if not found
-        var_type: Type to convert to (str, int, float, bool)
-
-    Returns:
-        Converted and validated value
-    """
-    # First check environment variables
+    """Get environment variable with type conversion and validation."""
     value = os.environ.get(key)
 
-    # If not found in environment, try the config file
     if value is None:
         try:
-            # Import here to avoid circular imports
             from .cli.config import get_config_value
-            # Convert key format: LOCALLAB_ENABLE_QUANTIZATION -> enable_quantization
-            if key.startswith("LOCALLAB_"):
-                config_key = key[9:].lower()
-            else:
-                config_key = key.lower()
+            config_key = key.lower()
+            if key == NGROK_TOKEN_ENV:
+                config_key = "ngrok_auth_token"
+            elif key == HF_TOKEN_ENV:
+                config_key = "huggingface_token"
 
             config_value = get_config_value(config_key)
             if config_value is not None:
                 value = config_value
-        except (ImportError, ModuleNotFoundError):
-            # If the config module isn't available yet, just use the environment variable
+        except:
             pass
 
-    # If still not found, use default
     if value is None:
         return default
 
@@ -54,6 +40,50 @@ def get_env_var(key: str, *, default: Any = None, var_type: Type = str) -> Any:
         logging.warning(f"Invalid value for {key}, using default: {default}")
         return default
 
+def set_env_var(name: str, value: str):
+    """Set environment variable with proper string handling"""
+    os.environ[name] = str(value).strip()
+
+def get_hf_token(interactive: bool = False) -> Optional[str]:
+    """Get HuggingFace token from environment or config"""
+    token = get_env_var(HF_TOKEN_ENV)
+    if token:
+        set_env_var(HF_TOKEN_ENV, token)  # Ensure it's in env vars
+    return token
+
+def get_ngrok_token() -> Optional[str]:
+    """Get ngrok token from environment or config"""
+    token = get_env_var(NGROK_TOKEN_ENV)
+    if token:
+        set_env_var(NGROK_TOKEN_ENV, token)  # Ensure it's in env vars
+    return token
+
+def save_config(config: Dict[str, Any]):
+    """Save configuration to file"""
+    from .cli.config import ensure_config_dir, CONFIG_FILE
+    
+    ensure_config_dir()
+    
+    # Ensure tokens are stored as proper strings
+    if "ngrok_auth_token" in config:
+        token = str(config["ngrok_auth_token"]).strip()
+        config["ngrok_auth_token"] = token
+        set_env_var(NGROK_TOKEN_ENV, token)
+        
+    if "huggingface_token" in config:
+        token = str(config["huggingface_token"]).strip()
+        config["huggingface_token"] = token
+        set_env_var(HF_TOKEN_ENV, token)
+    
+    try:
+        with open(CONFIG_FILE, "w") as f:
+            json.dump(config, f, indent=2)
+    except Exception as e:
+        logging.error(f"Error saving config: {e}")
+
+# Standard environment variable names - use official names
+NGROK_TOKEN_ENV = "NGROK_AUTHTOKEN"  # Official ngrok env var name
+HF_TOKEN_ENV = "HUGGINGFACE_TOKEN"   # Standard HF token env var
 
 # Server settings
 SERVER_HOST = get_env_var("SERVER_HOST", default="0.0.0.0")

locallab/utils/system.py

@@ -15,11 +15,13 @@
 from typing import Optional, Tuple, Dict, Any, List
 
 from ..logger import get_logger
-from ..config import MIN_FREE_MEMORY
 
 # Get logger
 logger = get_logger("locallab.utils.system")
 
+# System constants
+MIN_FREE_MEMORY = 2000  # Minimum required free memory in MB
+MIN_GPU_MEMORY = 4000  # Minimum required GPU memory in MB
 
 def get_system_memory() -> Tuple[int, int]:
     """Get system memory information in MB"""
@@ -28,84 +30,41 @@ def get_system_memory() -> Tuple[int, int]:
     free_memory = vm.available // (1024 * 1024)  # Convert to MB
     return total_memory, free_memory
 
-
 def get_gpu_memory() -> Optional[Tuple[int, int]]:
-    """Get GPU memory information in MB if available"""
+    """Get GPU memory information in MB"""
     if not TORCH_AVAILABLE or not torch.cuda.is_available():
         return None
-    
-    # First try nvidia-ml-py3 (nvidia_smi)
-    try:
-        import nvidia_smi
-        nvidia_smi.nvmlInit()
-        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
-        info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
 
-        total_memory = info.total // (1024 * 1024)  # Convert to MB
-        free_memory = info.free // (1024 * 1024)  # Convert to MB
-        
-        nvidia_smi.nvmlShutdown()
+    try:
+        device = torch.cuda.current_device()
+        total_memory = torch.cuda.get_device_properties(device).total_memory // (1024 * 1024)  # Convert to MB
+        free_memory = total_memory - (torch.cuda.memory_allocated() + torch.cuda.memory_reserved()) // (1024 * 1024)
         return total_memory, free_memory
-    except ImportError:
-        # If nvidia_smi not available, log at debug level to avoid noise
-        logger.debug("nvidia-ml-py3 not installed, falling back to torch for GPU info")
-        # Fall back to torch for basic info
-        try:
-            # Get basic info from torch
-            device = torch.cuda.current_device()
-            total_memory = torch.cuda.get_device_properties(device).total_memory // (1024 * 1024)
-            # Note: torch doesn't provide free memory info easily, so we estimate
-            # by allocating a tensor and seeing what's available
-            torch.cuda.empty_cache()
-            free_memory = total_memory  # Optimistic starting point
-            
-            # Rough estimate - we can't get exact free memory from torch easily
-            return total_memory, free_memory
-        except Exception as torch_error:
-            logger.debug(f"Torch GPU memory check also failed: {str(torch_error)}")
-            return None
     except Exception as e:
-        logger.debug(f"Failed to get detailed GPU memory info: {str(e)}")
-        # Fall back to torch for basic info (same as ImportError case)
-        try:
-            device = torch.cuda.current_device()
-            total_memory = torch.cuda.get_device_properties(device).total_memory // (1024 * 1024)
-            torch.cuda.empty_cache()
-            free_memory = total_memory  # Optimistic estimate
-            return total_memory, free_memory
-        except Exception:
-            return None
-
+        logger.warning(f"Failed to get GPU memory info: {e}")
+        return None
 
-def check_resource_availability(required_memory: int) -> bool:
-    """Check if system has enough resources for the requested operation"""
+def check_resource_availability(required_memory: int = MIN_FREE_MEMORY) -> bool:
+    """Check if system has enough resources"""
     _, free_memory = get_system_memory()
-    
-    # Check system memory
-    if free_memory < MIN_FREE_MEMORY:
-        logger.warning(f"Low system memory: {free_memory}MB available")
+    if free_memory < required_memory:
         return False
-    
-    # If GPU is available, check GPU memory
-    if TORCH_AVAILABLE and torch.cuda.is_available():
-        gpu_memory = get_gpu_memory()
-        if gpu_memory:
-            total_gpu, free_gpu = gpu_memory
-            if free_gpu < required_memory:
-                logger.warning(f"Insufficient GPU memory: {free_gpu}MB available, {required_memory}MB required")
-                return False
-    
+        
+    gpu_mem = get_gpu_memory()
+    if gpu_mem is not None:
+        _, free_gpu = gpu_mem
+        if free_gpu < MIN_GPU_MEMORY:
+            return False
+            
     return True
 
-
 def get_device() -> str:
     """Get the device to use for computations."""
     if TORCH_AVAILABLE and torch.cuda.is_available():
         return "cuda"
     else:
         return "cpu"
 
-
 def format_model_size(size_in_bytes: int) -> str:
     """Format model size in human-readable format"""
     for unit in ['B', 'KB', 'MB', 'GB']:
@@ -114,7 +73,6 @@ def format_model_size(size_in_bytes: int) -> str:
         size_in_bytes /= 1024
     return f"{size_in_bytes:.2f} TB"
 
-
 def get_system_resources() -> Dict[str, Any]:
     """Get system resource information"""
     resources = {
@@ -146,7 +104,6 @@ def get_system_resources() -> Dict[str, Any]:
 
     return resources
 
-
 def get_cpu_info() -> Dict[str, Any]:
     """Get information about the CPU."""
     return {
@@ -155,7 +112,6 @@ def get_cpu_info() -> Dict[str, Any]:
         "usage": psutil.cpu_percent(interval=0.1)
     }
 
-
 def get_gpu_info() -> List[Dict[str, Any]]:
     """Get detailed information about all available GPUs.
     
@@ -231,7 +187,6 @@ def get_gpu_info() -> List[Dict[str, Any]]:
 
     return gpu_info
 
-
 def get_memory_info() -> Dict[str, Any]:
     """Get information about the system memory."""
     mem = psutil.virtual_memory()
@@ -242,8 +197,7 @@ def get_memory_info() -> Dict[str, Any]:
         "percent": mem.percent
     }
 
-
 # Add this function for backward compatibility
 def get_system_info() -> Dict[str, Any]:
     """Get system resource information (alias for get_system_resources)"""
-    return get_system_resources() 
+    return get_system_resources()

setup.py

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