nvidia-smi command in bash vs in terminal for maxi

问题:

Overall goal: I'm trying to assign my python program to the GPU with the most available memory.

Current issue: When I use this command:

nvidia-smi --query-gpu=memory.free --format=csv

I get the following output:

memory.free [MiB]
4800 MiB
5332 MiB
5346 MiB

Of course, that changes very frequently, so perhaps this is not the best way to accomplish my overall goal. For the moment, I'm trying to use this to figure out which GPU in the set of three has the most free memory. In this case, it's clearly GPU2 (they're labelled 0, 1, 2).

I created a bash script:

#/bin/bash

myarr=( $(nvidia-smi --query-gpu=memory.free --format=csv) )
echo $myarr

which outputs this to the screen:

memory.free

It used to output in this format:

memory.free [MiB] 4800 MiB 5332 MiB 5346 MiB

I then did this (with thanks to the first answer here):

myarr2=${myarr[2,4,6]}
echo $myarr2
IC=(`tr ' ' '\n' <<<$myarr2 | cat -n | sort -k2,2nr | head -n1`)
echo $IC
Ival=${IC[0]}
Cval=${IC[1]}
echo $Ival $Cval

However, it doesn't seem to matter what I do, I always get the index of the maximum to be at position 1. In the example I gave above, this was correct. In general, it is not correct.

The complete script:

#/bin/bash

myarr=( $(nvidia-smi --query-gpu=memory.free --format=csv) )
echo $myarr
myarr2=${myarr[2,4,6]}
echo $myarr2
IC=(`tr ' ' '\n' <<<$myarr2 | cat -n | sort -k2,2nr | head -n1`)
echo $IC
Ival=${IC[0]}
Cval=${IC[1]}
echo $Ival $Cval

What is wrong with how I'm searching for the position of the maximum?

回答1:

The following is verbose, but careful to follow good practices; it can't be mislead by surprising inputs to process local filenames as data (as could happen if a future version of the software had *s in the header, for instance), and runs all processing internal to the shell (which can be slower for very large inputs, but for cases where only a small amount of data is being processed -- as here -- will be much faster on account of avoiding process startup overhead for sort, tr, cat or other tools external to the shell).

#!/usr/bin/env bash

max_idx=0
max_mem=0
idx=0

{
  read _;                         # discard first line (header)
  while read -r mem _; do         # for each subsequent line, read first word into mem
    if (( mem > max_mem )); then  # compare against maximum mem value seen
      max_mem=$mem                # ...if greater, then update both that max value
      max_idx=$idx                # ...and our stored index value.
    fi
    ((++idx))
  done
} < <(nvidia-smi --query-gpu=memory.free --format=csv)

echo "Maximum memory seen is $max_mem, at processor $idx"

If I define nvidia-smi as follows (note that due to the - in the name, this definition isn't portable to all versions of bash):

nvidia-smi() { printf '%s\n' 'memory.free [MiB]' '4800 MiB' '5332 MiB' '5346 MiB'; }

...this emits as output:

Maximum memory seen is 5346, at processor 3