DOCSEARCHbyAlgolia
Submit Docs

llama-swap

repository·main·Indexed 26 days ago

https://github.com/mostlygeek/llama-swap

A high-performance Go-based wrapper that allows running multiple generative AI models on a single machine with on-demand hot-swapping. It acts as a proxy for OpenAI, Anthropic, and llama.cpp compatible servers, automatically managing model loading and unloading based on API requests. Includes utilities like vllm-wrapper for vLLM daemon lifecycle management and wol-proxy for Wake-on-LAN server activation.

Tokens
11.8K
Snippets
35
Records
76
Agent score
89%

What's inside llama-swap

  1. Use wol-proxy to wake up llama-swap via Wake-on-LAN

    main

    wol-proxy is a utility that automatically wakes up a suspended llama-swap server using Wake-on-LAN (WOL) when requests are received.

    When a request arrives and the upstream llama-swap server is unavailable, wol-proxy sends a WOL packet and holds the request until the server becomes available. If the server does not respond within the specified timeout period (default: 60 seconds), the request is dropped. This is useful for energy conservation on GPU-heavy servers.

  2. Understand the llama-swap Router architecture

    main

    The llama-swap router is composed of three decoupled concerns that manage how models are served. This separation allows you to replace scheduling or eviction logic without affecting process management:

    1. Process machinery (baseRouter): Manages OS processes, lifecycle (start/stop), health checks, and HTTP request routing. Located in internal/router/base.go.
    2. Scheduling strategy (scheduler.Scheduler): Manages the request queue, in-flight bookkeeping, and the decision tree (e.g., whether to serve now or start a swap). The current implementation is FIFO.
    3. Eviction policy (scheduler.Swapper): A pure function that determines which currently running models must be stopped to make room for a target model. Implementations include groupSwapper and matrixSwapper.

  3. Use the Unified Docker Container

    main

    The Unified Docker Container is a custom llama-swap image that bundles multiple AI services into a single environment. It includes:

    • llama-server: Supports LLMs, rerank models, and embedding models.
    • sd-server (stable-diffusion.cpp): For image generation.
    • whisper.cpp: For Automatic Speech Recognition (ASR).

  5. Test a new Scheduler

    main

    Schedulers should be tested as pure state machines within the scheduler package.

    1. Unit Testing: Drive the On* methods directly against a fakeEffects instance. Assert on recorded grants, starts, and stops. Avoid using goroutines or sleeps in these tests. Follow the TestSchedulerName_<scenario> naming convention.
    2. Integration Testing: The baseRouter mechanism is tested in base_test.go. Use the testProcessed channel to wait for events to be fully processed instead of using sleeps.

    Commands:

    To run specific scheduler tests:

    go test -v -run TestMyScheduler_... ./internal/router/scheduler/

    To run a quick test and staticcheck pass over the proxy:

    make test-dev

  6. Manage in-flight request counts in the Scheduler

    main

    The scheduler tracks whether a model is "busy" by counting grants out and ServeDone events in. To prevent the in-flight counter from becoming permanently stuck (which would prevent a model from ever being evicted), you must follow this contract when using GrantServe:

    Only increment the inFlight counter if GrantServe returns true.

    If GrantServe returns false, it means the caller's Respond channel was unable to receive (likely because the HTTP client disconnected). In this case, trackedServe will not run, and no ServeDoneEvent will ever be sent. Incrementing the counter on a false return will strand the counter above zero, making the model un-evictable.

  8. Use Macros for Reusable Configuration Snippets

    main

    Macros allow you to define reusable string substitutions. They can be defined globally or per-model.

    • Global Macros: Available in any configuration value.
    • Model Macros: Override global macros for that specific model.
    • Environment Variables: Reference using ${env.VAR_NAME} syntax. If the variable is missing, loading fails.
    • Reserved Names: Do not use PID, PORT, or MODEL_ID as macro names.
    • Special Macros:
      • ${PORT}: Assigned automatically for each model using it in cmd.
      • ${PID}: Substituted when a model's cmdStop runs.
      • ${MODEL_ID}: The current model's ID.

  10. Understand Macro Substitution Rules and Ordering

    main

    Llama-swap uses a hierarchical, order-dependent macro substitution system. Macros are substituted in reverse definition order (LIFO - Last In, First Out) within their respective levels. This allows a macro defined later in the configuration to reference a macro defined earlier.

    Substitution Hierarchy

    1. Reserved macros (Highest priority, substituted last): ${PORT}, ${MODEL_ID}.
    2. Model-level macros (Middle priority): Defined within a specific model configuration; these override global macros of the same name.
    3. Global macros (Lowest priority): Defined at the root level of the configuration.

    Reference Rules

    • Allowed:
      • Referencing any macro defined before the current one in the file.
      • Model macros referencing global macros.
      • Referencing reserved macros ${PORT} and ${MODEL_ID}.
    • Prohibited:
      • Self-references: A macro cannot contain its own name (e.g., foo: "${foo}" is invalid).
      • Forward references: A macro cannot reference a macro defined after it in the file.
      • Circular references: The system uses single-pass substitution to prevent infinite loops.

  11. Implement a custom Swapper for custom eviction logic

    main

    A Swapper defines the eviction policy for the router. To implement a new one, follow these steps:

    1. Define your type: Create a struct that holds any necessary configuration. It does not need access to the process map, as the scheduler provides the running set.
    2. Implement EvictionFor(target string, running []string) []string: This must be a pure decision function.
      • running is the complete set of non-stopped processes plus targets of in-flight swaps.
      • Return a slice of model IDs that must stop for the target to run. Return nil or an empty slice if no eviction is needed.
      • CRITICAL: Do not mutate state and do not log inside this method. It is called frequently for speculative queries; logging here will cause duplicate or misleading log entries.
    3. Implement OnSwapStart(target string, running []string): This is called exactly once when a swap is actually committed. This is the correct place to perform logging.
    4. Wire it in: Instantiate your swapper and capture it in the Factory closure passed to newBaseRouter during construction.

  12. Use vllm-wrapper sleep as a model's cmdStop

    main

    The sleep subcommand is used as a model's cmdStop in llama-swap. It sends a sleep request to the vLLM daemon to free VRAM while keeping the process alive, allowing for near-instant wake-ups later.

    models:
  my-vllm-model:
    cmdStop: vllm-wrapper sleep --vllm-url http://127.0.0.1:8000