# GPU/GQE overlay for the clean-slate release compose.
#
#   RVBBIT_VERSION=1.0.0 docker compose \
#       -f docker-compose.release.yml -f docker-compose.release-gqe.yml up -d
#
# That PULLS the prebuilt runtime image (ghcr.io/ryrobes/rvbbit-postgres-gqe,
# ~9GB) — no on-box compilation. One image covers every CUDA CC 8.0+ GPU:
# A100 / RTX 30- and 40-series (sm_8x runs sm_80 SASS), H100/H200 (90),
# B100/B200/GB200 (100), RTX 50-series / RTX PRO Blackwell / DGX Spark
# (120 SASS + PTX forward-compat). Turing and older are not supported
# (cudf 26.x dropped them).
#
# BUILD-ON-BOX (optional, for custom GQE/cudf refs or arch lists): the fat
# dev image builds from the published base with docker/Dockerfile.rvbbit-gqe
# (one-time toolchain build, ~hours), then docker/Dockerfile.rvbbit-gqe-runtime
# extracts a slim publishable runtime image from it:
#
#   RVBBIT_VERSION=1.0.0 docker compose \
#       -f docker-compose.release.yml -f docker-compose.release-gqe.yml \
#       build postgres
#
# Requirements: NVIDIA GPU + driver + nvidia-container-toolkit (docker info
# must list the 'nvidia' runtime). Without them, use the plain release
# compose — the extension's GQE routing gate is on by default but self-gates
# on runtime availability, so a non-GPU box behaves identically either way.
services:
  # On every `up`, persist the GQE router-prior GUC to match the overlay env.
  # The env var alone reaches the Rust router, but rvbbit.warm_gpu_gqe() is
  # plpgsql and reads only the GUC — without this, heartbeat warm stays
  # "disabled" even though the overlay asked for it. Fully self-contained so
  # the overlay works over BOTH the release compose (which has a migrate
  # service) and the uber compose (which doesn't). Each -c is its own
  # transaction: ALTER SYSTEM refuses to run inside one.
  migrate:
    image: ${RVBBIT_POSTGRES_GQE_IMAGE:-ghcr.io/ryrobes/rvbbit-postgres-gqe:${RVBBIT_VERSION:-3.2.0}}
    container_name: rvbbit-migrate
    depends_on:
      postgres:
        condition: service_healthy
    environment:
      PGPASSWORD: ${POSTGRES_PASSWORD:-rvbbit}
    entrypoint: ["psql"]
    command:
      - -h
      - postgres
      - -U
      - postgres
      - -d
      - ${POSTGRES_DB:-rvbbit}
      - -v
      - ON_ERROR_STOP=1
      - -c
      - CREATE EXTENSION IF NOT EXISTS pg_rvbbit; ALTER EXTENSION pg_rvbbit UPDATE; SELECT rvbbit.migrate();
      - -c
      - ALTER SYSTEM SET rvbbit.route_gpu_gqe_prior = '${RVBBIT_ROUTE_GPU_GQE_PRIOR:-on}'
      - -c
      - SELECT pg_reload_conf()
    restart: "no"

  postgres:
    image: ${RVBBIT_POSTGRES_GQE_IMAGE:-ghcr.io/ryrobes/rvbbit-postgres-gqe:${RVBBIT_VERSION:-3.2.0}}
    build:
      context: .
      dockerfile: Dockerfile.rvbbit-gqe
      args:
        RVBBIT_BASE_IMAGE: ${RVBBIT_POSTGRES_IMAGE:-ghcr.io/ryrobes/rvbbit-postgres:${RVBBIT_VERSION:-3.2.0}}
        GQE_REPO: ${GQE_REPO:-https://github.com/rapidsai/gqe.git}
        GQE_REF: ${GQE_REF:-main}
        CUDF_REF: ${CUDF_REF:-v26.06.01}
        PARALLEL_LEVEL: ${PARALLEL_LEVEL:-}
        SCCACHE_VERSION: ${SCCACHE_VERSION:-0.15.0}
    gpus: all
    shm_size: ${RVBBIT_GQE_SHM_SIZE:-8gb}
    cap_add:
      - IPC_LOCK
    ulimits:
      memlock:
        soft: -1
        hard: -1
      stack: 67108864
    environment:
      NVIDIA_VISIBLE_DEVICES: ${NVIDIA_VISIBLE_DEVICES:-all}
      NVIDIA_DRIVER_CAPABILITIES: ${NVIDIA_DRIVER_CAPABILITIES:-compute,utility}
      # NVSHMEM VMM: =1 (disable) is required on some older cards, but BREAKS
      # Blackwell (InitNvshmem: pgas_memory_resource allocation failed ->
      # every GQE query times out + falls back). On Blackwell-class GPUs set
      # NVSHMEM_DISABLE_CUDA_VMM=0 (validated: RTX PRO 6000 / driver 580).
      NVSHMEM_DISABLE_CUDA_VMM: ${NVSHMEM_DISABLE_CUDA_VMM:-1}
      NVSHMEM_SYMMETRIC_SIZE: ${NVSHMEM_SYMMETRIC_SIZE:-6G}
      GQE_MAX_QUERY_MEMORY: ${GQE_MAX_QUERY_MEMORY:-6442450944}
      RVBBIT_GQE_CLI: ${RVBBIT_GQE_CLI:-/opt/gqe/rust/target/release/gqe-cli}
      RVBBIT_GQE_NODE_MANAGER: ${RVBBIT_GQE_NODE_MANAGER:-/opt/gqe/build/src/node_manager/gqe_node_manager}
      RVBBIT_GQE_TASK_MANAGER: ${RVBBIT_GQE_TASK_MANAGER:-/opt/gqe/build/src/task_manager/gqe_task_manager}
      RVBBIT_GQE_SERVER_URL: ${RVBBIT_GQE_SERVER_URL:-http://127.0.0.1:50051}
      RVBBIT_GQE_AUTO_START: ${RVBBIT_GQE_AUTO_START:-true}
      # This overlay IS the explicit GPU opt-in, so enable the GQE router prior:
      # the accel_tick heartbeat keeps the engine warm (rvbbit.warm_gpu_gqe) and
      # the router may pick GQE for large shapes without a profile — never cold
      # (gated on fresh gqe_warm_state). Set to "off" to revert to lazy start.
      RVBBIT_ROUTE_GPU_GQE_PRIOR: ${RVBBIT_ROUTE_GPU_GQE_PRIOR:-on}
      RVBBIT_GQE_CLIENT_MODE: ${RVBBIT_GQE_CLIENT_MODE:-flight}
      RVBBIT_GQE_FLIGHT_FALLBACK: ${RVBBIT_GQE_FLIGHT_FALLBACK:-true}
      # Routing gate (rvbbit.route_gpu_gqe) defaults ON in the extension and
      # self-gates on runtime availability — no override needed here.