GeezerVison is showing github.com — fix: Fixed intermittent failures in get_historical_features · Marcus-Rosti/feast@c335ec7
File treeExpand file treeCollapse file tree Open diff view settings Expand file treeCollapse file tree Open diff view settings Collapse file Copy file name to clipboardExpand all lines: sdk/python/feast/arrow_error_handler.py +36-7Lines changed: 36 additions & 7 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 1 | 1 | import logging | | 2 | +import time | 2 | 3 | from functools import wraps | 3 | 4 |
| 4 | 5 | import pyarrow.flight as fl | | 7 | 8 |
| 8 | 9 | logger = logging.getLogger(__name__) | 9 | 10 |
| | 11 | +BACKOFF_FACTOR = 0.5 | | 12 | + | 10 | 13 |
| 11 | 14 | def arrow_client_error_handling_decorator(func): | 12 | 15 | @wraps(func) | 13 | 16 | def wrapper(*args, **kwargs): | 14 | | - try: | 15 | | - return func(*args, **kwargs) | 16 | | - except Exception as e: | 17 | | - mapped_error = FeastError.from_error_detail(_get_exception_data(e.args[0])) | 18 | | - if mapped_error is not None: | 19 | | - raise mapped_error | 20 | | - raise e | | 17 | + # Retry only applies to FeastFlightClient methods where args[0] (self) | | 18 | + # carries _connection_retries from RemoteOfflineStoreConfig. | | 19 | + # Standalone stream functions (write_table, read_all) get 0 retries: | | 20 | + # broken streams can't be reused and retrying risks duplicate writes. | | 21 | + max_retries = max(0, getattr(args[0], "_connection_retries", 0)) if args else 0 | | 22 | + | | 23 | + for attempt in range(max_retries + 1): | | 24 | + try: | | 25 | + return func(*args, **kwargs) | | 26 | + except fl.FlightUnavailableError as e: | | 27 | + if attempt < max_retries: | | 28 | + wait_time = BACKOFF_FACTOR * (2**attempt) | | 29 | + logger.warning( | | 30 | + "Transient Arrow Flight error on attempt %d/%d, " | | 31 | + "retrying in %.1fs: %s", | | 32 | + attempt + 1, | | 33 | + max_retries + 1, | | 34 | + wait_time, | | 35 | + e, | | 36 | + ) | | 37 | + time.sleep(wait_time) | | 38 | + continue | | 39 | + mapped_error = FeastError.from_error_detail(_get_exception_data(str(e))) | | 40 | + if mapped_error is not None: | | 41 | + raise mapped_error | | 42 | + raise e | | 43 | + except Exception as e: | | 44 | + mapped_error = FeastError.from_error_detail( | | 45 | + _get_exception_data(e.args[0]) | | 46 | + ) | | 47 | + if mapped_error is not None: | | 48 | + raise mapped_error | | 49 | + raise e | 21 | 50 |
| 22 | 51 | return wrapper | 23 | 52 |
| |
Collapse file Copy file name to clipboardExpand all lines: sdk/python/feast/infra/offline_stores/remote.py +30-4Lines changed: 30 additions & 4 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 12 | 12 | import pyarrow.parquet | 13 | 13 | from pyarrow import Schema | 14 | 14 | from pyarrow._flight import FlightCallOptions, FlightDescriptor, Ticket | 15 | | -from pydantic import StrictInt, StrictStr | | 15 | +from pydantic import Field, StrictInt, StrictStr | 16 | 16 |
| 17 | 17 | from feast import OnDemandFeatureView | 18 | 18 | from feast.arrow_error_handler import arrow_client_error_handling_decorator | | 42 | 42 |
| 43 | 43 |
| 44 | 44 | class FeastFlightClient(fl.FlightClient): | | 45 | + def __init__(self, *args, connection_retries: int = 3, **kwargs): | | 46 | + super().__init__(*args, **kwargs) | | 47 | + self._connection_retries = max(0, connection_retries) | | 48 | + | 45 | 49 | @arrow_client_error_handling_decorator | 46 | 50 | def get_flight_info( | 47 | 51 | self, descriptor: FlightDescriptor, options: FlightCallOptions = None | @@ -71,7 +75,12 @@ def list_actions(self, options: FlightCallOptions = None): | 71 | 75 |
| 72 | 76 |
| 73 | 77 | def build_arrow_flight_client( | 74 | | - scheme: str, host: str, port, auth_config: AuthConfig, cert: str = "" | | 78 | + scheme: str, | | 79 | + host: str, | | 80 | + port, | | 81 | + auth_config: AuthConfig, | | 82 | + cert: str = "", | | 83 | + connection_retries: int = 3, | 75 | 84 | ): | 76 | 85 | arrow_scheme = "grpc+tcp" | 77 | 86 | if scheme == "https": | @@ -88,10 +97,17 @@ def build_arrow_flight_client( | 88 | 97 | if auth_config.type != AuthType.NONE.value: | 89 | 98 | middlewares = [FlightAuthInterceptorFactory(auth_config)] | 90 | 99 | return FeastFlightClient( | 91 | | - f"{arrow_scheme}://{host}:{port}", middleware=middlewares, **kwargs | | 100 | + f"{arrow_scheme}://{host}:{port}", | | 101 | + middleware=middlewares, | | 102 | + connection_retries=connection_retries, | | 103 | + **kwargs, | 92 | 104 | ) | 93 | 105 |
| 94 | | - return FeastFlightClient(f"{arrow_scheme}://{host}:{port}", **kwargs) | | 106 | + return FeastFlightClient( | | 107 | + f"{arrow_scheme}://{host}:{port}", | | 108 | + connection_retries=connection_retries, | | 109 | + **kwargs, | | 110 | + ) | 95 | 111 |
| 96 | 112 |
| 97 | 113 | class RemoteOfflineStoreConfig(FeastConfigBaseModel): | @@ -109,6 +125,9 @@ class RemoteOfflineStoreConfig(FeastConfigBaseModel): | 109 | 125 | """ str: Path to the public certificate when the offline server starts in TLS(SSL) mode. This may be needed if the offline server started with a self-signed certificate, typically this file ends with `*.crt`, `*.cer`, or `*.pem`. | 110 | 126 | If type is 'remote', then this configuration is needed to connect to remote offline server in TLS mode. """ | 111 | 127 |
| | 128 | + connection_retries: int = Field(default=3, ge=0) | | 129 | + """ int: Number of retries for transient Arrow Flight errors with exponential backoff (default 3). """ | | 130 | + | 112 | 131 |
| 113 | 132 | class RemoteRetrievalJob(RetrievalJob): | 114 | 133 | def __init__( | @@ -207,6 +226,7 @@ def get_historical_features( | 207 | 226 | port=config.offline_store.port, | 208 | 227 | auth_config=config.auth_config, | 209 | 228 | cert=config.offline_store.cert, | | 229 | + connection_retries=config.offline_store.connection_retries, | 210 | 230 | ) | 211 | 231 |
| 212 | 232 | feature_view_names = [fv.name for fv in feature_views] | @@ -257,6 +277,7 @@ def pull_all_from_table_or_query( | 257 | 277 | port=config.offline_store.port, | 258 | 278 | auth_config=config.auth_config, | 259 | 279 | cert=config.offline_store.cert, | | 280 | + connection_retries=config.offline_store.connection_retries, | 260 | 281 | ) | 261 | 282 |
| 262 | 283 | api_parameters = { | @@ -295,6 +316,7 @@ def pull_latest_from_table_or_query( | 295 | 316 | config.offline_store.port, | 296 | 317 | config.auth_config, | 297 | 318 | cert=config.offline_store.cert, | | 319 | + connection_retries=config.offline_store.connection_retries, | 298 | 320 | ) | 299 | 321 |
| 300 | 322 | api_parameters = { | @@ -334,6 +356,7 @@ def write_logged_features( | 334 | 356 | config.offline_store.port, | 335 | 357 | config.auth_config, | 336 | 358 | config.offline_store.cert, | | 359 | + connection_retries=config.offline_store.connection_retries, | 337 | 360 | ) | 338 | 361 |
| 339 | 362 | api_parameters = { | @@ -364,6 +387,7 @@ def offline_write_batch( | 364 | 387 | config.offline_store.port, | 365 | 388 | config.auth_config, | 366 | 389 | config.offline_store.cert, | | 390 | + connection_retries=config.offline_store.connection_retries, | 367 | 391 | ) | 368 | 392 |
| 369 | 393 | feature_view_names = [feature_view.name] | @@ -396,6 +420,7 @@ def validate_data_source( | 396 | 420 | config.offline_store.port, | 397 | 421 | config.auth_config, | 398 | 422 | config.offline_store.cert, | | 423 | + connection_retries=config.offline_store.connection_retries, | 399 | 424 | ) | 400 | 425 |
| 401 | 426 | api_parameters = { | @@ -421,6 +446,7 @@ def get_table_column_names_and_types_from_data_source( | 421 | 446 | config.offline_store.port, | 422 | 447 | config.auth_config, | 423 | 448 | config.offline_store.cert, | | 449 | + connection_retries=config.offline_store.connection_retries, | 424 | 450 | ) | 425 | 451 |
| 426 | 452 | api_parameters = { | |
Collapse file Copy file name to clipboardExpand all lines: sdk/python/tests/unit/test_arrow_error_decorator.py +148Lines changed: 148 additions & 0 deletions | Original file line number | Diff line number | Diff line change |
|---|
| | 1 | +from unittest.mock import MagicMock, patch | | 2 | + | 1 | 3 | import pyarrow.flight as fl | 2 | 4 | import pytest | | 5 | +from pydantic import ValidationError | 3 | 6 |
| 4 | 7 | from feast.arrow_error_handler import arrow_client_error_handling_decorator | 5 | 8 | from feast.errors import PermissionNotFoundException | | 9 | +from feast.infra.offline_stores.remote import RemoteOfflineStoreConfig | 6 | 10 |
| 7 | 11 | permissionError = PermissionNotFoundException("dummy_name", "dummy_project") | 8 | 12 |
| @@ -31,3 +35,147 @@ def test_rest_error_handling_with_feast_exception(error, expected_raised_error): | 31 | 35 | match=str(expected_raised_error), | 32 | 36 | ): | 33 | 37 | decorated_method(error) | | 38 | + | | 39 | + | | 40 | +class TestArrowClientRetry: | | 41 | + @patch("feast.arrow_error_handler.time.sleep") | | 42 | + def test_retries_on_flight_unavailable_error(self, mock_sleep): | | 43 | + client = MagicMock() | | 44 | + client._connection_retries = 3 | | 45 | + call_count = 0 | | 46 | + | | 47 | + @arrow_client_error_handling_decorator | | 48 | + def flaky_method(self_arg): | | 49 | + nonlocal call_count | | 50 | + call_count += 1 | | 51 | + if call_count < 3: | | 52 | + raise fl.FlightUnavailableError("Connection refused") | | 53 | + return "success" | | 54 | + | | 55 | + result = flaky_method(client) | | 56 | + assert result == "success" | | 57 | + assert call_count == 3 | | 58 | + assert mock_sleep.call_count == 2 | | 59 | + | | 60 | + @patch("feast.arrow_error_handler.time.sleep") | | 61 | + def test_raises_after_max_retries_exhausted(self, mock_sleep): | | 62 | + client = MagicMock() | | 63 | + client._connection_retries = 3 | | 64 | + | | 65 | + @arrow_client_error_handling_decorator | | 66 | + def always_unavailable(self_arg): | | 67 | + raise fl.FlightUnavailableError("Connection refused") | | 68 | + | | 69 | + with pytest.raises(fl.FlightUnavailableError, match="Connection refused"): | | 70 | + always_unavailable(client) | | 71 | + assert mock_sleep.call_count == 3 | | 72 | + | | 73 | + @patch("feast.arrow_error_handler.time.sleep") | | 74 | + def test_respects_connection_retries_from_client(self, mock_sleep): | | 75 | + client = MagicMock() | | 76 | + client._connection_retries = 1 | | 77 | + call_count = 0 | | 78 | + | | 79 | + @arrow_client_error_handling_decorator | | 80 | + def method_on_client(self_arg): | | 81 | + nonlocal call_count | | 82 | + call_count += 1 | | 83 | + raise fl.FlightUnavailableError("Connection refused") | | 84 | + | | 85 | + with pytest.raises(fl.FlightUnavailableError): | | 86 | + method_on_client(client) | | 87 | + | | 88 | + assert call_count == 2 # 1 initial + 1 retry | | 89 | + assert mock_sleep.call_count == 1 | | 90 | + | | 91 | + @patch("feast.arrow_error_handler.time.sleep") | | 92 | + def test_no_retry_on_non_transient_errors(self, mock_sleep): | | 93 | + client = MagicMock() | | 94 | + client._connection_retries = 3 | | 95 | + call_count = 0 | | 96 | + | | 97 | + @arrow_client_error_handling_decorator | | 98 | + def method_with_error(self_arg): | | 99 | + nonlocal call_count | | 100 | + call_count += 1 | | 101 | + raise fl.FlightError("Permanent error") | | 102 | + | | 103 | + with pytest.raises(fl.FlightError, match="Permanent error"): | | 104 | + method_with_error(client) | | 105 | + | | 106 | + assert call_count == 1 | | 107 | + mock_sleep.assert_not_called() | | 108 | + | | 109 | + @patch("feast.arrow_error_handler.time.sleep") | | 110 | + def test_exponential_backoff_timing(self, mock_sleep): | | 111 | + client = MagicMock() | | 112 | + client._connection_retries = 3 | | 113 | + | | 114 | + @arrow_client_error_handling_decorator | | 115 | + def always_unavailable(self_arg): | | 116 | + raise fl.FlightUnavailableError("Connection refused") | | 117 | + | | 118 | + with pytest.raises(fl.FlightUnavailableError): | | 119 | + always_unavailable(client) | | 120 | + | | 121 | + wait_times = [call.args[0] for call in mock_sleep.call_args_list] | | 122 | + assert wait_times == [0.5, 1.0, 2.0] | | 123 | + | | 124 | + @patch("feast.arrow_error_handler.time.sleep") | | 125 | + def test_zero_retries_disables_retry(self, mock_sleep): | | 126 | + client = MagicMock() | | 127 | + client._connection_retries = 0 | | 128 | + call_count = 0 | | 129 | + | | 130 | + @arrow_client_error_handling_decorator | | 131 | + def method_on_client(self_arg): | | 132 | + nonlocal call_count | | 133 | + call_count += 1 | | 134 | + raise fl.FlightUnavailableError("Connection refused") | | 135 | + | | 136 | + with pytest.raises(fl.FlightUnavailableError): | | 137 | + method_on_client(client) | | 138 | + | | 139 | + assert call_count == 1 | | 140 | + mock_sleep.assert_not_called() | | 141 | + | | 142 | + @patch("feast.arrow_error_handler.time.sleep") | | 143 | + def test_no_retry_for_standalone_stream_functions(self, mock_sleep): | | 144 | + """Standalone functions (write_table, read_all) where args[0] is a | | 145 | + writer/reader should not retry since broken streams can't be reused.""" | | 146 | + writer = MagicMock(spec=[]) # no _connection_retries attribute | | 147 | + call_count = 0 | | 148 | + | | 149 | + @arrow_client_error_handling_decorator | | 150 | + def write_table(w): | | 151 | + nonlocal call_count | | 152 | + call_count += 1 | | 153 | + raise fl.FlightUnavailableError("stream broken") | | 154 | + | | 155 | + with pytest.raises(fl.FlightUnavailableError, match="stream broken"): | | 156 | + write_table(writer) | | 157 | + | | 158 | + assert call_count == 1 | | 159 | + mock_sleep.assert_not_called() | | 160 | + | | 161 | + @patch("feast.arrow_error_handler.time.sleep") | | 162 | + def test_negative_connection_retries_treated_as_zero(self, mock_sleep): | | 163 | + """Negative _connection_retries must not skip function execution.""" | | 164 | + client = MagicMock() | | 165 | + client._connection_retries = -1 | | 166 | + call_count = 0 | | 167 | + | | 168 | + @arrow_client_error_handling_decorator | | 169 | + def method_on_client(self_arg): | | 170 | + nonlocal call_count | | 171 | + call_count += 1 | | 172 | + return "ok" | | 173 | + | | 174 | + result = method_on_client(client) | | 175 | + assert result == "ok" | | 176 | + assert call_count == 1 | | 177 | + mock_sleep.assert_not_called() | | 178 | + | | 179 | + def test_config_rejects_negative_connection_retries(self): | | 180 | + with pytest.raises(ValidationError): | | 181 | + RemoteOfflineStoreConfig(host="localhost", connection_retries=-1) |
|