# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0

services:
  redis:
    image: "redis/redis-stack"
    ports:
      - "6379:6379"

  page-elements:
    image: ${YOLOX_IMAGE:-nvcr.io/nim/nvidia/nemoretriever-page-elements-v2}:${YOLOX_TAG:-1.3.0}
    ports:
      - "8000:8000"
      - "8001:8001"
      - "8002:8002"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
      - NIM_TRITON_MODEL_BATCH_SIZE=${PAGE_ELEMENTS_BATCH_SIZE:-1}
      # NIM OpenTelemetry Settings
      - NIM_ENABLE_OTEL=0
      - NIM_OTEL_SERVICE_NAME=page-elements
      - NIM_OTEL_TRACES_EXPORTER=otlp
      - NIM_OTEL_METRICS_EXPORTER=console
      - NIM_OTEL_EXPORTER_OTLP_ENDPOINT=http://otel-collector:4318
      - NIM_ENABLE_OTEL=true
      # Triton OpenTelemetry Settings
      - TRITON_OTEL_URL=http://otel-collector:4318/v1/traces
      - TRITON_OTEL_RATE=1
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia

  graphic-elements:
    image: ${YOLOX_GRAPHIC_ELEMENTS_IMAGE:-nvcr.io/nim/nvidia/nemoretriever-graphic-elements-v1}:${YOLOX_GRAPHIC_ELEMENTS_TAG:-1.3.0}
    ports:
      - "8003:8000"
      - "8004:8001"
      - "8005:8002"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
      - NIM_TRITON_MODEL_BATCH_SIZE=${GRAPHIC_ELEMENTS_BATCH_SIZE:-1}
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia

  table-structure:
    image: ${YOLOX_TABLE_STRUCTURE_IMAGE:-nvcr.io/nim/nvidia/nemoretriever-table-structure-v1}:${YOLOX_TABLE_STRUCTURE_TAG:-1.3.0}
    ports:
      - "8006:8000"
      - "8007:8001"
      - "8008:8002"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
      - NIM_TRITON_MODEL_BATCH_SIZE=${TABLE_STRUCTURE_BATCH_SIZE:-1}
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia
    profiles:
      - table-structure

  paddle:
    image: ${PADDLE_IMAGE:-nvcr.io/nim/baidu/paddleocr}:${PADDLE_TAG:-1.3.0}
    shm_size: 2gb
    ports:
      - "8009:8000"
      - "8010:8001"
      - "8011:8002"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia

  embedding:
    # NIM ON
    image: ${EMBEDDING_IMAGE:-nvcr.io/nim/nvidia/llama-3.2-nv-embedqa-1b-v2}:${EMBEDDING_TAG:-1.6.0}
    shm_size: 16gb
    ports:
      - "8012:8000"
      - "8013:8001"
      - "8014:8002"
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
      - NIM_TRITON_MAX_BATCH_SIZE=${EMBEDDER_BATCH_SIZE:-1}
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia

  reranker:
    # NIM ON
    image: ${RERANKER_IMAGE:-nvcr.io/nim/nvidia/llama-3.2-nv-rerankqa-1b-v2}:${RERANKER_TAG:-1.5.0}
    shm_size: 16gb
    ports:
      - "8020:8000"
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia
    profiles:
      - reranker

  nemoretriever-parse:
    image: ${NEMORETRIEVER_PARSE_IMAGE:-nvcr.io/nvidia/nemo-microservices/nemoretriever-parse}:${NEMORETRIEVER_PARSE_TAG:-1.2.0ea}
    ports:
      - "8015:8000"
      - "8016:8001"
      - "8017:8002"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia
    profiles:
      - nemoretriever-parse

  vlm:
    image: ${VLM_IMAGE:-nvcr.io/nim/meta/llama-3.2-11b-vision-instruct}:${VLM_TAG:-1.1.1}
    ports:
      - "8018:8000"
    user: root
    environment:
      - NIM_HTTP_API_PORT=8000
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
      # VLM will use all available VRAM on device
      # For more info
      # https://docs.nvidia.com/nim/vision-language-models/latest/configuration.html
      #- NIM_KVCACHE_PERCENT=.25
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia
    profiles:
      - vlm

  audio:
    image: ${AUDIO_IMAGE:-nvcr.io/nim/nvidia/riva-asr}:${AUDIO_TAG:-1.3.0}
    shm_size: 2gb
    ports:
      - "8021:50051"  # grpc
      - "8022:9000"  # http
    user: root
    environment:
      - NIM_TAGS_SELECTOR=name=parakeet-1-1b-ctc-riva-en-us,mode=ofl
      - NIM_TRITON_LOG_VERBOSE=1
      - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    runtime: nvidia
    profiles:
      - audio

  nv-ingest-ms-runtime:
    image: nvcr.io/nvidia/nemo-microservices/nv-ingest:25.6.2
    shm_size: 40gb # Should be at minimum 30% of assigned memory per Ray documentation
    build:
      context: ${NV_INGEST_ROOT:-.}
      dockerfile: "./Dockerfile"
      target: runtime
      args:
        DOWNLOAD_LLAMA_TOKENIZER: ${DOWNLOAD_LLAMA_TOKENIZER:-False}
        HF_ACCESS_TOKEN: ${HF_ACCESS_TOKEN:-hfaccesstoken}
        MODEL_PREDOWNLOAD_PATH: ${MODEL_PREDOWNLOAD_PATH:-/workspace/models/}
    volumes:
      - ${DATASET_ROOT:-./data}:/workspace/data
    ports:
      # HTTP API
      - "7670:7670"
      # Simple Broker
      - "7671:7671"
      # Ray Dashboards
      - "8265:8265"
    cap_add:
      - sys_nice
    environment:
      - ARROW_DEFAULT_MEMORY_POOL=system
      - AUDIO_GRPC_ENDPOINT=audio:50051
      - AUDIO_INFER_PROTOCOL=grpc
      - CUDA_VISIBLE_DEVICES=-1
      - MAX_INGEST_PROCESS_WORKERS=${MAX_INGEST_PROCESS_WORKERS:-16}
      - EMBEDDING_NIM_ENDPOINT=${EMBEDDING_NIM_ENDPOINT:-http://embedding:8000/v1}
      - EMBEDDING_NIM_MODEL_NAME=${EMBEDDING_NIM_MODEL_NAME:-nvidia/llama-3.2-nv-embedqa-1b-v2}
      - INGEST_LOG_LEVEL=DEFAULT
      - INGEST_EDGE_BUFFER_SIZE=64
      - INGEST_DYNAMIC_MEMORY_THRESHOLD=0.80
      #- INGEST_DISABLE_DYNAMIC_SCALING=true # Disable dynamic scaling and use static worker count
      #- INSTALL_AUDIO_EXTRACTION_DEPS=true
      # Message client for development
      #- MESSAGE_CLIENT_HOST=0.0.0.0
      #- MESSAGE_CLIENT_PORT=7671
      #- MESSAGE_CLIENT_TYPE=simple # Configure the ingest service to use the simple broker
      # Message client for production
      - MESSAGE_CLIENT_HOST=redis
      - MESSAGE_CLIENT_PORT=6379
      - MESSAGE_CLIENT_TYPE=redis
      - MINIO_BUCKET=${MINIO_BUCKET:-nv-ingest}
      # build.nvidia.com hosted nemoretriever-parse
      # - NEMORETRIEVER_PARSE_HTTP_ENDPOINT=https://integrate.api.nvidia.com/v1/chat/completions
      - NEMORETRIEVER_PARSE_HTTP_ENDPOINT=http://nemoretriever-parse:8000/v1/chat/completions
      - NEMORETRIEVER_PARSE_INFER_PROTOCOL=http
      #- NEMORETRIEVER_PARSE_MODEL_NAME=nvidia/nemoretriever-parse
      - NGC_API_KEY=${NGC_API_KEY:-ngcapikey}
      - NVIDIA_BUILD_API_KEY=${NVIDIA_BUILD_API_KEY:-${NGC_API_KEY:-ngcapikey}}
      - NV_INGEST_MAX_UTIL=${NV_INGEST_MAX_UTIL:-48}
      - OTEL_EXPORTER_OTLP_ENDPOINT=otel-collector:4317
      # Self-hosted paddle endpoints.
      - PADDLE_GRPC_ENDPOINT=paddle:8001
      - PADDLE_HTTP_ENDPOINT=http://paddle:8000/v1/infer
      - PADDLE_INFER_PROTOCOL=grpc
      # build.nvidia.com hosted paddle endpoints.
      #- PADDLE_HTTP_ENDPOINT=https://ai.api.nvidia.com/v1/cv/baidu/paddleocr
      #- PADDLE_INFER_PROTOCOL=http
      - REDIS_INGEST_TASK_QUEUE=ingest_task_queue
      # Self-hosted redis endpoints.
      - YOLOX_GRPC_ENDPOINT=page-elements:8001
      - YOLOX_HTTP_ENDPOINT=http://page-elements:8000/v1/infer
      - YOLOX_INFER_PROTOCOL=grpc
      # build.nvidia.com hosted endpoints.
      #- YOLOX_HTTP_ENDPOINT=https://ai.api.nvidia.com/v1/cv/nvidia/nv-yolox-page-elements-v1
      #- YOLOX_INFER_PROTOCOL=http
      - YOLOX_GRAPHIC_ELEMENTS_GRPC_ENDPOINT=graphic-elements:8001
      - YOLOX_GRAPHIC_ELEMENTS_HTTP_ENDPOINT=http://graphic-elements:8000/v1/infer
      - YOLOX_GRAPHIC_ELEMENTS_INFER_PROTOCOL=grpc
      - YOLOX_TABLE_STRUCTURE_GRPC_ENDPOINT=table-structure:8001
      - YOLOX_TABLE_STRUCTURE_HTTP_ENDPOINT=http://table-structure:8000/v1/infer
      - YOLOX_TABLE_STRUCTURE_INFER_PROTOCOL=grpc
      - VLM_CAPTION_ENDPOINT=http://vlm:8000/v1/chat/completions
      - VLM_CAPTION_MODEL_NAME=meta/llama-3.2-11b-vision-instruct
      - MODEL_PREDOWNLOAD_PATH=${MODEL_PREDOWNLOAD_PATH:-/workspace/models/}
      # "ready" check configuration.
      # 1. COMPONENTS_TO_READY_CHECK= to disable and readiness checking
      # 2. COMPONENTS_TO_READY_CHECK=ALL for checking all services
      # 3. COMPONENTS_TO_READY_CHECK=YOLOX_HTTP_ENDPOINT, PADDLE_HTTP_ENDPOINT
      #    comma separated list of HTTP environment variables for specific services to check for ready
      - COMPONENTS_TO_READY_CHECK=ALL
    healthcheck:
      test: curl --fail http://nv-ingest-ms-runtime:7670/v1/health/ready || exit 1
      interval: 10s
      timeout: 5s
      retries: 20
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]

  otel-collector:
    image: otel/opentelemetry-collector-contrib:0.91.0
    hostname: otel-collector
    command: ["--config=/etc/otel-collector-config.yaml"]
    user: "${UID}:${GID}"
    volumes:
      - ./config/otel-collector-config.yaml:/etc/otel-collector-config.yaml
    ports:
      - "9988:9988" # Prometheus metrics exposed by the collector
      - "8889:8889" # Prometheus exporter metrics
      - "13133:13133" # health_check extension
      - "9411" # Zipkin receiver
      - "4317:4317" # OTLP gRPC receiver
      - "4318:4318" # OTLP/HTTP receiver
      - "55680:55679" # zpages extension
    depends_on:
      - zipkin

  zipkin:
    image: openzipkin/zipkin
    environment:
      JAVA_OPTS: "-Xms4g -Xmx8g -XX:+ExitOnOutOfMemoryError"
    ports:
      - "9411:9411" # Zipkin UI and API

  prometheus:
    image: prom/prometheus:latest
    command:
      - --web.console.templates=/etc/prometheus/consoles
      - --web.console.libraries=/etc/prometheus/console_libraries
      - --storage.tsdb.retention.time=1h
      - --config.file=/etc/prometheus/prometheus-config.yaml
      - --storage.tsdb.path=/prometheus
      - --web.enable-lifecycle
      - --web.route-prefix=/
      - --enable-feature=exemplar-storage
      - --enable-feature=otlp-write-receiver
    volumes:
      - ./config/prometheus.yaml:/etc/prometheus/prometheus-config.yaml
    ports:
      - "9090:9090"

  grafana:
    container_name: grafana-service
    image: grafana/grafana
    ports:
      - "3000:3000"

  etcd:
    # Turn on to leverage the `vdb_upload` task
    restart: always
    container_name: milvus-etcd
    image: quay.io/coreos/etcd:v3.5.5
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    volumes:
      - ./.volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    healthcheck:
      test: [ "CMD", "etcdctl", "endpoint", "health" ]
      interval: 30s
      timeout: 20s
      retries: 3
    profiles:
      - retrieval

  minio:
    # Turn on to leverage the `store` and `vdb_upload` task
    restart: always
    container_name: minio
    hostname: minio
    image: minio/minio:RELEASE.2023-03-20T20-16-18Z
    environment:
      MINIO_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minioadmin}
      MINIO_SECRET_KEY: ${MINIO_SECRET_KEY:-minioadmin}
    ports:
      - "9001:9001"
      - "9000:9000"
    volumes:
      - ./.volumes/minio:/minio_data
    command: minio server /minio_data --console-address ":9001"
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:9000/minio/health/live" ]
      interval: 30s
      timeout: 20s
      retries: 3
    profiles:
      - retrieval

  milvus:
    # Turn on to leverage the `vdb_upload` task
    restart: always
    container_name: milvus-standalone
    image: milvusdb/milvus:v2.5.3-gpu
    command: [ "milvus", "run", "standalone" ]
    hostname: milvus
    security_opt:
      - seccomp:unconfined
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ./.volumes/milvus:/var/lib/milvus
    healthcheck:
      test: [ "CMD", "curl", "-f", "http://localhost:9091/healthz" ]
      interval: 30s
      start_period: 90s
      timeout: 20s
      retries: 3
    ports:
      - "19530:19530"
      - "9091:9091"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    depends_on:
      - "etcd"
      - "minio"
    profiles:
      - retrieval

  attu:
    # Turn on to leverage the `vdb_upload` task
    restart: always
    container_name: milvus-attu
    image: zilliz/attu:v2.3.5
    hostname: attu
    environment:
      MILVUS_URL: http://milvus:19530
    ports:
      - "3001:3000"
    depends_on:
      - "milvus"
    profiles:
      - retrieval