lemonade-sdk 9.1.1__py3-none-any.whl

lemonade/tools/server/flm.py

@@ -0,0 +1,133 @@
+import os
+import logging
+import subprocess
+import time
+import threading
+
+import requests
+
+from lemonade_server.pydantic_models import (
+    PullConfig,
+    ChatCompletionRequest,
+)
+
+from lemonade.tools.server.wrapped_server import WrappedServerTelemetry, WrappedServer
+from lemonade.tools.flm.utils import install_flm, download_flm_model
+
+
+class FlmTelemetry(WrappedServerTelemetry):
+    """
+    Manages telemetry data collection and display for FLM server.
+    """
+
+    def parse_telemetry_line(self, line: str):
+        """
+        Parse telemetry data from FLM server output lines.
+
+        Note: as of FLM 0.9.10, no telemetry data is provided by the server CLI.
+                This function is required to be implemented, so we leave it empty
+                as a placeholder for now.
+        """
+
+        return
+
+
+class FlmServer(WrappedServer):
+    """
+    Routes OpenAI API requests to an FLM server instance and returns the result
+    back to Lemonade Server.
+    """
+
+    def __init__(self):
+        self.flm_model_name = None
+        super().__init__(server_name="flm-server", telemetry=FlmTelemetry())
+
+    def address(self):
+        return f"http://localhost:{self.port}/v1"
+
+    def install_server(self):
+        """
+        Check if FLM is installed and at minimum version.
+        If not, download and run the GUI installer, then wait for completion.
+        """
+        install_flm()
+
+    def download_model(
+        self, config_checkpoint, config_mmproj=None, do_not_upgrade=False
+    ) -> dict:
+        download_flm_model(config_checkpoint, config_mmproj, do_not_upgrade)
+
+    def _launch_server_subprocess(
+        self,
+        model_config: PullConfig,
+        snapshot_files: dict,
+        ctx_size: int,
+        supports_embeddings: bool = False,
+        supports_reranking: bool = False,
+    ):
+
+        self._choose_port()
+
+        # Keep track of the FLM model name so that we can use it later
+        self.flm_model_name = model_config.checkpoint
+
+        command = [
+            "flm",
+            "serve",
+            f"{self.flm_model_name}",
+            "--ctx-len",
+            str(ctx_size),
+            "--port",
+            str(self.port),
+        ]
+
+        # Set up environment with library path for Linux
+        env = os.environ.copy()
+
+        self.process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+            bufsize=1,
+            env=env,
+        )
+
+        # Start background thread to log subprocess output
+        threading.Thread(
+            target=self._log_subprocess_output,
+            args=("FLM SERVER",),
+            daemon=True,
+        ).start()
+
+    def _wait_for_load(self):
+        """
+        FLM doesn't seem to have a health API, so we'll use the "list local models"
+        API to check if the server is up.
+        """
+        status_code = None
+        while not self.process.poll() and status_code != 200:
+            health_url = f"http://localhost:{self.port}/api/tags"
+            try:
+                health_response = requests.get(health_url)
+            except requests.exceptions.ConnectionError:
+                logging.debug(
+                    "Not able to connect to %s yet, will retry", self.server_name
+                )
+            else:
+                status_code = health_response.status_code
+                logging.debug(
+                    "Testing %s readiness (will retry until ready), result: %s",
+                    self.server_name,
+                    health_response.json(),
+                )
+            time.sleep(1)
+
+    def chat_completion(self, chat_completion_request: ChatCompletionRequest):
+        # FLM requires the correct model name to be in the request
+        # (whereas llama-server ignores the model name field in the request)
+        chat_completion_request.model = self.flm_model_name
+
+        return super().chat_completion(chat_completion_request)

lemonade/tools/server/llamacpp.py

