target_architecture in acquire_jobs

pandajedi/jedibrokerage/AtlasBrokerUtils.py

@@ -8,14 +8,17 @@
 import traceback
 from typing import Any
 
-from packaging import version
 from pandacommon.pandautils.PandaUtils import naive_utcnow
 
 from pandajedi.jedicore import Interaction
 from pandajedi.jediddm.DDMInterface import DDMInterface
 from pandaserver.brokerage.SiteMapper import SiteMapper
 from pandaserver.dataservice import DataServiceUtils
 from pandaserver.dataservice.DataServiceUtils import select_scope
+from pandaserver.srvcore.hardware_matching import (
+    compare_version_string,
+    match_gpu_spec,
+)
 from pandaserver.taskbuffer import JobUtils, ProcessGroups, SiteSpec
 from pandaserver.taskbuffer.DdmSpec import DOWNTIME_STATUSES
 
@@ -916,49 +919,6 @@ def getAnalySitesClass(tbIF, fresher_than_minutes_ago=60):
     return ret_val, ret_map
 
 
-def compare_version_string(version_string, comparison_string):
-    """
-    Compares a version string with another string composed of a comparison operator and a version string.
-
-    Args:
-        version_string (str): The version string to compare.
-        comparison_string (str): The string containing the comparison operator and version string (e.g., ">=2.0").
-
-    Returns:
-        bool or None: True if the version string satisfies the comparison, False if it doesn't,
-                       or None if the comparison string is invalid.
-    """
-    match = re.match(r"([=><]+)(.+)", comparison_string)
-    if not match:
-        return None
-
-    operator = match.group(1).strip()
-    if operator == "=":
-        operator = "=="
-    version_to_compare = match.group(2).strip()
-
-    try:
-        version1 = version.parse(version_string)
-        version2 = version.parse(version_to_compare)
-    except version.InvalidVersion:
-        return None
-
-    if operator == "==":
-        return version1 == version2
-    elif operator == "!=":
-        return version1 != version2
-    elif operator == ">=":
-        return version1 >= version2
-    elif operator == "<=":
-        return version1 <= version2
-    elif operator == ">":
-        return version1 > version2
-    elif operator == "<":
-        return version1 < version2
-    else:
-        return None
-
-
 # check SW with json
 class JsonSoftwareCheck:
     # constructor
