1+ # -*- DISCLAIMER: this file contains code derived from gputil (https://github.com/anderskm/gputil)
2+ # and therefore is distributed under to the following license:
3+ #
4+ # MIT License
5+ #
6+ # Copyright (c) 2017 anderskm
7+ #
8+ # Permission is hereby granted, free of charge, to any person obtaining a copy
9+ # of this software and associated documentation files (the "Software"), to deal
10+ # in the Software without restriction, including without limitation the rights
11+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12+ # copies of the Software, and to permit persons to whom the Software is
13+ # furnished to do so, subject to the following conditions:
14+ #
15+ # The above copyright notice and this permission notice shall be included in all
16+ # copies or substantial portions of the Software.
17+ #
18+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24+ # SOFTWARE.
25+ 26+ import platform
27+ import shutil
28+ from subprocess import Popen , PIPE
29+ import os
30+ 31+ def gpu_count ():
32+ try :
33+ if platform .system () == "Windows" :
34+ nvidia_smi = shutil .which ('nvidia-smi' )
35+ if nvidia_smi is None :
36+ nvidia_smi = "%s\\ Program Files\\ NVIDIA Corporation\\ NVSMI\\ nvidia-smi.exe" % os .environ ['systemdrive' ]
37+ else :
38+ nvidia_smi = "nvidia-smi"
39+ 40+ p = Popen ([nvidia_smi ,
41+ "--query-gpu=name" ,
42+ "--format=csv,noheader,nounits" ], stdout = PIPE )
43+ stdout , stderror = p .communicate ()
44+ 45+ output = stdout .decode ('UTF-8' )
46+ lines = output .split (os .linesep )
47+ num_devices = len (lines ) - 1
48+ return num_devices
49+ except :
50+ return 0
51+
0 commit comments