From eee37802b4b39733cf7882cc9dd0533edd2032cc Mon Sep 17 00:00:00 2001 From: River Date: Fri, 14 Aug 2026 19:38:40 +0000 Subject: [PATCH] Bound orchestrator storage waits Add a timeout-aware GCS storage provider using google-cloud-storage's official per-call timeout parameter, wire GCS launchers through it, and bound the orchestrator retry helper with a configurable wall-clock deadline. Co-authored-by: Morgan Wowk --- .../launchers/google_kubernetes_launchers.py | 18 +- .../launchers/kubernetes_launchers.py | 12 +- cloud_pipelines_backend/orchestrator_sql.py | 61 +++- .../google_cloud_storage_with_timeout.py | 302 ++++++++++++++++++ .../test_google_cloud_storage_with_timeout.py | 93 ++++++ tests/test_orchestrator_retry_deadline.py | 82 +++++ 6 files changed, 557 insertions(+), 11 deletions(-) create mode 100644 cloud_pipelines_backend/storage_providers/google_cloud_storage_with_timeout.py create mode 100644 tests/test_google_cloud_storage_with_timeout.py create mode 100644 tests/test_orchestrator_retry_deadline.py diff --git a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py b/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py index 57f959f8..3788d8c6 100644 --- a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py @@ -2,7 +2,7 @@ from kubernetes import client as k8s_client_lib -from cloud_pipelines.orchestration.storage_providers import google_cloud_storage +from cloud_pipelines_backend.storage_providers import google_cloud_storage_with_timeout from . import kubernetes_launchers @@ -23,6 +23,9 @@ def __init__( service_account_name: str | None = None, request_timeout: int | tuple[int, int] = 10, gcs_client: "storage.Client | None" = None, + gcs_request_timeout: ( + google_cloud_storage_with_timeout.RequestTimeout | None + ) = None, pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: kubernetes_launchers.PodPostProcessor | None = None, @@ -44,8 +47,9 @@ def __init__( pod_labels=pod_labels, pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}), pod_postprocessor=final_pod_postporocessor, - _storage_provider=google_cloud_storage.GoogleCloudStorageProvider( - gcs_client + _storage_provider=google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + gcs_client, + request_timeout=gcs_request_timeout, ), _create_volume_and_volume_mount=kubernetes_launchers._create_volume_and_volume_mount_google_cloud_storage, ) @@ -64,6 +68,9 @@ def __init__( service_account_name: str | None = None, request_timeout: int | tuple[int, int] = 10, gcs_client: "storage.Client | None" = None, + gcs_request_timeout: ( + google_cloud_storage_with_timeout.RequestTimeout | None + ) = None, pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: kubernetes_launchers.PodPostProcessor | None = None, @@ -87,8 +94,9 @@ def __init__( pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}), pod_postprocessor=final_pod_postporocessor, always_launch_jobs=always_launch_jobs, - _storage_provider=google_cloud_storage.GoogleCloudStorageProvider( - gcs_client + _storage_provider=google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + gcs_client, + request_timeout=gcs_request_timeout, ), _create_volume_and_volume_mount=kubernetes_launchers._create_volume_and_volume_mount_google_cloud_storage, ) diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 9b1a661a..2857163e 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -715,6 +715,9 @@ def __init__( request_timeout: int | tuple[int, int] = 10, pod_name_prefix: str = "task-pod-", gcs_client: "storage.Client | None" = None, + gcs_request_timeout: ( + "google_cloud_storage_with_timeout.RequestTimeout | None" + ) = None, pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: PodPostProcessor | None = None, @@ -724,7 +727,9 @@ def __init__( pod_postprocessors.append(pod_postprocessor) final_pod_postporocessor = _create_pod_postprocessor_stack(pod_postprocessors) - from cloud_pipelines.orchestration.storage_providers import google_cloud_storage + from cloud_pipelines_backend.storage_providers import ( + google_cloud_storage_with_timeout, + ) super().__init__( namespace=namespace, @@ -732,8 +737,9 @@ def __init__( api_client=api_client, request_timeout=request_timeout, pod_name_prefix=pod_name_prefix, - _storage_provider=google_cloud_storage.GoogleCloudStorageProvider( - gcs_client + _storage_provider=google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + gcs_client, + request_timeout=gcs_request_timeout, ), pod_labels=pod_labels, pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}), diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index 8dfe13b7..58788122 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -33,6 +33,9 @@ DYNAMIC_DATA_SECRET_KEY = "secret" DYNAMIC_DATA_SECRET_NAME_KEY = "name" +ORCHESTRATOR_RETRY_DEADLINE_ENV = "TANGLE_ORCHESTRATOR_RETRY_DEADLINE_SECONDS" +DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS = 180.0 + class OrchestratorError(RuntimeError): pass @@ -1225,20 +1228,72 @@ def _update_dict_recursive(d1: dict, d2: dict): d1[k] = v2 +class RetryDeadlineExceededError(TimeoutError): + pass + + def _retry( - func: typing.Callable[[], _T], max_retries: int = 5, wait_seconds: float = 1.0 + func: typing.Callable[[], _T], + max_retries: int = 5, + wait_seconds: float = 1.0, + max_elapsed_seconds: float | None = None, ) -> _T: + if max_elapsed_seconds is None: + max_elapsed_seconds = _configured_retry_deadline_seconds() + deadline_at = time.monotonic() + max_elapsed_seconds + last_exception: Exception | None = None + for i in range(max_retries): + if time.monotonic() >= deadline_at: + raise RetryDeadlineExceededError( + f"Retry deadline expired before attempt {i + 1} of {max_retries} " + f"for {func}." + ) from last_exception + try: return func() - except Exception: + except Exception as ex: + last_exception = ex _logger.exception(f"Exception calling {func}.") - time.sleep(wait_seconds) if i == max_retries - 1: raise + + remaining_seconds = deadline_at - time.monotonic() + if remaining_seconds <= 0: + raise RetryDeadlineExceededError( + f"Retry deadline expired after attempt {i + 1} of " + f"{max_retries} for {func}." + ) from ex + time.sleep(min(wait_seconds, remaining_seconds)) raise +def _configured_retry_deadline_seconds() -> float: + raw_value = os.environ.get( + ORCHESTRATOR_RETRY_DEADLINE_ENV, + str(DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS), + ) + try: + timeout = float(raw_value) + except (TypeError, ValueError): + _logger.warning( + "Invalid %s=%r; using default %.1fs", + ORCHESTRATOR_RETRY_DEADLINE_ENV, + raw_value, + DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS, + ) + return DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS + if timeout <= 0: + _logger.warning( + "Invalid %s=%r; using default %.1fs", + ORCHESTRATOR_RETRY_DEADLINE_ENV, + raw_value, + DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS, + ) + return DEFAULT_ORCHESTRATOR_RETRY_DEADLINE_SECONDS + return timeout + + def record_system_error_exception(execution: bts.ExecutionNode, exception: Exception): app_metrics.execution_system_errors.add(1) bugsnag_instrumentation.notify( diff --git a/cloud_pipelines_backend/storage_providers/google_cloud_storage_with_timeout.py b/cloud_pipelines_backend/storage_providers/google_cloud_storage_with_timeout.py new file mode 100644 index 00000000..65c39fb1 --- /dev/null +++ b/cloud_pipelines_backend/storage_providers/google_cloud_storage_with_timeout.py @@ -0,0 +1,302 @@ +"""Google Cloud Storage provider with explicit per-request timeouts. + +The official google-cloud-storage Python client exposes timeouts as per-call +``timeout=`` arguments. This provider keeps Tangle's StorageProvider interface +but passes a configured timeout to each GCS request used by the upstream +GoogleCloudStorageProvider implementation. +""" + +from __future__ import annotations + +import base64 +import logging +import os +import typing +from typing import Optional, TypeAlias + +from cloud_pipelines.orchestration.storage_providers import google_cloud_storage +from cloud_pipelines.orchestration.storage_providers import interfaces + +_LOGGER = logging.getLogger(name=__name__) + +GCS_REQUEST_TIMEOUT_ENV = "TANGLE_GCS_REQUEST_TIMEOUT_SECONDS" +DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS = 60.0 + +RequestTimeout: TypeAlias = float | tuple[float, float] + +if typing.TYPE_CHECKING: + from google.cloud import storage + + +def _storage_module(): + from google.cloud import storage + + return storage + + +class GoogleCloudStorageProviderWithTimeout( + google_cloud_storage.GoogleCloudStorageProvider +): + """GoogleCloudStorageProvider that passes an explicit timeout to GCS calls.""" + + def __init__( + self, + client: Optional["storage.Client"] = None, + *, + request_timeout: RequestTimeout | None = None, + ) -> None: + """Construct a GCS provider whose every request is bounded by a timeout. + + Instantiated by the Kubernetes launchers, so any consumer that runs + Tangle on GCP gets a storage provider where a hung GCS call can never + block the orchestrator's poll loop indefinitely. The timeout defaults + from ``TANGLE_GCS_REQUEST_TIMEOUT_SECONDS``, letting a consumer tune GCS + patience per deployment without code changes. + """ + super().__init__(client=client) + self._request_timeout = request_timeout or _configured_request_timeout() + + def _upload_file(self, source_file_path: str, destination_blob_uri: str): + """Upload a single local file to a GCS object. + + Use cases: + - Staging a leaf output artifact (a single file produced by a completed + container) into its artifact URI. + - Writing a container's captured log file to its log URI. + """ + storage = _storage_module() + destination_blob = storage.Blob.from_string( + uri=destination_blob_uri, client=self._client + ) + destination_blob.upload_from_filename( + filename=source_file_path, + checksum="md5", + timeout=self._request_timeout, + ) + + def _upload_dir(self, source_dir_path: str, destination_dir_uri: str): + """Upload a local directory tree (recursively) to GCS. + + Used when a container output artifact is a directory rather than a + single file (e.g. a model checkpoint directory or a multi-file dataset): + each entry is uploaded under the destination prefix and a zero-byte + marker object represents the directory itself. + """ + # Creating the directory object (zero-byte object with name ending in slash) + storage = _storage_module() + storage.Blob.from_string( + uri=destination_dir_uri.rstrip("/") + "/", client=self._client + ).upload_from_string( + data="", + checksum="md5", + timeout=self._request_timeout, + ) + + for dir_entry_name in os.listdir(source_dir_path): + source_path = os.path.join(source_dir_path, dir_entry_name) + destination_uri = destination_dir_uri.rstrip("/") + "/" + dir_entry_name + self._upload_to_uri( + source_path=source_path, + destination_uri=destination_uri, + ) + + def upload_bytes( + self, data: bytes, destination_uri: google_cloud_storage.GoogleCloudStorageUri + ): + """Upload in-memory bytes directly to a GCS object (no local temp file). + + Use cases: + - Staging a small inline input-argument value into its staging URI + before a container is launched, so the container can consume it as a + file. + - Persisting small artifact or log payloads that a consumer already + holds in memory. + """ + storage = _storage_module() + destination_uri_str = destination_uri.uri + destination_blob = storage.Blob.from_string( + uri=destination_uri_str, client=self._client + ) + _LOGGER.debug(f"Uploading data to {destination_uri_str}") + destination_blob.upload_from_string( + data=data, + checksum="md5", + timeout=self._request_timeout, + ) + + def _download_from_uri(self, source_uri: str, destination_path: str): + """Download a GCS object — or every object under a directory prefix — to + a local path. + + Used by the launchers to materialize a container's input artifacts on + local disk before execution; handles both single-file and directory + artifacts. + """ + storage = _storage_module() + source_blob_or_dir = storage.Blob.from_string( + uri=source_uri, client=self._client + ) + if source_blob_or_dir.exists(timeout=self._request_timeout): + return _download_blob_to_filename_with_timeout( + blob=source_blob_or_dir, + destination_path=destination_path, + request_timeout=self._request_timeout, + ) + + source_dir_prefix = source_blob_or_dir.name.rstrip("/") + "/" + for source_blob in self._client.list_blobs( + bucket_or_name=source_blob_or_dir.bucket, + prefix=source_dir_prefix, + timeout=self._request_timeout, + ): + assert source_blob.name.startswith(source_dir_prefix) + relative_source_blob_name = source_blob.name[len(source_dir_prefix) :] + destination_file_path = os.path.join( + destination_path, relative_source_blob_name + ) + if source_blob.name.endswith("/"): + # It's a zero-size object that represents a directory + assert source_blob.size == 0 + os.makedirs(destination_file_path, exist_ok=True) + else: + _download_blob_to_filename_with_timeout( + blob=source_blob, + destination_path=destination_file_path, + request_timeout=self._request_timeout, + ) + + def download_bytes( + self, source_uri: google_cloud_storage.GoogleCloudStorageUri + ) -> bytes: + """Download a GCS object as raw bytes. + + Use cases: + - Launchers: reading an input artifact so its value can be inlined as + text and passed to a container. + - Orchestrator: preloading small (<=255 byte) output artifact values for + preservation after a container completes. + - API server: fetching stored container log text to serve to a consumer. + """ + storage = _storage_module() + source_uri_str = source_uri.uri + source_blob = storage.Blob.from_string(uri=source_uri_str, client=self._client) + _LOGGER.debug(f"Downloading data from {source_uri_str}") + return source_blob.download_as_bytes(timeout=self._request_timeout) + + def exists(self, uri: google_cloud_storage.GoogleCloudStorageUri) -> bool: + """Check whether a GCS object or directory exists. + + Used by the orchestrator after a container reports success to verify + every declared output artifact was actually produced; any missing output + marks the execution FAILED and skips its downstream nodes. + """ + storage = _storage_module() + blob_uri = uri.uri.rstrip("/") + file_blob = storage.Blob.from_string(uri=blob_uri, client=self._client) + # The "directory objects" are expected to exist for directories + dir_blob = storage.Blob.from_string(uri=blob_uri + "/", client=self._client) + return file_blob.exists(timeout=self._request_timeout) or dir_blob.exists( + timeout=self._request_timeout + ) + + def _get_info_from_uri(self, uri: str) -> interfaces.DataInfo: + """Return size, directory flag, and content hashes for a GCS object or + directory. + + Used by the orchestrator to record each produced output's ArtifactData + (total size, is_dir, hash). Those hashes drive artifact caching and + execution reuse and are surfaced to downstream consumers. + """ + storage = _storage_module() + file_info_list = [] + blob_or_dir = storage.Blob.from_string(uri=uri, client=self._client) + if blob_or_dir.exists(timeout=self._request_timeout): + blob = blob_or_dir + blob.reload(timeout=self._request_timeout) + return interfaces.DataInfo( + total_size=blob.size, + is_dir=False, + hashes=_get_gcs_blob_hashes(blob), + ) + + dir_prefix = blob_or_dir.name.rstrip("/") + "/" + for blob in self._client.list_blobs( + bucket_or_name=blob_or_dir.bucket, + prefix=dir_prefix, + timeout=self._request_timeout, + ): + blob.reload(timeout=self._request_timeout) + assert blob.name.startswith(dir_prefix) + relative_source_blob_name = blob.name[len(dir_prefix) :] + file_info_list.append( + interfaces._FileInfo( + path=relative_source_blob_name, + size=blob.size, + hashes=_get_gcs_blob_hashes(blob), + ) + ) + data_info = interfaces._make_data_info_for_dir(file_info_list) + data_info._file_info_list = file_info_list + return data_info + + +def _configured_request_timeout() -> RequestTimeout: + """Resolve the per-request timeout from ``TANGLE_GCS_REQUEST_TIMEOUT_SECONDS``. + + Falls back to the default when the variable is unset or invalid, letting a + consumer tune GCS patience per deployment via environment configuration. + """ + raw_value = os.environ.get( + GCS_REQUEST_TIMEOUT_ENV, + str(DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS), + ) + try: + timeout = float(raw_value) + except (TypeError, ValueError): + _LOGGER.warning( + "Invalid %s=%r; using default %.1fs", + GCS_REQUEST_TIMEOUT_ENV, + raw_value, + DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS, + ) + return DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS + if timeout <= 0: + _LOGGER.warning( + "Invalid %s=%r; using default %.1fs", + GCS_REQUEST_TIMEOUT_ENV, + raw_value, + DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS, + ) + return DEFAULT_GCS_REQUEST_TIMEOUT_SECONDS + return timeout + + +def _download_blob_to_filename_with_timeout( + *, blob: "storage.Blob", destination_path: str, request_timeout: RequestTimeout +) -> None: + """Download a single blob to a local file under the configured timeout. + + Creates parent directories first. Shared by the single-file and per-entry + directory download paths so both honor the same request timeout. + """ + os.makedirs(os.path.dirname(destination_path), exist_ok=True) + blob.download_to_filename(filename=destination_path, timeout=request_timeout) + + +def _get_gcs_blob_hashes(blob: "storage.Blob") -> dict[str, str]: + """Extract md5/crc32c hashes from a GCS blob's metadata. + + Feeds artifact integrity checks and cache-key / directory-hash computation + in ``get_info``. (Composite GCS objects have no md5 hash, so it may be + absent.) + """ + hashes = {} + # Note: Composite GCS objects do not have MD5 hash metadata. + # See: https://docs.cloud.google.com/storage/docs/composite-objects#metadata + if blob.md5_hash: + # blob.md5_hash is a base64-encoded hash digest byte array. E.g. "1B2M2Y8AsgTpgAmY7PhCfg==" + hashes["md5"] = base64.decodebytes(blob.md5_hash.encode("ascii")).hex() + if blob.crc32c: + # blob.crc32c is a base64-encoded hash digest byte array. E.g. "4gcgLQ==" + hashes["crc32c"] = base64.decodebytes(blob.crc32c.encode("ascii")).hex() + return hashes diff --git a/tests/test_google_cloud_storage_with_timeout.py b/tests/test_google_cloud_storage_with_timeout.py new file mode 100644 index 00000000..415b68ee --- /dev/null +++ b/tests/test_google_cloud_storage_with_timeout.py @@ -0,0 +1,93 @@ +import sys +import types +from unittest import mock + +# cloud-pipelines-backend keeps GCS support optional. Stub google.cloud.storage +# before importing the timeout provider so these unit tests do not require the +# optional Google client package. +_google_mod = types.ModuleType("google") +_cloud_mod = types.ModuleType("google.cloud") +_storage_mod = types.ModuleType("google.cloud.storage") +_storage_mod.Blob = mock.MagicMock() +_storage_mod.Client = mock.MagicMock() +_cloud_mod.storage = _storage_mod +_google_mod.cloud = _cloud_mod +sys.modules.setdefault("google", _google_mod) +sys.modules.setdefault("google.cloud", _cloud_mod) +sys.modules.setdefault("google.cloud.storage", _storage_mod) + +from cloud_pipelines_backend.storage_providers import google_cloud_storage_with_timeout + + +class TestGoogleCloudStorageProviderWithTimeout: + def test_upload_bytes_passes_timeout(self) -> None: + storage = mock.MagicMock() + blob = mock.MagicMock() + storage.Blob.from_string.return_value = blob + provider = ( + google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + client=mock.MagicMock(), request_timeout=12 + ) + ) + uri = mock.MagicMock(uri="gs://bucket/object") + + with mock.patch.object( + google_cloud_storage_with_timeout, + "_storage_module", + return_value=storage, + ): + provider.upload_bytes(b"data", uri) + + blob.upload_from_string.assert_called_once_with( + data=b"data", checksum="md5", timeout=12 + ) + + def test_exists_passes_timeout_to_file_and_dir_blobs(self) -> None: + storage = mock.MagicMock() + file_blob = mock.MagicMock() + file_blob.exists.return_value = False + dir_blob = mock.MagicMock() + dir_blob.exists.return_value = True + storage.Blob.from_string.side_effect = [file_blob, dir_blob] + provider = ( + google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + client=mock.MagicMock(), request_timeout=12 + ) + ) + uri = mock.MagicMock(uri="gs://bucket/path") + + with mock.patch.object( + google_cloud_storage_with_timeout, + "_storage_module", + return_value=storage, + ): + assert provider.exists(uri) is True + + file_blob.exists.assert_called_once_with(timeout=12) + dir_blob.exists.assert_called_once_with(timeout=12) + + def test_file_get_info_passes_timeout_to_exists_and_reload(self) -> None: + storage = mock.MagicMock() + blob = mock.MagicMock() + blob.exists.return_value = True + blob.size = 123 + blob.md5_hash = None + blob.crc32c = None + storage.Blob.from_string.return_value = blob + provider = ( + google_cloud_storage_with_timeout.GoogleCloudStorageProviderWithTimeout( + client=mock.MagicMock(), request_timeout=12 + ) + ) + + with mock.patch.object( + google_cloud_storage_with_timeout, + "_storage_module", + return_value=storage, + ): + data_info = provider._get_info_from_uri("gs://bucket/object") + + blob.exists.assert_called_once_with(timeout=12) + blob.reload.assert_called_once_with(timeout=12) + assert data_info.total_size == 123 + assert data_info.is_dir is False diff --git a/tests/test_orchestrator_retry_deadline.py b/tests/test_orchestrator_retry_deadline.py new file mode 100644 index 00000000..14861b19 --- /dev/null +++ b/tests/test_orchestrator_retry_deadline.py @@ -0,0 +1,82 @@ +from unittest import mock + +import pytest + +from cloud_pipelines_backend import orchestrator_sql + + +def test_retry_succeeds_before_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + calls = 0 + + def func() -> str: + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("try again") + return "ok" + + monkeypatch.setattr(orchestrator_sql.time, "sleep", lambda _seconds: None) + + assert orchestrator_sql._retry(func, max_retries=3, max_elapsed_seconds=10) == "ok" + assert calls == 2 + + +def test_retry_deadline_caps_elapsed_time(monkeypatch: pytest.MonkeyPatch) -> None: + now = 100.0 + calls = 0 + + def monotonic() -> float: + return now + + def sleep(seconds: float) -> None: + nonlocal now + now += seconds + + def func() -> None: + nonlocal calls, now + calls += 1 + now += 0.9 + raise RuntimeError("still failing") + + monkeypatch.setattr(orchestrator_sql.time, "monotonic", monotonic) + monkeypatch.setattr(orchestrator_sql.time, "sleep", sleep) + + with pytest.raises(orchestrator_sql.RetryDeadlineExceededError): + orchestrator_sql._retry( + func, max_retries=5, wait_seconds=1.0, max_elapsed_seconds=2 + ) + + assert calls == 2 + + +def test_max_retries_still_wins_before_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = 0 + + def func() -> None: + nonlocal calls + calls += 1 + raise RuntimeError("boom") + + monkeypatch.setattr(orchestrator_sql.time, "sleep", lambda _seconds: None) + + with pytest.raises(RuntimeError, match="boom"): + orchestrator_sql._retry(func, max_retries=2, max_elapsed_seconds=100) + + assert calls == 2 + + +def test_retry_deadline_uses_environment_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(orchestrator_sql.ORCHESTRATOR_RETRY_DEADLINE_ENV, "2") + + with mock.patch.object( + orchestrator_sql, + "_configured_retry_deadline_seconds", + wraps=orchestrator_sql._configured_retry_deadline_seconds, + ) as configured_retry_deadline_seconds: + orchestrator_sql._retry(lambda: "ok") + + configured_retry_deadline_seconds.assert_called_once_with()