services:
  ray_head:
    build:
      context: .
      dockerfile: images/argus-ray-node/Dockerfile
      args:
        BASE_IMAGE: verlai/verl:sgl055.latest
    image: argus/argus-ray-node:v2.5
    container_name: argus-ray-head
    ports:
      - "8265:8265"
      - "8080:8080"
    volumes:
      # NOTE: this compose file is intended for the dev env layout like:
      #   /home2/argus/infra/mvp/{shared,verl,src/mvp}
      # so from `src/mvp/` the shared/verl dirs are `../../shared` and `../../verl`.
      - ../../verl:/workspace/verl
      - ../../shared:/private
      - .:/workspace/mvp
    shm_size: "10g"
    ulimits:
      nofile:
        soft: 65536
        hard: 65536
    cap_add:
      - SYS_ADMIN
      - SYS_PTRACE
    networks:
      - argus-ray-net
    environment:
      ARGUS_ROLE: "head"
      ARGUS_SHARED_ROOT: "/private"
      ARGUS_CLUSTER_NAME: "argus-ray"
      ARGUS_LOG_DIR: "/private/common/logs"
      HF_HOME: "/private/hf"
      HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
      TRANSFORMERS_CACHE: "/private/hf/transformers"
      HF_ENDPOINT: "https://hf-mirror.com"
      PYTHONUNBUFFERED: "1"

  ray_worker_0:
    image: argus/argus-ray-node:v2.5
    container_name: argus-ray-worker-0
    volumes:
      - ../../verl:/workspace/verl
      - ../../shared:/private
      - .:/workspace/mvp
    shm_size: "10g"
    ulimits:
      nofile:
        soft: 65536
        hard: 65536
    cap_add:
      - SYS_ADMIN
      - SYS_PTRACE
    networks:
      - argus-ray-net
    runtime: nvidia
    environment:
      ARGUS_ROLE: "worker"
      ARGUS_SHARED_ROOT: "/private"
      ARGUS_CLUSTER_NAME: "argus-ray"
      ARGUS_LOG_DIR: "/private/common/logs"
      ARGUS_WORKER_RESOURCES_KV: "worker_node=100"
      NVIDIA_VISIBLE_DEVICES: "0,1,2,3"
      NVIDIA_DRIVER_CAPABILITIES: "all"
      HF_HOME: "/private/hf"
      HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
      TRANSFORMERS_CACHE: "/private/hf/transformers"
      HF_ENDPOINT: "https://hf-mirror.com"
      PYTHONUNBUFFERED: "1"

  ray_worker_1:
    image: argus/argus-ray-node:v2.5
    container_name: argus-ray-worker-1
    volumes:
      - ../../verl:/workspace/verl
      - ../../shared:/private
      - .:/workspace/mvp
    shm_size: "10g"
    ulimits:
      nofile:
        soft: 65536
        hard: 65536
    cap_add:
      - SYS_ADMIN
      - SYS_PTRACE
    networks:
      - argus-ray-net
    runtime: nvidia
    environment:
      ARGUS_ROLE: "worker"
      ARGUS_SHARED_ROOT: "/private"
      ARGUS_CLUSTER_NAME: "argus-ray"
      ARGUS_LOG_DIR: "/private/common/logs"
      ARGUS_WORKER_RESOURCES_KV: "worker_node=100"
      NVIDIA_VISIBLE_DEVICES: "4,5,6,7"
      NVIDIA_DRIVER_CAPABILITIES: "all"
      HF_HOME: "/private/hf"
      HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
      TRANSFORMERS_CACHE: "/private/hf/transformers"
      HF_ENDPOINT: "https://hf-mirror.com"
      PYTHONUNBUFFERED: "1"

networks:
  argus-ray-net:
    driver: bridge
