herd/manage.py at master · ndcorder/herd

658 lines (547 loc) · 23.7 KB
#!/usr/bin/env python3
Herd — Model registry and lifecycle manager for local GGUF models.
    herd build          Generate config.yaml for llama-swap
    herd status         Show model status dashboard
    herd list           List models with filters
    herd add            Interactive model addition
    herd download       Download models from HuggingFace
    herd enable <id>    Enable a model
    herd disable <id>   Disable a model
    herd validate       Check for problems
    herd cleanup        Remove orphaned files from disk
    herd monitor        Live TUI dashboard for llama-swap server
import argparse
from pathlib import Path
from lib.registry import Registry
from lib.builder import build_config, build_cmd
from lib.status import get_model_status, validate_registry, get_orphaned_files, check_context_lengths, check_mmproj_files, ModelState
from lib.mutator import RegistryMutator
from lib.downloader import generate_download_commands, run_downloads, smart_download
from lib.scanner import scan_orphaned
from lib.hf import (
    validate_hf_repo,
    fetch_hf_config,
    fetch_context_length,
    extract_context_length,
    list_gguf_files,
    find_mmproj_files,
    pick_best_mmproj,
    detect_model_type,
    sanitize_model_name,
    detect_special_flags,
    infer_category_dir,
    extract_model_size,
DEFAULT_REGISTRY = Path(__file__).parent / "models"
def cmd_build(args, registry):
    """Generate llama-swap config.yaml."""
    categories = args.only.split(",") if args.only else None
    tags = args.tags.split(",") if args.tags else None
    exclude = args.exclude.split(",") if args.exclude else None
    output = build_config(registry, categories=categories, tags=tags, exclude_categories=exclude)
    output_path = Path(args.output)
    output_path.write_text(output)
    print(f"Generated {output_path} ({len(registry.enabled_models(categories=categories, tags=tags, exclude_categories=exclude))} models)")
def _fmt_size(size_bytes: int) -> str:
    """Format bytes as human-readable size."""
    if size_bytes == 0:
        return "-"
    if size_bytes >= 1 << 30:
        return f"{size_bytes / (1 << 30):.1f}G"
    if size_bytes >= 1 << 20:
        return f"{size_bytes / (1 << 20):.0f}M"
    return f"{size_bytes / (1 << 10):.0f}K"
def cmd_info(args, registry):
    """Show resolved config and command for a single model."""
    model_id = args.model_id
    alias_source = None
    # Check if it's an alias
    target = registry.resolve_alias(model_id)
    if target:
        alias_source = model_id
        model_id = target
        model = registry.resolve(model_id)
    except KeyError:
        print(f"Error: '{args.model_id}' not found in registry or aliases")
        sys.exit(1)
    # Disk status
    full_path = registry.base_path / model["path"]
    if full_path.exists():
        size = full_path.stat().st_size
        if size >= 1 << 30:
            size_str = f"{size / (1 << 30):.1f}G"
        elif size >= 1 << 20:
            size_str = f"{size / (1 << 20):.0f}M"
        else:
            size_str = f"{size / (1 << 10):.0f}K"
        status_str = f"downloaded ({size_str})"
        status_str = "missing"
    if alias_source:
        print(f"  (resolved from alias '{alias_source}')")
    print(f"{model['id']} ({model['category']})")
    print(f"  path:       {registry.base_path / model['path']}")
    print(f"  ctx:        {model['ctx']}")
    print(f"  gpu_layers: {model['gpu_layers']}")
    flags_str = ", ".join(model["flags"]) if model["flags"] else "(none)"
    print(f"  flags:      {flags_str}")
    if model.get("mmproj"):
        print(f"  mmproj:     {registry.base_path / model['mmproj']}")
    if model.get("repo"):
        print(f"  repo:       {model['repo']}")
    if model.get("tags"):
        print(f"  tags:       {', '.join(model['tags'])}")
    print(f"  enabled:    {model['enabled']}")
    print(f"  status:     {status_str}")
    if model.get("note"):
        print(f"\n  {model['note']}")
    # Show aliases pointing to this model
    aliases_for = [a for a, t in registry.aliases.items() if t == model_id]
    if aliases_for:
        print(f"  aliases:    {', '.join(sorted(aliases_for))}")
    # Print the llama-server command
    cmd = build_cmd(model, str(registry.base_path))
    print(f"\n  command:\n    {cmd.replace(chr(10), chr(10) + '    ')}")
def cmd_status(args, registry):
    """Show model status dashboard."""
    statuses = get_model_status(registry)
    # Group by category
    by_category = {}
    for s in statuses:
        if not args.all and s["state"] == ModelState.DISABLED:
            continue
        cat = s["category"]
        by_category.setdefault(cat, []).append(s)
    # Apply filters
    if args.only:
        cats = args.only.split(",")
        by_category = {k: v for k, v in by_category.items() if k in cats}
    counts = {"downloaded": 0, "missing": 0, "disabled": 0}
    total_size = 0
    for cat, models in sorted(by_category.items()):
        cat_size = sum(m["size"] for m in models)
        print(f"\n{cat.upper()} ({len(models)} models, {_fmt_size(cat_size)})")
        for m in models:
            state = m["state"]
            if state == ModelState.DOWNLOADED:
                icon = "+"
                counts["downloaded"] += 1
            elif state == ModelState.MISSING:
                icon = "x"
                counts["missing"] += 1
            else:
                icon = "o"
                counts["disabled"] += 1
            total_size += m["size"]
            size_display = _fmt_size(m["size"])
            ctx_display = f"{m['ctx'] // 1024}K" if m['ctx'] >= 1024 else str(m['ctx'])
            tags_display = f"  [{', '.join(m['tags'])}]" if m['tags'] else ""
            print(f"  {icon} {m['id']:<40} {size_display:>6}  {ctx_display:<8} {state.value}{tags_display}")
    print(f"\nSummary: {counts['downloaded']} downloaded, {counts['disabled']} disabled, {counts['missing']} missing")
    print(f"Disk usage: {_fmt_size(total_size)}")
def cmd_list(args, registry):
    """List models with filters."""
    categories = args.only.split(",") if args.only else None
    tags = args.tags.split(",") if args.tags else None
    if args.all:
        models = registry.all_models()
        if categories:
            models = [m for m in models if m["category"] in categories]
        if tags:
            models = [m for m in models if any(t in m["tags"] for t in tags)]
        models = registry.enabled_models(categories=categories, tags=tags)
    for m in models:
        enabled_str = "" if m["enabled"] else " (disabled)"
        tags_str = f"  [{', '.join(m['tags'])}]" if m["tags"] else ""
        print(f"  {m['id']:<40} {m['category']:<12}{enabled_str}{tags_str}")
    print(f"\n{len(models)} model(s)")
def cmd_add(args, registry):
    """Interactive model addition."""
    repo = args.repo
    if not repo:
        repo = input("HuggingFace repo (author/model-name): ").strip()
    if not repo or "/" not in repo:
        print("Error: Invalid repo format. Expected 'author/model-name'")
        sys.exit(1)
    print(f"Validating {repo}...")
    if not validate_hf_repo(repo):
        print(f"Warning: Could not validate repository {repo}")
        if input("Continue anyway? (y/n): ").strip().lower() != "y":
            sys.exit(1)
        print("  Repository found")
    # Detect context length (GGUF metadata → config.json → base model)
    print("Detecting context length...")
    detected_ctx = fetch_context_length(repo)
    if detected_ctx:
        print(f"  Context length: {detected_ctx}")
        print("  Could not detect context length")
    # List GGUF files (filter out sub-4-bit quants)
    print("Listing GGUF files...")
    gguf_files = list_gguf_files(repo)
    sub4_patterns = re.compile(r'[_-](IQ[12]|IQ3_XXS|Q[123]_|Q[123]K|Q2_K|Q3_K)', re.IGNORECASE)
    filtered = [f for f in gguf_files if not sub4_patterns.search(f)]
    hidden = len(gguf_files) - len(filtered)
    if filtered:
        for i, f in enumerate(filtered, 1):
            print(f"  {i}. {f}")
        if hidden:
            print(f"  ({hidden} sub-4-bit quants hidden)")
    # Get filename
    file = input("\nGGUF filename or pattern (e.g., model.gguf or *.gguf): ").strip()
    if not file:
        print("Error: Filename required")
        sys.exit(1)
    # Auto-detect
    category = detect_model_type(file, repo)
    model_id = sanitize_model_name(file)
    flags = detect_special_flags(model_id, category, repo)
    cat_dir = infer_category_dir(model_id, category)
    # Determine config filename vs download pattern
    config_file = file
    if "*" in file or "?" in file:
        config_file = input(f"Exact filename for config (not wildcard): ").strip()
        if not config_file:
            print("Error: Config filename required for wildcard patterns")
            sys.exit(1)
        model_id = sanitize_model_name(config_file)
    model_path = f"{cat_dir}/{config_file}"
    # Context size
    default_ctx = detected_ctx or 32768
    ctx_input = input(f"Context size [{default_ctx}]: ").strip()
    ctx = int(ctx_input) if ctx_input else default_ctx
    # Check for mmproj (vision model projector)
    mmproj_path = None
    mmproj_files = find_mmproj_files(repo)
    if mmproj_files:
        best = pick_best_mmproj(mmproj_files)
        print(f"\n  Vision model detected — mmproj files available:")
        for f in mmproj_files:
            marker = " (recommended)" if f == best else ""
            print(f"    {f}{marker}")
        mmproj_input = input(f"  Use mmproj [{best}] (Enter to accept, 'n' to skip): ").strip()
        if mmproj_input.lower() != "n":
            mmproj_file = mmproj_input if mmproj_input and mmproj_input != "" else best
            mmproj_path = f"{cat_dir}/{mmproj_file}"
    tags_input = input("Tags (comma-separated, or Enter for none): ").strip()
    tags = [t.strip() for t in tags_input.split(",") if t.strip()] if tags_input else None
    # Preview
    print(f"\n  ID:       {model_id}")
    print(f"  Category: {category}")
    print(f"  Path:     {model_path}")
    if mmproj_path:
        print(f"  mmproj:   {mmproj_path}")
    print(f"  Repo:     {repo}")
    print(f"  File:     {file}")
    print(f"  Context:  {ctx}")
    print(f"  Flags:    {flags or 'none'}")
    print(f"  Tags:     {tags or 'none'}")
    if input("\nAdd to registry? (y/n): ").strip().lower() != "y":
        print("Aborted.")
        return
    mutator = RegistryMutator(registry.registry_path)
    mutator.add_model(
        model_id=model_id,
        category=category,
        model_path=model_path,
        repo=repo,
        file=file,
        ctx=ctx if ctx != default_ctx else None,
        flags=flags or None,
        tags=tags,
        mmproj=mmproj_path,
    print(f"Added '{model_id}' to registry. Run 'herd build' to regenerate config.")
def cmd_download(args, registry):
    """Download models from HuggingFace with smart lock file checks."""
    model_ids = [args.model] if args.model else None
    categories = args.only.split(",") if args.only else None
    lock_path = registry.registry_path.parent / "models.lock"
    print("Checking models...")
    results = smart_download(
        registry,
        lock_path,
        model_ids=model_ids,
        categories=categories,
        dry_run=args.dry_run,
        fast=args.fast,
        throttle=args.throttle,
    if not results:
        print("No models to download.")
        return
    # Summary
    up_to_date = sum(1 for r in results if r["status"] in ("up_to_date", "locked"))
    ok = sum(1 for r in results if r["status"] == "ok")
    errors = sum(1 for r in results if r["status"] == "error")
    needs_download = sum(1 for r in results if r["status"] in ("not_downloaded", "update_available"))
    no_meta = sum(1 for r in results if r["status"] == "no_metadata")
    print()
    if args.dry_run:
        parts = []
        if up_to_date:
            parts.append(f"{up_to_date} up to date")
        if needs_download:
            parts.append(f"{needs_download} to download")
        if no_meta:
            parts.append(f"{no_meta} no metadata")
        print(f"Dry run: {', '.join(parts)}")
        parts = []
        if up_to_date:
            parts.append(f"{up_to_date} up to date")
        if ok:
            parts.append(f"{ok} downloaded")
        if errors:
            parts.append(f"{errors} errors")
        if no_meta:
            parts.append(f"{no_meta} no metadata")
        print(f"Done: {', '.join(parts)}")
    for r in results:
        if r["status"] == "error":
            print(f"  FAILED: {r['model_id']}: {r.get('error', '')}")
def cmd_enable(args, registry):
    """Enable a model."""
    mutator = RegistryMutator(registry.registry_path)
    mutator.set_enabled(args.model_id, True)
    print(f"Enabled '{args.model_id}'")
def cmd_disable(args, registry):
    """Disable a model."""
    mutator = RegistryMutator(registry.registry_path)
    mutator.set_enabled(args.model_id, False)
    print(f"Disabled '{args.model_id}'")
def _fmt_ctx(ctx: int) -> str:
    """Format context length as human-readable (e.g. 32768 -> 32K, 131072 -> 128K)."""
    if ctx >= 1024 and ctx % 1024 == 0:
        return f"{ctx // 1024}K"
    return str(ctx)
def cmd_validate(args, registry):
    """Validate the registry."""
    report = validate_registry(registry)
    if report["errors"]:
        print(f"\nERRORS ({len(report['errors'])}):")
        for e in report["errors"]:
            print(f"  x {e}")
    if report["warnings"]:
        print(f"\nWARNINGS ({len(report['warnings'])}):")
        for w in report["warnings"]:
            print(f"  ! {w}")
    if report["orphaned"]:
        print(f"\nORPHANED FILES ({len(report['orphaned'])}):")
        for o in report["orphaned"]:
            print(f"  ? {o}")
    if args.check_ctx:
        print("\nChecking context lengths against HuggingFace...")
        ctx_results = check_context_lengths(registry)
        mismatches = [r for r in ctx_results if r["status"] == "mismatch"]
        failures = [r for r in ctx_results if r["status"] == "fetch_failed"]
        if mismatches:
            if args.fix:
                print(f"\nFIXING CONTEXT MISMATCHES ({len(mismatches)}):")
                mutator = RegistryMutator(registry.registry_path)
                for r in mismatches:
                    mutator.set_ctx(r["id"], r["actual"])
                    print(f"  + {r['id']}: {_fmt_ctx(r['configured'])} -> {_fmt_ctx(r['actual'])}")
                print(f"\nUpdated {len(mismatches)} model(s) in registry")
            else:
                print(f"\nCONTEXT MISMATCHES ({len(mismatches)}):")
                for r in mismatches:
                    print(f"  ! {r['id']}: configured {_fmt_ctx(r['configured'])}, actual {_fmt_ctx(r['actual'])}  ({r['repo']})")
                print(f"\nRun with --fix to update the registry")
        if failures:
            print(f"\nCONTEXT CHECK FAILED ({len(failures)}):")
            for r in failures:
                print(f"  ? {r['id']}: could not fetch config  ({r['repo']})")
        if not mismatches and not failures:
            print("  All context lengths match HuggingFace configs")
    if args.check_mmproj:
        print("\nChecking mmproj filenames against HuggingFace repos...")
        mmproj_results = check_mmproj_files(registry)
        not_found = [r for r in mmproj_results if r["status"] == "not_found"]
        fetch_fails = [r for r in mmproj_results if r["status"] == "fetch_failed"]
        if not_found:
            print(f"\nMMPROJ NOT FOUND IN REPO ({len(not_found)}):")
            for r in not_found:
                print(f"  x {r['id']}: {r['mmproj']} not in {r['repo']}")
                if r["available"]:
                    print(f"    available: {', '.join(r['available'])}")
                else:
                    print(f"    (no mmproj files found in repo)")
        if fetch_fails:
            print(f"\nMMPROJ CHECK FAILED ({len(fetch_fails)}):")
            for r in fetch_fails:
                print(f"  ? {r['id']}: could not list files  ({r['repo']})")
        if not not_found and not fetch_fails:
            print("  All mmproj filenames match HuggingFace repos")
    has_remote_check = args.check_ctx or args.check_mmproj
    if not has_remote_check and not report["errors"] and not report["warnings"] and not report["orphaned"]:
        print("All models validated OK")
def cmd_cleanup(args, registry):
    """Remove orphaned files and optionally disabled models' files."""
    import shutil
    to_remove = []
    # Orphaned files
    orphaned = get_orphaned_files(registry)
    for f in orphaned:
        full = registry.base_path / f
        to_remove.append(("orphaned", f, full.stat().st_size if full.exists() else 0))
    # Disabled models' files (optional)
    if args.include_disabled:
        for model in registry.all_models():
            if not model["enabled"]:
                full = registry.base_path / model["path"]
                if full.exists():
                    to_remove.append(("disabled", model["path"], full.stat().st_size))
    if not to_remove:
        print("Nothing to clean up.")
        return
    total_size = sum(s for _, _, s in to_remove)
    print(f"\nFiles to remove ({_fmt_size(total_size)}):\n")
    for kind, path, size in to_remove:
        print(f"  [{kind}] {path}  ({_fmt_size(size)})")
    if args.dry_run:
        print(f"\n{len(to_remove)} file(s) would be removed (dry run)")
        return
    if input(f"\nRemove {len(to_remove)} file(s)? (y/n): ").strip().lower() != "y":
        print("Aborted.")
        return
    removed = 0
    freed = 0
    for kind, path, size in to_remove:
        full = registry.base_path / path
        try:
            full.unlink()
            removed += 1
            freed += size
            # Remove parent dir if now empty
            parent = full.parent
            if parent.exists() and not any(parent.iterdir()):
                parent.rmdir()
        except Exception as e:
            print(f"  Failed to remove {path}: {e}")
    print(f"\nRemoved {removed} file(s), freed {_fmt_size(freed)}")
def cmd_monitor(args, registry):
    """Launch the live TUI dashboard for llama-swap server."""
    from lib.monitor import run_monitor
    base_url = f"http://{args.host}:{args.port}"
    run_monitor(
        base_url=base_url,
        registry_path=str(registry.registry_path),
        server_pid=args.server_pid,
def cmd_scan(args, registry):
    """Scan for orphaned GGUFs and propose registry entries."""
    proposals = scan_orphaned(registry)
    if not proposals:
        print("No orphaned GGUF files found.")
        return
    print(f"Found {len(proposals)} orphaned GGUF file(s):\n")
    for p in proposals:
        print(f"  {p['id']}")
        print(f"    category: {p['category']}")
        print(f"    path:     {p['path']}")
        print(f"    ctx:      {p['ctx']}")
        print()
    if args.dry_run:
        print(f"{len(proposals)} model(s) would be added (dry run)")
        return
    if input(f"Add {len(proposals)} model(s) to registry? (y/n): ").strip().lower() != "y":
        print("Aborted.")
        return
    mutator = RegistryMutator(registry.registry_path)
    added = 0
    for p in proposals:
        try:
            mutator.add_model(
                model_id=p["id"],
                category=p["category"],
                model_path=p["path"],
                repo=None,
                file=p["file"],
            added += 1
        except Exception as e:
            print(f"  Failed to add {p['id']}: {e}")
    print(f"Added {added} model(s). Run 'herd build' to regenerate config.")
def main():
    parser = argparse.ArgumentParser(description="Herd — model registry and lifecycle manager for local GGUF models")
    parser.add_argument("--registry", type=Path, default=DEFAULT_REGISTRY, help="Path to models directory")
    sub = parser.add_subparsers(dest="command", required=True)
    # build
    p_build = sub.add_parser("build", help="Generate llama-swap config.yaml")
    p_build.add_argument("--output", "-o", default="config.yaml", help="Output path (default: config.yaml)")
    p_build.add_argument("--only", help="Only include these categories (comma-separated)")
    p_build.add_argument("--tags", help="Only include models with these tags (comma-separated)")
    p_build.add_argument("--exclude", help="Exclude these categories (comma-separated)")
    p_info = sub.add_parser("info", help="Show resolved config for a model")
    p_info.add_argument("model_id", help="Model ID or alias")
    # status
    p_status = sub.add_parser("status", help="Show model status dashboard")
    p_status.add_argument("--only", help="Filter by categories")
    p_status.add_argument("--all", action="store_true", help="Include disabled models")
    p_list = sub.add_parser("list", help="List models")
    p_list.add_argument("--only", help="Filter by categories")
    p_list.add_argument("--tags", help="Filter by tags")
    p_list.add_argument("--all", action="store_true", help="Include disabled models")
    p_add = sub.add_parser("add", help="Add a model interactively")
    p_add.add_argument("repo", nargs="?", help="HuggingFace repo (author/model-name)")
    # download
    p_dl = sub.add_parser("download", help="Download models from HuggingFace")
    p_dl.add_argument("model", nargs="?", help="Specific model ID to download")
    p_dl.add_argument("--only", help="Filter by categories")
    p_dl.add_argument("--dry-run", action="store_true", help="Show commands without running")
    p_dl.add_argument("--fast", action="store_true", help="Use hf-transfer for speed (no resume on interrupt)")
    p_dl.add_argument("--throttle", type=int, metavar="MBIT", help="Limit download speed in Mbit/s (e.g. 900)")
    # enable/disable
    p_en = sub.add_parser("enable", help="Enable a model")
    p_en.add_argument("model_id", help="Model ID to enable")
    p_dis = sub.add_parser("disable", help="Disable a model")
    p_dis.add_argument("model_id", help="Model ID to disable")
    # validate
    p_val = sub.add_parser("validate", help="Validate registry")
    p_val.add_argument("--check-ctx", action="store_true", help="Verify context lengths against HuggingFace")
    p_val.add_argument("--check-mmproj", action="store_true", help="Verify mmproj filenames exist in HuggingFace repos")
    p_val.add_argument("--fix", action="store_true", help="Auto-fix context length mismatches (requires --check-ctx)")
    p_scan = sub.add_parser("scan", help="Scan for orphaned GGUFs and propose registry entries")
    p_scan.add_argument("--dry-run", action="store_true", help="Show proposals without adding")
    # monitor
    p_mon = sub.add_parser("monitor", help="Live TUI dashboard for llama-swap server")
    p_mon.add_argument("--host", default="localhost", help="llama-swap host (default: localhost)")
    p_mon.add_argument("--port", type=int, default=8080, help="llama-swap port (default: 8080)")
    p_mon.add_argument("--server-pid", type=int, default=None, help="PID of llama-swap process to manage")
    # cleanup
    p_clean = sub.add_parser("cleanup", help="Remove orphaned files from disk")
    p_clean.add_argument("--include-disabled", action="store_true", help="Also remove disabled models' files")
    p_clean.add_argument("--dry-run", action="store_true", help="Show what would be removed without deleting")
    args = parser.parse_args()
    registry = Registry(args.registry)
    commands = {
        "build": cmd_build,
        "info": cmd_info,
        "status": cmd_status,
        "list": cmd_list,
        "add": cmd_add,
        "download": cmd_download,
        "enable": cmd_enable,
        "disable": cmd_disable,
        "validate": cmd_validate,
        "monitor": cmd_monitor,
        "scan": cmd_scan,
        "cleanup": cmd_cleanup,
    commands[args.command](args, registry)
if __name__ == "__main__":
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

manage.py

Latest commit

History

manage.py

File metadata and controls
