Skip to content

Commit c335ec7

Browse files
ntkatholentkathole
committed
fix: Fixed intermittent failures in get_historical_features
Signed-off-by: ntkathole <[email protected]>
1 parent 3639570 commit c335ec7

3 files changed

Lines changed: 214 additions & 11 deletions

File tree

‎sdk/python/feast/arrow_error_handler.py‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import logging
2+
import time
23
from functools import wraps
34

45
import pyarrow.flight as fl
@@ -7,17 +8,45 @@
78

89
logger = logging.getLogger(__name__)
910

11+
BACKOFF_FACTOR = 0.5
12+
1013

1114
def arrow_client_error_handling_decorator(func):
1215
@wraps(func)
1316
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
2150

2251
return wrapper
2352

‎sdk/python/feast/infra/offline_stores/remote.py‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import pyarrow.parquet
1313
from pyarrow import Schema
1414
from pyarrow._flight import FlightCallOptions, FlightDescriptor, Ticket
15-
from pydantic import StrictInt, StrictStr
15+
from pydantic import Field, StrictInt, StrictStr
1616

1717
from feast import OnDemandFeatureView
1818
from feast.arrow_error_handler import arrow_client_error_handling_decorator
@@ -42,6 +42,10 @@
4242

4343

4444
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+
4549
@arrow_client_error_handling_decorator
4650
def get_flight_info(
4751
self, descriptor: FlightDescriptor, options: FlightCallOptions = None
@@ -71,7 +75,12 @@ def list_actions(self, options: FlightCallOptions = None):
7175

7276

7377
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,
7584
):
7685
arrow_scheme = "grpc+tcp"
7786
if scheme == "https":
@@ -88,10 +97,17 @@ def build_arrow_flight_client(
8897
if auth_config.type != AuthType.NONE.value:
8998
middlewares = [FlightAuthInterceptorFactory(auth_config)]
9099
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,
92104
)
93105

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+
)
95111

96112

97113
class RemoteOfflineStoreConfig(FeastConfigBaseModel):
@@ -109,6 +125,9 @@ class RemoteOfflineStoreConfig(FeastConfigBaseModel):
109125
""" 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`.
110126
If type is 'remote', then this configuration is needed to connect to remote offline server in TLS mode. """
111127

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+
112131

113132
class RemoteRetrievalJob(RetrievalJob):
114133
def __init__(
@@ -207,6 +226,7 @@ def get_historical_features(
207226
port=config.offline_store.port,
208227
auth_config=config.auth_config,
209228
cert=config.offline_store.cert,
229+
connection_retries=config.offline_store.connection_retries,
210230
)
211231

212232
feature_view_names = [fv.name for fv in feature_views]
@@ -257,6 +277,7 @@ def pull_all_from_table_or_query(
257277
port=config.offline_store.port,
258278
auth_config=config.auth_config,
259279
cert=config.offline_store.cert,
280+
connection_retries=config.offline_store.connection_retries,
260281
)
261282

262283
api_parameters = {
@@ -295,6 +316,7 @@ def pull_latest_from_table_or_query(
295316
config.offline_store.port,
296317
config.auth_config,
297318
cert=config.offline_store.cert,
319+
connection_retries=config.offline_store.connection_retries,
298320
)
299321

300322
api_parameters = {
@@ -334,6 +356,7 @@ def write_logged_features(
334356
config.offline_store.port,
335357
config.auth_config,
336358
config.offline_store.cert,
359+
connection_retries=config.offline_store.connection_retries,
337360
)
338361

339362
api_parameters = {
@@ -364,6 +387,7 @@ def offline_write_batch(
364387
config.offline_store.port,
365388
config.auth_config,
366389
config.offline_store.cert,
390+
connection_retries=config.offline_store.connection_retries,
367391
)
368392

369393
feature_view_names = [feature_view.name]
@@ -396,6 +420,7 @@ def validate_data_source(
396420
config.offline_store.port,
397421
config.auth_config,
398422
config.offline_store.cert,
423+
connection_retries=config.offline_store.connection_retries,
399424
)
400425

401426
api_parameters = {
@@ -421,6 +446,7 @@ def get_table_column_names_and_types_from_data_source(
421446
config.offline_store.port,
422447
config.auth_config,
423448
config.offline_store.cert,
449+
connection_retries=config.offline_store.connection_retries,
424450
)
425451

426452
api_parameters = {

‎sdk/python/tests/unit/test_arrow_error_decorator.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1+
from unittest.mock import MagicMock, patch
2+
13
import pyarrow.flight as fl
24
import pytest
5+
from pydantic import ValidationError
36

47
from feast.arrow_error_handler import arrow_client_error_handling_decorator
58
from feast.errors import PermissionNotFoundException
9+
from feast.infra.offline_stores.remote import RemoteOfflineStoreConfig
610

711
permissionError = PermissionNotFoundException("dummy_name", "dummy_project")
812

@@ -31,3 +35,147 @@ def test_rest_error_handling_with_feast_exception(error, expected_raised_error):
3135
match=str(expected_raised_error),
3236
):
3337
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)

0 commit comments

Comments
 (0)