From 28d18ec836fe42bebddaa56c41f8432caa0add79 Mon Sep 17 00:00:00 2001 From: mehrdadn Date: Sat, 4 Jul 2020 13:20:23 -0700 Subject: [PATCH] [windows] Improve GPU detection (#9300) Co-authored-by: Mehrdad --- python/ray/resource_spec.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/python/ray/resource_spec.py b/python/ray/resource_spec.py index 9da487988..0399952ae 100644 --- a/python/ray/resource_spec.py +++ b/python/ray/resource_spec.py @@ -3,6 +3,8 @@ from collections import namedtuple import logging import multiprocessing import os +import subprocess +import sys import ray import ray.ray_constants as ray_constants @@ -229,12 +231,23 @@ class ResourceSpec( def _autodetect_num_gpus(): """Attempt to detect the number of GPUs on this machine. - TODO(rkn): This currently assumes Nvidia GPUs and Linux. + TODO(rkn): This currently assumes NVIDIA GPUs on Linux. + TODO(mehrdadn): This currently does not work on macOS. + TODO(mehrdadn): Use a better mechanism for Windows. + + Possibly useful: tensorflow.config.list_physical_devices() Returns: The number of GPUs if any were detected, otherwise 0. """ - proc_gpus_path = "/proc/driver/nvidia/gpus" - if os.path.isdir(proc_gpus_path): - return len(os.listdir(proc_gpus_path)) - return 0 + result = 0 + if sys.platform.startswith("linux"): + proc_gpus_path = "/proc/driver/nvidia/gpus" + if os.path.isdir(proc_gpus_path): + result = len(os.listdir(proc_gpus_path)) + elif sys.platform == "win32": + props = "AdapterCompatibility" + cmdargs = ["WMIC", "PATH", "Win32_VideoController", "GET", props] + lines = subprocess.check_output(cmdargs).splitlines()[1:] + result = len([l.rstrip() for l in lines if l.startswith(b"NVIDIA")]) + return result