pandaserver/srvcore/hardware_matching.py

0001 """
0002 matching of hardware requirements against actual hardware
0003
0004 The requirement side comes from the task architecture, e.g. JediTaskSpec.get_host_gpu_spec().
0005 The hardware side is a list of GPU dictionaries using the key names of the worker node GPU
0006 monitoring (ATLAS_PANDA.worker_node_gpus / MV_WORKER_NODE_GPU_SUMMARY), i.e. vendor, model,
0007 vram, architecture, framework_version, and driver_version. Those dictionaries either describe
0008 all worker nodes of a PanDA queue, when brokering tasks to queues, or the GPUs of a single
0009 worker node, when dispatching jobs to a pilot.
0010 """
0011
0012 import re
0013
0014 from packaging import version
0015
0016
0017 def compare_version_string(version_string, comparison_string):
0018     """
0019     Compares a version string with another string composed of a comparison operator and a version string.
0020
0021     Args:
0022         version_string (str): The version string to compare.
0023         comparison_string (str): The string containing the comparison operator and version string (e.g., ">=2.0").
0024
0025     Returns:
0026         bool or None: True if the version string satisfies the comparison, False if it doesn't,
0027                        or None if the comparison string is invalid.
0028     """
0029     match = re.match(r"([=><!]+)(.+)", comparison_string)
0030     if not match:
0031         return None
0032
0033     operator = match.group(1).strip()
0034     if operator == "=":
0035         operator = "=="
0036     version_to_compare = match.group(2).strip()
0037
0038     try:
0039         version1 = version.parse(version_string)
0040         version2 = version.parse(version_to_compare)
0041     except version.InvalidVersion:
0042         return None
0043
0044     if operator == "==":
0045         return version1 == version2
0046     elif operator == "!=":
0047         return version1 != version2
0048     elif operator == ">=":
0049         return version1 >= version2
0050     elif operator == "<=":
0051         return version1 <= version2
0052     elif operator == ">":
0053         return version1 > version2
0054     elif operator == "<":
0055         return version1 < version2
0056     else:
0057         return None
0058
0059
0060 def match_gpu_spec(required_gpu_spec, gpus):
0061     """
0062     Checks whether GPUs satisfy the GPU requirement of a task.
0063
0064     Selection attributes (vendor, model, microarchitecture) use an any match, i.e. it is enough that
0065     one GPU is of the requested type. Minimum-requirement attributes (vram, version, driver_version)
0066     use an all match, i.e. every GPU has to satisfy the constraint, so that a job cannot end up on a
0067     non-compliant GPU of a mixed set.
0068
0069     Args:
0070         required_gpu_spec (dict): The GPU requirement of the task, with the keys vendor, model, vram,
0071                                   microarchitecture, version, and driver_version. Only vendor and model
0072                                   are mandatory and `*` is the wildcard for them. The model is either a
0073                                   regular expression for inclusion or a dictionary with pattern and excl
0074                                   keys for exclusion. The version, driver_version, and vram are
0075                                   operator-prefixed strings, e.g. `>=12.0`.
0076         gpus (list): List of dictionaries describing the actual GPUs, with the keys vendor, model, vram,
0077                      architecture, framework_version, and driver_version.
0078
0079     Returns:
0080         bool: True if the GPUs satisfy the requirement.
0081     """
0082     # check vendor
0083     required_vendor = required_gpu_spec.get("vendor", "*")
0084     if required_vendor != "*":
0085         if not gpus or not any(gpu.get("vendor") and re.match(required_vendor, gpu["vendor"], re.IGNORECASE) for gpu in gpus):
0086             return False
0087
0088     # check model (include or exclude pattern)
0089     required_model = required_gpu_spec.get("model", "*")
0090     if required_model != "*":
0091         if isinstance(required_model, dict):
0092             model_pattern = required_model["pattern"]
0093             model_excl = required_model.get("excl", False)
0094         else:
0095             model_pattern = required_model
0096             model_excl = False
0097         if not gpus:
0098             return False
0099         matches = any(gpu.get("model") and re.match(model_pattern, gpu["model"], re.IGNORECASE) for gpu in gpus)
0100         if matches == model_excl:
0101             return False
0102
0103     # check VRAM (in MB); supports operators: ==, >=, <=, >, <, != (e.g. ">=40960")
0104     if "vram" in required_gpu_spec:
0105         if not gpus or not all(gpu.get("vram") and compare_version_string(str(gpu["vram"]), required_gpu_spec["vram"]) for gpu in gpus):
0106             return False
0107
0108     # check GPU microarchitecture generation (e.g. Ampere, Hopper, Ada Lovelace)
0109     if "microarchitecture" in required_gpu_spec:
0110         req_arch = required_gpu_spec["microarchitecture"]
0111         if isinstance(req_arch, str):
0112             req_arch = [req_arch]
0113         if not gpus or not any(gpu.get("architecture") in req_arch for gpu in gpus):
0114             return False
0115
0116     # check CUDA toolkit version
0117     if "version" in required_gpu_spec:
0118         if not gpus or not all(gpu.get("framework_version") and compare_version_string(gpu["framework_version"], required_gpu_spec["version"]) for gpu in gpus):
0119             return False
0120
0121     # check GPU kernel driver version (e.g. 575.57.08)
0122     if "driver_version" in required_gpu_spec:
0123         if not gpus or not all(
0124             gpu.get("driver_version") and compare_version_string(gpu["driver_version"], required_gpu_spec["driver_version"]) for gpu in gpus
0125         ):
0126             return False
0127
0128     return True