diff --git a/job_creator/jobcreator/job_creator.py b/job_creator/jobcreator/job_creator.py index d8424d7..0ec4700 100644 --- a/job_creator/jobcreator/job_creator.py +++ b/job_creator/jobcreator/job_creator.py @@ -5,6 +5,8 @@ from typing import Any from kubernetes import client # type: ignore[import-untyped] +from kubernetes.client.exceptions import ApiException # type: ignore[import-untyped] +from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential from jobcreator.utils import load_kubernetes_config, logger @@ -123,6 +125,20 @@ def _setup_extras_pv(job_name: str, secret_namespace: str, manila_share_id: str, return pv_name +def _is_retryable_k8s_error(exc: ApiException) -> bool: + """Only retry on transient K8s API errors, not on conflicts or client errors.""" + if isinstance(exc, ApiException): + return exc.status in (429, 500, 502, 503, 504) + # Also retry on connection-level errors + return isinstance(exc, (ConnectionError, TimeoutError)) + + +@retry( + retry=retry_if_exception(_is_retryable_k8s_error), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=1, max=10), + reraise=True, +) def _setup_ceph_pv( pv_name: str, ceph_creds_k8s_secret_name: str, @@ -299,179 +315,198 @@ def spawn_job( # noqa: PLR0913 pv_names = [] pvc_names = [] - # Setup Archive PV and PVC - archive_pv_name = f"{job_name}-archive-pv-smb" - _setup_smb_pv( - archive_pv_name, - "archive-creds", - job_namespace, - "//isisdatar55.isis.cclrc.ac.uk/inst$/", - ["noserverino", "_netdev", "vers=2.1"], - ) - pv_names.append(archive_pv_name) - - archive_pvc_name = f"{job_name}-archive-pvc" - _setup_pvc(archive_pvc_name, archive_pv_name, job_namespace) - pvc_names.append(archive_pvc_name) - - # Setup Extras PV and PVC - extras_pv_name = _setup_extras_pv( - job_name=job_name, - secret_namespace=job_namespace, - manila_share_id=manila_share_id, - manila_share_access_id=manila_share_access_id, - ) - pv_names.append(extras_pv_name) - - extras_pvc_name = f"{job_name}-extras-pvc" - _setup_pvc(extras_pvc_name, extras_pv_name, job_namespace) - pvc_names.append(extras_pvc_name) - - # Setup ceph PV and PVC - if not self.dev_mode: - ceph_pv_name = f"{job_name}-ceph-pv" - ( - _setup_ceph_pv( - ceph_pv_name, - ceph_creds_k8s_secret_name, - ceph_creds_k8s_namespace, - cluster_id, - fs_name, - ceph_mount_path, - ), - ) - pv_names.append(ceph_pv_name) - - ceph_pvc_name = f"{job_name}-ceph-pvc" - _setup_pvc(ceph_pvc_name, ceph_pv_name, job_namespace, access_mode="ReadWriteMany") - pvc_names.append(ceph_pvc_name) - - ceph_volume = client.V1Volume( - name="ceph-mount", - persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( - claim_name=ceph_pvc_name, - read_only=False, - ), + try: + # Setup Archive PV and PVC + archive_pv_name = f"{job_name}-archive-pv-smb" + _setup_smb_pv( + archive_pv_name, + "archive-creds", + job_namespace, + "//isisdatar55.isis.cclrc.ac.uk/inst$/", + ["noserverino", "_netdev", "vers=2.1"], ) - else: - ceph_volume = client.V1Volume( - name="ceph-mount", - empty_dir=client.V1EmptyDirVolumeSource(size_limit="100Gi"), + pv_names.append(archive_pv_name) + + archive_pvc_name = f"{job_name}-archive-pvc" + _setup_pvc(archive_pvc_name, archive_pv_name, job_namespace) + pvc_names.append(archive_pvc_name) + + # Setup Extras PV and PVC + extras_pv_name = _setup_extras_pv( + job_name=job_name, + secret_namespace=job_namespace, + manila_share_id=manila_share_id, + manila_share_access_id=manila_share_access_id, ) + pv_names.append(extras_pv_name) + + extras_pvc_name = f"{job_name}-extras-pvc" + _setup_pvc(extras_pvc_name, extras_pv_name, job_namespace) + pvc_names.append(extras_pvc_name) + + # Setup ceph PV and PVC + if not self.dev_mode: + ceph_pv_name = f"{job_name}-ceph-pv" + ( + _setup_ceph_pv( + ceph_pv_name, + ceph_creds_k8s_secret_name, + ceph_creds_k8s_namespace, + cluster_id, + fs_name, + ceph_mount_path, + ), + ) + pv_names.append(ceph_pv_name) + + ceph_pvc_name = f"{job_name}-ceph-pvc" + _setup_pvc(ceph_pvc_name, ceph_pv_name, job_namespace, access_mode="ReadWriteMany") + pvc_names.append(ceph_pvc_name) + + ceph_volume = client.V1Volume( + name="ceph-mount", + persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( + claim_name=ceph_pvc_name, + read_only=False, + ), + ) + else: + ceph_volume = client.V1Volume( + name="ceph-mount", + empty_dir=client.V1EmptyDirVolumeSource(size_limit="100Gi"), + ) - # Create the Job - logger.info("Spawning job: %s", job_name) + # Create the Job + logger.info("Spawning job: %s", job_name) - volumes = [ - client.V1Volume( - name="archive-mount", - persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( - claim_name=archive_pvc_name, - read_only=True, - ), - ), - ceph_volume, - client.V1Volume( - name="extras-mount", - persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( - claim_name=extras_pvc_name, - read_only=True, + volumes = [ + client.V1Volume( + name="archive-mount", + persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( + claim_name=archive_pvc_name, + read_only=True, + ), ), - ), - ] - volumes_mounts = [ - client.V1VolumeMount(name="archive-mount", mount_path="/archive"), - client.V1VolumeMount(name="ceph-mount", mount_path="/output"), - client.V1VolumeMount(name="extras-mount", mount_path="/extras"), - ] - # Setup special PVs and add them to the volume mounts - if "imat" in special_pvs: - _setup_imat_pv_and_pvcs(job_name, job_namespace, pv_names, pvc_names) - imat_pvc_source = client.V1PersistentVolumeClaimVolumeSource( - claim_name=f"{job_name}-ndximat-pvc", read_only=True - ) - volumes.append(client.V1Volume(name="imat-mount", persistent_volume_claim=imat_pvc_source)) - volumes_mounts.append(client.V1VolumeMount(name="imat-mount", mount_path="/imat")) - # Because imat is special and uses mantid imaging to load large .tiff files, we need to ensure the /dev/shm - # is larger than 64mb. We do however have a soft-ish limit of around 32GiB on the size of datasets when - # doing this. - volumes.append( + ceph_volume, client.V1Volume( - name="dev-shm", empty_dir=client.V1EmptyDirVolumeSource(size_limit="32Gi", medium="Memory") + name="extras-mount", + persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( + claim_name=extras_pvc_name, + read_only=True, + ), + ), + ] + volumes_mounts = [ + client.V1VolumeMount(name="archive-mount", mount_path="/archive"), + client.V1VolumeMount(name="ceph-mount", mount_path="/output"), + client.V1VolumeMount(name="extras-mount", mount_path="/extras"), + ] + # Setup special PVs and add them to the volume mounts + if "imat" in special_pvs: + _setup_imat_pv_and_pvcs(job_name, job_namespace, pv_names, pvc_names) + imat_pvc_source = client.V1PersistentVolumeClaimVolumeSource( + claim_name=f"{job_name}-ndximat-pvc", read_only=True ) + volumes.append(client.V1Volume(name="imat-mount", persistent_volume_claim=imat_pvc_source)) + volumes_mounts.append(client.V1VolumeMount(name="imat-mount", mount_path="/imat")) + # Because imat is special and uses mantid imaging to load large .tiff files, we need to + # ensure the /dev/shm is larger than 64mb. We do however have a soft-ish limit of around + # 32GiB on the size of datasets when doing this. + volumes.append( + client.V1Volume( + name="dev-shm", empty_dir=client.V1EmptyDirVolumeSource(size_limit="32Gi", medium="Memory") + ) + ) + volumes_mounts.append(client.V1VolumeMount(name="dev-shm", mount_path="/dev/shm")) # noqa: S108 + + # Decide whether this is a GPU workload. IMAT jobs run mantid imaging on GPU nodes and need + # the NVIDIA runtime + a GPU resource request so the GPU Operator injects the matching + # userspace driver libraries (libcuda.so.*) into the container. + gpu_job = "imat" in special_pvs + + main_container = client.V1Container( + name=job_name, + image=runner_image, + args=[script], + env=[client.V1EnvVar(name="PYTHONUNBUFFERED", value="1")], + volume_mounts=volumes_mounts, + resources=client.V1ResourceRequirements( + limits={"nvidia.com/gpu": "1"}, + ) + if gpu_job + else None, ) - volumes_mounts.append(client.V1VolumeMount(name="dev-shm", mount_path="/dev/shm")) # noqa: S108 - - # Decide whether this is a GPU workload. IMAT jobs run mantid imaging on GPU nodes and need - # the NVIDIA runtime + a GPU resource request so the GPU Operator injects the matching - # userspace driver libraries (libcuda.so.*) into the container. - gpu_job = "imat" in special_pvs - - main_container = client.V1Container( - name=job_name, - image=runner_image, - args=[script], - env=[client.V1EnvVar(name="PYTHONUNBUFFERED", value="1")], - volume_mounts=volumes_mounts, - resources=client.V1ResourceRequirements( - limits={"nvidia.com/gpu": "1"}, - ) - if gpu_job - else None, - ) - watcher_container = client.V1Container( - name="job-watcher", - image=f"ghcr.io/fiaisis/jobwatcher@sha256:{self.watcher_sha}", - env=[ - client.V1EnvVar(name="FIA_API_HOST", value=fia_api_host), - client.V1EnvVar(name="FIA_API_API_KEY", value=fia_api_api_key), - client.V1EnvVar(name="MAX_TIME_TO_COMPLETE_JOB", value=str(max_time_to_complete_job)), - client.V1EnvVar(name="CONTAINER_NAME", value=job_name), - client.V1EnvVar(name="JOB_NAME", value=job_name), - client.V1EnvVar(name="POD_NAME", value=job_name), - ], - ) + watcher_container = client.V1Container( + name="job-watcher", + image=f"ghcr.io/fiaisis/jobwatcher@sha256:{self.watcher_sha}", + env=[ + client.V1EnvVar(name="FIA_API_HOST", value=fia_api_host), + client.V1EnvVar(name="FIA_API_API_KEY", value=fia_api_api_key), + client.V1EnvVar(name="MAX_TIME_TO_COMPLETE_JOB", value=str(max_time_to_complete_job)), + client.V1EnvVar(name="CONTAINER_NAME", value=job_name), + client.V1EnvVar(name="JOB_NAME", value=job_name), + client.V1EnvVar(name="POD_NAME", value=job_name), + ], + ) - affinity = _generate_affinities(node_affinity_dict=affinity) - tolerations = _generate_tolerations_from_taints(taints) - - pod_spec = client.V1PodSpec( - affinity=affinity, - service_account_name="jobwatcher", - containers=[main_container, watcher_container], - restart_policy="Never", - tolerations=tolerations, - volumes=volumes, - runtime_class_name="nvidia" if gpu_job else None, - ) + affinity = _generate_affinities(node_affinity_dict=affinity) + tolerations = _generate_tolerations_from_taints(taints) + + pod_spec = client.V1PodSpec( + affinity=affinity, + service_account_name="jobwatcher", + containers=[main_container, watcher_container], + restart_policy="Never", + tolerations=tolerations, + volumes=volumes, + runtime_class_name="nvidia" if gpu_job else None, + ) - pod_metadata = client.V1ObjectMeta( - labels={"reduce.isis.cclrc.ac.uk/job-source": "automated-reduction"}, - ) + pod_metadata = client.V1ObjectMeta( + labels={"reduce.isis.cclrc.ac.uk/job-source": "automated-reduction"}, + ) - template = client.V1PodTemplateSpec(spec=pod_spec, metadata=pod_metadata) + template = client.V1PodTemplateSpec(spec=pod_spec, metadata=pod_metadata) - spec = client.V1JobSpec( - template=template, - backoff_limit=0, - ttl_seconds_after_finished=21600, # 6 hours - ) + spec = client.V1JobSpec( + template=template, + backoff_limit=0, + ttl_seconds_after_finished=21600, # 6 hours + ) - job_metadata = client.V1ObjectMeta( - name=job_name, - annotations={ - "job-id": str(job_id), - "pvs": str(pv_names), - "pvcs": str(pvc_names), - "kubectl.kubernetes.io/default-container": main_container.name, - }, - ) + job_metadata = client.V1ObjectMeta( + name=job_name, + annotations={ + "job-id": str(job_id), + "pvs": str(pv_names), + "pvcs": str(pvc_names), + "kubectl.kubernetes.io/default-container": main_container.name, + }, + ) - job = client.V1Job( - api_version="batch/v1", - kind="Job", - metadata=job_metadata, - spec=spec, - ) - client.BatchV1Api().create_namespaced_job(namespace=job_namespace, body=job) + job = client.V1Job( + api_version="batch/v1", + kind="Job", + metadata=job_metadata, + spec=spec, + ) + client.BatchV1Api().create_namespaced_job(namespace=job_namespace, body=job) + + except Exception: + logger.exception("Failed to spawn job %s, cleaning up resources", job_name) + self._cleanup_resources(pv_names, pvc_names, job_namespace) + raise + + def _cleanup_resources(self, pv_names: list[str], pvc_names: list[str], namespace: str) -> None: + """Best-effort cleanup of created K8s resources.""" + for pvc_name in pvc_names: + try: + client.CoreV1Api().delete_namespaced_persistent_volume_claim(name=pvc_name, namespace=namespace) + except client.ApiException: + logger.warning("Failed to cleanup PVC: %s", pvc_name) + for pv_name in pv_names: + try: + client.CoreV1Api().delete_persistent_volume(name=pv_name) + except client.ApiException: + logger.warning("Failed to cleanup PV: %s", pv_name) diff --git a/job_creator/jobcreator/queue_consumer.py b/job_creator/jobcreator/queue_consumer.py index c600839..d60c9b5 100644 --- a/job_creator/jobcreator/queue_consumer.py +++ b/job_creator/jobcreator/queue_consumer.py @@ -6,7 +6,8 @@ import time from collections.abc import Callable -from pika import BlockingConnection, ConnectionParameters, PlainCredentials # type: ignore[import-untyped] +import pika # type: ignore[import-untyped] +from pika import BlockingConnection, ConnectionParameters, PlainCredentials from jobcreator.utils import logger @@ -34,24 +35,35 @@ def __init__( self.channel = None self.connect_to_broker() - def connect_to_broker(self) -> None: + def connect_to_broker(self, max_retries: int = 10) -> None: """ Use this to connect to the broker :return: None """ - self.connection = BlockingConnection(self.connection_parameters) - self.channel = self.connection.channel() # type: ignore[attr-defined] - self.channel.exchange_declare( # type: ignore[attr-defined] - self.queue_name, - exchange_type="direct", - durable=True, - ) - self.channel.queue_declare( # type: ignore[attr-defined] - self.queue_name, - durable=True, - arguments={"x-queue-type": "quorum"}, - ) - self.channel.queue_bind(self.queue_name, self.queue_name, routing_key="") # type: ignore[attr-defined] + for attempt in range(max_retries): + try: + self.connection = BlockingConnection(self.connection_parameters) + self.channel = self.connection.channel() # type: ignore[attr-defined] + self.channel.exchange_declare( # type: ignore[attr-defined] + self.queue_name, + exchange_type="direct", + durable=True, + ) + self.channel.queue_declare( # type: ignore[attr-defined] + self.queue_name, + durable=True, + arguments={"x-queue-type": "quorum"}, + ) + self.channel.queue_bind(self.queue_name, self.queue_name, routing_key="") # type: ignore[attr-defined] + logger.info("Connected to broker") + return + except pika.exceptions.AMQPConnectionError: + wait_time = min(2**attempt, 30) + logger.warning( + "Broker unavailable (attempt %d/%d), retrying in %ds", attempt + 1, max_retries, wait_time + ) + time.sleep(wait_time) + raise RuntimeError(f"Failed to connect to message broker after {max_retries} attempts") def _message_handler(self, msg: str) -> None: """ @@ -77,19 +89,23 @@ def start_consuming(self, callback_func: Callable[[], None], run_once: bool = Fa while run: if run_once: run = False - callback_func() - for header, _, body in self.channel.consume( # type: ignore[attr-defined] - self.queue_name, - inactivity_timeout=5, - ): - try: - self._message_handler(body.decode()) - self.channel.basic_ack(header.delivery_tag) # type: ignore[attr-defined] - except AttributeError: - # If the message frame or body is missing attributes required e.g. the delivery tag - pass - except Exception: - logger.warning("Problem processing message: %s", body) - break - - time.sleep(0.1) + try: + callback_func() + for header, _, body in self.channel.consume( # type: ignore[attr-defined] + self.queue_name, + inactivity_timeout=5, + ): + try: + self._message_handler(body.decode()) + self.channel.basic_ack(header.delivery_tag) # type: ignore[attr-defined] + except AttributeError: + # If the message frame or body is missing attributes required e.g. the delivery tag + pass + except Exception: + logger.warning("Problem processing message: %s", body) + break + time.sleep(0.1) + except (pika.exceptions.AMQPConnectionError, pika.exceptions.ChannelClosedByBroker) as e: + logger.warning("Lost connection to broker: %s. Reconnecting...", e) + time.sleep(5) + self.connect_to_broker() diff --git a/job_creator/jobcreator/utils.py b/job_creator/jobcreator/utils.py index ed3109c..884e2d1 100644 --- a/job_creator/jobcreator/utils.py +++ b/job_creator/jobcreator/utils.py @@ -11,6 +11,7 @@ import requests from kubernetes import config # type: ignore[import-untyped] from kubernetes.config import ConfigException # type: ignore[import-untyped] +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential stdout_handler = logging.StreamHandler(stream=sys.stdout) logging.basicConfig( @@ -59,6 +60,12 @@ def load_kubernetes_config() -> None: config.load_kube_config() +@retry( + retry=retry_if_exception_type(OSError), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=15), + reraise=True, +) def create_ceph_mount_path_simple( user_number: str | None = None, experiment_number: str | None = None, @@ -90,6 +97,12 @@ def create_ceph_mount_path_simple( return Path(mount_path) / ceph_path +@retry( + retry=retry_if_exception_type(OSError), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=15), + reraise=True, +) def ensure_ceph_path_exists_autoreduction(ceph_path: Path) -> Path: """ Takes a path that is intended to be on ceph and ensures that it will be correct for what we should mount and @@ -144,6 +157,12 @@ def get_org_image_name_and_version_from_image_path(image_path: str) -> tuple[str return org_name, image_name, version # Use organisation and image name without ghcr.io +@retry( + retry=retry_if_exception_type((requests.ConnectionError, requests.Timeout)), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=1, max=5), + reraise=True, +) def get_sha256_using_image_from_ghcr(user_image: str, version: str = "") -> str: """ Take the user image and request from the github api the sha256 of the image tag diff --git a/job_creator/pyproject.toml b/job_creator/pyproject.toml index b1cdac3..d7bbcde 100644 --- a/job_creator/pyproject.toml +++ b/job_creator/pyproject.toml @@ -6,6 +6,7 @@ requires-python = ">= 3.11" dependencies = [ "kubernetes==36.0.2", "pika==1.4.1", + "tenacity==9.1.4", ] [project.urls] diff --git a/job_creator/test/test_job_creator.py b/job_creator/test/test_job_creator.py index c5a9397..43d64b4 100644 --- a/job_creator/test/test_job_creator.py +++ b/job_creator/test/test_job_creator.py @@ -2,8 +2,12 @@ from unittest import mock from unittest.mock import call +import pytest +from kubernetes.client.exceptions import ApiException + from jobcreator.job_creator import ( JobCreator, + _is_retryable_k8s_error, _setup_ceph_pv, _setup_extras_pv, _setup_extras_pvc, @@ -669,3 +673,152 @@ def test_jobcreator_spawn_job_dev_mode_false( assert client.V1PersistentVolumeClaimVolumeSource.call_count == 2 # noqa: PLR2004 setup_ceph_pv.assert_not_called() + + +# --- Tests for _is_retryable_k8s_error --- + + +@pytest.mark.parametrize("status_code", [429, 500, 502, 503, 504]) +def test_is_retryable_k8s_error_returns_true_for_transient_status_codes(status_code): + exc = ApiException(status=status_code) + assert _is_retryable_k8s_error(exc) is True + + +@pytest.mark.parametrize("status_code", [400, 401, 403, 404, 409, 422]) +def test_is_retryable_k8s_error_returns_false_for_non_transient_status_codes(status_code): + exc = ApiException(status=status_code) + assert _is_retryable_k8s_error(exc) is False + + +def test_is_retryable_k8s_error_returns_true_for_connection_error(): + assert _is_retryable_k8s_error(ConnectionError("connection refused")) is True + + +def test_is_retryable_k8s_error_returns_true_for_timeout_error(): + assert _is_retryable_k8s_error(TimeoutError("timed out")) is True + + +def test_is_retryable_k8s_error_returns_false_for_unrelated_exception(): + assert _is_retryable_k8s_error(ValueError("something else")) is False + + +def test_is_retryable_k8s_error_returns_false_for_runtime_error(): + assert _is_retryable_k8s_error(RuntimeError("unexpected")) is False + + +# --- Tests for _cleanup_resources --- + + +@mock.patch("jobcreator.job_creator.client") +@mock.patch("jobcreator.job_creator.load_kubernetes_config") +def test_cleanup_resources_deletes_pvcs_and_pvs(mock_load_k8s_config, client): + job_creator = JobCreator("sha", False) + pv_names = ["pv-1", "pv-2"] + pvc_names = ["pvc-1", "pvc-2"] + namespace = "test-ns" + + job_creator._cleanup_resources(pv_names, pvc_names, namespace) + + core_api = client.CoreV1Api.return_value + assert core_api.delete_namespaced_persistent_volume_claim.call_count == 2 # noqa: PLR2004 + core_api.delete_namespaced_persistent_volume_claim.assert_any_call(name="pvc-1", namespace="test-ns") + core_api.delete_namespaced_persistent_volume_claim.assert_any_call(name="pvc-2", namespace="test-ns") + assert core_api.delete_persistent_volume.call_count == 2 # noqa: PLR2004 + core_api.delete_persistent_volume.assert_any_call(name="pv-1") + core_api.delete_persistent_volume.assert_any_call(name="pv-2") + + +@mock.patch("jobcreator.job_creator.logger") +@mock.patch("jobcreator.job_creator.client") +@mock.patch("jobcreator.job_creator.load_kubernetes_config") +def test_cleanup_resources_logs_warning_on_pvc_deletion_failure(mock_load_k8s_config, client, mock_logger): + client.ApiException = ApiException + core_api = client.CoreV1Api.return_value + core_api.delete_namespaced_persistent_volume_claim.side_effect = ApiException(status=404) + + job_creator = JobCreator("sha", False) + job_creator._cleanup_resources(["pv-1"], ["pvc-1"], "test-ns") + + mock_logger.warning.assert_any_call("Failed to cleanup PVC: %s", "pvc-1") + # PV cleanup should still proceed + core_api.delete_persistent_volume.assert_called_once_with(name="pv-1") + + +@mock.patch("jobcreator.job_creator.logger") +@mock.patch("jobcreator.job_creator.client") +@mock.patch("jobcreator.job_creator.load_kubernetes_config") +def test_cleanup_resources_logs_warning_on_pv_deletion_failure(mock_load_k8s_config, client, mock_logger): + client.ApiException = ApiException + core_api = client.CoreV1Api.return_value + core_api.delete_persistent_volume.side_effect = ApiException(status=404) + + job_creator = JobCreator("sha", False) + job_creator._cleanup_resources(["pv-1"], [], "test-ns") + + mock_logger.warning.assert_any_call("Failed to cleanup PV: %s", "pv-1") + + +@mock.patch("jobcreator.job_creator.client") +@mock.patch("jobcreator.job_creator.load_kubernetes_config") +def test_cleanup_resources_handles_empty_lists(mock_load_k8s_config, client): + job_creator = JobCreator("sha", False) + job_creator._cleanup_resources([], [], "test-ns") + + core_api = client.CoreV1Api.return_value + core_api.delete_namespaced_persistent_volume_claim.assert_not_called() + core_api.delete_persistent_volume.assert_not_called() + + +# --- Tests for spawn_job cleanup-on-failure --- + + +@mock.patch("jobcreator.job_creator._setup_extras_pv") +@mock.patch("jobcreator.job_creator._setup_extras_pvc") +@mock.patch("jobcreator.job_creator._setup_smb_pv") +@mock.patch("jobcreator.job_creator._setup_pvc") +@mock.patch("jobcreator.job_creator._setup_ceph_pv") +@mock.patch("jobcreator.job_creator.load_kubernetes_config") +@mock.patch("jobcreator.job_creator.client") +def test_spawn_job_calls_cleanup_and_reraises_on_failure( + client, + _, # noqa: PT019 + setup_ceph_pv, + setup_pvc, + setup_smb_pv, + setup_extras_pvc, + setup_extras_pv, +): + """When job creation fails, cleanup_resources is called and the exception is re-raised.""" + client.BatchV1Api.return_value.create_namespaced_job.side_effect = ApiException(status=500) + job_creator = JobCreator("sha", False) + job_creator._cleanup_resources = mock.MagicMock() + job_name = "test-job" + + with pytest.raises(ApiException): + job_creator.spawn_job( + job_name, + "script", + "namespace", + "ceph-secret", + "ceph-ns", + "cluster-id", + "fs-name", + "/ceph/path", + 1, + 3600, + "fia-api-host", + "api-key", + "runner-image", + "manila-share-id", + "manila-access-id", + [], + [], + None, + ) + + job_creator._cleanup_resources.assert_called_once() + call_args = job_creator._cleanup_resources.call_args + # Verify pv_names and pvc_names lists were passed (non-empty since setup steps succeeded) + assert len(call_args[0][0]) > 0 # pv_names + assert len(call_args[0][1]) > 0 # pvc_names + assert call_args[0][2] == "namespace" # namespace diff --git a/job_creator/test/test_queue_consumer.py b/job_creator/test/test_queue_consumer.py index fe772f9..0be4efb 100644 --- a/job_creator/test/test_queue_consumer.py +++ b/job_creator/test/test_queue_consumer.py @@ -1,5 +1,6 @@ from unittest import mock +import pika.exceptions import pytest from pika import PlainCredentials @@ -12,7 +13,7 @@ QUEUE_NAME = mock.MagicMock() -@pytest.fixture(autouse=True, scope="module") +@pytest.fixture(autouse=True) def setup_queue_consumer(): with ( mock.patch("jobcreator.queue_consumer.ConnectionParameters") as connection_parameters, @@ -107,3 +108,117 @@ def raise_exception(): quc.start_consuming(mock.MagicMock(), run_once=True) logger.warning.assert_called_once_with("Problem processing message: %s", body) + + +# --- Tests for connect_to_broker retry logic --- + + +@mock.patch("jobcreator.queue_consumer.time") +@mock.patch("jobcreator.queue_consumer.ConnectionParameters") +@mock.patch("jobcreator.queue_consumer.BlockingConnection") +def test_connect_to_broker_retries_on_amqp_error_then_succeeds(blocking_connection, mock_conn_params, mock_time): + """connect_to_broker retries on AMQPConnectionError and succeeds on a subsequent attempt.""" + + blocking_connection.side_effect = [ + pika.exceptions.AMQPConnectionError("refused"), + pika.exceptions.AMQPConnectionError("refused"), + mock.MagicMock(), # Third attempt succeeds + ] + + quc = QueueConsumer(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + assert quc.connection is not None + assert quc.channel is not None + assert blocking_connection.call_count == 3 # noqa: PLR2004 + # Verify sleep was called for the two failed attempts + assert mock_time.sleep.call_count >= 2 # noqa: PLR2004 + + +@mock.patch("jobcreator.queue_consumer.time") +@mock.patch("jobcreator.queue_consumer.ConnectionParameters") +@mock.patch("jobcreator.queue_consumer.BlockingConnection") +def test_connect_to_broker_raises_runtime_error_after_max_retries(blocking_connection, mock_conn_params, mock_time): + """connect_to_broker raises RuntimeError after exhausting all retry attempts.""" + + blocking_connection.side_effect = pika.exceptions.AMQPConnectionError("refused") + + with pytest.raises(RuntimeError, match="Failed to connect to message broker after 10 attempts"): + QueueConsumer(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + +@mock.patch("jobcreator.queue_consumer.time") +@mock.patch("jobcreator.queue_consumer.ConnectionParameters") +@mock.patch("jobcreator.queue_consumer.BlockingConnection") +def test_connect_to_broker_uses_exponential_backoff(blocking_connection, mock_conn_params, mock_time): + """connect_to_broker sleeps with exponentially increasing wait times.""" + + blocking_connection.side_effect = pika.exceptions.AMQPConnectionError("refused") + + with pytest.raises(RuntimeError): + QueueConsumer(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + + # Verify exponential backoff: min(2^0, 30)=1, min(2^1, 30)=2, min(2^2, 30)=4, ... + sleep_calls = [c[0][0] for c in mock_time.sleep.call_args_list] + for i, wait in enumerate(sleep_calls): + assert wait == min(2**i, 30) + + +# --- Tests for start_consuming reconnection logic --- + + +@mock.patch("jobcreator.queue_consumer.time") +@mock.patch("jobcreator.queue_consumer.ConnectionParameters") +@mock.patch("jobcreator.queue_consumer.BlockingConnection") +def test_start_consuming_reconnects_on_amqp_connection_error(blocking_connection, mock_conn_params, mock_time): + """start_consuming reconnects to the broker when AMQPConnectionError is raised during consuming.""" + + quc = QueueConsumer(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + quc.connect_to_broker = mock.MagicMock() + + call_count = 0 + + def callback_side_effect(): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise pika.exceptions.AMQPConnectionError("lost connection") + + callback = mock.MagicMock(side_effect=callback_side_effect) + + # run_once=True won't help here since we need the while loop to iterate. + # Instead, have the second callback call succeed and use channel.consume normally. + # The reconnect path sets run=False on second iteration via run_once logic. + # Actually, let's just test that connect_to_broker is called after the error. + # We need to break the loop, so after reconnection, raise KeyboardInterrupt. + quc.connect_to_broker.side_effect = KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + quc.start_consuming(callback, run_once=False) + + quc.connect_to_broker.assert_called_once() + + +@mock.patch("jobcreator.queue_consumer.time") +@mock.patch("jobcreator.queue_consumer.ConnectionParameters") +@mock.patch("jobcreator.queue_consumer.BlockingConnection") +def test_start_consuming_reconnects_on_channel_closed_by_broker(blocking_connection, mock_conn_params, mock_time): + """start_consuming reconnects when ChannelClosedByBroker is raised.""" + + quc = QueueConsumer(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()) + quc.connect_to_broker = mock.MagicMock() + + call_count = 0 + + def callback_side_effect(): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise pika.exceptions.ChannelClosedByBroker(406, "PRECONDITION_FAILED") + + callback = mock.MagicMock(side_effect=callback_side_effect) + quc.connect_to_broker.side_effect = KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + quc.start_consuming(callback, run_once=False) + + quc.connect_to_broker.assert_called_once()