# =====================================================================
# Ollama, plus a watchdog that restarts it when the GPU goes away.
#
# This is the model host. Both the assistant and the VM sub-agent talk to
# one Ollama with one model loaded, which is the whole reason the
# sub-agent's session is a single long block rather than many short ones:
# a handover evicts the other model's prefix cache and forces a cold
# re-prefill of its entire context.
#
# WHAT TO CHANGE
#   * the port binding, if 172.17.0.1 is not your docker bridge
#   * mem_limit, to suit the machine
#   * the ntfy block at the bottom, or delete it - alerting is optional
#
# WHAT NOT TO CHANGE
#   OLLAMA_NUM_PARALLEL and OLLAMA_MAX_LOADED_MODELS are both 1 on
#   purpose. The assistant runs at a 160k context, and the KV cache costs
#   roughly 85 KB of VRAM per token on a 27B Q4_K_M model - about 13.4 GB
#   on top of 16.5 GB of weights. A second slot of that size does not fit
#   in 32 GB, and letting Ollama try means it unloads and reloads the
#   model mid-conversation.
#
# WHY IT BINDS TO THE DOCKER BRIDGE AND NOT 0.0.0.0
#   Ollama has NO AUTHENTICATION and serves DELETE /api/delete,
#   POST /api/pull and POST /api/create. Anything that can reach it can
#   delete your model. Nothing in this project talks to it directly: the
#   sub-agent reaches it through owuitools-vmgw, an nginx that exposes
#   exactly /api/chat and /api/ps and 404s everything else.
# =====================================================================
# Pinned so the project, its volume and its network have the same names
# wherever this file is run from - see the note in open-webui.compose.yml.
name: ollama

