version: '3.8'

services:
  ubuntu-mineru-start-mpi-complie:
    build: 
      context: .
      dockerfile: Dockerfile.mineru310_mpich_complie
      args: 
        CACHEBUST: 1
        CONDA_ENV_NAME: "mineru"
        PYTHON_VERSION: "3.10"
        TAG_VERSION: "12.4.1"
        NV_DRIVER_VERSION: "535"
        GO_VERSION: "1.21.13"
    # env_file:
    #   - .env_mpich
    volumes:
      - ../src:/work
      - /mnt/beegfs/pdf_clean/zly_project/pdf_clean_agent/PDF-Extract-Kit_bak2/models:/models
      - /mnt/beegfs/downloads:/mnt/beegfs/downloads
      - ./results:/results
      - /mnt/beegfs/pdf_clean/yanbaodata:/data
      - ./start_mineru_update_mpi_print.py:/start_mineru_update_mpi_print.py
      - ./start_mineru_update_mpi.py:/start_mineru_update_mpi.py
      - ./magic-pdf.json:/root/magic-pdf.json
      - ./hostfile:/hostfile
      - ./hostfile_mpich:/hostfile_mpich
      - ./mpi_test.py:/mpi_test.py
      - /root/test:/root/test
    container_name: ubuntu-mineru-start-mpi-complie
    pull_policy: if_not_present
    ulimits:
      memlock:
        soft: -1
        hard: -1
    restart: unless-stopped
    image: hotwa/mpi:testbuild
    privileged: true
    cap_add:
      - ALL
      - CAP_SYS_PTRACE
    shm_size: '2000gb'
    devices:
      - /dev/infiniband/rdma_cm
      - /dev/infiniband/uverbs0
      - /dev/infiniband/uverbs1
      - /dev/infiniband/uverbs2
      - /dev/infiniband/uverbs3
      - /dev/infiniband/uverbs4
      - /dev/infiniband/uverbs5
      - /dev/infiniband/uverbs6
      - /dev/infiniband/uverbs7
      - /dev/infiniband/uverbs8
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - NVIDIA_DRIVER_CAPABILITIES=compute,utility
      - TMPDIR=/var/tmp
      - MAGIC_PDF_METHOD=auto
      - MAGIC_PDF_MODEL_MODE=full
      - MAGIC_PDF_INSIDE_MODEL=true
      - MAX_PROCESSES_PER_GPU=10
      - PDF_DIR=/data
      # - UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1,mlx5_8:1
      - UCX_NET_DEVICES=mlx5_3:1
    network_mode: host
    command: ["/usr/sbin/sshd", "-D"]
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]