@@ -1090,60 +1050,12 @@ def check(
                                 continue
 
                             # All attribute checks use WN GPU monitoring (MV_WORKER_NODE_GPU_SUMMARY)
-                            # which has richer per-host data (vram, architecture, driver version)
-                            wn_gpus = self.wn_gpu_map.get(tmp_site_name, [])
-
-                            # check vendor
-                            if host_gpu_spec["vendor"] != "*":
-                                if not wn_gpus or not any(g.get("vendor") and re.match(host_gpu_spec["vendor"], g["vendor"], re.IGNORECASE) for g in wn_gpus):
-                                    continue
-
-                            # check model (include or exclude pattern)
-                            if host_gpu_spec["model"] != "*":
-                                if isinstance(host_gpu_spec["model"], dict):
-                                    model_pattern = host_gpu_spec["model"]["pattern"]
-                                    model_excl = host_gpu_spec["model"].get("excl", False)
-                                else:
-                                    model_pattern = host_gpu_spec["model"]
-                                    model_excl = False
-                                if not wn_gpus:
-                                    continue
-                                matches = any(g.get("model") and re.match(model_pattern, g["model"], re.IGNORECASE) for g in wn_gpus)
-                                if matches == model_excl:
-                                    continue
-
-                            # check VRAM (in MB); supports operators: ==, >=, <=, >, <, != (e.g. ">=40960")
-                            # all() ensures every GPU entry in the queue meets the minimum — prevents brokering to
-                            # mixed sites where some nodes fall below the requirement
-                            if "vram" in host_gpu_spec:
-                                if not wn_gpus or not all(g.get("vram") and compare_version_string(str(g["vram"]), host_gpu_spec["vram"]) for g in wn_gpus):
-                                    continue
-
-                            # check GPU microarchitecture generation (e.g. Ampere, Hopper, Ada Lovelace)
-                            if "microarchitecture" in host_gpu_spec:
-                                req_arch = host_gpu_spec["microarchitecture"]
-                                if isinstance(req_arch, str):
-                                    req_arch = [req_arch]
-                                if not wn_gpus or not any(g.get("architecture") in req_arch for g in wn_gpus):
-                                    continue
-
-                            # check minimum CUDA version
-                            # all() ensures every GPU entry in the queue meets the minimum — prevents brokering to
-                            # mixed sites where some nodes fall below the requirement
-                            if "version" in host_gpu_spec:
-                                if not wn_gpus or not all(
-                                    g.get("framework_version") and compare_version_string(g["framework_version"], host_gpu_spec["version"]) for g in wn_gpus
-                                ):
-                                    continue
-
-                            # check minimum GPU driver version (kernel driver, e.g. 575.57.08)
-                            # all() ensures every GPU entry in the queue meets the minimum — prevents brokering to
-                            # mixed sites where some nodes fall below the requirement
-                            if "driver_version" in host_gpu_spec:
-                                if not wn_gpus or not all(
-                                    g.get("driver_version") and compare_version_string(g["driver_version"], host_gpu_spec["driver_version"]) for g in wn_gpus
-                                ):
-                                    continue
+                            # which has richer per-host data (vram, architecture, driver version).
+                            # The minimum-requirement attributes are checked against all GPU entries of
+                            # the queue, to prevent brokering to mixed sites where some nodes fall below
+                            # the requirement
+                            if not match_gpu_spec(host_gpu_spec, self.wn_gpu_map.get(tmp_site_name, [])):
+                                continue
                     go_ahead = True
                 except Exception as e:
                     if log_stream:

pandaserver/api/v1/pilot_api.py

@@ -1,4 +1,5 @@
 import datetime
+import json
 import os
 import sys
 import time
@@ -72,6 +73,7 @@ def acquire_jobs(
     job_type: str = None,
     via_topic: bool = None,
     remaining_time=None,
+    target_architecture: dict | str = None,
 ) -> dict:
     """
     Acquire jobs
@@ -104,6 +106,14 @@ def acquire_jobs(
                                  to disambiguate the cases of test jobs that can be production or analysis. Optional and defaults to `None`.
         via_topic(bool, optional): Topic for message broker. Optional and defaults to `None`.
         remaining_time(int, optional): Remaining walltime. Optional and defaults to `None`.
+        target_architecture(dict or str, optional): Hardware of the worker node, either as a dictionary or as a JSON-encoded string.
+                                 Only jobs of tasks whose hardware requirements are satisfied by the worker node are returned.
+                                 The `gpus` key contains the list of GPUs, using the same key names as `update_worker_node_gpu`, e.g.
+                                 ``{"gpus": [{"vendor": "NVIDIA", "model": "NVIDIA A100-SXM4-40GB", "vram": 40960,
+                                 "architecture": "Ampere", "framework_version": "12.4", "driver_version": "575.57.08"}]}``.
+                                 An empty `gpus` list means that the worker node has no GPU, while an absent `gpus` key means that
+                                 the worker node doesn't report GPU information, in which case GPU requirements are not checked.
+                                 Optional and defaults to `None`.
 
     Returns:
         dict: The system response `{"success": success, "message": message, "data": data}`. The data is a list of job dictionaries.
@@ -153,6 +163,23 @@ def acquire_jobs(
     except (ValueError, TypeError):
         remaining_time = 0
 
+    # convert target architecture. Bad values are rejected instead of being ignored,
+    # since ignoring them would dispatch jobs to unsuitable hardware
+    if target_architecture:
+        if isinstance(target_architecture, str):
+            try:
+                target_architecture = json.loads(target_architecture)
+            except Exception as e:
+                message = f"failed to parse target_architecture with {str(e)}"
+                tmp_logger.error(message)
+                return generate_response(False, message=message)
+        if not isinstance(target_architecture, dict):
+            message = "target_architecture must be a JSON object"
+            tmp_logger.error(message)
+            return generate_response(False, message=message)
+    else:
+        target_architecture = None
+
     # harvester ID was not set, but we haver the scheduler ID, which should be the same
     if not harvester_id and scheduler_id:
         harvester_id = scheduler_id
@@ -162,7 +189,8 @@ def acquire_jobs(
         f"node={node}, ce={computing_element}, user={prod_user_id}, proxy={get_proxy_key}, "
         f"task_id={task_id}, DN={real_dn}, role={is_production_manager}, "
         f"bg={background}, rt={resource_type}, harvester_id={harvester_id}, worker_id={worker_id}, "
-        f"scheduler_id={scheduler_id}, job_type={job_type}, via_topic={via_topic} remaining_time={remaining_time}"
+        f"scheduler_id={scheduler_id}, job_type={job_type}, via_topic={via_topic} remaining_time={remaining_time}, "
+        f"target_architecture={target_architecture}"
     )
 
     # log the acquire_jobs as it's used for site activity metrics
@@ -204,6 +232,7 @@ def acquire_jobs(
         is_grandly_unified,
         via_topic,
         remaining_time,
+        target_architecture,
     )
 
     # Time-out
@@ -267,7 +296,7 @@ def acquire_jobs(
             tmp_logger.error(f"{tmp_msg}\n{traceback.format_exc()}")
             raise
 
-    tmp_logger.debug(f"Done for {site_name} {node}")
+    tmp_logger.debug(f"Sent {len(response_list)} jobs for {site_name} {node}")
 
     t_end = time.time()
     t_delta = t_end - t_start

pandaserver/api/v1/tests/pilot_api_tests.py

@@ -32,37 +32,19 @@ def test_acquire_jobs(self):
             "scheduler_id": "imaginary_scheduler",
             "job_type": "user",
             "via_topic": False,
-        }
-
-        status, output = self.http_client.post(url, data)
-        print(output)
-        output["status"] = status
-
-        expected_response = {"status": 0, "success": False, "data": 2, "message": ""}
-        self.assertEqual(output, expected_response)
-
-    def test_acquire_jobs(self):
-        url = f"{api_url_ssl}/pilot/acquire_jobs"
-        print(f"Testing URL: {url}")
-        data = {
-            "site_name": "CERN",
-            "timeout": 60,
-            "memory": 999999999,
-            "disk_space": 999999999,
-            "prod_source_label": "managed",
-            "node": "aipanda120.cern.ch",
-            "computing_element": "CERN",
-            "prod_user_id": None,
-            "get_proxy_key": None,
-            "task_id": None,
-            "n_jobs": 1,
-            "background": False,
-            "resource_type": "SCORE",
-            "harvester_id": "imaginary_harvester",
-            "worker_id": 12345,
-            "scheduler_id": "imaginary_scheduler",
-            "job_type": "user",
-            "via_topic": False,
+            "remaining_time": 3600,
+            "target_architecture": {
+                "gpus": [
+                    {
+                        "vendor": "NVIDIA",
+                        "model": "NVIDIA A100-SXM4-40GB",
+                        "vram": 40960,
+                        "architecture": "Ampere",
+                        "framework_version": "12.4",
+                        "driver_version": "575.57.08",
+                    }
+                ]
+            },
         }
 
         status, output = self.http_client.post(url, data)

pandaserver/srvcore/hardware_matching.py

@@ -0,0 +1,128 @@
+"""
+matching of hardware requirements against actual hardware
+
+The requirement side comes from the task architecture, e.g. JediTaskSpec.get_host_gpu_spec().
+The hardware side is a list of GPU dictionaries using the key names of the worker node GPU
+monitoring (ATLAS_PANDA.worker_node_gpus / MV_WORKER_NODE_GPU_SUMMARY), i.e. vendor, model,
+vram, architecture, framework_version, and driver_version. Those dictionaries either describe
+all worker nodes of a PanDA queue, when brokering tasks to queues, or the GPUs of a single
+worker node, when dispatching jobs to a pilot.
+"""
+
+import re
+
+from packaging import version
+
+
+def compare_version_string(version_string, comparison_string):
+    """
+    Compares a version string with another string composed of a comparison operator and a version string.
+
+    Args:
+        version_string (str): The version string to compare.
+        comparison_string (str): The string containing the comparison operator and version string (e.g., ">=2.0").
+
+    Returns:
+        bool or None: True if the version string satisfies the comparison, False if it doesn't,
+                       or None if the comparison string is invalid.
+    """
+    match = re.match(r"([=><!]+)(.+)", comparison_string)
+    if not match:
+        return None
+
+    operator = match.group(1).strip()
+    if operator == "=":
+        operator = "=="
+    version_to_compare = match.group(2).strip()
+
+    try:
+        version1 = version.parse(version_string)
+        version2 = version.parse(version_to_compare)
+    except version.InvalidVersion:
+        return None
+
+    if operator == "==":
+        return version1 == version2
+    elif operator == "!=":
+        return version1 != version2
+    elif operator == ">=":
+        return version1 >= version2
+    elif operator == "<=":
+        return version1 <= version2
+    elif operator == ">":
+        return version1 > version2
+    elif operator == "<":
+        return version1 < version2
+    else:
+        return None
+
+
+def match_gpu_spec(required_gpu_spec, gpus):
+    """
+    Checks whether GPUs satisfy the GPU requirement of a task.
+
+    Selection attributes (vendor, model, microarchitecture) use an any match, i.e. it is enough that
+    one GPU is of the requested type. Minimum-requirement attributes (vram, version, driver_version)
+    use an all match, i.e. every GPU has to satisfy the constraint, so that a job cannot end up on a
+    non-compliant GPU of a mixed set.
+
+    Args:
+        required_gpu_spec (dict): The GPU requirement of the task, with the keys vendor, model, vram,
+                                  microarchitecture, version, and driver_version. Only vendor and model
+                                  are mandatory and `*` is the wildcard for them. The model is either a
+                                  regular expression for inclusion or a dictionary with pattern and excl
+                                  keys for exclusion. The version, driver_version, and vram are
+                                  operator-prefixed strings, e.g. `>=12.0`.
+        gpus (list): List of dictionaries describing the actual GPUs, with the keys vendor, model, vram,
+                     architecture, framework_version, and driver_version.
+
+    Returns:
+        bool: True if the GPUs satisfy the requirement.
+    """
+    # check vendor
+    required_vendor = required_gpu_spec.get("vendor", "*")
+    if required_vendor != "*":
+        if not gpus or not any(gpu.get("vendor") and re.match(required_vendor, gpu["vendor"], re.IGNORECASE) for gpu in gpus):
+            return False
+
+    # check model (include or exclude pattern)
+    required_model = required_gpu_spec.get("model", "*")
+    if required_model != "*":
+        if isinstance(required_model, dict):
+            model_pattern = required_model["pattern"]
+            model_excl = required_model.get("excl", False)
+        else:
+            model_pattern = required_model
+            model_excl = False
+        if not gpus:
+            return False
+        matches = any(gpu.get("model") and re.match(model_pattern, gpu["model"], re.IGNORECASE) for gpu in gpus)
+        if matches == model_excl:
+            return False
+
+    # check VRAM (in MB); supports operators: ==, >=, <=, >, <, != (e.g. ">=40960")
+    if "vram" in required_gpu_spec:
+        if not gpus or not all(gpu.get("vram") and compare_version_string(str(gpu["vram"]), required_gpu_spec["vram"]) for gpu in gpus):
+            return False
+
+    # check GPU microarchitecture generation (e.g. Ampere, Hopper, Ada Lovelace)
+    if "microarchitecture" in required_gpu_spec:
+        req_arch = required_gpu_spec["microarchitecture"]
+        if isinstance(req_arch, str):
+            req_arch = [req_arch]
+        if not gpus or not any(gpu.get("architecture") in req_arch for gpu in gpus):
+            return False
+
+    # check CUDA toolkit version
+    if "version" in required_gpu_spec:
+        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):
+            return False
+
+    # check GPU kernel driver version (e.g. 575.57.08)
+    if "driver_version" in required_gpu_spec:
+        if not gpus or not all(
+            gpu.get("driver_version") and compare_version_string(gpu["driver_version"], required_gpu_spec["driver_version"]) for gpu in gpus
+        ):
+            return False
+
+    return True