diff --git a/cloud_pipelines_backend/launchers/interfaces.py b/cloud_pipelines_backend/launchers/interfaces.py index 8844f7d..fa333aa 100644 --- a/cloud_pipelines_backend/launchers/interfaces.py +++ b/cloud_pipelines_backend/launchers/interfaces.py @@ -26,7 +26,15 @@ class LauncherError(RuntimeError): - pass + """Base class for errors raised by a launcher. + + `is_retriable` marks a failure worth trying again (e.g. the backing platform + was unavailable) rather than a definitive error. + """ + + def __init__(self, *args: object, is_retriable: bool = False) -> None: + super().__init__(*args) + self.is_retriable = is_retriable @dataclasses.dataclass(kw_only=True) diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 0007c12..88886d8 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -595,7 +595,12 @@ def get_refreshed_launched_container_from_dict( launched_container = LaunchedKubernetesContainer.from_dict( launched_container_dict, launcher=self ) - return launched_container.get_refreshed() + try: + return launched_container.get_refreshed() + except kubernetes.client.exceptions.ApiException as ex: + raise _launcher_error_from_api_exception( + ex, message="Failed to refresh pod status" + ) from ex def deserialize_launched_container_from_dict( self, launched_container_dict: dict @@ -1257,7 +1262,12 @@ def get_refreshed_launched_container_from_dict( launched_container = LaunchedKubernetesJob.from_dict( launched_container_dict, launcher=self ) - return launched_container.get_refreshed() + try: + return launched_container.get_refreshed() + except kubernetes.client.exceptions.ApiException as ex: + raise _launcher_error_from_api_exception( + ex, message="Failed to refresh job status" + ) from ex def deserialize_launched_container_from_dict( self, launched_container_dict: dict @@ -1848,3 +1858,20 @@ def _remove_keys_with_none_values(d: dict): del d[k] if isinstance(v, dict): _remove_keys_with_none_values(v) + + +def _launcher_error_from_api_exception( + exception: kubernetes.client.exceptions.ApiException, + *, + message: str, +) -> interfaces.LauncherError: + """Translate a Kubernetes API error into a launcher error. + + A 5xx is retriable (the API server broke or shed load); any other status is + a definitive failure. + """ + status = exception.status + is_retriable = isinstance(status, int) and 500 <= status < 600 + return interfaces.LauncherError( + f"{message}: {exception!r}", is_retriable=is_retriable + ) diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index 3e53516..4297b7f 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -2,6 +2,7 @@ import json import datetime import logging +import os import time import traceback import typing @@ -52,6 +53,7 @@ def __init__( output_data_purge_duration: datetime.timedelta = None, *, # Internal/experimental: + _max_container_execution_refresh_error_retries: int = 3, _max_queue_batch_size: int = 1, _max_queue_batch_duration: datetime.timedelta = datetime.timedelta(), ): @@ -65,6 +67,10 @@ def __init__( self._queued_executions_queue_idle = False self._running_executions_queue_idle = False self._output_data_purge_duration = output_data_purge_duration + self._max_container_execution_refresh_error_retries = ( + _max_container_execution_refresh_error_retries + ) + self._container_execution_refresh_error_counts: dict[Any, int] = {} self._max_queue_batch_size = _max_queue_batch_size self._max_queue_batch_duration = _max_queue_batch_duration @@ -213,6 +219,38 @@ def internal_process_running_executions_queue(self, session: orm.Session): except Exception as ex: _logger.exception("Error processing running container execution") session.rollback() + + is_retriable = ( + isinstance(ex, launcher_interfaces.LauncherError) + and ex.is_retriable + ) + error_count = ( + self._container_execution_refresh_error_counts.get( + running_container_execution.id, 0 + ) + + 1 + ) + if ( + is_retriable + and error_count + < self._max_container_execution_refresh_error_retries + ): + self._container_execution_refresh_error_counts[ + running_container_execution.id + ] = error_count + _logger.warning( + "Could not refresh running container execution " + f"({error_count} of " + f"{self._max_container_execution_refresh_error_retries}" + f" consecutive retriable failures): {ex!r}. Leaving " + f"it in {running_container_execution.status} and " + "refreshing it on a later sweep." + ) + return True + + self._container_execution_refresh_error_counts.pop( + running_container_execution.id, None + ) running_container_execution.status = ( bts.ContainerExecutionStatus.SYSTEM_ERROR ) @@ -238,6 +276,10 @@ def internal_process_running_executions_queue(self, session: orm.Session): session=session, execution=execution_node ) session.commit() + else: + self._container_execution_refresh_error_counts.pop( + running_container_execution.id, None + ) finally: duration_ms = (time.monotonic_ns() - start_timestamp) / 1_000_000 _logger.info( diff --git a/tests/test_container_execution_refresh_retries.py b/tests/test_container_execution_refresh_retries.py new file mode 100644 index 0000000..9d56414 --- /dev/null +++ b/tests/test_container_execution_refresh_retries.py @@ -0,0 +1,265 @@ +"""Tests for the refresh-failure budget in the running-executions queue.""" + +from typing import Any, Callable +from unittest import mock + +from sqlalchemy import orm +from sqlalchemy import sql + +from cloud_pipelines_backend import api_server_sql +from cloud_pipelines_backend import backend_types_sql as bts +from cloud_pipelines_backend import component_structures as structures +from cloud_pipelines_backend import database_ops +from cloud_pipelines_backend import orchestrator_sql +from cloud_pipelines_backend.launchers import interfaces as launcher_interfaces + + +def _retriable_error() -> launcher_interfaces.LauncherError: + return launcher_interfaces.LauncherError( + "The platform was unavailable", is_retriable=True + ) + + +def _non_retriable_error() -> launcher_interfaces.LauncherError: + return launcher_interfaces.LauncherError("Something went definitively wrong") + + +def _create_session_factory() -> Callable[[], orm.Session]: + db_engine = database_ops.create_db_engine_and_migrate_db(database_uri="sqlite://") + return lambda: orm.Session(bind=db_engine) + + +def _make_launched_container(launcher_data: dict[str, Any]) -> mock.MagicMock: + return mock.MagicMock( + status=launcher_interfaces.ContainerStatus.PENDING, + to_dict=lambda: dict(launcher_data), + ) + + +def _container_task(image: str = "python") -> structures.TaskSpec: + return structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + implementation=structures.ContainerImplementation( + container=structures.ContainerSpec(image=image) + ) + ) + ) + ) + + +def _create_launched_container_executions( + task_count: int = 1, +) -> Callable[[], orm.Session]: + """A pipeline run with ``task_count`` container tasks, launched and PENDING. + + Each launch gets distinct ``launcher_data`` so tests can tell which + execution the orchestrator refreshed. The tasks use distinct images so that + each gets its own ``ContainerExecution`` instead of a cache hit on the + first one. + """ + pipeline_spec = structures.ComponentSpec( + implementation=structures.GraphImplementation( + graph=structures.GraphSpec( + tasks={ + f"task{i}": _container_task(image=f"python:3.{i}") + for i in range(task_count) + } + ) + ), + ) + session_factory = _create_session_factory() + api_server_sql.PipelineRunsApiService_Sql().create( + session=session_factory(), + root_task=structures.TaskSpec( + component_ref=structures.ComponentReference(spec=pipeline_spec) + ), + created_by="user1", + ) + + launch_count = iter(range(task_count)) + launch_orchestrator = orchestrator_sql.OrchestratorService_Sql( + session_factory=session_factory, + launcher=mock.MagicMock( + launch_container_task=mock.MagicMock( + side_effect=lambda *a, **kw: _make_launched_container( + {"pod": f"pod-{next(launch_count)}"} + ) + ) + ), + storage_provider=mock.MagicMock(), + data_root_uri="file:///tmp/artifacts", + logs_root_uri="file:///tmp/logs", + ) + session = session_factory() + for _ in range(20 * task_count): + if not launch_orchestrator.internal_process_queued_executions_queue( + session=session + ): + break + return session_factory + + +def _make_orchestrator( + session_factory: Callable[[], orm.Session], + get_refreshed: Callable[..., Any], + max_failures: int = 3, +) -> orchestrator_sql.OrchestratorService_Sql: + launcher = mock.MagicMock( + deserialize_launched_container_from_dict=mock.MagicMock( + side_effect=_make_launched_container + ), + get_refreshed_launched_container_from_dict=mock.MagicMock( + side_effect=get_refreshed + ), + ) + return orchestrator_sql.OrchestratorService_Sql( + session_factory=session_factory, + launcher=launcher, + storage_provider=mock.MagicMock(), + data_root_uri="file:///tmp/artifacts", + logs_root_uri="file:///tmp/logs", + _max_container_execution_refresh_error_retries=max_failures, + ) + + +def _statuses( + session_factory: Callable[[], orm.Session], +) -> list[bts.ContainerExecutionStatus]: + return list( + session_factory().scalars(sql.select(bts.ContainerExecution.status)).all() + ) + + +def _only_status( + session_factory: Callable[[], orm.Session], +) -> bts.ContainerExecutionStatus: + statuses = _statuses(session_factory) + assert len(statuses) == 1 + return statuses[0] + + +class TestContainerExecutionRefreshRetries: + def test_non_retriable_error_terminalizes_immediately(self) -> None: + session_factory = _create_launched_container_executions() + orchestrator = _make_orchestrator( + session_factory, + mock.MagicMock(side_effect=_non_retriable_error()), + max_failures=3, + ) + + orchestrator.internal_process_running_executions_queue( + session=session_factory() + ) + + assert _only_status(session_factory) == ( + bts.ContainerExecutionStatus.SYSTEM_ERROR + ) + + def test_retriable_errors_below_budget_leave_execution_running(self) -> None: + session_factory = _create_launched_container_executions() + orchestrator = _make_orchestrator( + session_factory, + mock.MagicMock(side_effect=_retriable_error()), + max_failures=3, + ) + session = session_factory() + + for _ in range(2): + orchestrator.internal_process_running_executions_queue(session=session) + assert _only_status(session_factory) == ( + bts.ContainerExecutionStatus.PENDING + ) + + def test_retriable_errors_at_budget_terminalize_execution(self) -> None: + session_factory = _create_launched_container_executions() + orchestrator = _make_orchestrator( + session_factory, + mock.MagicMock(side_effect=_retriable_error()), + max_failures=3, + ) + session = session_factory() + + for _ in range(3): + orchestrator.internal_process_running_executions_queue(session=session) + + assert _only_status(session_factory) == ( + bts.ContainerExecutionStatus.SYSTEM_ERROR + ) + + def test_successful_refresh_resets_the_budget(self) -> None: + session_factory = _create_launched_container_executions() + outcomes = [ + _retriable_error(), + _retriable_error(), + None, # refreshed successfully + _retriable_error(), + _retriable_error(), + ] + + def get_refreshed(launcher_data: dict[str, Any]) -> Any: + outcome = outcomes.pop(0) + if outcome is not None: + raise outcome + return _make_launched_container(launcher_data) + + orchestrator = _make_orchestrator( + session_factory, get_refreshed, max_failures=3 + ) + session = session_factory() + + for _ in range(5): + orchestrator.internal_process_running_executions_queue(session=session) + + # Five sweeps, but never three consecutive failures. + assert _only_status(session_factory) == bts.ContainerExecutionStatus.PENDING + + def test_retry_moves_execution_to_the_back_of_the_queue(self) -> None: + session_factory = _create_launched_container_executions() + orchestrator = _make_orchestrator( + session_factory, + mock.MagicMock(side_effect=_retriable_error()), + max_failures=3, + ) + before = session_factory().scalar( + sql.select(bts.ContainerExecution.last_processed_at) + ) + + orchestrator.internal_process_running_executions_queue( + session=session_factory() + ) + + after = session_factory().scalar( + sql.select(bts.ContainerExecution.last_processed_at) + ) + assert before is not None and after is not None + assert after > before + + def test_failing_execution_does_not_hold_up_the_queue(self) -> None: + """A retriable-failing execution must not be refreshed ahead of others.""" + session_factory = _create_launched_container_executions(task_count=2) + assert len(_statuses(session_factory)) == 2 + refreshed: list[dict[str, Any]] = [] + + def get_refreshed(launcher_data: dict[str, Any]) -> Any: + refreshed.append(dict(launcher_data)) + if len(refreshed) == 1: + raise _retriable_error() + return _make_launched_container(launcher_data) + + orchestrator = _make_orchestrator( + session_factory, get_refreshed, max_failures=3 + ) + session = session_factory() + + orchestrator.internal_process_running_executions_queue(session=session) + orchestrator.internal_process_running_executions_queue(session=session) + + assert len(refreshed) == 2 + # The second sweep moved on to the other execution instead of retrying + # the failed one, so the orchestrator keeps making progress. + assert refreshed[0] != refreshed[1] + assert _statuses(session_factory) == [ + bts.ContainerExecutionStatus.PENDING, + bts.ContainerExecutionStatus.PENDING, + ] diff --git a/tests/test_kubernetes_launcher_error_classification.py b/tests/test_kubernetes_launcher_error_classification.py new file mode 100644 index 0000000..c72beb9 --- /dev/null +++ b/tests/test_kubernetes_launcher_error_classification.py @@ -0,0 +1,37 @@ +"""Tests for translating Kubernetes API errors into launcher errors.""" + +import kubernetes.client.exceptions + +from cloud_pipelines_backend.launchers import interfaces +from cloud_pipelines_backend.launchers import kubernetes_launchers + + +def _api_exception(status: int) -> kubernetes.client.exceptions.ApiException: + return kubernetes.client.exceptions.ApiException(status=status, reason="test") + + +class TestLauncherErrorFromApiException: + def test_server_error_is_retriable(self) -> None: + error = kubernetes_launchers._launcher_error_from_api_exception( + _api_exception(500), message="Failed to refresh pod status" + ) + assert isinstance(error, interfaces.LauncherError) + assert error.is_retriable + + def test_service_unavailable_is_retriable(self) -> None: + error = kubernetes_launchers._launcher_error_from_api_exception( + _api_exception(503), message="Failed to refresh pod status" + ) + assert error.is_retriable + + def test_not_found_is_not_retriable(self) -> None: + error = kubernetes_launchers._launcher_error_from_api_exception( + _api_exception(404), message="Failed to refresh pod status" + ) + assert not error.is_retriable + + def test_client_error_is_not_retriable(self) -> None: + error = kubernetes_launchers._launcher_error_from_api_exception( + _api_exception(403), message="Failed to refresh pod status" + ) + assert not error.is_retriable