@@ -0,0 +1,320 @@
+import os
+import logging
+import subprocess
+import re
+import threading
+import platform
+
+from dotenv import load_dotenv
+from fastapi import HTTPException, status
+
+from lemonade_server.pydantic_models import (
+    PullConfig,
+)
+from lemonade.tools.llamacpp.utils import (
+    get_llama_server_exe_path,
+    install_llamacpp,
+    download_gguf,
+    resolve_local_gguf_model,
+    parse_checkpoint,
+)
+from lemonade.tools.server.wrapped_server import WrappedServerTelemetry, WrappedServer
+
+# Embedding model batch configuration set to 8192 as default
+EMBEDDING_CTX_SIZE = 8192
+EMBEDDING_BATCH_SIZE = 8192
+EMBEDDING_UBATCH_SIZE = 8192
+
+
+class LlamaTelemetry(WrappedServerTelemetry):
+    """
+    Manages telemetry data collection and display for llama server.
+    """
+
+    def parse_telemetry_line(self, line: str):
+        """
+        Parse telemetry data from llama server output lines.
+        """
+
+        if "vk::PhysicalDevice::createDevice: ErrorExtensionNotPresent" in line:
+            msg = (
+                "Your AMD GPU driver version is not compatible with this software.\n"
+                "Please update and try again: "
+                "https://www.amd.com/en/support/download/drivers.html"
+            )
+            logging.error(msg)
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=msg,
+            )
+        elif "error" in line.lower():
+            logging.error(line)
+
+        # Parse Vulkan device detection
+        vulkan_match = re.search(r"ggml_vulkan: Found (\d+) Vulkan devices?:", line)
+        if vulkan_match:
+            device_count = int(vulkan_match.group(1))
+            if device_count > 0:
+                logging.info(
+                    f"GPU acceleration active: {device_count} device(s) "
+                    "detected by llama-server"
+                )
+            return
+
+        # Parse prompt evaluation line
+        prompt_match = re.search(
+            r"prompt eval time\s*=\s*([\d.]+)\s*ms\s*/\s*(\d+)\s*tokens.*?"
+            r"([\d.]+)\s*tokens per second",
+            line,
+        )
+        if prompt_match:
+            prompt_time_ms = float(prompt_match.group(1))
+            input_tokens = int(prompt_match.group(2))
+
+            self.prompt_eval_time = prompt_time_ms / 1000.0
+            self.input_tokens = input_tokens
+            self.time_to_first_token = prompt_time_ms / 1000.0
+            return
+
+        # Parse generation evaluation line
+        eval_match = re.search(
+            r"eval time\s*=\s*([\d.]+)\s*ms\s*/\s*(\d+)\s*tokens.*?"
+            r"([\d.]+)\s*tokens per second",
+            line,
+        )
+        if eval_match:
+            eval_time_ms = float(eval_match.group(1))
+            output_tokens = int(eval_match.group(2))
+            tokens_per_second = float(eval_match.group(3))
+
+            self.eval_time = eval_time_ms / 1000.0
+            self.output_tokens = output_tokens
+            self.tokens_per_second = tokens_per_second
+            return
+
+
+class LlamaServer(WrappedServer):
+    def __init__(self, backend: str):
+        self.backend = backend
+        super().__init__(server_name="llama-server", telemetry=LlamaTelemetry())
+
+    def install_server(self, backend=None):
+        """
+        Install the wrapped server
+        """
+        install_llamacpp(self.backend)
+
+    def download_model(
+        self, config_checkpoint, config_mmproj=None, do_not_upgrade=False
+    ) -> dict:
+        """
+        Download a model for the wrapper server.
+        First checks local cache, then downloads from internet if needed.
+        """
+        # If it's a direct file path, just return it
+
+        if os.path.exists(config_checkpoint):
+            result = {"variant": config_checkpoint}
+            if config_mmproj:
+                result["mmproj"] = config_mmproj
+            return result
+
+        # Try to resolve from local cache first to avoid unnecessary downloads
+        checkpoint, variant = parse_checkpoint(config_checkpoint)
+        local_result = resolve_local_gguf_model(checkpoint, variant, config_mmproj)
+
+        if local_result:
+            return local_result
+
+        # Not found locally - download from internet
+        return download_gguf(
+            config_checkpoint=config_checkpoint,
+            config_mmproj=config_mmproj,
+            do_not_upgrade=do_not_upgrade,
+        )
+
+    def _launch_device_backend_subprocess(
+        self,
+        snapshot_files: dict,
+        use_gpu: bool,
+        ctx_size: int,
+        supports_embeddings: bool = False,
+        supports_reranking: bool = False,
+    ) -> subprocess.Popen:
+        """
+        Launch llama server subprocess with appropriate configuration.
+
+        Args:
+            snapshot_files: Dictionary of model files to load
+            use_gpu: Whether to use GPU acceleration
+            telemetry: Telemetry object for tracking performance metrics
+            backend: Backend to use (e.g., 'vulkan', 'rocm', 'cpu')
+            supports_embeddings: Whether the model supports embeddings
+            supports_reranking: Whether the model supports reranking
+
+        Returns:
+            Subprocess handle for the llama server
+        """
+
+        # Get the current executable path (handles both Windows and Ubuntu structures)
+        exe_path = get_llama_server_exe_path(self.backend)
+
+        # For embedding models, use a larger context size to support longer individual
+        # strings. Embedding requests can include multiple strings in a batch, and each
+        # string needs to fit within the context window.
+        if supports_embeddings and ctx_size < EMBEDDING_CTX_SIZE:
+            ctx_size = EMBEDDING_CTX_SIZE
+
+        # Build the base command
+        base_command = [
+            exe_path,
+            "-m",
+            snapshot_files["variant"],
+            "--ctx-size",
+            str(ctx_size),
+        ]
+
+        # Lock random seed for deterministic behavior in CI
+        if os.environ.get("LEMONADE_CI_MODE"):
+            base_command.extend(["--seed", "42"])
+            logging.info(f"Seed applied to base command: {base_command}")
+
+        if "mmproj" in snapshot_files:
+            base_command.extend(["--mmproj", snapshot_files["mmproj"]])
+            if not use_gpu:
+                base_command.extend(["--no-mmproj-offload"])
+
+        # Find a port, and save it in the telemetry object for future reference
+        # by other functions
+        self._choose_port()
+
+        # Add port and jinja to enable tool use
+        base_command.extend(["--port", str(self.port), "--jinja"])
+
+        # Enable context shift and avoid attention sink issues by preserving the initial tokens
+        # Note: --context-shift is not supported on all backends (e.g., Metal on macOS)
+        # Only add context-shift for backends that support it
+        context_shift_supported_backends = ["vulkan", "rocm"]
+        if self.backend in context_shift_supported_backends:
+            base_command.extend(["--context-shift", "--keep", "16"])
+        else:
+            # For backends that don't support context-shift (e.g., Metal), just use keep
+            base_command.extend(["--keep", "16"])
+            logging.debug(
+                f"Skipped --context-shift for backend: {self.backend} (not supported)"
+            )
+
+        # Use legacy reasoning formatting, since not all apps support the new
+        # reasoning_content field
+        base_command.extend(["--reasoning-format", "auto"])
+
+        # Add embeddings support if the model supports it
+        if supports_embeddings:
+            # For embedding models, set batch sizes to handle multiple documents in a single request
+            # batch-size: logical batch size (total tokens across all sequences)
+            # ubatch-size: physical batch size (tokens processed in a single forward pass)
+            base_command.extend(
+                [
+                    "--embeddings",
+                    "--batch-size",
+                    str(EMBEDDING_BATCH_SIZE),
+                    "--ubatch-size",
+                    str(EMBEDDING_UBATCH_SIZE),
+                ]
+            )
+
+        # Add reranking support if the model supports it
+        if supports_reranking:
+            base_command.append("--reranking")
+
+        # Configure GPU layers: 99 for GPU, 0 for CPU-only
+        ngl_value = "99" if use_gpu else "0"
+        command = base_command + ["-ngl", ngl_value]
+
+        # Set up environment with library path for Linux
+        env = os.environ.copy()
+
+        # Load environment variables from .env file in the executable directory
+        exe_dir = os.path.dirname(exe_path)
+        env_file_path = os.path.join(exe_dir, ".env")
+        if os.path.exists(env_file_path):
+            load_dotenv(env_file_path, override=False)
+            env.update(os.environ)
+            logging.debug(f"Loaded environment variables from {env_file_path}")
+
+        system = platform.system().lower()
+        if system == "linux":
+            lib_dir = os.path.dirname(exe_path)  # Same directory as the executable
+            current_ld_path = env.get("LD_LIBRARY_PATH", "")
+            if current_ld_path:
+                env["LD_LIBRARY_PATH"] = f"{lib_dir}:{current_ld_path}"
+            else:
+                env["LD_LIBRARY_PATH"] = lib_dir
+            logging.debug(f"Set LD_LIBRARY_PATH to {env['LD_LIBRARY_PATH']}")
+        elif system == "darwin":
+            lib_dir = os.path.dirname(exe_path)
+            current_dyld_path = env.get("DYLD_LIBRARY_PATH", "")
+            if current_dyld_path:
+                env["DYLD_LIBRARY_PATH"] = f"{lib_dir}:{current_dyld_path}"
+            else:
+                env["DYLD_LIBRARY_PATH"] = lib_dir
+            logging.debug(f"Set DYLD_LIBRARY_PATH to {env['DYLD_LIBRARY_PATH']}")
+
+        # Start subprocess with output capture
+        self.process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+            bufsize=1,
+            env=env,
+        )
+
+        # Start background thread to log subprocess output
+        device_type = "GPU" if use_gpu else "CPU"
+        threading.Thread(
+            target=self._log_subprocess_output,
+            args=(f"LLAMA SERVER {device_type}",),
+            daemon=True,
+        ).start()
+
+    def _launch_server_subprocess(
+        self,
+        model_config: PullConfig,
+        snapshot_files: dict,
+        ctx_size: int,
+        supports_embeddings: bool = False,
+        supports_reranking: bool = False,
+    ):
+
+        # Attempt loading on GPU first
+        self._launch_device_backend_subprocess(
+            snapshot_files,
+            use_gpu=True,
+            ctx_size=ctx_size,
+            supports_embeddings=supports_embeddings,
+            supports_reranking=supports_reranking,
+        )
+
+        # Check the /health endpoint until GPU server is ready
+        self._wait_for_load()
+
+        # If loading on GPU failed, try loading on CPU
+        if self.process.poll():
+            logging.warning(
+                f"Loading {model_config.model_name} on GPU didn't work, re-attempting on CPU"
+            )
+
+            if os.environ.get("LEMONADE_LLAMACPP_NO_FALLBACK"):
+                # Used for testing, when the test should fail if GPU didn't work
+                raise Exception("llamacpp GPU loading failed")
+
+            self._launch_device_backend_subprocess(
+                snapshot_files,
+                use_gpu=False,
+                ctx_size=ctx_size,
+                supports_embeddings=supports_embeddings,
+                supports_reranking=supports_reranking,
+            )