cc-condor-sync/scripts/condor_status_to_gpu_map.py
2022-12-16 13:42:04 +01:00

34 lines
1.1 KiB
Python

#!/usr/bin/python3
import json
if __name__ == "__main__":
import argparse
about = """This script parses a `condor_status -json` dump file and spits out a GPU hash ID to PCIe address map.
"""
parser = argparse.ArgumentParser(description=about)
parser.add_argument("condor_status_file",
help="`condor_status -json` dump", default="condor_status.json")
args = parser.parse_args()
# open dump file
with open(args.condor_status_file, 'r', encoding='utf-8') as f:
condor_status = json.load(f)
slot_gpu_map = {}
for slot in condor_status:
machine = slot["Machine"].split('.')[0]
gpu_map = {}
if machine in slot_gpu_map:
gpu_map = slot_gpu_map[machine]
else:
slot_gpu_map[machine] = {}
if not "AssignedGPUs" in slot:
continue
gpus = slot["AssignedGPUs"].split(',')
for gpu_id in gpus:
gpu = slot["GPUs_" + gpu_id.strip().replace("-", "_")]
gpu_map[gpu["Id"]] = gpu["DevicePciBusId"]
slot_gpu_map[machine] = gpu_map
print(json.dumps(slot_gpu_map))