diff --git a/roles/gpu/files/scripts/nvidia-xconf-gen.py b/roles/gpu/files/scripts/nvidia-xconf-gen.py index f26446971f04c022e19897edf5d9d05cefe61da7..7cd9cb551f348d608b583466322b1acc137e9b8d 100755 --- a/roles/gpu/files/scripts/nvidia-xconf-gen.py +++ b/roles/gpu/files/scripts/nvidia-xconf-gen.py @@ -11,14 +11,19 @@ from subprocess import call import re import json +def check_nvidia_smi(): + try: + smi = subprocess.check_output(["which","nvidia-smi"]) + except subprocess.CalledProcessError: + print("nvidia-smi binary not found!") + exit(1) + + def grab_card_ids(): - # This method runs nvidia-smi to grab the card ids, then returns a list + # This method runs nvidia-smi to grab the card ids, then returns a list + check_nvidia_smi() - if not os.path.isfile("/bin/nvidia-smi"): - print("nvidia-smi binary not found!") - exit(1) - - cmd = ["/bin/nvidia-smi", "--query-gpu=pci.bus_id","--format=csv,noheader"] + cmd = ["nvidia-smi", "--query-gpu=pci.bus_id","--format=csv,noheader"] p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) cards = [] @@ -27,76 +32,43 @@ def grab_card_ids(): line = line.rstrip().split(":")[2] pcibus_num = int(re.sub('[.:]', '', line).rstrip("0"),16) card = "PCI:0:{}:0".format(str(pcibus_num)) - cards.append(card) + cards.append(card) return cards def grab_card_boardname(): - if not os.path.isfile("/bin/nvidia-smi"): - print("nvidia-smi binary not found!") - exit(1) + check_nvidia_smi() - cmd = ["/bin/nvidia-smi", "--query-gpu=name","--format=csv,noheader"] - cards = [] - p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) - for line in p.stdout.readlines(): - line = line.rstrip() - cards.append(line) - return cards + cmd = ["nvidia-smi", "--query-gpu=name","--format=csv,noheader"] + cards = [] + p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + for line in p.stdout.readlines(): + line = line.rstrip() + cards.append(line) + return cards def write_xorg_conf(cards): - num_of_cards = len(cards) + 1 - boardname = (grab_card_boardname())[0] - gpus = [] + num_of_cards = len(cards) + 1 + boardname = (grab_card_boardname())[0] + gpus = [] - for i in range(1, num_of_cards): - monitors = [] - screens = [] - res = list(itertools.combinations(cards,i)) + for i in range(1, num_of_cards): + monitors = [] + screens = [] + res = list(itertools.combinations(cards,i)) - for j in range(i): - monitors.append("Monitor"+str(j)) - screens.append("Screen"+str(j)) + for j in range(i): + monitors.append("Monitor"+str(j)) + screens.append("Screen"+str(j)) - for card in res: - filename = "/etc/X11/xorg.conf." + str(i) + str(res.index(card)) - #gpus[filename] = filename - - gpus.append({'filename':filename,'boardname':boardname,'monitors':monitors,'screens':screens,'devices':card}) + for card in res: + filename = "/etc/X11/xorg.conf." + str(i) + str(res.index(card)) + + gpus.append({'filename':filename,'boardname':boardname,'monitors':monitors,'screens':screens,'devices':card}) - print(json.dumps(gpus)) + print(json.dumps(gpus)) if __name__ == "__main__": -# cards = grab_card_ids() -# write_xorg_conf(cards) - cards = grab_card_ids() - write_xorg_conf(cards) - - -''' with open("xorgconftest.j2") as f: - template = Template(f.read()) - -num_of_monitors = 2 - -monitors = [] -screens = [] -busid = ["PCI:0:6:0","PCI:0:7:0","PCI:0:8:0","PCI:0:9:0"] - -for i in range(num_of_monitors): - monitors.append("Monitor" + str(i)) - screens.append("Screen" + str(i)) - -#output = template.render(monitors=monitors,devices=busid) -#print(output) - -res = itertools.combinations(busid,2) -listres = list(res) - -for i in listres: - output = template.render(monitors=monitors,devices=i,screens=screens) - filename = "xorg.conf." + str(listres.index(i)) - print(filename) - f = open(filename,"w") - f.write(output) - f.close() ''' + cards = grab_card_ids() + write_xorg_conf(cards)