diff --git a/cloud_pipelines_backend/launchers/interfaces.py b/cloud_pipelines_backend/launchers/interfaces.py index 2c1309c..bfb5a72 100644 --- a/cloud_pipelines_backend/launchers/interfaces.py +++ b/cloud_pipelines_backend/launchers/interfaces.py @@ -71,6 +71,25 @@ class ContainerStatus(str, enum.Enum): ERROR = "ERROR" +class LauncherResourceNotFound(LauncherError): + def __init__( + self, + *, + kind: str, + namespace: str, + name: str, + cached_status: ContainerStatus, + ): + self.kind = kind + self.namespace = namespace + self.name = name + self.cached_status = cached_status + super().__init__( + f"Kubernetes {kind} {namespace}/{name} was not found while refreshing " + f"cached {cached_status.value} state." + ) + + class LaunchedContainer(abc.ABC): # @classmethod diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 25cd267..3c0d595 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -6,6 +6,7 @@ import logging import os import pathlib +import time import typing from typing import Any, Optional @@ -72,6 +73,21 @@ _T = typing.TypeVar("_T") +_KUBERNETES_DIAGNOSTIC_MAX_ATTEMPTS = 3 +_KUBERNETES_DIAGNOSTIC_RETRY_WAIT_SECONDS = 0.1 + + +def _retry_kubernetes_diagnostic_call(func: typing.Callable[[], _T]) -> _T: + for attempt in range(_KUBERNETES_DIAGNOSTIC_MAX_ATTEMPTS): + try: + return func() + except Exception: + if attempt == _KUBERNETES_DIAGNOSTIC_MAX_ATTEMPTS - 1: + raise + time.sleep(_KUBERNETES_DIAGNOSTIC_RETRY_WAIT_SECONDS) + raise RuntimeError("Unreachable") + + _CONTAINER_FILE_NAME = "data" @@ -910,11 +926,21 @@ def from_dict( def get_refreshed(self) -> "LaunchedKubernetesContainer": launcher = self._get_launcher() core_api_client = k8s_client_lib.CoreV1Api(api_client=launcher._api_client) - pod: k8s_client_lib.V1Pod = core_api_client.read_namespaced_pod( - name=self._pod_name, - namespace=self._namespace, - _request_timeout=launcher._request_timeout, - ) + try: + pod: k8s_client_lib.V1Pod = core_api_client.read_namespaced_pod( + name=self._pod_name, + namespace=self._namespace, + _request_timeout=launcher._request_timeout, + ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status == 404: + raise interfaces.LauncherResourceNotFound( + kind="Pod", + namespace=self._namespace, + name=self._pod_name, + cached_status=self.status, + ) from None + raise new_launched_container = copy.copy(self) new_launched_container._debug_pod = pod return new_launched_container @@ -975,11 +1001,19 @@ def __str__(self) -> str: def terminate(self): launcher = self._get_launcher() core_api_client = k8s_client_lib.CoreV1Api(api_client=launcher._api_client) - core_api_client.delete_namespaced_pod( - name=self._pod_name, - namespace=self._namespace, - grace_period_seconds=10, - ) + try: + core_api_client.delete_namespaced_pod( + name=self._pod_name, + namespace=self._namespace, + grace_period_seconds=10, + ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status != 404: + raise + _logger.info( + f"Pod {self._pod_name} in namespace {self._namespace} was already absent." + ) + return _logger.info(f"Terminated pod {self._pod_name} in namespace {self._namespace}") @@ -1449,28 +1483,50 @@ def from_dict( def get_refreshed(self) -> "LaunchedKubernetesJob": launcher = self._get_launcher() batch_api_client = k8s_client_lib.BatchV1Api(api_client=launcher._api_client) - job: k8s_client_lib.V1Job = batch_api_client.read_namespaced_job( - name=self._job_name, - namespace=self._namespace, - _request_timeout=launcher._request_timeout, - ) - # Refreshing the job pods. We do not strictly need them. - # But this information is useful for debugging and it will also allow slightly better status reporting. - core_api_client = k8s_client_lib.CoreV1Api(launcher._api_client) - pod_list_response: k8s_client_lib.V1PodList = ( - core_api_client.list_namespaced_pod( + try: + job: k8s_client_lib.V1Job = batch_api_client.read_namespaced_job( + name=self._job_name, namespace=self._namespace, - label_selector=f"job-name={self._job_name}", - watch=False, _request_timeout=launcher._request_timeout, ) - ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status == 404: + raise interfaces.LauncherResourceNotFound( + kind="Job", + namespace=self._namespace, + name=self._job_name, + cached_status=self.status, + ) from None + raise + + new_launched_container = copy.copy(self) + new_launched_container._debug_job = job + + # Child Pods are diagnostics. Once the parent Job is terminal, failures to + # collect them must not replace the Job's authoritative terminal state. + core_api_client = k8s_client_lib.CoreV1Api(launcher._api_client) + try: + pod_list_response: k8s_client_lib.V1PodList = ( + _retry_kubernetes_diagnostic_call( + lambda: core_api_client.list_namespaced_pod( + namespace=self._namespace, + label_selector=f"job-name={self._job_name}", + watch=False, + _request_timeout=launcher._request_timeout, + ) + ) + ) + except Exception: + if not new_launched_container.has_ended: + raise + _logger.exception( + f"Unable to collect child Pod diagnostics for terminal Job {self._job_name} in namespace {self._namespace}." + ) + return new_launched_container + pod_map: dict[str, k8s_client_lib.V1Pod] = {} if job.spec.completion_mode == "Indexed": - # There can be multiple pods for each index. - # This can happen when Job gets suspended and resumed multiple times and Pods can get stuck in "Terminating" phase. - # In this case we want to get the latest Pods. - # Alternatively, we could store all pods, but this can complicate the routines that determine status and get logs. + # A suspended and resumed indexed Job can have multiple Pods per index. pods_sorted_by_creation_time = sorted( pod_list_response.items, key=lambda pod: pod.metadata.creation_timestamp or "", @@ -1492,8 +1548,7 @@ def get_refreshed(self) -> "LaunchedKubernetesJob": for pod in pod_list_response.items: index_str: str = pod.metadata.name pod_map[index_str] = pod - new_launched_container = copy.copy(self) - new_launched_container._debug_job = job + new_launched_container._debug_pods = pod_map return new_launched_container @@ -1531,7 +1586,22 @@ def _get_log_by_pod_key(self, pod_name: str) -> str | None: def _get_all_logs(self) -> dict[str, str]: logs = {} for pod_key, pod in self._debug_pods.items(): - log = self._get_log_by_pod_key(pod.metadata.name) + if not pod.metadata or not pod.metadata.name: + raise ValueError( + f"Child Pod {pod_key} of Job {self._job_name} does not have a name." + ) + pod_name = pod.metadata.name + try: + log = _retry_kubernetes_diagnostic_call( + lambda: self._get_log_by_pod_key(pod_name) + ) + except Exception: + if not self.has_ended: + raise + _logger.exception( + f"Unable to collect log diagnostics for Pod {pod_name} of terminal Job {self._job_name}." + ) + continue if log: logs[pod_key] = log return logs @@ -1565,6 +1635,11 @@ def get_log(self) -> str: def upload_log(self): all_logs = self._get_all_logs() + if not all_logs: + _logger.info( + f"No child Pod logs were available for Job {self._job_name}; preserving any existing uploaded log." + ) + return merged_log = self._merge_logs(all_logs) # Uploading the merged log @@ -1610,12 +1685,20 @@ def __str__(self) -> str: def terminate(self): launcher = self._get_launcher() batch_api_client = k8s_client_lib.BatchV1Api(api_client=launcher._api_client) - batch_api_client.delete_namespaced_job( - name=self._job_name, - namespace=self._namespace, - grace_period_seconds=10, - propagation_policy="Foreground", - ) + try: + batch_api_client.delete_namespaced_job( + name=self._job_name, + namespace=self._namespace, + grace_period_seconds=10, + propagation_policy="Foreground", + ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status != 404: + raise + _logger.info( + f"Job {self._job_name} in namespace {self._namespace} was already absent." + ) + return _logger.info(f"Terminated job {self._job_name} in namespace {self._namespace}") diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index 1a81f45..0db7b21 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -212,8 +212,13 @@ def internal_process_running_executions_queue(self, session: orm.Session): # Doing an intermediate commit here because it's most important to mark the problematic execution as SYSTEM_ERROR. session.commit() - # Mark our ExecutionNode as SYSTEM_ERROR - execution_nodes = running_container_execution.execution_nodes + # Mark our active ExecutionNodes as SYSTEM_ERROR. + execution_nodes = [ + execution_node + for execution_node in running_container_execution.execution_nodes + if execution_node.container_execution_status + != bts.ContainerExecutionStatus.CANCELLED + ] for execution_node in execution_nodes: execution_node.container_execution_status = ( bts.ContainerExecutionStatus.SYSTEM_ERROR @@ -689,70 +694,77 @@ def internal_process_one_running_execution( launched_container: launcher_interfaces.LaunchedContainer = ( self._launcher.deserialize_launched_container_from_dict(launcher_data) ) - previous_status = launched_container.status + previous_status = launcher_interfaces.ContainerStatus( + container_execution.status.value + ) if previous_status not in ( - bts.ContainerExecutionStatus.PENDING, - bts.ContainerExecutionStatus.RUNNING, + launcher_interfaces.ContainerStatus.PENDING, + launcher_interfaces.ContainerStatus.RUNNING, ): raise OrchestratorError( f"Unexpected running container status: {previous_status=}, {launched_container=}" ) - # Handling cancellation - votes_to_terminate = [] - votes_to_not_terminate = [] - # TODO: Get the desired state from the pipeline runs, not execution nodes - execution_nodes = container_execution.execution_nodes - for execution_node in execution_nodes: - should_terminate = (execution_node.extra_data or {}).get( - "desired_state" - ) == "TERMINATED" - if should_terminate: - votes_to_terminate.append(execution_node) - else: - votes_to_not_terminate.append(execution_node) - - if votes_to_terminate: - # Terminate the container execution only when all execution nodes pointing to it are asked to terminate. - terminated = False - if votes_to_not_terminate: - _logger.info( - f"Not terminating container execution since some other executions ({[execution_node.id for execution_node in votes_to_not_terminate]}) are still using it." - ) - else: - _logger.info("Terminating container execution.") - # We should preserve the logs before terminating/deleting the container - try: - _retry(lambda: launched_container.upload_log()) - except Exception: - _logger.exception("Error uploading logs before termination.") - # Requesting container termination. - # Termination might not happen immediately (e.g. Kubernetes has grace period). - launched_container.terminate() - container_execution.ended_at = _get_current_time() - # We need to mark the execution as CANCELLED otherwise orchestrator will continue polling it. - container_execution.status = bts.ContainerExecutionStatus.CANCELLED - terminated = True - - # Mark the execution nodes as cancelled only after the launched container is successfully terminated (if needed) - for execution_node in votes_to_terminate: - _logger.info( - f"Cancelling execution {execution_node.id} and skipping all downstream executions." - ) - execution_node.container_execution_status = ( - bts.ContainerExecutionStatus.CANCELLED - ) - _mark_all_downstream_executions_as_skipped( - session=session, execution=execution_node - ) - session.commit() - if terminated: - return - - # Asking the launcher to refresh the container state. - reloaded_launched_container: launcher_interfaces.LaunchedContainer = ( - self._launcher.get_refreshed_launched_container_from_dict(launcher_data) + cached_snapshot_is_terminal = launched_container.status in ( + launcher_interfaces.ContainerStatus.SUCCEEDED, + launcher_interfaces.ContainerStatus.FAILED, ) + + if not cached_snapshot_is_terminal: + # Handling cancellation + votes_to_terminate = [] + votes_to_not_terminate = [] + # TODO: Get the desired state from the pipeline runs, not execution nodes + execution_nodes = container_execution.execution_nodes + for execution_node in execution_nodes: + should_terminate = (execution_node.extra_data or {}).get( + "desired_state" + ) == "TERMINATED" + if should_terminate: + votes_to_terminate.append(execution_node) + else: + votes_to_not_terminate.append(execution_node) + + if votes_to_terminate: + # Terminate the container execution only when all execution nodes pointing to it are asked to terminate. + terminated = False + if votes_to_not_terminate: + _logger.info( + f"Not terminating container execution since some other executions ({[execution_node.id for execution_node in votes_to_not_terminate]}) are still using it." + ) + else: + _logger.info("Terminating container execution.") + # We should preserve the logs before terminating/deleting the container + _upload_log_best_effort(launched_container) + # Requesting container termination. + # Termination might not happen immediately (e.g. Kubernetes has grace period). + launched_container.terminate() + container_execution.ended_at = _get_current_time() + # We need to mark the execution as CANCELLED otherwise orchestrator will continue polling it. + container_execution.status = bts.ContainerExecutionStatus.CANCELLED + terminated = True + + # Mark the execution nodes as cancelled only after the launched container is successfully terminated (if needed) + for execution_node in votes_to_terminate: + _logger.info( + f"Cancelling execution {execution_node.id} and skipping all downstream executions." + ) + execution_node.container_execution_status = ( + bts.ContainerExecutionStatus.CANCELLED + ) + _mark_all_downstream_executions_as_skipped( + session=session, execution=execution_node + ) + session.commit() + if terminated: + return + + # Asking the launcher to refresh the container state. + reloaded_launched_container: launcher_interfaces.LaunchedContainer = ( + self._launcher.get_refreshed_launched_container_from_dict(launcher_data) + ) + else: + reloaded_launched_container = launched_container current_time = _get_current_time() # Saving the updated launcher data reloaded_launcher_data = reloaded_launched_container.to_dict() @@ -772,7 +784,12 @@ def internal_process_one_running_execution( ) session.rollback() container_execution.updated_at = current_time - execution_nodes = container_execution.execution_nodes + execution_nodes = [ + execution_node + for execution_node in container_execution.execution_nodes + if execution_node.container_execution_status + != bts.ContainerExecutionStatus.CANCELLED + ] if not execution_nodes: raise OrchestratorError( "Could not find ExecutionNode associated with ContainerExecution." @@ -796,14 +813,8 @@ def internal_process_one_running_execution( container_execution.started_at = reloaded_launched_container.started_at container_execution.ended_at = reloaded_launched_container.ended_at - # Don't fail the execution if log upload fails. - # Logs are important, but not so important that we should fail a successfully completed container execution. - try: - _retry(reloaded_launched_container.upload_log) - except Exception as ex: - _logger.exception( - f"! Error during `LaunchedContainer.upload_log` call: {ex}." - ) + # Logs are useful diagnostics, but do not determine execution success. + _upload_log_best_effort(reloaded_launched_container) _MAX_PRELOAD_VALUE_SIZE = 255 @@ -939,7 +950,7 @@ def _maybe_preload_value( message=orchestration_error_message, ) - _retry(reloaded_launched_container.upload_log) + _upload_log_best_effort(reloaded_launched_container) # Skip downstream executions for execution_node in execution_nodes: execution_node.container_execution_status = ( @@ -1092,6 +1103,15 @@ def _retry( raise +def _upload_log_best_effort( + launched_container: launcher_interfaces.LaunchedContainer, +) -> None: + try: + _retry(launched_container.upload_log) + except Exception: + _logger.exception("Error during `LaunchedContainer.upload_log` call.") + + def record_system_error_exception(execution: bts.ExecutionNode, exception: Exception): app_metrics.execution_system_errors.add(1) bugsnag_instrumentation.notify( diff --git a/tests/test_kubernetes_terminal_recovery.py b/tests/test_kubernetes_terminal_recovery.py new file mode 100644 index 0000000..1e8f4a5 --- /dev/null +++ b/tests/test_kubernetes_terminal_recovery.py @@ -0,0 +1,717 @@ +from __future__ import annotations + +import datetime +from types import SimpleNamespace +from unittest import mock + +import pytest +from kubernetes import client as k8s_client_lib +from kubernetes.client import exceptions as k8s_exceptions +from sqlalchemy import orm + +from cloud_pipelines_backend import backend_types_sql as bts +from cloud_pipelines_backend import database_ops +from cloud_pipelines_backend import orchestrator_sql +from cloud_pipelines_backend.launchers import interfaces +from cloud_pipelines_backend.launchers import kubernetes_launchers + +_NAMESPACE = "terminal-recovery" +_RESOURCE_NAME = "execution-resource" +_LOG_URI = "memory://execution.log" +_NOW = datetime.datetime.now(datetime.timezone.utc) + + +def _make_pod(status: interfaces.ContainerStatus, *, name: str = _RESOURCE_NAME): + phase = { + interfaces.ContainerStatus.PENDING: "Pending", + interfaces.ContainerStatus.RUNNING: "Running", + interfaces.ContainerStatus.SUCCEEDED: "Succeeded", + interfaces.ContainerStatus.FAILED: "Failed", + }[status] + return k8s_client_lib.V1Pod( + metadata=k8s_client_lib.V1ObjectMeta(name=name, namespace=_NAMESPACE), + spec=k8s_client_lib.V1PodSpec( + containers=[k8s_client_lib.V1Container(name="main")], + restart_policy="Never", + ), + status=k8s_client_lib.V1PodStatus(phase=phase), + ) + + +def _make_job(status: interfaces.ContainerStatus): + conditions = [] + if status == interfaces.ContainerStatus.SUCCEEDED: + conditions.append( + k8s_client_lib.V1JobCondition( + type="Complete", status="True", last_transition_time=_NOW + ) + ) + elif status == interfaces.ContainerStatus.FAILED: + conditions.append( + k8s_client_lib.V1JobCondition( + type="Failed", status="True", last_transition_time=_NOW + ) + ) + return k8s_client_lib.V1Job( + metadata=k8s_client_lib.V1ObjectMeta(name=_RESOURCE_NAME, namespace=_NAMESPACE), + spec=k8s_client_lib.V1JobSpec( + completions=1, + completion_mode="Indexed", + template=k8s_client_lib.V1PodTemplateSpec( + spec=k8s_client_lib.V1PodSpec( + containers=[k8s_client_lib.V1Container(name="main")], + restart_policy="Never", + ) + ), + ), + status=k8s_client_lib.V1JobStatus( + conditions=conditions, + active=(1 if status == interfaces.ContainerStatus.RUNNING else None), + ), + ) + + +def _make_kubernetes_resource( + kind: str, + status: interfaces.ContainerStatus, + *, + storage_provider=None, +): + launcher = SimpleNamespace( + _api_client=object(), + _request_timeout=10, + _storage_provider=storage_provider or mock.MagicMock(), + ) + if kind == "Pod": + return kubernetes_launchers.LaunchedKubernetesContainer( + pod_name=_RESOURCE_NAME, + namespace=_NAMESPACE, + output_uris={}, + log_uri=_LOG_URI, + debug_pod=_make_pod(status), + launcher=launcher, + ) + return kubernetes_launchers.LaunchedKubernetesJob( + job_name=_RESOURCE_NAME, + namespace=_NAMESPACE, + output_uris={}, + log_uri=_LOG_URI, + debug_job=_make_job(status), + debug_pods={}, + launcher=launcher, + ) + + +def _install_kubernetes_api_mocks(monkeypatch, *, core_api=None, batch_api=None): + core_api = core_api or mock.MagicMock() + batch_api = batch_api or mock.MagicMock() + monkeypatch.setattr(k8s_client_lib, "CoreV1Api", lambda *args, **kwargs: core_api) + monkeypatch.setattr(k8s_client_lib, "BatchV1Api", lambda *args, **kwargs: batch_api) + return core_api, batch_api + + +@pytest.mark.parametrize( + ("kind", "cached_status"), + [ + ("Pod", interfaces.ContainerStatus.RUNNING), + ("Pod", interfaces.ContainerStatus.SUCCEEDED), + ("Job", interfaces.ContainerStatus.RUNNING), + ("Job", interfaces.ContainerStatus.FAILED), + ], +) +def test_authoritative_get_404_raises_typed_launcher_error( + monkeypatch, kind, cached_status +): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + api_error = k8s_exceptions.ApiException(status=404, reason="Not Found") + api_error.body = "sensitive response body" + if kind == "Pod": + core_api.read_namespaced_pod.side_effect = api_error + else: + batch_api.read_namespaced_job.side_effect = api_error + + resource = _make_kubernetes_resource(kind, cached_status) + with pytest.raises(interfaces.LauncherResourceNotFound) as exc_info: + resource.get_refreshed() + + error = exc_info.value + assert error.kind == kind + assert error.namespace == _NAMESPACE + assert error.name == _RESOURCE_NAME + assert error.cached_status == cached_status + assert str(error) + assert "sensitive response body" not in str(error) + assert error.__cause__ is None + assert error.__suppress_context__ + + +@pytest.mark.parametrize("kind", ["Pod", "Job"]) +def test_authoritative_get_non_404_is_not_reclassified(monkeypatch, kind): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + api_error = k8s_exceptions.ApiException(status=500, reason="Server Error") + if kind == "Pod": + core_api.read_namespaced_pod.side_effect = api_error + else: + batch_api.read_namespaced_job.side_effect = api_error + + with pytest.raises(k8s_exceptions.ApiException) as exc_info: + _make_kubernetes_resource( + kind, interfaces.ContainerStatus.RUNNING + ).get_refreshed() + + assert exc_info.value is api_error + + +@pytest.mark.parametrize( + ("parent_status", "child_result"), + [ + (interfaces.ContainerStatus.SUCCEEDED, []), + (interfaces.ContainerStatus.FAILED, []), + ( + interfaces.ContainerStatus.SUCCEEDED, + k8s_exceptions.ApiException(status=404, reason="Not Found"), + ), + ( + interfaces.ContainerStatus.FAILED, + k8s_exceptions.ApiException(status=404, reason="Not Found"), + ), + ( + interfaces.ContainerStatus.SUCCEEDED, + k8s_exceptions.ApiException(status=500, reason="Server Error"), + ), + ( + interfaces.ContainerStatus.FAILED, + k8s_exceptions.ApiException(status=500, reason="Server Error"), + ), + ], +) +def test_terminal_job_parent_truth_survives_child_list_diagnostics( + monkeypatch, parent_status, child_result +): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + batch_api.read_namespaced_job.return_value = _make_job(parent_status) + if isinstance(child_result, Exception): + core_api.list_namespaced_pod.side_effect = child_result + else: + core_api.list_namespaced_pod.return_value = k8s_client_lib.V1PodList( + items=child_result + ) + + refreshed = _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.RUNNING + ).get_refreshed() + + assert refreshed.status == parent_status + assert refreshed.exit_code == ( + 0 if parent_status == interfaces.ContainerStatus.SUCCEEDED else None + ) + assert refreshed.launcher_error_message is None + + +def test_terminal_job_retries_child_list_diagnostics(monkeypatch): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + batch_api.read_namespaced_job.return_value = _make_job( + interfaces.ContainerStatus.SUCCEEDED + ) + core_api.list_namespaced_pod.side_effect = [ + k8s_exceptions.ApiException(status=500, reason="Server Error"), + k8s_client_lib.V1PodList(items=[]), + ] + + refreshed = _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.RUNNING + ).get_refreshed() + + assert refreshed.status == interfaces.ContainerStatus.SUCCEEDED + assert core_api.list_namespaced_pod.call_count == 2 + + +def test_nonterminal_job_child_list_failure_propagates(monkeypatch): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + batch_api.read_namespaced_job.return_value = _make_job( + interfaces.ContainerStatus.RUNNING + ) + api_error = k8s_exceptions.ApiException(status=404, reason="Not Found") + core_api.list_namespaced_pod.side_effect = api_error + + with pytest.raises(k8s_exceptions.ApiException) as exc_info: + _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.RUNNING + ).get_refreshed() + + assert exc_info.value is api_error + + +class _MemoryStorageProvider: + def __init__(self): + self.uploads: dict[str, list[str]] = {} + + def make_uri(self, uri): + uploads = self.uploads + + class _Writer: + def upload_from_text(self, text): + uploads.setdefault(uri, []).append(text) + + return SimpleNamespace(get_writer=lambda: _Writer()) + + +def _log_response(text: str): + return SimpleNamespace(data=text.encode("utf-8"), release_conn=mock.MagicMock()) + + +def test_terminal_job_retries_child_log_diagnostics(monkeypatch): + storage = _MemoryStorageProvider() + resource = _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.SUCCEEDED, storage_provider=storage + ) + resource._debug_pods = { + "0": _make_pod(interfaces.ContainerStatus.SUCCEEDED, name="pod-0") + } + core_api, _ = _install_kubernetes_api_mocks(monkeypatch) + core_api.read_namespaced_pod_log.side_effect = [ + k8s_exceptions.ApiException(status=500, reason="Server Error"), + _log_response("2026-01-01T00:00:00Z completed\n"), + ] + + resource.upload_log() + + assert core_api.read_namespaced_pod_log.call_count == 2 + assert storage.uploads[_LOG_URI] == ["2026-01-01T00:00:00Z completed\n"] + + +def test_terminal_job_does_not_suppress_invalid_cached_child_metadata(monkeypatch): + resource = _make_kubernetes_resource("Job", interfaces.ContainerStatus.SUCCEEDED) + malformed_pod = _make_pod(interfaces.ContainerStatus.SUCCEEDED) + malformed_pod.metadata = None + resource._debug_pods = {"0": malformed_pod} + core_api, _ = _install_kubernetes_api_mocks(monkeypatch) + + with pytest.raises(ValueError, match="does not have a name"): + resource.upload_log() + + core_api.read_namespaced_pod_log.assert_not_called() + + +def test_terminal_job_uploads_available_child_logs_when_another_log_fails( + monkeypatch, +): + storage = _MemoryStorageProvider() + resource = _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.FAILED, storage_provider=storage + ) + resource._debug_pods = { + "0": _make_pod(interfaces.ContainerStatus.FAILED, name="pod-0"), + "1": _make_pod(interfaces.ContainerStatus.SUCCEEDED, name="pod-1"), + } + core_api, _ = _install_kubernetes_api_mocks(monkeypatch) + + def read_log(*, name, **kwargs): + if name == "pod-0": + raise k8s_exceptions.ApiException(status=500, reason="Server Error") + return _log_response("2026-01-01T00:00:00Z completed\n") + + core_api.read_namespaced_pod_log.side_effect = read_log + + resource.upload_log() + + assert storage.uploads[_LOG_URI] == ["2026-01-01T00:00:00Z completed\n"] + assert storage.uploads[f"{_LOG_URI}.1"] == ["2026-01-01T00:00:00Z completed\n"] + + +@pytest.mark.parametrize("with_child", [False, True]) +def test_job_upload_does_not_overwrite_log_when_no_child_log_is_read( + monkeypatch, with_child +): + storage = _MemoryStorageProvider() + resource = _make_kubernetes_resource( + "Job", interfaces.ContainerStatus.FAILED, storage_provider=storage + ) + if with_child: + resource._debug_pods = { + "0": _make_pod(interfaces.ContainerStatus.FAILED, name="pod-0") + } + core_api, _ = _install_kubernetes_api_mocks(monkeypatch) + core_api.read_namespaced_pod_log.side_effect = k8s_exceptions.ApiException( + status=404, reason="Not Found" + ) + + resource.upload_log() + + assert storage.uploads == {} + + +@pytest.mark.parametrize("kind", ["Pod", "Job"]) +def test_delete_404_is_idempotent_but_other_errors_propagate(monkeypatch, kind): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + delete = ( + core_api.delete_namespaced_pod + if kind == "Pod" + else batch_api.delete_namespaced_job + ) + resource = _make_kubernetes_resource(kind, interfaces.ContainerStatus.RUNNING) + + delete.side_effect = k8s_exceptions.ApiException(status=404, reason="Not Found") + resource.terminate() + + delete.side_effect = k8s_exceptions.ApiException(status=403, reason="Forbidden") + with pytest.raises(k8s_exceptions.ApiException): + resource.terminate() + + +class _FakeLaunchedContainer: + def __init__( + self, + status: interfaces.ContainerStatus, + *, + upload_error: Exception | None = None, + ): + self.status = status + self.exit_code = 0 if status == interfaces.ContainerStatus.SUCCEEDED else None + self.started_at = _NOW + self.ended_at = ( + _NOW + if status + in { + interfaces.ContainerStatus.SUCCEEDED, + interfaces.ContainerStatus.FAILED, + } + else None + ) + self.launcher_error_message = None + self.upload_error = upload_error + self.upload_calls = 0 + self.terminate_calls = 0 + + def to_dict(self): + return {"fake": {"status": self.status.value}} + + def upload_log(self): + self.upload_calls += 1 + if self.upload_error: + raise self.upload_error + + def terminate(self): + self.terminate_calls += 1 + + +class _FakeLauncher: + def __init__(self, cached, *, refreshed=None, refresh_error=None): + self.cached = cached + self.refreshed = refreshed or cached + self.refresh_error = refresh_error + self.refresh_calls = 0 + + def deserialize_launched_container_from_dict(self, launcher_data): + return self.cached + + def get_refreshed_launched_container_from_dict(self, launcher_data): + self.refresh_calls += 1 + if self.refresh_error: + raise self.refresh_error + return self.refreshed + + +class _ResourceLauncher(_FakeLauncher): + def get_refreshed_launched_container_from_dict(self, launcher_data): + self.refresh_calls += 1 + return self.cached.get_refreshed() + + +def _new_session(): + engine = database_ops.create_db_engine_and_migrate_db( + database_uri="sqlite://", do_skip_backfill=True + ) + return orm.Session(engine) + + +def _persist_running_execution( + session, + snapshot, + *, + db_status=bts.ContainerExecutionStatus.RUNNING, + desired_states=(None,), + output_artifact_data_map=None, +): + container_execution = bts.ContainerExecution( + status=db_status, + last_processed_at=_NOW, + created_at=_NOW, + launcher_data=snapshot.to_dict(), + input_artifact_data_map={}, + output_artifact_data_map=output_artifact_data_map or {}, + log_uri=_LOG_URI, + ) + execution_nodes = [] + for index, desired_state in enumerate(desired_states): + extra_data = ( + {"desired_state": desired_state} if desired_state is not None else None + ) + execution_nodes.append( + bts.ExecutionNode( + task_spec={}, + container_execution_status=db_status, + container_execution_cache_key=f"cache-{index}", + container_execution=container_execution, + extra_data=extra_data, + ) + ) + session.add_all(execution_nodes) + session.commit() + return container_execution.id, [node.id for node in execution_nodes] + + +def _make_orchestrator(session, launcher, *, storage_provider=None): + return orchestrator_sql.OrchestratorService_Sql( + session_factory=lambda: session, + launcher=launcher, + storage_provider=storage_provider or mock.MagicMock(), + data_root_uri="memory://artifacts", + logs_root_uri="memory://logs", + ) + + +def _process_running(session, launcher, *, storage_provider=None): + _make_orchestrator( + session, launcher, storage_provider=storage_provider + ).internal_process_running_executions_queue(session) + + +@pytest.mark.parametrize( + "db_status", + [ + bts.ContainerExecutionStatus.PENDING, + bts.ContainerExecutionStatus.RUNNING, + ], +) +def test_cached_terminal_snapshot_resumes_without_refresh(db_status): + session = _new_session() + cached = _FakeLaunchedContainer(interfaces.ContainerStatus.SUCCEEDED) + launcher = _FakeLauncher(cached, refresh_error=AssertionError("must not refresh")) + execution_id, node_ids = _persist_running_execution( + session, cached, db_status=db_status + ) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.SUCCEEDED + ) + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + bts.ContainerExecutionStatus.SUCCEEDED + ) + assert launcher.refresh_calls == 0 + + +@pytest.mark.parametrize( + ("kind", "terminal_status", "expected_status"), + [ + ( + "Pod", + interfaces.ContainerStatus.SUCCEEDED, + bts.ContainerExecutionStatus.SUCCEEDED, + ), + ( + "Job", + interfaces.ContainerStatus.FAILED, + bts.ContainerExecutionStatus.FAILED, + ), + ], +) +def test_cached_terminal_snapshot_wins_over_later_cancellation( + monkeypatch, kind, terminal_status, expected_status +): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + resource = _make_kubernetes_resource(kind, terminal_status) + resource.upload_log = mock.MagicMock() + launcher = _ResourceLauncher(resource) + session = _new_session() + execution_id, node_ids = _persist_running_execution( + session, resource, desired_states=("TERMINATED",) + ) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == expected_status + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + expected_status + ) + assert launcher.refresh_calls == 0 + core_api.delete_namespaced_pod.assert_not_called() + batch_api.delete_namespaced_job.assert_not_called() + + +@pytest.mark.parametrize("kind", ["Pod", "Job"]) +def test_nonterminal_cancellation_precedes_refresh_and_delete_404_is_idempotent( + monkeypatch, kind +): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + delete = ( + core_api.delete_namespaced_pod + if kind == "Pod" + else batch_api.delete_namespaced_job + ) + delete.side_effect = k8s_exceptions.ApiException(status=404, reason="Not Found") + resource = _make_kubernetes_resource(kind, interfaces.ContainerStatus.RUNNING) + resource.upload_log = mock.MagicMock() + launcher = _ResourceLauncher(resource) + session = _new_session() + execution_id, node_ids = _persist_running_execution( + session, resource, desired_states=("TERMINATED",) + ) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.CANCELLED + ) + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + bts.ContainerExecutionStatus.CANCELLED + ) + assert launcher.refresh_calls == 0 + delete.assert_called_once() + + +@pytest.mark.parametrize("kind", ["Pod", "Job"]) +def test_nonterminal_authoritative_404_becomes_system_error(monkeypatch, kind): + core_api, batch_api = _install_kubernetes_api_mocks(monkeypatch) + api_error = k8s_exceptions.ApiException(status=404, reason="Not Found") + api_error.body = "sensitive response body" + if kind == "Pod": + core_api.read_namespaced_pod.side_effect = api_error + else: + batch_api.read_namespaced_job.side_effect = api_error + resource = _make_kubernetes_resource(kind, interfaces.ContainerStatus.RUNNING) + launcher = _ResourceLauncher(resource) + session = _new_session() + execution_id, node_ids = _persist_running_execution(session, resource) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.SYSTEM_ERROR + ) + node = session.get(bts.ExecutionNode, node_ids[0]) + assert node.container_execution_status == bts.ContainerExecutionStatus.SYSTEM_ERROR + assert ( + "LauncherResourceNotFound" + in node.extra_data[ + bts.EXECUTION_NODE_EXTRA_DATA_SYSTEM_ERROR_EXCEPTION_MESSAGE_KEY + ] + ) + assert ( + "sensitive response body" + not in node.extra_data[ + bts.EXECUTION_NODE_EXTRA_DATA_SYSTEM_ERROR_EXCEPTION_FULL_KEY + ] + ) + + +def test_db_status_drives_pending_to_running_transition(): + session = _new_session() + cached = _FakeLaunchedContainer(interfaces.ContainerStatus.RUNNING) + refreshed = _FakeLaunchedContainer(interfaces.ContainerStatus.RUNNING) + launcher = _FakeLauncher(cached, refreshed=refreshed) + execution_id, node_ids = _persist_running_execution( + session, cached, db_status=bts.ContainerExecutionStatus.PENDING + ) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.RUNNING + ) + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + bts.ContainerExecutionStatus.RUNNING + ) + + +@pytest.mark.parametrize("refresh_fails", [False, True]) +def test_shared_cancellation_is_not_overwritten_by_refresh(refresh_fails): + session = _new_session() + cached = _FakeLaunchedContainer(interfaces.ContainerStatus.RUNNING) + refreshed = _FakeLaunchedContainer(interfaces.ContainerStatus.SUCCEEDED) + refresh_error = None + if refresh_fails: + refresh_error = interfaces.LauncherResourceNotFound( + kind="Job", + namespace=_NAMESPACE, + name=_RESOURCE_NAME, + cached_status=interfaces.ContainerStatus.RUNNING, + ) + launcher = _FakeLauncher(cached, refreshed=refreshed, refresh_error=refresh_error) + execution_id, node_ids = _persist_running_execution( + session, cached, desired_states=("TERMINATED", None) + ) + + _process_running(session, launcher) + + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + bts.ContainerExecutionStatus.CANCELLED + ) + expected_active_status = ( + bts.ContainerExecutionStatus.SYSTEM_ERROR + if refresh_fails + else bts.ContainerExecutionStatus.SUCCEEDED + ) + assert session.get(bts.ExecutionNode, node_ids[1]).container_execution_status == ( + expected_active_status + ) + assert session.get(bts.ContainerExecution, execution_id).status == ( + expected_active_status + ) + + +def test_failed_log_upload_remains_failed(monkeypatch): + monkeypatch.setattr( + orchestrator_sql, "_retry", lambda function, **kwargs: function() + ) + session = _new_session() + cached = _FakeLaunchedContainer( + interfaces.ContainerStatus.FAILED, + upload_error=RuntimeError("log storage unavailable"), + ) + launcher = _FakeLauncher(cached) + execution_id, node_ids = _persist_running_execution(session, cached) + + _process_running(session, launcher) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.FAILED + ) + assert session.get(bts.ExecutionNode, node_ids[0]).container_execution_status == ( + bts.ContainerExecutionStatus.FAILED + ) + + +def test_complete_job_with_missing_outputs_still_fails(monkeypatch): + monkeypatch.setattr( + orchestrator_sql, "_retry", lambda function, **kwargs: function() + ) + session = _new_session() + cached = _FakeLaunchedContainer(interfaces.ContainerStatus.SUCCEEDED) + launcher = _FakeLauncher(cached) + storage_provider = mock.MagicMock() + storage_provider.make_uri.return_value.get_reader.return_value.exists.return_value = ( + False + ) + execution_id, node_ids = _persist_running_execution( + session, + cached, + output_artifact_data_map={"model": {"uri": "memory://missing-model"}}, + ) + + _process_running(session, launcher, storage_provider=storage_provider) + + assert session.get(bts.ContainerExecution, execution_id).status == ( + bts.ContainerExecutionStatus.FAILED + ) + node = session.get(bts.ExecutionNode, node_ids[0]) + assert node.container_execution_status == bts.ContainerExecutionStatus.FAILED + assert ( + "missing outputs" + in node.extra_data[ + bts.EXECUTION_NODE_EXTRA_DATA_ORCHESTRATION_ERROR_MESSAGE_KEY + ] + ) + assert launcher.refresh_calls == 0