GPU Info

[907Resident]

TL;DR

Overall, I wanted an argument in watermark to allow the user to view GPU information.

Details

• I stumbled upon this issue Would love the option to get the GPU and Nvidia revision info #31 while trying to view my GPU capabilities
• Since the issue has been open since 2017, I have taken a stab by submitting a PR to the argument gpu to watermark
• Note: This will only work with GPUs, where NVIDIA drivers are installed already. In the issue, @dartdog mentions, and I agree with, most people are not using AMD GPUs for data science. Therefore, I think that only accounting for NVIDIA drivers should be fine
• Lastly, I added a new file in tests/ called test_gpu_info.py, which should allow the maintainers of this repo to quickly run a pytest to ensure that the gpu argument works as the maintainers see fit

[rasbt]

Awesome! This is a very exciting PR! Thanks!

[rasbt]

I just tried it on a machine with NVIDIA GPUs, and it couldn't find these. How did it look on your machine, am I missing any requirements @907Resident ?

[907Resident]

Hi @rasbt,

I think I accidentally provided the incorrect pythonic interface for the gpu. As I mentioned in commit 87e7201, the correct package is py3nvml.

I have made the correction, and the change can be seen in this Colab notebook (assuming the user has Colab Pro with GPU). The output can be viewed below.

[907Resident]

I also noticed that as of commit ab5cf05, that the automated checks have failed. But my local execution of pytest appears to be working:

But for some reason that I do no understand, the automated CI process is unable to install py3nvml.

Nonetheless, I think once py3nvml is installed, I believe that you will be able to get the name of the GPU being used with my additions. If there is more functionality that you want to see, just let me know.

[rasbt]

Thanks for the update! I did some fixes in the formatting and the CI works now! I'll also try it on a non-Colab GPU machine to make sure.

[rasbt]

Seems to work (see screenshot below). I think if there's a single GPU,

GPU Info: GPU 0: NVIDIA A100-SXM4-40GB

would look fine. But if there are more than one, perhaps

GPU Info:
  GPU 0: NVIDIA A100-SXM4-40GB
  GPU 1: NVIDIA A100-SXM4-40GB
  GPU 2: NVIDIA A100-SXM4-40GB
  GPU 3: NVIDIA A100-SXM4-40GB
  GPU 4: NVIDIA A100-SXM4-40GB
  GPU 5: NVIDIA A100-SXM4-40GB
  GPU 6: NVIDIA A100-SXM4-40GB
  GPU 7: NVIDIA A100-SXM4-40GB

[rasbt]

I changed it to always use the indentation now. In hindsight, I think this should look better since it avoids the double-colon.

I.e., it's now

GPU Info:
  GPU 0: NVIDIA A100-SXM4-40GB
  GPU 1: NVIDIA A100-SXM4-40GB
  GPU 2: NVIDIA A100-SXM4-40GB
  GPU 3: NVIDIA A100-SXM4-40GB
  GPU 4: NVIDIA A100-SXM4-40GB
  GPU 5: NVIDIA A100-SXM4-40GB
  GPU 6: NVIDIA A100-SXM4-40GB
  GPU 7: NVIDIA A100-SXM4-40GB

or

GPU Info:
  GPU 0: NVIDIA A100-SXM4-40GB

[907Resident]

Thanks for the updates and integrating my changes.

I changed it to always use the indentation now. In hindsight, I think this should look better since it avoids the double-colon.

I.e., it's now

GPU Info:
  GPU 0: NVIDIA A100-SXM4-40GB
  GPU 1: NVIDIA A100-SXM4-40GB
  GPU 2: NVIDIA A100-SXM4-40GB
  GPU 3: NVIDIA A100-SXM4-40GB
  GPU 4: NVIDIA A100-SXM4-40GB
  GPU 5: NVIDIA A100-SXM4-40GB
  GPU 6: NVIDIA A100-SXM4-40GB
  GPU 7: NVIDIA A100-SXM4-40GB

or

GPU Info:
  GPU 0: NVIDIA A100-SXM4-40GB

Yes, I agree. I only have access to a machine with one GPU; so, I did not think about formatting for multiple GPUs. But your formatting looks much better.

@rasbt, at this point what do we need to do to merge this PR? Anything else you need me to add?

[rasbt]

@907Resident It looks all good to me know! Just added a changelog entry and am happy to merge (and make version release then)! Thanks so much for the PR. It's a super nice contribution!

[907Resident]

Cool, sounds good!