#!/bin/bash
# poll.sh <container> <out.csv>: unix time, that container's process GPU memory (MiB, nvidia-smi per process), card total
pid=$(docker inspect -f '{{.State.Pid}}' "$1")
gpu=$(nvidia-smi --query-compute-apps=pid,gpu_uuid --format=csv,noheader | awk -F', ' -v p="$pid" '$1==p{print $2}')
echo "# pid=$pid gpu=$gpu" > "$2"
while [ -d /proc/$pid ]; do
  procmem=$(nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits | awk -F', ' -v p="$pid" '$1==p{print $2}')
  card=$(nvidia-smi --id="$gpu" --query-gpu=memory.used --format=csv,noheader,nounits)
  printf '%s,%s,%s\n' "$(date +%s.%N)" "${procmem:-0}" "$card" >> "$2"
done