diff --git a/alts/scheduler/scheduling.py b/alts/scheduler/scheduling.py index c3cfd40..86acc53 100644 --- a/alts/scheduler/scheduling.py +++ b/alts/scheduler/scheduling.py @@ -68,6 +68,38 @@ def get_excluded_packages(self, distro: str) -> Optional[dict]: pkgs = self.__cached_config.get('excluded_packages', {}) return pkgs.get(distro, {}) + @staticmethod + def get_queue_arch( + runner_class, + payload: TaskRequestPayload, + ) -> Optional[str]: + """ + Returns the architecture part of the queue name for the task, + applying the first matching override from the config. + """ + queue_arch = None + for ( + arch, + supported_arches, + ) in runner_class.ARCHITECTURES_MAPPING.items(): + if payload.dist_arch in supported_arches: + queue_arch = arch + if not queue_arch: + return None + dist_version = str(payload.dist_version) + major_version = dist_version.split('.')[0] + for override in CONFIG.queue_arch_overrides: + if ( + override.dist_name == payload.dist_name + and override.arch == queue_arch + and ( + dist_version in override.dist_versions + or major_version in override.dist_versions + ) + ): + return override.queue_arch + return queue_arch + def schedule_test_task(self, payload: TaskRequestPayload): """ Schedules new tasks in Test System. @@ -114,14 +146,7 @@ def schedule_test_task(self, payload: TaskRequestPayload): # TODO: Make decision on what queue to use for particular task based on # queues load - queue_arch = None - for ( - arch, - supported_arches, - ) in runner_class.ARCHITECTURES_MAPPING.items(): - if payload.dist_arch in supported_arches: - queue_arch = arch - + queue_arch = self.get_queue_arch(runner_class, payload) if not queue_arch: self.logger.error( 'Cannot map requested architecture to any ' diff --git a/alts/shared/constants.py b/alts/shared/constants.py index 9ae2083..cbcccba 100644 --- a/alts/shared/constants.py +++ b/alts/shared/constants.py @@ -25,6 +25,9 @@ ARCHITECTURES = ( 'x86_64', 'x86_64_v2', + # Queue-only architecture for x86_64 workers whose CPU supports + # x86-64-v3, see SchedulerConfig.queue_arch_overrides + 'x86_64_v3', 'aarch64', 'ppc64le', 's390x', diff --git a/alts/shared/models.py b/alts/shared/models.py index 4181647..4ea9a4c 100644 --- a/alts/shared/models.py +++ b/alts/shared/models.py @@ -11,7 +11,7 @@ Set, ) -from pydantic import BaseModel, ConfigDict, computed_field +from pydantic import BaseModel, ConfigDict, computed_field, field_validator from alts.shared import constants from alts.shared.utils.path_utils import get_abspath @@ -170,6 +170,27 @@ class DockerImageOverride(BaseModel): image: str +class QueueArchOverride(BaseModel): + # Sends matching tasks to a different queue architecture, e.g. to + # schedule AlmaLinux 10 x86_64 tasks only on x86-64-v3 capable workers + dist_name: str + # Matched against the full version and its major part + dist_versions: List[str] + # Queue architecture the task would be scheduled to otherwise + arch: str + queue_arch: str + + @field_validator('arch', 'queue_arch') + @classmethod + def check_queue_arch(cls, value: str) -> str: + # Tasks sent to an undeclared queue are never picked up + if value not in constants.ARCHITECTURES: + raise ValueError( + f'{value} is not one of {", ".join(constants.ARCHITECTURES)}' + ) + return value + + class RabbitmqBrokerConfig(BaseBrokerConfig): use_ssl: bool = False ssl_config: Optional[SslConfig] = None @@ -431,3 +452,4 @@ class SchedulerConfig(CeleryConfig): excluded_tests_url: str = 'https://git.almalinux.org/almalinux/alts-exclusions/raw/branch/main/skipped_tests.json' excluded_tests_cache_size: str = 1 excluded_tests_cache_update_interval: int = 600 + queue_arch_overrides: List[QueueArchOverride] = [] diff --git a/tests/scheduler/__init__.py b/tests/scheduler/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/scheduler/test_scheduling.py b/tests/scheduler/test_scheduling.py new file mode 100644 index 0000000..ff0f780 --- /dev/null +++ b/tests/scheduler/test_scheduling.py @@ -0,0 +1,100 @@ +import pytest +from pydantic import ValidationError + +from alts.scheduler import CONFIG +from alts.scheduler import scheduling +from alts.shared.models import QueueArchOverride, TaskRequestPayload +from alts.worker.runners import DockerRunner + +get_queue_arch = scheduling.TestsScheduler.get_queue_arch + + +def make_payload(dist_name: str, dist_version, dist_arch: str): + return TaskRequestPayload( + bs_task_id=1, + runner_type='docker', + dist_name=dist_name, + dist_version=dist_version, + dist_arch=dist_arch, + package_name='bash', + ) + + +class TestGetQueueArch: + @pytest.mark.parametrize( + 'dist_name, dist_version, dist_arch, expected', + [ + pytest.param( + 'almalinux', '10', 'x86_64', 'x86_64_v3', + id='major-version', + ), + pytest.param( + 'almalinux', '10.1', 'x86_64', 'x86_64_v3', + id='full-version', + ), + pytest.param( + 'almalinux', 10, 'x86_64', 'x86_64_v3', + id='int-version', + ), + pytest.param( + 'almalinux', '10', 'i686', 'x86_64_v3', + id='mapped-arch', + ), + pytest.param( + 'almalinux', '10', 'x86_64_v2', 'x86_64_v2', + id='other-arch', + ), + pytest.param( + 'almalinux', '9', 'x86_64', 'x86_64', + id='other-version', + ), + pytest.param( + 'centos', '10', 'x86_64', 'x86_64', + id='other-dist', + ), + pytest.param( + 'almalinux', '10', 'aarch64', 'aarch64', + id='no-override-for-arch', + ), + ], + ) + def test_get_queue_arch( + self, + monkeypatch: pytest.MonkeyPatch, + dist_name: str, + dist_version, + dist_arch: str, + expected: str, + ): + monkeypatch.setattr( + CONFIG, + 'queue_arch_overrides', + [ + QueueArchOverride( + dist_name='almalinux', + dist_versions=['10'], + arch='x86_64', + queue_arch='x86_64_v3', + ), + ], + ) + payload = make_payload(dist_name, dist_version, dist_arch) + assert get_queue_arch(DockerRunner, payload) == expected + + def test_no_overrides(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(CONFIG, 'queue_arch_overrides', []) + payload = make_payload('almalinux', '10', 'x86_64') + assert get_queue_arch(DockerRunner, payload) == 'x86_64' + + def test_unknown_arch(self): + payload = make_payload('almalinux', '10', 'riscv64') + assert get_queue_arch(DockerRunner, payload) is None + + def test_undeclared_queue_arch_rejected(self): + with pytest.raises(ValidationError): + QueueArchOverride( + dist_name='almalinux', + dist_versions=['10'], + arch='x86_64', + queue_arch='x86-64-v3', + )