services:
  ollama:
    image: ollama/ollama:latest
    container_name: ollama
    restart: unless-stopped

    ports:
      # The docker bridge address, not 0.0.0.0. See the note above.
      - "172.17.0.1:11434:11434"

    volumes:
      - ollama:/root/.ollama

    environment:
      OLLAMA_MAX_LOADED_MODELS: "1"
      OLLAMA_NUM_PARALLEL: "1"
      NVIDIA_VISIBLE_DEVICES: "all"
      NVIDIA_DRIVER_CAPABILITIES: "compute,utility"

    mem_limit: 24g
    memswap_limit: 24g
    gpus: all

    deploy:
      resources:
        limits:
          memory: 24g
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

  # -------------------------------------------------------------------
  # The watchdog. Optional, and worth having.
  #
  # It exists because of a failure mode that does not look like a
  # failure: the container keeps running and answering, and silently
  # falls back to CPU, so every reply still arrives and takes twenty
  # times as long. Three probes catch it - NVML through nvidia-smi, an
  # HTTP GET that catches a hung server, and a grep of the logs for the
  # CPU-fallback strings.
  #
  # Restarts are rate-limited, and alerts are throttled to one a day, so
  # a persistent fault does not become a persistent notification.
  # -------------------------------------------------------------------
  ollama-gpu-watchdog:
    container_name: ollama-gpu-watchdog
    depends_on:
      - ollama
    image: docker:cli
    restart: unless-stopped
    init: true
    mem_limit: 128m

    volumes:
      - /var/run/docker.sock:/var/run/docker.sock

    healthcheck:
      test: ["CMD-SHELL", "test $$(( $$(date +%s) - $$(stat -c %Y /tmp/heartbeat) )) -lt 300"]
      interval: 60s
      timeout: 5s
      retries: 3
      start_period: 30s

    entrypoint: ["/bin/sh", "-c"]
    command:
      - |
        set -u
        : "$${TARGET:=ollama}"
        : "$${INSTANCE:=$$TARGET}"
        : "$${CHECK_INTERVAL:=30}"
        : "$${RESTART_COOLDOWN:=600}"
        : "$${STARTUP_GRACE:=60}"
        : "$${GPU_FAIL_LIMIT:=2}"
        : "$${API_FAIL_LIMIT:=3}"
        : "$${LOG_WINDOW:=60s}"
        : "$${PROBE_TIMEOUT:=20}"
        : "$${GPU_CHECK_INTERVAL:=600}"
        : "$${GPU_RETRY_INTERVAL:=60}"
        : "$${GPU_PROBE_OFFSET:=0}"
        : "$${API_URL:=http://$$TARGET:11434/api/tags}"
        : "$${NOTIFY_COOLDOWN:=86400}"
        : "$${NTFY_URL:=}"
        : "$${NTFY_TOPIC:=}"
        : "$${NTFY_TOKEN:=}"

        log() { echo "$$(date -u +%Y-%m-%dT%H:%M:%SZ) [$$INSTANCE] $$*"; }

        LAST_RESTART=0
        LAST_NOTIFY=0
        GPU_FAILS=0
        API_FAILS=0

        notify() {
          [ -n "$$NTFY_URL" ] && [ -n "$$NTFY_TOKEN" ] || return 0
          now="$$(date +%s)"
          if [ "$$((now - LAST_NOTIFY))" -lt "$$NOTIFY_COOLDOWN" ]; then
            log "ntfy suppressed (24h throttle): $$*"
            return 0
          fi
          if wget -q -T 10 -O /dev/null \
               --header "Authorization: Bearer $$NTFY_TOKEN" \
               --post-data "$$(date -u +%Y-%m-%dT%H:%M:%SZ) $$*" \
               "$$NTFY_URL/$$NTFY_TOPIC?title=Ollama+watchdog:+$$INSTANCE&priority=high&tags=rotating_light"; then
            LAST_NOTIFY="$$(date +%s)"
            log "ntfy alert sent: $$*"
          else
            log "ERROR: ntfy publish failed (will retry on next event)"
          fi
        }

        restart_target() {
          reason="$$1"
          notify "$$reason - watchdog intervening on $$INSTANCE"
          now="$$(date +%s)"
          if [ "$$((now - LAST_RESTART))" -lt "$$RESTART_COOLDOWN" ]; then
            log "restart wanted ($$reason) but in cooldown"
            return 0
          fi
          log "RESTARTING $$TARGET: $$reason"
          if timeout 180 docker restart "$$TARGET" >/dev/null 2>&1; then
            LAST_RESTART="$$(date +%s)"
            GPU_FAILS=0
            API_FAILS=0
            log "restart issued OK; grace sleep $${STARTUP_GRACE}s"
            sleep "$$STARTUP_GRACE"
          else
            log "ERROR: docker restart $$TARGET failed/timed out"
            notify "docker restart of $$INSTANCE FAILED after: $$reason"
          fi
        }

        log "watchdog up"
        NEXT_GPU_PROBE="$$(( $$(date +%s) + GPU_PROBE_OFFSET ))"
        while true; do
          touch /tmp/heartbeat

          state="$$(timeout 10 docker inspect -f '{{.State.Status}}' "$$TARGET" 2>/dev/null || echo unknown)"
          if [ "$$state" != "running" ]; then
            log "target state=$$state - waiting (docker restart policy owns crash recovery)"
            sleep "$$CHECK_INTERVAL"
            continue
          fi

          # 1. GPU / NVML, only every GPU_CHECK_INTERVAL: without a
          # persistence client every nvidia-smi is a full driver
          # init/teardown. After a failure, re-probe sooner.
          now="$$(date +%s)"
          if [ "$$now" -ge "$$NEXT_GPU_PROBE" ]; then
            if timeout "$$PROBE_TIMEOUT" docker exec "$$TARGET" nvidia-smi >/dev/null 2>&1; then
              [ "$$GPU_FAILS" -gt 0 ] && log "GPU probe recovered"
              GPU_FAILS=0
              NEXT_GPU_PROBE="$$((now + GPU_CHECK_INTERVAL))"
            else
              GPU_FAILS="$$((GPU_FAILS + 1))"
              log "GPU probe FAILED ($$GPU_FAILS/$$GPU_FAIL_LIMIT)"
              NEXT_GPU_PROBE="$$((now + GPU_RETRY_INTERVAL))"
              if [ "$$GPU_FAILS" -ge "$$GPU_FAIL_LIMIT" ]; then
                restart_target "container lost GPU (NVML dead)"
                NEXT_GPU_PROBE="$$(( $$(date +%s) + GPU_RETRY_INTERVAL ))"
              fi
            fi
          fi

          # 2. the API itself - catches a hung server with no docker exec
          if timeout "$$PROBE_TIMEOUT" wget -q -T "$$PROBE_TIMEOUT" -O /dev/null "$$API_URL"; then
            [ "$$API_FAILS" -gt 0 ] && log "API probe recovered"
            API_FAILS=0
          else
            API_FAILS="$$((API_FAILS + 1))"
            log "API probe FAILED ($$API_FAILS/$$API_FAIL_LIMIT)"
            if [ "$$API_FAILS" -ge "$$API_FAIL_LIMIT" ]; then
              restart_target "API unresponsive"
            fi
          fi

          # 3. the silent one: still answering, but on the CPU
          if timeout 30 docker logs --since "$$LOG_WINDOW" "$$TARGET" 2>&1 | grep -qE \
            'msg="inference compute".*id=cpu|library=cpu|total_vram="0 B"|no compatible GPUs|failed to discover.*GPU|CUDA.*error|CUDA.*failure|unable to load CUDA'; then
            restart_target "GPU fallback / CUDA error in logs"
          fi

          sleep "$$CHECK_INTERVAL"
        done

    environment:
      TARGET: "ollama"
      INSTANCE: "ollama"
      CHECK_INTERVAL: "30"
      RESTART_COOLDOWN: "600"
      STARTUP_GRACE: "60"
      GPU_FAIL_LIMIT: "2"
      API_FAIL_LIMIT: "3"
      LOG_WINDOW: "60s"
      PROBE_TIMEOUT: "20"
      GPU_CHECK_INTERVAL: "600"
      GPU_RETRY_INTERVAL: "60"
      GPU_PROBE_OFFSET: "0"
      NOTIFY_COOLDOWN: "86400"
      # Alerting is optional: leave these empty and the watchdog still
      # restarts, it just does not tell anyone. Fill them in to get a
      # push when it intervenes. Use a write-only token.
      NTFY_URL: ""
      NTFY_TOPIC: ""
      NTFY_TOKEN: ""

volumes:
  ollama: