100 lines
3.0 KiB
Python
Executable File
100 lines
3.0 KiB
Python
Executable File
#! /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()) |