From 93ad98b6dd9e1b195dea99a6460b50898bbbca82 Mon Sep 17 00:00:00 2001 From: tmaeno Date: Wed, 26 Aug 2026 13:48:19 +0200 Subject: [PATCH 1/2] target_architecture in acquire_jobs --- pandajedi/jedibrokerage/AtlasBrokerUtils.py | 108 +---------- pandaserver/api/v1/pilot_api.py | 31 ++- pandaserver/api/v1/tests/pilot_api_tests.py | 44 ++--- pandaserver/srvcore/hardware_matching.py | 128 +++++++++++++ pandaserver/srvcore/tests/__init__.py | 0 .../srvcore/tests/hardware_matching_tests.py | 144 ++++++++++++++ pandaserver/taskbuffer/TaskBuffer.py | 2 + .../db_proxy_mods/job_complex_module.py | 177 +++++++++++++++--- 8 files changed, 476 insertions(+), 158 deletions(-) create mode 100644 pandaserver/srvcore/hardware_matching.py create mode 100644 pandaserver/srvcore/tests/__init__.py create mode 100644 pandaserver/srvcore/tests/hardware_matching_tests.py diff --git a/pandajedi/jedibrokerage/AtlasBrokerUtils.py b/pandajedi/jedibrokerage/AtlasBrokerUtils.py index 4044546a1..6eacb2dec 100644 --- a/pandajedi/jedibrokerage/AtlasBrokerUtils.py +++ b/pandajedi/jedibrokerage/AtlasBrokerUtils.py @@ -8,7 +8,6 @@ import traceback from typing import Any -from packaging import version from pandacommon.pandautils.PandaUtils import naive_utcnow from pandajedi.jedicore import Interaction @@ -16,6 +15,10 @@ 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: diff --git a/pandaserver/api/v1/pilot_api.py b/pandaserver/api/v1/pilot_api.py index 422a8f70e..9085875fd 100644 --- a/pandaserver/api/v1/pilot_api.py +++ b/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 diff --git a/pandaserver/api/v1/tests/pilot_api_tests.py b/pandaserver/api/v1/tests/pilot_api_tests.py index 88df82b92..6e53465ad 100644 --- a/pandaserver/api/v1/tests/pilot_api_tests.py +++ b/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) diff --git a/pandaserver/srvcore/hardware_matching.py b/pandaserver/srvcore/hardware_matching.py new file mode 100644 index 000000000..ca68dc7e1 --- /dev/null +++ b/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"([=>=": + 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 diff --git a/pandaserver/srvcore/tests/__init__.py b/pandaserver/srvcore/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pandaserver/srvcore/tests/hardware_matching_tests.py b/pandaserver/srvcore/tests/hardware_matching_tests.py new file mode 100644 index 000000000..4a8ee73a5 --- /dev/null +++ b/pandaserver/srvcore/tests/hardware_matching_tests.py @@ -0,0 +1,144 @@ +# Description: Unit tests for the hardware matching functions +import unittest + +from pandaserver.srvcore.hardware_matching import ( + compare_version_string, + match_gpu_spec, +) + +A100 = { + "vendor": "NVIDIA", + "model": "NVIDIA A100-SXM4-40GB", + "vram": 40960, + "architecture": "Ampere", + "framework_version": "12.4", + "driver_version": "575.57.08", +} + +V100 = { + "vendor": "NVIDIA", + "model": "Tesla V100-SXM2-16GB", + "vram": 16384, + "architecture": "Volta", + "framework_version": "12.2", + "driver_version": "535.104.05", +} + +ANY_GPU = {"vendor": "*", "model": "*"} + + +class TestCompareVersionString(unittest.TestCase): + def test_operators(self): + self.assertTrue(compare_version_string("12.4", ">=12.0")) + self.assertFalse(compare_version_string("11.8", ">=12.0")) + self.assertTrue(compare_version_string("11.8", "<=12.0")) + self.assertTrue(compare_version_string("12.4", ">12.0")) + self.assertFalse(compare_version_string("12.0", ">12.0")) + self.assertTrue(compare_version_string("11.8", "<12.0")) + self.assertTrue(compare_version_string("12.0", "==12.0")) + self.assertFalse(compare_version_string("12.4", "==12.0")) + + def test_single_equal_is_equality(self): + self.assertTrue(compare_version_string("12.0", "=12.0")) + self.assertFalse(compare_version_string("12.4", "=12.0")) + + def test_not_equal(self): + # the != operator used to be unreachable since the operator pattern didn't accept ! + self.assertTrue(compare_version_string("12.4", "!=12.0")) + self.assertFalse(compare_version_string("12.0", "!=12.0")) + + def test_multi_component_version(self): + self.assertTrue(compare_version_string("575.57.08", ">=575.0")) + self.assertFalse(compare_version_string("535.104.05", ">=575.0")) + + def test_invalid_input(self): + # no operator + self.assertIsNone(compare_version_string("12.0", "12.0")) + # unparsable versions + self.assertIsNone(compare_version_string("Ampere", ">=12.0")) + self.assertIsNone(compare_version_string("12.0", ">=Ampere")) + + +class TestMatchGpuSpec(unittest.TestCase): + def test_wildcard_requirement(self): + self.assertTrue(match_gpu_spec(ANY_GPU, [A100])) + self.assertTrue(match_gpu_spec(ANY_GPU, [A100, V100])) + # no GPU information is not a rejection for a wildcard requirement + self.assertTrue(match_gpu_spec(ANY_GPU, [])) + + def test_specific_requirement_without_gpus(self): + self.assertFalse(match_gpu_spec({"vendor": "NVIDIA", "model": "*"}, [])) + self.assertFalse(match_gpu_spec({"vendor": "*", "model": ".*A100.*"}, [])) + self.assertFalse(match_gpu_spec({"vendor": "*", "model": "*", "vram": ">=40960"}, [])) + self.assertFalse(match_gpu_spec({"vendor": "*", "model": "*", "microarchitecture": "Ampere"}, [])) + + def test_vendor(self): + self.assertTrue(match_gpu_spec({"vendor": "NVIDIA", "model": "*"}, [A100])) + self.assertTrue(match_gpu_spec({"vendor": "nvidia", "model": "*"}, [A100])) + self.assertFalse(match_gpu_spec({"vendor": "AMD", "model": "*"}, [A100])) + # any match, one of the GPUs is enough + self.assertTrue(match_gpu_spec({"vendor": "NVIDIA", "model": "*"}, [{"vendor": "AMD", "model": "MI250"}, A100])) + + def test_model_inclusion(self): + self.assertTrue(match_gpu_spec({"vendor": "*", "model": ".*A100.*"}, [A100])) + # matching is case-insensitive + self.assertTrue(match_gpu_spec({"vendor": "*", "model": ".*a100.*"}, [A100])) + self.assertFalse(match_gpu_spec({"vendor": "*", "model": ".*A100.*"}, [V100])) + # any match + self.assertTrue(match_gpu_spec({"vendor": "*", "model": ".*A100.*"}, [V100, A100])) + + def test_model_exclusion(self): + excl_p100 = {"vendor": "*", "model": {"pattern": ".*P100.*", "excl": True}} + self.assertTrue(match_gpu_spec(excl_p100, [A100])) + P100 = dict(A100, model="Tesla P100-PCIE-16GB") + self.assertFalse(match_gpu_spec(excl_p100, [P100])) + # excluded when any of the GPUs matches the pattern + self.assertFalse(match_gpu_spec(excl_p100, [A100, P100])) + + def test_vram(self): + self.assertTrue(match_gpu_spec(dict(ANY_GPU, vram=">=40960"), [A100])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, vram=">=40960"), [V100])) + # all match, every GPU has to meet the minimum + self.assertFalse(match_gpu_spec(dict(ANY_GPU, vram=">=40960"), [A100, V100])) + self.assertTrue(match_gpu_spec(dict(ANY_GPU, vram=">=16384"), [A100, V100])) + self.assertTrue(match_gpu_spec(dict(ANY_GPU, vram="==40960"), [A100])) + + def test_microarchitecture(self): + self.assertTrue(match_gpu_spec(dict(ANY_GPU, microarchitecture="Ampere"), [A100])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, microarchitecture="Ampere"), [V100])) + # a list of generations is accepted + self.assertTrue(match_gpu_spec(dict(ANY_GPU, microarchitecture=["Ampere", "Hopper"]), [A100])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, microarchitecture=["Ampere", "Hopper"]), [V100])) + # any match, one of the GPUs is enough + self.assertTrue(match_gpu_spec(dict(ANY_GPU, microarchitecture="Ampere"), [V100, A100])) + + def test_framework_version(self): + self.assertTrue(match_gpu_spec(dict(ANY_GPU, version=">=12.0"), [A100, V100])) + # all match, a single old GPU excludes the whole set + self.assertFalse(match_gpu_spec(dict(ANY_GPU, version=">=12.3"), [A100, V100])) + self.assertTrue(match_gpu_spec(dict(ANY_GPU, version=">=12.3"), [A100])) + + def test_driver_version(self): + self.assertTrue(match_gpu_spec(dict(ANY_GPU, driver_version=">=575.0"), [A100])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, driver_version=">=575.0"), [V100])) + # all match + self.assertFalse(match_gpu_spec(dict(ANY_GPU, driver_version=">=575.0"), [A100, V100])) + + def test_missing_attributes_in_gpus(self): + bare = {"vendor": "NVIDIA", "model": "NVIDIA A100-SXM4-40GB"} + self.assertTrue(match_gpu_spec({"vendor": "NVIDIA", "model": ".*A100.*"}, [bare])) + # constraints on attributes which the GPU doesn't report are not satisfied + self.assertFalse(match_gpu_spec(dict(ANY_GPU, vram=">=40960"), [bare])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, version=">=12.0"), [bare])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, driver_version=">=575.0"), [bare])) + self.assertFalse(match_gpu_spec(dict(ANY_GPU, microarchitecture="Ampere"), [bare])) + + def test_combined_constraints(self): + spec = {"vendor": "NVIDIA", "model": ".*A100.*", "vram": ">=40960", "microarchitecture": "Ampere", "version": ">=12.0", "driver_version": ">=575.0"} + self.assertTrue(match_gpu_spec(spec, [A100])) + self.assertFalse(match_gpu_spec(spec, [V100])) + self.assertFalse(match_gpu_spec(spec, [A100, V100])) + + +if __name__ == "__main__": + unittest.main() diff --git a/pandaserver/taskbuffer/TaskBuffer.py b/pandaserver/taskbuffer/TaskBuffer.py index f767847ac..8a1b856a4 100755 --- a/pandaserver/taskbuffer/TaskBuffer.py +++ b/pandaserver/taskbuffer/TaskBuffer.py @@ -902,6 +902,7 @@ def getJobs( is_gu, via_topic, remaining_time, + target_architecture, ): # get DBproxy with self.proxyPool.get() as proxy: @@ -927,6 +928,7 @@ def getJobs( is_gu, via_topic, remaining_time, + target_architecture, ) t_after = time.time() t_total = t_after - t_before diff --git a/pandaserver/taskbuffer/db_proxy_mods/job_complex_module.py b/pandaserver/taskbuffer/db_proxy_mods/job_complex_module.py index 6d9a6509a..3ceec2844 100644 --- a/pandaserver/taskbuffer/db_proxy_mods/job_complex_module.py +++ b/pandaserver/taskbuffer/db_proxy_mods/job_complex_module.py @@ -11,6 +11,7 @@ from pandaserver.config import panda_config from pandaserver.srvcore import CoreUtils, srv_msg_utils +from pandaserver.srvcore.hardware_matching import match_gpu_spec from pandaserver.taskbuffer import ( ErrorCode, EventServiceUtils, @@ -28,8 +29,15 @@ from pandaserver.taskbuffer.db_proxy_mods.task_event_module import get_task_event_module from pandaserver.taskbuffer.db_proxy_mods.worker_module import get_worker_module from pandaserver.taskbuffer.FileSpec import FileSpec +from pandaserver.taskbuffer.JediTaskSpec import JediTaskSpec from pandaserver.taskbuffer.JobSpec import JobSpec, get_task_queued_time +# maximum number of task IDs excluded from job dispatch due to hardware mismatch +MAX_EXCLUDED_TASK_IDS = 500 + +# maximum number of attempts to get job candidates while excluding tasks with hardware mismatch +MAX_ARCHITECTURE_MATCHING_TRIES = 5 + # Module class to define job-related methods that use another module's methods or serve as their dependencies class JobComplexModule(BaseModule): @@ -2073,6 +2081,7 @@ def construct_where_clause( task_id, average_memory_limit, remaining_time, + excluded_task_ids=None, ): get_val_map = {":oldJobStatus": "activated", ":computingSite": site_name} @@ -2147,8 +2156,79 @@ def construct_where_clause( sql_where_clause += "AND minramcount / NVL(corecount, 1)<=:average_memory_limit " get_val_map[":average_memory_limit"] = average_memory_limit + # skip tasks whose hardware requirements are not satisfied by the worker node + if excluded_task_ids: + var_names_str, var_map = get_sql_IN_bind_variables(excluded_task_ids, prefix=":excludedTaskID") + sql_where_clause += f"AND (jediTaskID IS NULL OR jediTaskID NOT IN ({var_names_str})) " + get_val_map.update(var_map) + return sql_where_clause, get_val_map + # get architectures of tasks + def get_task_architectures(self, task_ids): + """ + Get the architecture of tasks. jediTaskID is the primary key of JEDI_Tasks, so that this is an index lookup. + Task IDs without a row in JEDI_Tasks are missing from the returned dictionary. + + Args: + task_ids (iterable): JEDI task IDs. + + Returns: + dict: Map of JEDI task ID and architecture. Empty when the lookup fails. + """ + comment = " /* DBProxy.get_task_architectures */" + tmp_log = self.create_tagged_logger(comment) + task_ids = list(task_ids) + if not task_ids: + return {} + try: + var_names_str, var_map = get_sql_IN_bind_variables(task_ids, prefix=":jediTaskID") + sql = f"SELECT jediTaskID,architecture FROM {panda_config.schemaJEDI}.JEDI_Tasks WHERE jediTaskID IN ({var_names_str}) " + # start transaction + self.conn.begin() + self.cur.arraysize = 100 + self.cur.execute(sql + comment, var_map) + res = self.cur.fetchall() + # commit + if not self._commit(): + raise RuntimeError("Commit error") + return {jedi_task_id: architecture for jedi_task_id, architecture in res} + except Exception: + # roll back + self._rollback() + # error + self.dump_error_message(tmp_log) + return {} + + # check if a task can run on the worker node + def check_task_architecture(self, architecture, target_architecture): + """ + Check if the hardware requirements of a task are satisfied by the hardware of the worker node. + Unknown requirements and unreported hardware are accepted, so that jobs are not withheld when + information is missing. + + Args: + architecture (str): The architecture of the task, i.e. JEDI_Tasks.architecture. + target_architecture (dict): The hardware of the worker node, with the `gpus` key containing + the list of GPUs. See pandaserver.srvcore.hardware_matching. + + Returns: + bool: True if the task can run on the worker node. + """ + if not architecture: + return True + # reuse the task spec parser which understands both the old and new architecture formats + task_spec = JediTaskSpec() + task_spec.architecture = architecture + gpu_spec = task_spec.get_host_gpu_spec() + # the task doesn't require any GPU + if not gpu_spec: + return True + # the worker node doesn't report GPU information + if "gpus" not in target_architecture: + return True + return match_gpu_spec(gpu_spec, target_architecture["gpus"]) + # get jobs def getJobs( self, @@ -2171,12 +2251,16 @@ def getJobs( is_gu, via_topic, remaining_time, + target_architecture, ): """ 1. Construct where clause (sql_where_clause) based on applicable filters for request 2. Select n jobs with the highest priorities and the lowest pandaids - 3. Update the jobs to status SENT - 4. Pack the files and if jobs are AES also the event ranges + 3. Skip candidates of tasks whose hardware requirements are not satisfied by the worker node, + and retry the selection while excluding those tasks, so that a high priority task requiring + other hardware doesn't starve the worker node + 4. Update the jobs to status SENT + 5. Pack the files and if jobs are AES also the event ranges """ comment = " /* DBProxy.getJobs */" timeStart = naive_utcnow() @@ -2220,32 +2304,42 @@ def getJobs( average_memory_limit = average_memory_target tmp_log.info(f"Queue {siteName} meanRSS will be throttled to jobs under {average_memory_limit}MB") - # generate the WHERE clauses based on the requirements for the job - sql_where_clause, getValMap = self.construct_where_clause( - site_name=siteName, - mem=mem, - disk_space=diskSpace, - background=background, - resource_type=resourceType, - prod_source_label=prodSourceLabel, - computing_element=computingElement, - is_gu=is_gu, - job_type=jobType, - prod_user_id=prodUserID, - task_id=taskID, - average_memory_limit=average_memory_limit, - remaining_time=remaining_time, - ) - # get the sorting criteria (global shares, age, etc.) sorting_sql, sorting_varmap = get_entity_module(self).getSortingCriteria(siteName, maxAttemptIDx) - if sorting_varmap: # copy the var map, but not the sql, since it has to be at the very end - for tmp_key in sorting_varmap: - getValMap[tmp_key] = sorting_varmap[tmp_key] + + # task IDs excluded since the worker node doesn't satisfy their hardware requirements, and the + # verdicts to avoid looking up the same task twice. They are local to this request and are reused + # over the iterations to get multiple jobs + excluded_task_ids = set() + task_match_verdict = {} + + # generate the WHERE clause based on the requirements for the job + def build_where_clause(): + tmp_where_clause, tmp_val_map = self.construct_where_clause( + site_name=siteName, + mem=mem, + disk_space=diskSpace, + background=background, + resource_type=resourceType, + prod_source_label=prodSourceLabel, + computing_element=computingElement, + is_gu=is_gu, + job_type=jobType, + prod_user_id=prodUserID, + task_id=taskID, + average_memory_limit=average_memory_limit, + remaining_time=remaining_time, + excluded_task_ids=sorted(excluded_task_ids), + ) + if sorting_varmap: # copy the var map, but not the sql, since it has to be at the very end + for tmp_key in sorting_varmap: + tmp_val_map[tmp_key] = sorting_varmap[tmp_key] + return tmp_where_clause, tmp_val_map + + sql_where_clause, getValMapOrig = build_where_clause() retJobs = [] nSent = 0 - getValMapOrig = copy.copy(getValMap) try: timeLimit = datetime.timedelta(seconds=timeout - 10) @@ -2255,10 +2349,13 @@ def getJobs( getValMap = copy.copy(getValMapOrig) pandaID = 0 - nTry = 1 + # retry with tasks excluded by the hardware matching removed from the candidates + nTry = MAX_ARCHITECTURE_MATCHING_TRIES if target_architecture else 1 for iTry in range(nTry): # set siteID tmpSiteID = siteName + # task IDs newly excluded in this attempt due to hardware mismatch + newly_excluded_task_ids = set() # get file lock tmp_log.debug("lock") if (naive_utcnow() - timeStart) < timeLimit: @@ -2268,7 +2365,7 @@ def getJobs( if toGetPandaIDs: # get PandaIDs - sqlP = "SELECT /*+ INDEX_RS_ASC(tab (PRODSOURCELABEL COMPUTINGSITE JOBSTATUS) ) */ PandaID,currentPriority,specialHandling FROM ATLAS_PANDA.jobsActive4 tab " + sqlP = "SELECT /*+ INDEX_RS_ASC(tab (PRODSOURCELABEL COMPUTINGSITE JOBSTATUS) ) */ PandaID,currentPriority,specialHandling,jediTaskID FROM ATLAS_PANDA.jobsActive4 tab " sqlP += sql_where_clause if sorting_sql: @@ -2286,14 +2383,35 @@ def getJobs( if not self._commit(): raise RuntimeError("Commit error") + # check the hardware requirements of the tasks which are not yet checked + if target_architecture: + unchecked_task_ids = {tmpRes[-1] for tmpRes in resIDs if tmpRes[-1] and tmpRes[-1] not in task_match_verdict} + if unchecked_task_ids: + architecture_map = self.get_task_architectures(unchecked_task_ids) + for tmpTaskID in unchecked_task_ids: + # tasks missing in JEDI_Tasks are treated as unconstrained + task_match_verdict[tmpTaskID] = self.check_task_architecture(architecture_map.get(tmpTaskID), target_architecture) + for ( tmpPandaID, tmpCurrentPriority, tmpSpecialHandling, + tmpJediTaskID, ) in resIDs: + # skip jobs of tasks which cannot run on the worker node + if target_architecture and tmpJediTaskID and not task_match_verdict.get(tmpJediTaskID, True): + if len(excluded_task_ids) < MAX_EXCLUDED_TASK_IDS and tmpJediTaskID not in excluded_task_ids: + excluded_task_ids.add(tmpJediTaskID) + newly_excluded_task_ids.add(tmpJediTaskID) + continue pandaIDs.append(tmpPandaID) specialHandlingMap[tmpPandaID] = tmpSpecialHandling + # rebuild the query so that the excluded tasks are skipped in the subsequent attempts + if newly_excluded_task_ids: + tmp_log.debug(f"excluding jediTaskIDs={sorted(newly_excluded_task_ids)} due to hardware mismatch") + sql_where_clause, getValMapOrig = build_where_clause() + if pandaIDs == []: tmp_log.debug("no PandaIDs") retU = 0 # retU: return from update @@ -2433,9 +2551,12 @@ def getJobs( # succeeded if retU != 0: break - if iTry + 1 < nTry: - # time.sleep(0.5) - pass + # retry only when tasks were newly excluded by the hardware matching, since otherwise + # the same candidates would be selected again. Matching jobs can be hidden behind + # the excluded tasks in the candidate list + if not newly_excluded_task_ids: + break + getValMap = copy.copy(getValMapOrig) # failed to UPDATE if retU == 0: # reset pandaID From 84390993a953dc6b35d2cb1bc1f8a0bf31263d23 Mon Sep 17 00:00:00 2001 From: tmaeno Date: Wed, 26 Aug 2026 14:10:00 +0200 Subject: [PATCH 2/2] Update debug log message in acquire_jobs to include job count --- pandaserver/api/v1/pilot_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandaserver/api/v1/pilot_api.py b/pandaserver/api/v1/pilot_api.py index 9085875fd..a99b57ab2 100644 --- a/pandaserver/api/v1/pilot_api.py +++ b/pandaserver/api/v1/pilot_api.py @@ -296,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