1

Add a script for extracting GPU info from deadline

This commit is contained in:
Evan Pratten 2023-12-01 20:14:02 -05:00
parent ac8f7cdc35
commit 4f8c5ac4f3

100
scripts/deadline-gpu-inventory Executable file
View File

@ -0,0 +1,100 @@
#! /usr/bin/env python
import argparse
import sys
import logging
import pyperclip
import re
logger = logging.getLogger(__name__)
GPU_FILTERS = [
re.compile(r"(Intel\(R\)[a-zA-Z\d ]+)"),
re.compile(r"(NVIDIA [a-zA-Z\d ]+)"),
]
MSRP_ESTIMATES = {
"NVIDIA RTX A5000": 5400,
"NVIDIA RTX A4000": 1600,
"NVIDIA Quadro M5000": 2000,
"NVIDIA Geforce RTX 2070": 800
}
def main() -> int:
# Handle program arguments
ap = argparse.ArgumentParser(prog='deadline-gpu-inventory', description='Reads Deadline machine info from the clipboard and displays the GPU inventory')
ap.add_argument("--format", "-f", help="Output format", choices=["pretty", "tsv", "csv"], default="pretty")
ap.add_argument("--guess-cost", "-g", help="Guess the cost of the GPUs", action="store_true")
ap.add_argument('-v', '--verbose', help='Enable verbose logging', action='store_true')
args = ap.parse_args()
# Configure logging
logging.basicConfig(
level=logging.DEBUG if args.verbose else logging.INFO,
format='%(levelname)s: %(message)s',
)
# Read from clipboard
clipboard = pyperclip.paste()
# Split clipboard into lines
lines = clipboard.splitlines()
# Search for GPUs
gpus = []
for line in lines:
for gpu_filter in GPU_FILTERS:
# Find any matches
match = gpu_filter.search(line)
if match:
# Add to list of GPUs
gpus.append(match.group(1))
# If we don't have any GPUs, exit
if not gpus:
logger.error("No GPUs found.")
logger.info("Are you sure you have copied the machine info to the clipboard?")
return 1
# Reshape the list into [(gpu, count)]
gpu_counts = {}
for gpu in gpus:
gpu_counts[gpu] = gpu_counts.get(gpu, 0) + 1
max_count = max(gpu_counts.values())
gpu_counts = list(gpu_counts.items())
# Sort the list by count
gpu_counts.sort(key=lambda x: x[1], reverse=True)
if args.format == "pretty":
# Figure out the width of the largest count
count_width = len(str(max_count))
# Print the list
for gpu, count in gpu_counts:
print(f"{count:>{count_width}} {gpu}")
elif args.format == "tsv":
# Print the list
for gpu, count in gpu_counts:
print(f"{count}\t{gpu}")
elif args.format == "csv":
# Print the list
print("Count,GPU")
for gpu, count in gpu_counts:
print(f"{count},{gpu}")
# If cost guessing is enabled, print the estimated cost
if args.guess_cost:
total_cost = 0
known_cost_gpu_count = 0
for gpu, count in gpu_counts:
if gpu in MSRP_ESTIMATES:
total_cost += MSRP_ESTIMATES[gpu] * count
known_cost_gpu_count += 1
print(f"Estimated cost: ${total_cost}. ({known_cost_gpu_count} models included in estimate)")
return 0
if __name__ == "__main__":
sys.exit(main())