#!/usr/bin/env bash

CONFIG=$1
TASK=$2
GPUS=$3
MORE_ARG=${*:4}

RANDOM_MASTER=$(shuf -i 2000-65000 -n 1)

_command="print(len('${GPUS}'.split(',')))"
GPU_COUNT=$(python3 -c "${_command}")

echo "GPU_COUNT:${GPU_COUNT}"

echo CUDA_VISIBLE_DEVICES=$GPUS \
  PYTHONPATH=.:$PYTHONPATH OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node="$GPU_COUNT" \
  main.py "$TASK" "$CONFIG" --backup_config --setup_output_dir --setup_random_seed "$MORE_ARG"

CUDA_VISIBLE_DEVICES=$GPUS \
  PYTHONPATH=.:$PYTHONPATH OMP_NUM_THREADS=1 python -m torch.distributed.launch --nproc_per_node="$GPU_COUNT" \
  --master_port=${RANDOM_MASTER} \
  main.py "$TASK" "$CONFIG" $MORE_ARG --backup_config --setup_output_dir --setup_random_seed