lemonade-sdk 8.1.4__py3-none-any.whl → 8.2.2__py3-none-any.whl

lemonade/tools/oga/migration.py

@@ -0,0 +1,403 @@
+"""
+Migration utilities for handling RyzenAI version upgrades.
+
+This module provides functionality to detect and clean up incompatible RyzenAI models
+when upgrading between major versions (e.g., 1.4/1.5 -> 1.6).
+"""
+
+import os
+import json
+import shutil
+import logging
+from typing import List, Dict, Optional, Tuple
+
+
+def get_directory_size(path: str) -> int:
+    """
+    Calculate the total size of a directory in bytes.
+
+    Args:
+        path: Path to the directory
+
+    Returns:
+        Total size in bytes
+    """
+    total_size = 0
+    try:
+        for dirpath, _, filenames in os.walk(path):
+            for filename in filenames:
+                filepath = os.path.join(dirpath, filename)
+                try:
+                    total_size += os.path.getsize(filepath)
+                except (OSError, FileNotFoundError):
+                    # Skip files that can't be accessed
+                    pass
+    except (OSError, FileNotFoundError):
+        pass
+    return total_size
+
+
+def format_size(size_bytes: int) -> str:
+    """
+    Format byte size to human-readable string.
+
+    Args:
+        size_bytes: Size in bytes
+
+    Returns:
+        Formatted string (e.g., "1.5 GB", "450 MB")
+    """
+    for unit in ["B", "KB", "MB"]:
+        if size_bytes < 1024.0:
+            return f"{size_bytes:.1f} {unit}"
+        size_bytes /= 1024.0
+    return f"{size_bytes:.1f} GB"
+
+
+def check_rai_config_version(model_path: str, required_version: str = "1.6.0") -> bool:
+    """
+    Check if a model's rai_config.json contains the required version.
+
+    Args:
+        model_path: Path to the model directory
+        required_version: Version string to check for (default: "1.6.0")
+
+    Returns:
+        True if model is compatible (has required version), False otherwise
+    """
+    rai_config_path = os.path.join(model_path, "rai_config.json")
+
+    # If no rai_config.json exists, it's not a RyzenAI model
+    if not os.path.exists(rai_config_path):
+        return True
+
+    try:
+        with open(rai_config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+
+        # Check if max_prompt_length exists and has the required version
+        if "max_prompt_length" in config:
+            max_prompt_length = config["max_prompt_length"]
+            if isinstance(max_prompt_length, dict):
+                # If it's a dict with version keys, check for required version
+                return required_version in max_prompt_length
+            # Fallback to True to avoid deleting models if format changes
+            return True
+
+        return True
+
+    except (json.JSONDecodeError, OSError) as e:
+        logging.warning(f"Could not read rai_config.json from {model_path}: {e}")
+        # If we can't read it, assume it's compatible to avoid false positives
+        return True
+
+
+def scan_oga_models_cache(cache_dir: str) -> List[Dict[str, any]]:
+    """
+    Scan the Lemonade OGA models cache for incompatible models.
+
+    Args:
+        cache_dir: Path to the Lemonade cache directory
+
+    Returns:
+        List of dicts with model info (path, name, size, compatible)
+    """
+    oga_models_path = os.path.join(cache_dir, "oga_models")
+    incompatible_models = []
+
+    if not os.path.exists(oga_models_path):
+        return incompatible_models
+
+    try:
+        # Iterate through model directories in oga_models
+        for model_name in os.listdir(oga_models_path):
+            model_dir = os.path.join(oga_models_path, model_name)
+
+            if not os.path.isdir(model_dir):
+                continue
+
+            # Check all subdirectories (e.g., npu-int4, hybrid-int4)
+            for subdir in os.listdir(model_dir):
+                subdir_path = os.path.join(model_dir, subdir)
+
+                if not os.path.isdir(subdir_path):
+                    continue
+
+                # Check if this model version is compatible
+                if not check_rai_config_version(subdir_path):
+                    size = get_directory_size(subdir_path)
+                    incompatible_models.append(
+                        {
+                            "path": subdir_path,
+                            "name": f"{model_name}/{subdir}",
+                            "size": size,
+                            "size_formatted": format_size(size),
+                            "cache_type": "lemonade",
+                        }
+                    )
+
+    except (OSError, PermissionError) as e:
+        logging.warning(f"Error scanning oga_models cache: {e}")
+
+    return incompatible_models
+
+
+def scan_huggingface_cache(hf_home: Optional[str] = None) -> List[Dict[str, any]]:
+    """
+    Scan the HuggingFace cache for incompatible RyzenAI models.
+
+    Args:
+        hf_home: Path to HuggingFace home directory (default: from env or ~/.cache/huggingface)
+
+    Returns:
+        List of dicts with model info (path, name, size, compatible)
+    """
+    if hf_home is None:
+        hf_home = os.environ.get(
+            "HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
+        )
+
+    hub_path = os.path.join(hf_home, "hub")
+    incompatible_models = []
+
+    if not os.path.exists(hub_path):
+        return incompatible_models
+
+    try:
+        # Iterate through model directories in HuggingFace cache
+        for item in os.listdir(hub_path):
+            if not item.startswith("models--"):
+                continue
+
+            model_dir = os.path.join(hub_path, item)
+            if not os.path.isdir(model_dir):
+                continue
+
+            # Look in snapshots subdirectory
+            snapshots_dir = os.path.join(model_dir, "snapshots")
+            if not os.path.exists(snapshots_dir):
+                continue
+
+            # Check each snapshot
+            for snapshot_hash in os.listdir(snapshots_dir):
+                snapshot_path = os.path.join(snapshots_dir, snapshot_hash)
+
+                if not os.path.isdir(snapshot_path):
+                    continue
+
+                # Check if this snapshot has incompatible RyzenAI model
+                if not check_rai_config_version(snapshot_path):
+                    # Extract readable model name from directory
+                    model_name = item.replace("models--", "").replace("--", "/")
+                    size = get_directory_size(
+                        model_dir
+                    )  # Size of entire model directory
+                    incompatible_models.append(
+                        {
+                            "path": model_dir,
+                            "name": model_name,
+                            "size": size,
+                            "size_formatted": format_size(size),
+                            "cache_type": "huggingface",
+                        }
+                    )
+                    break
+
+    except (OSError, PermissionError) as e:
+        logging.warning(f"Error scanning HuggingFace cache: {e}")
+
+    return incompatible_models
+
+
+def detect_incompatible_ryzenai_models(
+    cache_dir: str, hf_home: Optional[str] = None
+) -> Tuple[List[Dict[str, any]], int]:
+    """
+    Detect all incompatible RyzenAI models in both Lemonade and HuggingFace caches.
+
+    Args:
+        cache_dir: Path to the Lemonade cache directory
+        hf_home: Path to HuggingFace home directory (optional)
+
+    Returns:
+        Tuple of (list of incompatible models, total size in bytes)
+    """
+    incompatible_models = []
+
+    # Scan Lemonade cache
+    oga_models = scan_oga_models_cache(cache_dir)
+    incompatible_models.extend(oga_models)
+
+    # Scan HuggingFace cache
+    hf_models = scan_huggingface_cache(hf_home)
+    incompatible_models.extend(hf_models)
+
+    # Calculate total size
+    total_size = sum(model["size"] for model in incompatible_models)
+
+    logging.info(
+        f"Found {len(incompatible_models)} incompatible RyzenAI models "
+        f"({format_size(total_size)} total)"
+    )
+
+    return incompatible_models, total_size
+
+
+def delete_model_directory(model_path: str) -> bool:
+    """
+    Safely delete a model directory.
+
+    Args:
+        model_path: Path to the model directory to delete
+
+    Returns:
+        True if deletion successful, False otherwise
+    """
+    try:
+        if os.path.exists(model_path):
+            shutil.rmtree(model_path)
+            logging.info(f"Deleted model directory: {model_path}")
+            return True
+        else:
+            logging.warning(f"Model directory not found: {model_path}")
+            return False
+    except (OSError, PermissionError) as e:
+        logging.error(f"Failed to delete model directory {model_path}: {e}")
+        return False
+
+
+def _extract_checkpoint_from_path(path: str) -> Optional[str]:
+    """
+    Extract the checkpoint name from a model path.
+
+    Args:
+        path: Model directory path (either Lemonade cache or HuggingFace cache)
+
+    Returns:
+        Checkpoint name (e.g., "amd/Qwen2.5-1.5B-Instruct-awq") or None if not extractable
+    """
+    # Normalize path separators to handle both Unix and Windows paths
+    normalized_path = path.replace("\\", "/")
+    parts = normalized_path.split("/")
+
+    # Handle HuggingFace cache paths: models--{org}--{repo}
+    if "models--" in normalized_path:
+        for part in parts:
+            if part.startswith("models--"):
+                # Convert models--org--repo to org/repo
+                # Replace first two occurrences of -- with /
+                checkpoint = part.replace("models--", "", 1).replace("--", "/", 1)
+                return checkpoint
+        return None
+
+    # Handle Lemonade cache paths: oga_models/{model_name}/{device}-{dtype}
+    if "oga_models" in normalized_path:
+        try:
+            oga_models_idx = parts.index("oga_models")
+            if oga_models_idx + 1 < len(parts):
+                model_name = parts[oga_models_idx + 1]
+                # Convert model_name back to checkpoint (e.g., amd_model -> amd/model)
+                # This is a heuristic - we look for the pattern {org}_{model}
+                checkpoint = model_name.replace("_", "/", 1)
+                return checkpoint
+        except (ValueError, IndexError):
+            return None
+
+    return None
+
+
+def _cleanup_user_models_json(deleted_checkpoints: List[str], user_models_file: str):
+    """
+    Remove entries from user_models.json for models that have been deleted.
+
+    Args:
+        deleted_checkpoints: List of checkpoint names that were deleted
+        user_models_file: Path to user_models.json
+    """
+    if not deleted_checkpoints or not os.path.exists(user_models_file):
+        return
+
+    try:
+        with open(user_models_file, "r", encoding="utf-8") as f:
+            user_models = json.load(f)
+
+        # Track which models to remove
+        models_to_remove = []
+        for model_name, model_info in user_models.items():
+            checkpoint = model_info.get("checkpoint", "")
+            # Check if this checkpoint matches any deleted checkpoints
+            # We do a case-insensitive comparison since paths may have been lowercased
+            for deleted_checkpoint in deleted_checkpoints:
+                if checkpoint.lower() == deleted_checkpoint.lower():
+                    models_to_remove.append(model_name)
+                    break
+
+        # Remove the models
+        for model_name in models_to_remove:
+            del user_models[model_name]
+            logging.info(f"Removed {model_name} from user_models.json")
+
+        # Save the updated file only if we removed something
+        if models_to_remove:
+            with open(user_models_file, "w", encoding="utf-8") as f:
+                json.dump(user_models, f, indent=2)
+            logging.info(
+                f"Updated user_models.json - removed {len(models_to_remove)} entries"
+            )
+
+    except (json.JSONDecodeError, OSError) as e:
+        logging.warning(f"Could not update user_models.json: {e}")
+
+
+def delete_incompatible_models(
+    model_paths: List[str], user_models_file: Optional[str] = None
+) -> Dict[str, any]:
+    """
+    Delete multiple incompatible model directories and clean up user_models.json.
+
+    Args:
+        model_paths: List of paths to delete
+        user_models_file: Path to user_models.json (optional, will use default if not provided)
+
+    Returns:
+        Dict with deletion results (success_count, failed_count, freed_size, cleaned_user_models)
+    """
+    success_count = 0
+    failed_count = 0
+    freed_size = 0
+    deleted_checkpoints = []
+
+    for path in model_paths:
+        # Calculate size before deletion
+        size = get_directory_size(path)
+
+        # Extract checkpoint name before deleting
+        checkpoint = _extract_checkpoint_from_path(path)
+        if checkpoint:
+            deleted_checkpoints.append(checkpoint)
+
+        if delete_model_directory(path):
+            success_count += 1
+            freed_size += size
+        else:
+            failed_count += 1
+
+    # Clean up user_models.json if we deleted any models
+    cleaned_user_models = False
+    if deleted_checkpoints:
+        # Use default path if not provided
+        if user_models_file is None:
+            from lemonade.cache import DEFAULT_CACHE_DIR
+
+            user_models_file = os.path.join(DEFAULT_CACHE_DIR, "user_models.json")
+
+        _cleanup_user_models_json(deleted_checkpoints, user_models_file)
+        cleaned_user_models = True
+
+    return {
+        "success_count": success_count,
+        "failed_count": failed_count,
+        "freed_size": freed_size,
+        "freed_size_formatted": format_size(freed_size),
+        "cleaned_user_models": cleaned_user_models,
+    }

lemonade/tools/report/table.py

@@ -48,6 +48,18 @@ def _merge_join(str1, str2) -> str:
     return str1 + ("\n" if str1 and str2 else "") + str2
 
 
+def _window_sum(data: list, n_windows: int) -> list:
+    """Sums data into n_windows windows"""
+    if n_windows <= 0:
+        return data
+    window_size = max(1, len(data) // n_windows)
+    summed_data = []
+    for i in range(0, len(data), window_size):
+        window_sum = sum(data[i : i + window_size])
+        summed_data.append(window_sum)
+    return summed_data
+
+
 ################################################################################
 # CLASSES THAT DESCRIBE TEXT TABLE COLUMNS
 ################################################################################
@@ -88,10 +100,54 @@ class SimpleStat(TableColumn):
         if lean and self.omit_if_lean:
             return None
         data = build_stats.get(self.stat, None)
-        if data is None:
-            return ""
+        if data is None or (data == []):
+            return "-"
         if self.stat_fn:
             data = self.stat_fn(data)
+        cell_str = "\n".join(
+            [
+                _wrap("-" if x is None else f"{x:{self.format_str}}", self.wrap)
+                for x in _to_list(data)
+            ]
+        )
+        return cell_str
+
+
+class DependentStat(TableColumn):
+    """
+    These are for statistics already declared by the tool or basic build stats that
+    rely on one or more additional stats to compute their value.  The dependency is
+    embodied by the stat_fn function.
+    """
+
+    def __init__(
+        self,
+        column_header,
+        stats,
+        format_str,
+        align="center",
+        omit_if_lean=False,
+        wrap=None,
+        stat_fn=None,
+    ):
+        self.column_header = column_header
+        self.stats = stats
+        self.format_str = format_str
+        self.align = align
+        self.omit_if_lean = omit_if_lean
+        self.wrap = wrap or self.default_wrap
+        self.stat_fn = stat_fn
+
+    def get_str(self, build_stats, lean=False):
+        if lean and self.omit_if_lean:
+            return None
+        stats_data = [build_stats.get(stat, None) for stat in self.stats]
+        if self.stat_fn:
+            data = self.stat_fn(stats_data)
+        else:
+            data = stats_data[0]
+        if data is None or (data == []):
+            return "-"
         cell_str = "\n".join(
             [_wrap(f"{x:{self.format_str}}", self.wrap) for x in _to_list(data)]
         )
@@ -434,10 +490,12 @@ class Table(ABC):
             row = []
 
             # First columns
+            first_columns_count = 0
             for entry in first_columns:
                 entry_str = entry.get_str(build_stats, self.lean)
                 if entry_str is not None:
                     row.append(entry_str)
+                    first_columns_count += 1
 
             # Per tool columns
             for tool in tools:
@@ -460,22 +518,24 @@ class Table(ABC):
                     row.append(entry_str)
 
             # Final columns
+            last_columns_count = 0
             for entry in last_columns:
                 entry_str = entry.get_str(build_stats, self.lean)
                 if entry_str is not None:
                     row.append(entry_str)
+                    last_columns_count += 1
 
             # See if this row should be merged with the last row
             if last_build_stats and self.merge_test_fn(last_build_stats, build_stats):
                 # Merge with last row
-                for col in range(0, len(first_columns)):
+                for col in range(0, first_columns_count):
                     # If identical, don't duplicate
                     if last_row[col] != row[col]:
                         last_row[col] = _merge_join(last_row[col], row[col])
-                for col in range(len(first_columns), len(row) - len(last_columns)):
+                for col in range(first_columns_count, len(row) - last_columns_count):
                     # Allow duplicates
                     last_row[col] = _merge_join(last_row[col], row[col])
-                for col in range(len(row) - len(last_columns), len(row)):
+                for col in range(len(row) - last_columns_count, len(row)):
                     # If identical, don't duplicate
                     if last_row[col] != row[col]:
                         last_row[col] = _merge_join(last_row[col], row[col])
@@ -581,17 +641,25 @@ class LemonadePerfTable(Table):
                     Keys.STD_DEV_SECONDS_TO_FIRST_TOKEN,
                     ".2f",
                 ),
+                StatWithSD(
+                    _wrap("Prefill Tokens per Second", 8),
+                    Keys.PREFILL_TOKENS_PER_SECOND,
+                    Keys.STD_DEV_PREFILL_TOKENS_PER_SECOND,
+                    ".2f",
+                ),
                 StatWithSD(
                     _wrap("Tokens per Second", 8),
                     Keys.TOKEN_GENERATION_TOKENS_PER_SECOND,
                     Keys.STD_DEV_TOKENS_PER_SECOND,
                     ".2f",
                 ),
-                SimpleStat(
+                DependentStat(
                     _wrap("Total Generated Tokens", 9),
-                    Keys.RESPONSE_TOKENS,
+                    [Keys.RESPONSE_TOKENS, Keys.PROMPT_TOKENS],
                     "d",
-                    stat_fn=sum,
+                    stat_fn=lambda x: _window_sum(
+                        _to_list(x[0]), n_windows=len(_to_list(x[1]))
+                    ),
                 ),
                 SimpleStat(
                     _wrap("Memory Used (GB)", 8), Keys.MAX_MEMORY_USED_GBYTE, ".3f"

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)