# Docker Compose for LocalAI Distributed Mode
#
# Starts a full distributed stack: PostgreSQL, NATS, Garage (S3-compatible
# object storage), a LocalAI frontend, a chat worker, a model worker and an
# agent worker.
#
# Model files are transferred from the frontend to backend nodes via HTTP
# — no shared volumes needed between frontend and backends.
#
# The bytes of the chat files go to Garage, so every frontend and chat worker
# reads what another one wrote. All credentials below are examples: change
# them for anything beyond a local test.
#
# Usage:
#   docker compose -f docker-compose.distributed.yaml up
#
# See docs: https://localai.io/features/distributed-mode/

services:
  # --- Infrastructure ---

  postgres:
    image: quay.io/mudler/localrecall:v0.5.5-postgresql  # PostgreSQL with pgvector
    environment:
      POSTGRES_DB: localai
      POSTGRES_USER: localai
      POSTGRES_PASSWORD: localai
    volumes:
      - postgres_data:/var/lib/postgresql
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U localai"]
      interval: 5s
      timeout: 3s
      retries: 10

  nats:
    image: nats:2-alpine
    ports:
      - "4222:4222"   # Client connections
      - "8222:8222"   # HTTP monitoring (optional, useful for debugging)
    command: ["--js", "-m", "8222"]  # Enable JetStream + monitoring

  # Garage: a small S3-compatible object store. One node, no replication.
  # The chat files (and optionally model files) are stored here, so that every
  # frontend sees the same bytes. The node needs a layout before it takes
  # data: garage-init below assigns it and creates the bucket and the key.
  garage:
    image: dxflrs/garage:v2.1.0
    volumes:
      - ./examples/distributed/garage.toml:/etc/garage.toml:ro
      - garage_data:/var/lib/garage

  # One-shot setup of Garage: layout, bucket, access key. Safe to run again.
  # The Garage image has no shell, so this image copies the garage binary
  # next to one. The node ID is the public key that Garage keeps in its
  # metadata directory, which this container reads from the same volume.
  garage-init:
    build:
      dockerfile_inline: |
        FROM dxflrs/garage:v2.1.0 AS garage
        FROM alpine:3.20
        COPY --from=garage /garage /usr/local/bin/garage
    restart: "no"
    volumes:
      - ./examples/distributed/garage.toml:/etc/garage.toml:ro
      - garage_data:/var/lib/garage:ro
    depends_on:
      garage:
        condition: service_started
    entrypoint: ["/bin/sh", "-c"]
    command:
      - |
        set -e
        until [ -s /var/lib/garage/meta/node_key.pub ]; do sleep 1; done
        ID=$$(od -An -v -tx1 /var/lib/garage/meta/node_key.pub | tr -d ' \n')
        G="garage -c /etc/garage.toml -h $$ID@garage:3901"
        until $$G status >/dev/null 2>&1; do sleep 1; done
        # These steps fail harmlessly when a previous run already did them.
        $$G layout assign -z dc1 -c 10G $$ID || true
        $$G layout apply --version 1 || true
        until $$G bucket create localai >/dev/null 2>&1 || $$G bucket info localai >/dev/null 2>&1; do sleep 1; done
        $$G key import --yes -n localai GK0123456789abcdef01234567 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef || true
        $$G bucket allow --read --write --owner localai --key localai
        echo "garage ready: bucket localai"

  # --- LocalAI Frontend ---
  # Stateless API server that routes requests to backend nodes.
  # Add more replicas behind a load balancer for HA.

  localai:
    # image: localai/localai:latest-cpu
    build:
      context: .
      dockerfile: Dockerfile
      args:
        - IMAGE_TYPE=core
        - BASE_IMAGE=ubuntu:24.04
    ports:
      - "8080:8080"
    environment:
      # Distributed mode
      LOCALAI_DISTRIBUTED: "true"
      LOCALAI_NATS_URL: "nats://nats:4222"
      LOCALAI_COLLECTIONS_EMBEDDING_MODEL: "granite-embedding-107m-multilingual"
      LOCALAI_COLLECTIONS_VECTOR_ENGINE: "postgres"
      LOCALAI_COLLECTIONS_DATABASE_URL: "postgresql://localai:localai@postgres:5432/localai?sslmode=disable"
      LOCALAI_REGISTRATION_TOKEN: "changeme"  # Change this in production!
      # Object storage (Garage). Chat files need it in a cluster: without it
      # they are off. Path-style addressing is always used, which Garage needs.
      LOCALAI_STORAGE_URL: "http://garage:3900"
      LOCALAI_STORAGE_BUCKET: "localai"
      LOCALAI_STORAGE_REGION: "us-east-1"  # Must match s3_region in garage.toml
      LOCALAI_STORAGE_ACCESS_KEY: "GK0123456789abcdef01234567"
      LOCALAI_STORAGE_SECRET_KEY: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
      # Send the turns of triggers, agent jobs and connector messages to the
      # chat worker below. With no chat worker running, the frontend runs them.
      LOCALAI_DISTRIBUTED_CHAT_WORKERS: "true"
      # Shared-models mode (optional): set when every node mounts the SAME
      # models directory at the SAME path (see "Shared Volume Mode" below).
      # The router then skips gRPC file staging and workers load models
      # directly from the shared volume instead of re-downloading them.
      # LOCALAI_DISTRIBUTED_SHARED_MODELS: "true"
      # Auth (required for distributed mode — must use PostgreSQL)
      LOCALAI_AUTH: "true"
      LOCALAI_AUTH_DATABASE_URL: "postgresql://localai:localai@postgres:5432/localai?sslmode=disable"
      # Force pure-Go DNS resolver. The default cgo resolver follows the
      # container's nsswitch.conf and ends up forwarding to host
      # systemd-resolved (127.0.0.53), which isn't reachable from inside
      # the container — failing every postgres/nats hostname lookup at
      # boot. The pure-Go path reads /etc/resolv.conf directly and uses
      # Docker's embedded DNS at 127.0.0.11.
      GODEBUG: "netdns=go"
      # Paths
      MODELS_PATH: /models
      # Avoid probing remote gallery GGUF metadata during container startup.
      # Remove this line or set a positive limit to opt back into cache warming.
      LOCALAI_VRAM_WARM_LIMIT: "0"
    volumes:
      - frontend_models:/models
      - frontend_data:/data
    depends_on:
      postgres:
        condition: service_healthy
      nats:
        condition: service_started
      garage-init:
        condition: service_completed_successfully

  # --- Chat Worker ---
  # Runs the turns of triggers, agent jobs and connector messages away from
  # the frontend. It is the frontend configuration plus LOCALAI_CHAT_WORKER:
  # same database, NATS, storage, models path and data path. It starts no
  # triggers, connectors or jobs, and it does not take user traffic.

  chat-worker-1:
    # image: localai/localai:latest-cpu
    build:
      context: .
      dockerfile: Dockerfile
      args:
        - IMAGE_TYPE=core
        - BASE_IMAGE=ubuntu:24.04
    environment:
      LOCALAI_DISTRIBUTED: "true"
      LOCALAI_CHAT_WORKER: "true"
      LOCALAI_NATS_URL: "nats://nats:4222"
      LOCALAI_AUTH: "true"
      LOCALAI_AUTH_DATABASE_URL: "postgresql://localai:localai@postgres:5432/localai?sslmode=disable"
      LOCALAI_COLLECTIONS_EMBEDDING_MODEL: "granite-embedding-107m-multilingual"
      LOCALAI_COLLECTIONS_VECTOR_ENGINE: "postgres"
      LOCALAI_COLLECTIONS_DATABASE_URL: "postgresql://localai:localai@postgres:5432/localai?sslmode=disable"
      LOCALAI_REGISTRATION_TOKEN: "changeme"  # Must match frontend token
      LOCALAI_STORAGE_URL: "http://garage:3900"
      LOCALAI_STORAGE_BUCKET: "localai"
      LOCALAI_STORAGE_REGION: "us-east-1"
      LOCALAI_STORAGE_ACCESS_KEY: "GK0123456789abcdef01234567"
      LOCALAI_STORAGE_SECRET_KEY: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
      LOCALAI_CHAT_WORKER_MAX_TURNS: "32"
      GODEBUG: "netdns=go"  # See note in localai service
      MODELS_PATH: /models
      LOCALAI_VRAM_WARM_LIMIT: "0"
    volumes:
      - chat_worker_1_models:/models
      - chat_worker_1_data:/data
    depends_on:
      postgres:
        condition: service_healthy
      nats:
        condition: service_started
      garage-init:
        condition: service_completed_successfully

  # --- Worker Node ---
  # A generic worker that self-registers with the frontend.
  # The same LocalAI image is used — no separate image needed.
  # The SmartRouter dynamically tells workers which backend to install via NATS.
  #
  # Model files are transferred from the frontend via HTTP file staging.
  # The worker has its own independent models volume.

  worker-1:
    # image: localai/localai:latest-cpu
    build:
      context: .
      dockerfile: Dockerfile
      args:
        - IMAGE_TYPE=core
        - BASE_IMAGE=ubuntu:24.04
    command:
      - worker
    # No HEALTHCHECK_ENDPOINT override is needed: the image's healthcheck
    # detects worker mode and derives the port from LOCALAI_SERVE_ADDR below
    # (gRPC base port - 1 = 50050). The worker's /readyz reports 503 when its
    # NATS connection is down, and also when a backend process it is already
    # serving no longer answers a TCP dial on its gRPC address. So `unhealthy`
    # here means the worker genuinely cannot receive work, or cannot serve the
    # models it is holding.
    environment:
      LOCALAI_SERVE_ADDR: "0.0.0.0:50051"
      LOCALAI_ADVERTISE_ADDR: "worker-1:50051"
      LOCALAI_ADVERTISE_HTTP_ADDR: "worker-1:50050"
      DEBUG: "true"
      LOCALAI_REGISTER_TO: "http://localai:8080"
      LOCALAI_NODE_NAME: "worker-1"
      LOCALAI_REGISTRATION_TOKEN: "changeme"  # Must match frontend token
      LOCALAI_HEARTBEAT_INTERVAL: "10s"
      LOCALAI_NATS_URL: "nats://nats:4222"
      GODEBUG: "netdns=go"  # See note in localai service
      MODELS_PATH: /models
    volumes:
      - worker_1_models:/models
    depends_on:
      localai:
        condition: service_started
      nats:
        condition: service_started

    # --- GPU Support (NVIDIA) ---
    # Uncomment the following and change the image to a CUDA variant
    # (e.g., localai/localai:latest-gpu-nvidia-cuda-12) to enable GPU.
    #
    # NVIDIA_DRIVER_CAPABILITIES must include `utility` so nvidia-smi / NVML
    # are available inside the container; without it the worker cannot report
    # free VRAM and the Nodes page will show 0 free / total used.
    # `init: true` avoids zombie-reap races that make nvidia-smi flaky.
    #
    # init: true
    # environment:
    #   NVIDIA_DRIVER_CAPABILITIES: "compute,utility"
    # deploy:
    #   resources:
    #     reservations:
    #       devices:
    #         - driver: nvidia.com/gpu
    #           count: all
    #           capabilities: [gpu, utility]

  # --- Shared Volume Mode (optional) ---
  # If all services run on the same Docker host, you can skip gRPC file transfer
  # by sharing a single models volume. Replace the volumes above with:
  #
  #   localai:
  #     volumes:
  #       - shared_models:/models
  #       - frontend_data:/data
  #
  #   backend-llama-cpp:
  #     volumes:
  #       - shared_models:/models
  #
  # Then add to the volumes section:
  #   shared_models:
  #
  # With shared volumes the model files are already present on every worker at
  # the same path. Set LOCALAI_DISTRIBUTED_SHARED_MODELS=true on the frontend
  # (see its environment above) so the router skips gRPC file staging and the
  # worker loads the model directly from the shared path instead of
  # re-downloading it into a per-model subdirectory.

  # --- Adding More Workers ---
  # Copy the worker-1 service above and change:
  #   - Service name (e.g., worker-2)
  #   - LOCALAI_NODE_NAME (must be unique)
  #   - LOCALAI_ADVERTISE_ADDR (must match service name)
  #
  # Workers are generic — no backend type needed. The SmartRouter
  # will dynamically install the required backend via NATS when
  # a model request arrives.

  # --- Agent Worker ---
  # Dedicated process that runs the MCP servers of the models, such as stdio
  # servers that need docker. The frontend sends it tool calls and discovery
  # requests over NATS. No database access needed.

  agent-worker-1:
    # image: localai/localai:latest-cpu
    build:
      context: .
      dockerfile: Dockerfile
      args:
        - IMAGE_TYPE=core
        - BASE_IMAGE=ubuntu:24.04
    # Install Docker CLI and start agent-worker.
    # The Docker socket is mounted from the host so that MCP stdio servers
    # using "docker run" commands can spawn containers on the host Docker.
    entrypoint: ["/bin/sh", "-c"]
    command:
      - |
        apt-get update -qq && apt-get install -y -qq docker.io >/dev/null 2>&1
        exec /entrypoint.sh agent-worker
    # The agent worker is NATS-only — no HTTP server to probe. The image's
    # healthcheck detects that mode and reports healthy rather than probing a
    # port that will never bind, so no override is needed here.
    environment:
      LOCALAI_NATS_URL: "nats://nats:4222"
      LOCALAI_REGISTER_TO: "http://localai:8080"
      LOCALAI_NODE_NAME: "agent-worker-1"
      LOCALAI_REGISTRATION_TOKEN: "changeme"  # Must match frontend token
      GODEBUG: "netdns=go"  # See note in localai service
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
    depends_on:
      localai:
        condition: service_started
      nats:
        condition: service_started

volumes:
  postgres_data:
  frontend_models:
  frontend_data:
  worker_1_models:
  garage_data:
  chat_worker_1_models:
  chat_worker_1_data: