Merge pull request #3978 from SebastianAchilles/gpuinfo

add support for collecting GPU info (via rocm-smi)

Author: bartoldeman

easybuild/tools/systemtools.py

@@ -610,6 +610,29 @@ def get_gpu_info():
         except Exception as err:
             _log.debug("Exception was raised when running nvidia-smi: %s", err)
             _log.info("No NVIDIA GPUs detected")
+
+        try:
+            cmd = "rocm-smi --showdriverversion --csv"
+            _log.debug("Trying to determine AMD GPU driver on Linux via cmd '%s'", cmd)
+            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False)
+            if ec == 0:
+                amd_driver = out.strip().split('\n')[1].split(',')[1]
+
+            cmd = "rocm-smi --showproductname --csv"
+            _log.debug("Trying to determine AMD GPU info on Linux via cmd '%s'", cmd)
+            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False)
+            if ec == 0:
+                for line in out.strip().split('\n')[1:]:
+                    amd_card_model = line.split(',')[2]
+                    amd_gpu = ', '.join([amd_card_model, amd_driver])
+                    amd_gpu_info = gpu_info.setdefault('AMD', {})
+                    amd_gpu_info.setdefault(amd_gpu, 0)
+                    amd_gpu_info[amd_gpu] += 1
+            else:
+                _log.debug("None zero exit (%s) from rocm-smi: %s", ec, out)
+        except Exception as err:
+            _log.debug("Exception was raised when running rocm-smi: %s", err)
+            _log.info("No AMD GPUs detected")
     else:
         _log.info("Only know how to get GPU info on Linux, assuming no GPUs are present")