lpm_kernel/api/services/local_llm_service.py

import os
import json
import logging
import psutil
import time
import subprocess
import torch  # Add torch import for CUDA detection
import threading
import queue
from typing import Iterator, Any, Optional, Generator, Dict
from datetime import datetime
from flask import Response
from openai import OpenAI
from lpm_kernel.api.domains.kernel2.dto.server_dto import ServerStatus, ProcessInfo
from lpm_kernel.configs.config import Config
import uuid

logger = logging.getLogger(__name__)

class LocalLLMService:
    """Service for managing local LLM client and server"""
    
    def __init__(self):
        self._client = None
        self._stopping_server = False
        
    @property
    def client(self) -> OpenAI:
        config = Config.from_env()
        """Get the OpenAI client for local LLM server"""
        if self._client is None:
            base_url = config.get("LOCAL_LLM_SERVICE_URL")
            if not base_url:
                raise ValueError("LOCAL_LLM_SERVICE_URL environment variable is not set")
                
            self._client = OpenAI(
                base_url=base_url,
                api_key="sk-no-key-required"
            )
        return self._client

    def start_server(self, model_path: str, use_gpu: bool = True) -> bool:
        """
        Start the llama-server service with GPU acceleration when available
        
        Args:
            model_path: Path to the GGUF model file
            use_gpu: Whether to use GPU acceleration if available
            
        Returns:
            bool: True if server started successfully, False otherwise
        """
        try:
            # Check if server is already running
            status = self.get_server_status()
            if status.is_running:
                logger.info("LLama server is already running")
                return True

            # Check for CUDA availability if GPU was requested
            cuda_available = torch.cuda.is_available() if use_gpu else False
            cuda_available = False
            gpu_info = ""
            
            if use_gpu and cuda_available:
                gpu_device = torch.cuda.current_device()
                gpu_info = f" using GPU: {torch.cuda.get_device_name(gpu_device)}"
                gpu_memory = torch.cuda.get_device_properties(gpu_device).total_memory / (1024**3)
                
                logger.info(f"CUDA is available. Using GPU acceleration{gpu_info}")
                logger.info(f"CUDA device capabilities: {torch.cuda.get_device_capability(gpu_device)}")
                logger.info(f"CUDA memory: {gpu_memory:.2f} GB")
                
                # Pre-initialize CUDA to speed up first inference
                logger.info("Pre-initializing CUDA context to speed up first inference")
                torch.cuda.init()
                torch.cuda.empty_cache()
            elif use_gpu and not cuda_available:
                logger.warning("CUDA was requested but is not available. Using CPU instead.")
            else:
                logger.info("Using CPU for inference (GPU not requested)")

            # Check for GPU optimization marker
            gpu_optimized = False
            model_dir = os.path.dirname(model_path)
            gpu_marker_path = os.path.join(model_dir, "gpu_optimized.json")
            if os.path.exists(gpu_marker_path):
                try:
                    with open(gpu_marker_path, 'r') as f:
                        gpu_data = json.load(f)
                        if gpu_data.get("gpu_optimized", False):
                            gpu_optimized = True
                            logger.info(f"Found GPU optimization marker created on {gpu_data.get('optimized_on', 'unknown date')}")
                except Exception as e:
                    logger.warning(f"Error reading GPU marker file: {e}")

            # Get the correct path to the llama-server executable
            base_dir = os.getcwd()
            server_path = os.path.join(base_dir, "llama.cpp", "build", "bin", "llama-server")
            
            # For Windows, add .exe extension if needed
            if os.name == 'nt' and not server_path.endswith('.exe'):
                server_path += '.exe'
                
            # Verify executable exists
            if not os.path.exists(server_path):
                logger.error(f"llama-server executable not found at: {server_path}")
                return False
                
            # Start server with optimal parameters for faster startup
            cmd = [
                server_path,
                "-m", model_path,
                "--host", "0.0.0.0",
                "--port", "8080",
                "--ctx-size", "2048",     # Default context size (adjust based on needs)
                "--parallel", "2",        # Enable request parallelism
                "--cont-batching"         # Enable continuous batching
            ]
            
            # Set up environment with CUDA variables to ensure GPU detection
            env = os.environ.copy()
            env["CUDA_VISIBLE_DEVICES"] = ""
            
            # Add GPU-related parameters if CUDA is available
            if cuda_available and use_gpu:
                # Force GPU usage with optimal parameters for faster loads
                cmd.extend([
                    "--n-gpu-layers", "999",  # Use all layers on GPU
                    "--tensor-split", "0",    # Use the first GPU for all operations
                    "--main-gpu", "0",        # Use GPU 0 as the primary device
                    "--mlock"                 # Lock memory to prevent swapping during inference
                ])
                
                # Set CUDA environment variables to help with GPU detection
                env["CUDA_VISIBLE_DEVICES"] = "0"  # Force using first GPU
                
                # Ensure comprehensive library paths for CUDA
                cuda_lib_paths = [
                    "/usr/local/cuda/lib64",
                    "/usr/lib/cuda/lib64",
                    "/usr/local/lib",
                    "/usr/lib/x86_64-linux-gnu",
                    "/usr/lib/wsl/lib"  # For Windows WSL environments
                ]
                
                # Build a comprehensive LD_LIBRARY_PATH
                current_ld_path = env.get("LD_LIBRARY_PATH", "")
                for path in cuda_lib_paths:
                    if os.path.exists(path) and path not in current_ld_path:
                        current_ld_path = f"{path}:{current_ld_path}" if current_ld_path else path
                
                env["LD_LIBRARY_PATH"] = current_ld_path
                logger.info(f"Setting LD_LIBRARY_PATH to: {current_ld_path}")
                
                # If this is Windows, use different approach for CUDA libraries
                if os.name == 'nt':
                    # Windows typically has CUDA in PATH already if installed
                    logger.info("Windows system detected, using system CUDA libraries")
                else:
                    # On Linux, try to find CUDA libraries in common locations
                    for cuda_path in [
                        # Common CUDA paths
                        "/usr/local/cuda/lib64",
                        "/usr/lib/cuda/lib64",
                        "/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib",
                        "/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib",
                    ]:
                        if os.path.exists(cuda_path):
                            # Add CUDA path to library path
                            env["LD_LIBRARY_PATH"] = f"{cuda_path}:{env.get('LD_LIBRARY_PATH', '')}"
                            env["CUDA_HOME"] = os.path.dirname(cuda_path)
                            logger.info(f"Found CUDA at {cuda_path}, setting environment variables")
                            break

                # NOTE: CUDA support and rebuild should be handled at build/setup time (e.g., Docker build or setup script).
                # The runtime check and rebuild logic has been removed for efficiency and reliability.
                # Ensure llama.cpp is built with CUDA support before running the server if GPU is required.

                # Pre-heat GPU to ensure faster initial response
                if torch.cuda.is_available():
                    logger.info("Pre-warming GPU to reduce initial latency...")
                    dummy_tensor = torch.zeros(1, 1).cuda()
                    del dummy_tensor
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
                    logger.info("GPU warm-up complete")
                
                logger.info("Using GPU acceleration for inference with optimized settings")
            else:
                # If GPU isn't available or supported, optimize for CPU
                cmd.extend([
                    "--threads", str(max(1, os.cpu_count() - 1)),  # Use all CPU cores except one
                ])
                logger.info(f"Using CPU-only mode with {max(1, os.cpu_count() - 1)} threads")
            
            logger.info(f"Starting llama-server with command: {' '.join(cmd)}")
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
                env=env
            )
            
            # Wait for server to start (longer wait for GPU initialization)
            wait_time = 5 if cuda_available and use_gpu else 3
            logger.info(f"Waiting {wait_time} seconds for server to start...")
            time.sleep(wait_time)
            
            # Check if process is still running
            if process.poll() is None:
                # Log initialization success
                if cuda_available and use_gpu:
                    logger.info(f"✅ LLama server started successfully with GPU acceleration{gpu_info}")
                else:
                    logger.info("✅ LLama server started successfully in CPU-only mode")
                return True
            else:
                stdout, stderr = process.communicate()
                logger.error(f"Failed to start llama-server: {stderr}")
                return False
                
        except Exception as e:
            logger.error(f"Error starting llama-server: {str(e)}")
            return False

    def stop_server(self) -> ServerStatus:
        """
        Stop the llama-server service.
        Find and forcibly terminate all llama-server processes
        
        Returns:
            ServerStatus: Service status object containing information about whether processes are still running
        """
        try:
            if self._stopping_server:
                logger.info("Server is already in the process of stopping")
                return self.get_server_status()
            
            self._stopping_server = True
        
            try:
                # Find all possible llama-server processes and forcibly terminate them
                terminated_pids = []
                for proc in psutil.process_iter(["pid", "name", "cmdline"]):
                    try:
                        cmdline = proc.cmdline()
                        if any("llama-server" in cmd for cmd in cmdline):
                            pid = proc.pid
                            logger.info(f"Force terminating llama-server process, PID: {pid}")
                            
                            # Directly use kill signal to forcibly terminate
                            proc.kill()
                            
                            # Ensure the process has been terminated
                            try:
                                proc.wait(timeout=0.2)  # Slightly increase wait time to ensure process termination
                                terminated_pids.append(pid)
                                logger.info(f"Successfully terminated llama-server process {pid}")
                            except psutil.TimeoutExpired:
                                # If timeout, try to terminate again
                                logger.warning(f"Process {pid} still running, sending SIGKILL again")
                                try:
                                    import os
                                    import signal
                                    os.kill(pid, signal.SIGKILL)  # Use system-level SIGKILL signal
                                    terminated_pids.append(pid)
                                    logger.info(f"Successfully force killed llama-server process {pid} with SIGKILL")
                                except ProcessLookupError:
                                    # Process no longer exists
                                    terminated_pids.append(pid)
                                    logger.info(f"Process {pid} no longer exists after kill attempt")
                    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                        continue
                
                if terminated_pids:
                    logger.info(f"Terminated llama-server processes: {terminated_pids}")
                else:
                    logger.info("No running llama-server process found")
                
                # Check again if any llama-server processes are still running
                return self.get_server_status()
            
            finally:
                self._stopping_server = False
            
        except Exception as e:
            logger.error(f"Error stopping llama-server: {str(e)}")
            self._stopping_server = False
            return ServerStatus.not_running()

    def get_server_status(self) -> ServerStatus:
        """
        Get the current status of llama-server
        Returns: ServerStatus object
        """
        try:
            base_dir = os.getcwd()
            server_path = os.path.join(base_dir, "llama.cpp", "build", "bin", "llama-server")
            server_exec_name = os.path.basename(server_path)
            
            for proc in psutil.process_iter(["pid", "name", "cmdline"]):
                try:
                    cmdline = proc.cmdline()
                    # Check both for the executable name and the full path
                    if any(server_exec_name in cmd for cmd in cmdline) or any("llama-server" in cmd for cmd in cmdline):
                        with proc.oneshot():
                            process_info = ProcessInfo(
                                pid=proc.pid,
                                cpu_percent=proc.cpu_percent(),
                                memory_percent=proc.memory_percent(),
                                create_time=proc.create_time(),
                                cmdline=cmdline,
                            )
                            return ServerStatus.running(process_info)
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    continue
                    
            return ServerStatus.not_running()
            
        except Exception as e:
            logger.error(f"Error checking llama-server status: {str(e)}")
            return ServerStatus.not_running()

    def _parse_response_chunk(self, chunk):
        """Parse different response chunk formats into a standardized format."""
        try:
            if chunk is None:
                logger.warning("Received None chunk")
                return None
                
            # logger.info(f"Parsing response chunk: {chunk}")
            # Handle custom format
            if isinstance(chunk, dict) and "type" in chunk and chunk["type"] == "chat_response":
                logger.info(f"Processing custom format response: {chunk}")
                return {
                    "id": str(uuid.uuid4()),  # Generate a unique ID
                    "object": "chat.completion.chunk",
                    "created": int(datetime.now().timestamp()),
                    "model": "models/lpm",
                    "system_fingerprint": None,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {
                                "content": chunk.get("content", "")
                            },
                            "finish_reason": "stop" if chunk.get("done", False) else None
                        }
                    ]
                }
            
            # Handle OpenAI format
            if not hasattr(chunk, 'choices'):
                logger.warning(f"Chunk has no choices attribute: {chunk}")
                return None
                
            choices = getattr(chunk, 'choices', [])
            if not choices:
                logger.warning("Chunk has empty choices")
                return None
                
            # logger.info(f"Processing OpenAI format response: choices={choices}")
            delta = choices[0].delta
            
            # Create standard response structure
            response_data = {
                "id": chunk.id,
                "object": "chat.completion.chunk",
                "created": int(datetime.now().timestamp()),
                "model": "models/lpm",
                "system_fingerprint": chunk.system_fingerprint if hasattr(chunk, 'system_fingerprint') else None,
                "choices": [
                    {
                        "index": 0,
                        "delta": {
                            # Keep even if content is None, let the client handle it
                            "content": delta.content if hasattr(delta, 'content') else ""
                        },
                        "finish_reason": choices[0].finish_reason
                    }
                ]
            }
            
            # If there is neither content nor finish_reason, skip
            if not (hasattr(delta, 'content') or choices[0].finish_reason):
                logger.debug("Skipping chunk with no content and no finish_reason")
                return None
                
            return response_data
            
        except Exception as e:
            logger.error(f"Error parsing response chunk: {e}, chunk: {chunk}")
            return None

    def handle_stream_response(self, response_iter: Iterator[Any]) -> Response:
        """Handle streaming response from the LLM server"""
        # Create a queue for thread communication
        message_queue = queue.Queue()
        # Create an event flag to notify when model processing is complete
        completion_event = threading.Event()
        # Create a variable to track if heartbeat is needed after first response
        first_response_received = False
        
        def heartbeat_thread():
            """Thread function for sending heartbeats"""
            start_time = time.time()
            heartbeat_interval = 10  # Send heartbeat every 10 seconds
            heartbeat_count = 0
            
            logger.info("[STREAM_DEBUG] Heartbeat thread started")
            
            try:
                # Send initial heartbeat
                message_queue.put((b": initial heartbeat\n\n", "[INITIAL_HEARTBEAT]"))
                last_heartbeat_time = time.time()
                
                while not completion_event.is_set():
                    current_time = time.time()
                    
                    # Check if we need to send a heartbeat
                    if current_time - last_heartbeat_time >= heartbeat_interval:
                        heartbeat_count += 1
                        elapsed = current_time - start_time
                        logger.info(f"[STREAM_DEBUG] Sending heartbeat #{heartbeat_count} at {elapsed:.2f}s")
                        message_queue.put((f": heartbeat #{heartbeat_count}\n\n".encode('utf-8'), "[HEARTBEAT]"))
                        last_heartbeat_time = current_time
                    
                    # Short sleep to prevent CPU spinning
                    time.sleep(0.1)
                
                logger.info(f"[STREAM_DEBUG] Heartbeat thread stopping after {heartbeat_count} heartbeats")
            except Exception as e:
                logger.error(f"[STREAM_DEBUG] Error in heartbeat thread: {str(e)}", exc_info=True)
                message_queue.put((f"data: {{\"error\": \"Heartbeat error: {str(e)}\"}}\n\n".encode('utf-8'), "[ERROR]"))
        
        def model_response_thread():
            """Thread function for processing model responses"""
            chunk = None
            start_time = time.time()
            chunk_count = 0
            
            try:
                logger.info("[STREAM_DEBUG] Model response thread started")
                
                # Process model responses
                for chunk in response_iter:
                    current_time = time.time()
                    elapsed_time = current_time - start_time
                    chunk_count += 1
                    
                    logger.info(f"[STREAM_DEBUG] Received chunk #{chunk_count} after {elapsed_time:.2f}s")
                    
                    if chunk is None:
                        logger.warning("[STREAM_DEBUG] Received None chunk, skipping")
                        continue
                    
                    # Check if it's an end marker
                    if chunk == "[DONE]":
                        logger.info(f"[STREAM_DEBUG] Received [DONE] marker after {elapsed_time:.2f}s")
                        message_queue.put((b"data: [DONE]\n\n", "[DONE]"))
                        break
                    
                    # Handle error responses
                    if isinstance(chunk, dict) and "error" in chunk:
                        logger.warning(f"[STREAM_DEBUG] Received error response: {chunk}")
                        data_str = json.dumps(chunk)
                        message_queue.put((f"data: {data_str}\n\n".encode('utf-8'), "[ERROR]"))
                        message_queue.put((b"data: [DONE]\n\n", "[DONE]"))
                        break
                    
                    # Handle normal responses
                    response_data = self._parse_response_chunk(chunk)
                    if response_data:
                        data_str = json.dumps(response_data)
                        content = response_data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        content_length = len(content) if content else 0
                        logger.info(f"[STREAM_DEBUG] Sending chunk #{chunk_count}, content length: {content_length}, elapsed: {elapsed_time:.2f}s")
                        message_queue.put((f"data: {data_str}\n\n".encode('utf-8'), "[CONTENT]"))
                    else:
                        logger.warning(f"[STREAM_DEBUG] Parsed response data is None for chunk #{chunk_count}")
                
                # Handle the case where no responses were received
                if chunk_count == 0:
                    logger.info("[STREAM_DEBUG] No chunks received, sending empty message")
                    thinking_message = {
                        "id": str(uuid.uuid4()),
                        "object": "chat.completion.chunk",
                        "created": int(datetime.now().timestamp()),
                        "model": "models/lpm",
                        "system_fingerprint": None,
                        "choices": [
                            {
                                "index": 0,
                                "delta": {
                                    "content": ""  # Empty content won't affect frontend display
                                },
                                "finish_reason": None
                            }
                        ]
                    }
                    data_str = json.dumps(thinking_message)
                    message_queue.put((f"data: {data_str}\n\n".encode('utf-8'), "[THINKING]"))
                
                # Model processing is complete, send end marker
                if chunk != "[DONE]":
                    logger.info(f"[STREAM_DEBUG] Sending final [DONE] marker after {elapsed_time:.2f}s")
                    message_queue.put((b"data: [DONE]\n\n", "[DONE]"))
                
            except Exception as e:
                logger.error(f"[STREAM_DEBUG] Error processing model response: {str(e)}", exc_info=True)
                message_queue.put((f"data: {{\"error\": \"{str(e)}\"}}\n\n".encode('utf-8'), "[ERROR]"))
                message_queue.put((b"data: [DONE]\n\n", "[DONE]"))
            finally:
                # Set completion event to notify heartbeat thread to stop
                completion_event.set()
                logger.info(f"[STREAM_DEBUG] Model response thread completed with {chunk_count} chunks")
        
        def generate():
            """Main generator function for generating responses"""
            # Start heartbeat thread
            heart_thread = threading.Thread(target=heartbeat_thread, daemon=True)
            heart_thread.start()
            
            # Start model response processing thread
            model_thread = threading.Thread(target=model_response_thread, daemon=True)
            model_thread.start()
            
            try:
                # Get messages from queue and return to client
                while True:
                    try:
                        # Use short timeout to get message, prevent blocking
                        message, message_type = message_queue.get(timeout=0.1)
                        logger.debug(f"[STREAM_DEBUG] Yielding message type: {message_type}")
                        yield message
                        
                        # If end marker is received, exit loop
                        if message_type == "[DONE]":
                            logger.info("[STREAM_DEBUG] Received [DONE] marker, ending generator")
                            break
                    except queue.Empty:
                        # Queue is empty, continue trying to get message
                        # Check if model thread has completed but didn't send [DONE]
                        if completion_event.is_set() and not model_thread.is_alive():
                            logger.warning("[STREAM_DEBUG] Model thread completed without [DONE], ending generator")
                            yield b"data: [DONE]\n\n"
                            break
                        pass
            except GeneratorExit:
                # Client closed connection
                logger.info("[STREAM_DEBUG] Client closed connection (GeneratorExit)")
                completion_event.set()
            except Exception as e:
                logger.error(f"[STREAM_DEBUG] Error in generator: {str(e)}", exc_info=True)
                try:
                    yield f"data: {{\"error\": \"Generator error: {str(e)}\"}}\n\n".encode('utf-8')
                    yield b"data: [DONE]\n\n"
                except:
                    pass
                completion_event.set()
            finally:
                # Ensure completion event is set
                completion_event.set()
                # Wait for threads to complete
                if heart_thread.is_alive():
                    heart_thread.join(timeout=1.0)
                if model_thread.is_alive():
                    model_thread.join(timeout=1.0)
                logger.info("[STREAM_DEBUG] Generator completed")
        
        # Return response
        return Response(
            generate(),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache, no-transform',
                'X-Accel-Buffering': 'no',
                'Connection': 'keep-alive',
                'Transfer-Encoding': 'chunked'
            }
        )


# Global instance
local_llm_service = LocalLLMService()