From 9d2b34ab60ef88537d651fde19a0b8979deb23bc Mon Sep 17 00:00:00 2001
From: evonneng <evonneng@users.noreply.github.com>
Date: Mon, 30 May 2022 09:01:52 -0700
Subject: [PATCH] Added auto find gpu; added method name to arg list (#111)

* update benchmarking for better practice

* fixed logic for finding/distributing across GPUS

* added method to argument list
---
 scripts/benchmarking/launch_train_blender.sh | 38 +++++++++++++++++---
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh
index d523770c1..e5b679539 100755
--- a/scripts/benchmarking/launch_train_blender.sh
+++ b/scripts/benchmarking/launch_train_blender.sh
@@ -1,15 +1,44 @@
 #!/bin/bash
 
-DATASETS=("mic" "ficus" "chair" "hotdog" "materials" "drums" "ship" "lego")
-GPU_IDX=1
+METHOD=$1
+shift
 
+# Deal with gpu's. If passed in, use those.
+GPU_IDX=("$@")
+if [ -z "$GPU_IDX" ]; then
+    echo "no gpus set... finding available gpus"
+    # Find available devices
+    num_device=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+    START=0
+    END=${num_device}-1
+    GPU_IDX=()
+
+    for (( id=$START; id<$END; id++ )); do
+        free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo [0-9]+)
+        if [[ $free_mem -gt 10000 ]]; then
+            GPU_IDX+=( $id )
+        fi
+    done
+fi
+echo available gpus... ${GPU_IDX[@]}
+
+DATASETS=("mic" "ficus" "chair" "hotdog" "materials" "drums" "ship" "lego")
 date
 now=$(date)
 tag=$(date +'%m-%d-%Y')
+idx=0
+len=${#GPU_IDX[@]}
+((len=len-1))
 
 for dataset in ${DATASETS[@]}; do
-    export CUDA_VISIBLE_DEVICES=${!GPU_IDX}
+    export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]}
+    if [ $idx == $len ]; then
+        idx=0
+    else
+        ((idx=idx+1))
+    fi
     python scripts/run_train.py \
+           --config-name ${METHOD} \
            '~logging.writer.LocalWriter' \
            data.dataset.data_directory=data/blender/${dataset} \
            data.dataset.downscale_factor=1 \
@@ -18,6 +47,5 @@ for dataset in ${DATASETS[@]}; do
            graph.steps_per_save=25000 \
            graph.max_num_iterations=2000000 \
            logging.enable_profiler=False &
-    echo "Launched ${dataset} on gpu ${!GPU_IDX}, ${tag}"
-    GPU_IDX=$((GPU_IDX+1))
+    echo "Launched ${METHOD} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
 done
\ No newline at end of file