# NeMo Lens observability stack
#
# Usage:
#   docker compose -f docker-compose.otel.yml up -d
#
# Then bash into the Megatron container to run training:
#   docker compose -f docker-compose.otel.yml exec megatron bash
#
# UIs:
#   Jaeger         http://localhost:16686
#   Grafana        http://localhost:3000   (anonymous admin, no login required)
#   Kibana         http://localhost:5601
#   Prometheus     http://localhost:9090   (optional, for ad-hoc queries)
#   Elasticsearch  http://localhost:9200
#
# Architecture:
#   Megatron ──OTLP──► OTel Collector ──► (one of three modes, pick in `command:`)
#
# Modes (swap which `command:` line is uncommented under otel-collector):
#   collector.yaml           — local stack: Jaeger + Prometheus + Elasticsearch
#   collector-file.yaml      — write to /data volume (for offline inspection)
#   collector-weave.yaml     — forward traces to W&B Weave (needs WANDB_API_KEY)
#   collector-honeycomb.yaml — forward all signals to Honeycomb (needs HONEYCOMB_API_KEY)
#
# For hosted-backend modes (Weave, Honeycomb), copy .env.example → .env and
# fill in the relevant credentials. See .env.example for details.

services:

  # ── Megatron-LM training container ─────────────────────────────────────────
  # Interactive container for running training with OTel enabled.
  # Traces go to Jaeger, metrics go to the OTel Collector, logs go to ES.
  megatron:
    image: local-megatron-lm:otel-compose
    build:
      context: ..
      dockerfile_inline: |
        FROM local-megatron-lm:otel
        COPY lens /tmp/nemo-otel
        RUN python3 -m pip install --no-cache-dir 'opentelemetry-api>=1.40.0' '/tmp/nemo-otel[sdk]' \
         && /usr/bin/python -m pip install --no-cache-dir 'opentelemetry-api>=1.40.0' '/tmp/nemo-otel[sdk]' \
         && rm -rf /tmp/nemo-otel
    volumes:
      - ..:/workspace
    working_dir: /workspace/Megatron-LM
    # .env is gitignored; carries WANDB_ENTITY / WANDB_PROJECT used by Weave.
    # For non-Weave modes the values are harmless (just unused env).
    env_file:
      - .env
    environment:
      - MEGATRON_OTEL_ENABLED=1
      - MEGATRON_OTEL_SPAN_GROUPS=all
      - NEMO_LENS_LOGS_ENABLED=1
      - NEMO_LENS_RUN_ID=my-experiment-1   # optional: auto-generated from SLURM_JOB_ID or UUID
      - NEMO_LENS_USER_ID=team-alpha        # optional: user/team label
      - OTEL_METRIC_EXPORT_INTERVAL=1000
      - OTEL_EXPORTER_OTLP_ENDPOINT=http://otel-collector:4317
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    ipc: host
    ulimits:
      memlock: -1
      stack: 67108864
    stdin_open: true
    tty: true
    networks:
      - otel-net
    depends_on:
      jaeger:
        condition: service_started
      otel-collector:
        condition: service_started

  # ── Jaeger v2 ──────────────────────────────────────────────────────────────
  # Trace storage and search UI. Jaeger v2 natively accepts OTLP — no
  # separate collector needed for the trace pipeline.
  jaeger:
    image: jaegertracing/jaeger:2.15.0
    command: ["--config", "/etc/jaeger/config.yaml"]
    volumes:
      - ./observability/jaeger.yaml:/etc/jaeger/config.yaml:ro
    ports:
      - "16686:16686"   # Jaeger UI
    networks:
      - otel-net
    restart: unless-stopped

  # ── OpenTelemetry Collector ────────────────────────────────────────────────
  # Four modes — uncomment exactly ONE `command:` line below to pick a mode.
  #
  #   collector.yaml            Local stack: Jaeger + Prometheus + Elasticsearch.
  #                             Use when you want full local observability UIs.
  #
  #   collector-file.yaml       Write all signals to /data as JSONL.
  #                             Use for offline / reproducible debugging; no UIs.
  #
  #   collector-weave.yaml      Forward traces to W&B Weave. Requires
  #                             WANDB_API_KEY in .env. Weave accepts traces
  #                             only — metrics/logs are accepted-and-dropped.
  #
  #   collector-honeycomb.yaml  Forward all signals (traces, metrics, logs)
  #                             to Honeycomb. Requires HONEYCOMB_API_KEY and
  #                             HONEYCOMB_DATASET in .env.
  otel-collector:
    image: otel/opentelemetry-collector-contrib:0.146.0
    user: "0:0"  # file exporter needs write access to /data volume
    # .env supplies credentials to hosted-backend modes (Weave, Honeycomb).
    # Harmless in local/file modes (keys are simply unused).
    env_file:
      - .env

    # ── Pick exactly one mode ──
    # command: ["--config=/etc/otel/collector.yaml"]             # local stack: Jaeger + Prometheus + Elasticsearch
    # command: ["--config=/etc/otel/collector-file.yaml"]        # file: write to /data volume
    command: ["--config=/etc/otel/collector-weave.yaml"]         # W&B Weave: traces only (needs WANDB_API_KEY)
    # command: ["--config=/etc/otel/collector-honeycomb.yaml"]   # Honeycomb: all signals (needs HONEYCOMB_API_KEY, HONEYCOMB_DATASET)

    volumes:
      - ./observability/otel-collector.yaml:/etc/otel/collector.yaml:ro
      - ./observability/otel-collector-file.yaml:/etc/otel/collector-file.yaml:ro
      - ./observability/otel-collector-weave.yaml:/etc/otel/collector-weave.yaml:ro
      - ./observability/otel-collector-honeycomb.yaml:/etc/otel/collector-honeycomb.yaml:ro
      - /proc:/hostfs/proc:ro
      - /sys:/hostfs/sys:ro
      - otel-export:/data    # used by file mode
    ports:
      - "8889:8889"          # Prometheus scrape endpoint (local-stack mode)
    networks:
      - otel-net
    restart: unless-stopped
    depends_on:
      - dcgm-exporter
      - node-exporter

  # ── DCGM Exporter ─────────────────────────────────────────────────────────
  # GPU metrics: utilization, memory, temperature, power, NVLink, PCIe, ECC,
  # clock throttle reasons, SM activity — per GPU.
  dcgm-exporter:
    image: nvcr.io/nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    cap_add:
      - SYS_ADMIN     # required for DCGM to access GPU telemetry
    networks:
      - otel-net
    restart: unless-stopped

  # ── Node Exporter ──────────────────────────────────────────────────────────
  # InfiniBand port counters, RDMA stats, NIC bandwidth, TCP retransmits.
  node-exporter:
    image: prom/node-exporter:v1.9.1
    command:
      - "--path.procfs=/host/proc"
      - "--path.sysfs=/host/sys"
      - "--path.rootfs=/host/root"
      - "--collector.infiniband"
      - "--collector.netclass"
      - "--collector.tcpstat"
      - "--no-collector.arp"
      - "--no-collector.bcache"
      - "--no-collector.bonding"
      - "--no-collector.btrfs"
      - "--no-collector.entropy"
      - "--no-collector.hwmon"
      - "--no-collector.ipvs"
      - "--no-collector.mdadm"
      - "--no-collector.ntp"
      - "--no-collector.textfile"
      - "--no-collector.timex"
      - "--no-collector.xfs"
      - "--no-collector.zfs"
    volumes:
      - /proc:/host/proc:ro
      - /sys:/host/sys:ro
      - /:/host/root:ro
    pid: host
    networks:
      - otel-net
    restart: unless-stopped

  # ── Prometheus ─────────────────────────────────────────────────────────────
  # Scrapes the collector's Prometheus endpoint and stores time-series data.
  prometheus:
    image: prom/prometheus:v3.10.0
    command:
      - "--config.file=/etc/prometheus/prometheus.yml"
      - "--storage.tsdb.path=/prometheus"
      - "--storage.tsdb.retention.time=30d"
      - "--web.enable-lifecycle"          # allows hot-reload: POST /-/reload
    volumes:
      - ./observability/prometheus.yml:/etc/prometheus/prometheus.yml:ro
      - prometheus-data:/prometheus
    ports:
      - "9090:9090"
    networks:
      - otel-net
    restart: unless-stopped
    depends_on:
      - otel-collector

  # ── Elasticsearch ──────────────────────────────────────────────────────────
  elasticsearch:
    image: docker.elastic.co/elasticsearch/elasticsearch:9.2.0
    environment:
      - discovery.type=single-node
      - xpack.security.enabled=false
      - action.auto_create_index=true
      - ES_JAVA_OPTS=-Xms512m -Xmx512m
    ports:
      - "9200:9200"
    volumes:
      - es-data:/usr/share/elasticsearch/data
    healthcheck:
      test: ["CMD-SHELL", "curl -sf http://localhost:9200/_cluster/health || exit 1"]
      interval: 10s
      timeout: 5s
      retries: 10
    networks:
      - otel-net
    restart: unless-stopped

  # ── Kibana ─────────────────────────────────────────────────────────────────
  # Log exploration UI. Pre-configured with nemo-logs data view and dashboards.
  kibana:
    image: docker.elastic.co/kibana/kibana:9.2.0
    volumes:
      - ./observability/kibana/kibana.yml:/usr/share/kibana/config/kibana.yml:ro
      - ./observability/kibana/setup.sh:/usr/share/kibana/setup/setup.sh:ro
      - ./observability/kibana/saved-objects.ndjson:/usr/share/kibana/setup/saved-objects.ndjson:ro
    command: >
      bash -c "/usr/share/kibana/setup/setup.sh &
               /usr/local/bin/kibana-docker"
    ports:
      - "5601:5601"
    networks:
      - otel-net
    restart: unless-stopped
    depends_on:
      elasticsearch:
        condition: service_healthy


  # ── Grafana ────────────────────────────────────────────────────────────────
  # Unified dashboard — pre-wired to Prometheus (metrics) and Jaeger (traces).
  grafana:
    image: grafana/grafana:12.4.0
    environment:
      - GF_AUTH_ANONYMOUS_ENABLED=true
      - GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
      - GF_AUTH_DISABLE_LOGIN_FORM=true
      - GF_FEATURE_TOGGLES_ENABLE=traceToMetrics
    volumes:
      - ./observability/grafana/provisioning:/etc/grafana/provisioning:ro
      - ./observability/grafana/dashboards:/var/lib/grafana/dashboards:ro
      - grafana-data:/var/lib/grafana
    ports:
      - "3000:3000"
    networks:
      - otel-net
    restart: unless-stopped
    depends_on:
      - prometheus
      - jaeger

volumes:
  prometheus-data:
  grafana-data:
  es-data:
  otel-export:

networks:
  otel-net:
    driver: bridge