Use this custom Metaflow card to observe GPU utilization and memory consumption on any Metaflow task.
The uses the nvidia-smi command-line tool to retrieve the metrics. Make sure the tool is installed on the image
you use. Most CUDA-enabled images include it by default.
To see pretty charts, ensure that matplotlib is installed either in your @conda environment or in the image.
-
Copy
gpu_profile.pyin the same directory with your flow file. -
Add
from gpu_profile import gpu_profileat the top of the flow file. -
Include
@gpu_profile()(remember the parentheses) above every step that you want to monitor.
Run the flow as usual and observe the card, after a task has finished, either through command-line, notebook, or Metaflow UI.
See example/cudasum.py for an example that exercises GPU and monitors utilization with this card.