Fix gpu_id

[fatwir]

This Pull Request enhances the GPU selection mechanism based on memory utilization. Specifically, it implements the following changes:

1. GPU ID Selection:

   • The GPU ID will be selected based on memory utilization. If the memory utilization of a given GPU is less than 20%, the first MIG ID under that GPU will be selected.

2. Limitations:

   • While nvidia-smi does support a MIG option, I currently do not have the necessary permissions to access this feature.
   • Detailed information about the MIG-partitioned GPU utilization can be retrieved using dcgmi. However, this tool is not installed and cannot be utilized.

3. Current Approach:

   • For the time being, a manual search through the output of nvidia-smi -L is performed to select the appropriate GPU IDs.

[vianamp]

I did not test it.

[vianamp]

Wonder if we should use something like this instead: https://github.com/anderskm/gputil

[fatwir]

Thank you for this! I believe it would be beneficial to update it for the MIG partitioned GPUs as well. Additionally, I noticed that they are leveraging nvidia-smi --query-gpu to retrieve the statistics, which we are also using. The repo you shared offers a lot of stats about the GPUs, and I think incorporating that could be useful. I'll talk to ritvik about this!

[ritvikvasan]

this PR seems to be doing a lot more than this? did you mean to merge into the other nb-> python file branch?

[fatwir]

I had forked from this branch and updated the _setup_gpu function in the scripts!

[ritvikvasan]

can you merge into that branch then so its easier to know what you added?

[ritvikvasan]

I should have clarified here I didn't mean to drop this PR and just merge your changes, but to just base your PR off that branch

[ritvikvasan]

Much better, thanks!