Skip to content

Commit 5388522

Browse files
Reactor11Reactor11
authored
feat(bigtable): Make write batch size and concurrency configurable (#6847)
feat(bigtable): make write batch size and concurrency configurable The Bigtable online store hardcodes the mutations-per-write batch size (MUTATIONS_PER_OP = 50_000) and the write thread-pool size (BIGTABLE_CLIENT_CONNECTION_POOL_SIZE = 10). On a shared Bigtable instance, a large materialization issues its writes as an unthrottled burst that can saturate the instance and inflate read-path tail latency for other workloads sharing it. Expose both as optional BigtableOnlineStoreConfig fields, mutations_per_write and write_concurrency, defaulting to the existing constants so behavior is unchanged. Operators can lower either to reduce the write load a materialization places on the instance, at the cost of longer materialization time. Add unit tests covering the defaults, positive-int validation, batch chunking, the one-row-per-request floor for very wide feature views, and the configurable thread-pool size. Signed-off-by: Manas Bhardwaj <[email protected]>
1 parent f8883d0 commit 5388522

2 files changed

Lines changed: 141 additions & 7 deletions

File tree

‎sdk/python/feast/infra/online_stores/bigtable.py‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import google
88
from google.cloud import bigtable
99
from google.cloud.bigtable import row_filters
10-
from pydantic import StrictStr
10+
from pydantic import PositiveInt, StrictStr
1111

1212
from feast import Entity, FeatureView, utils
1313
from feast.feature_view import DUMMY_ENTITY_NAME
@@ -19,11 +19,13 @@
1919

2020
logger = logging.getLogger(__name__)
2121

22-
# Number of mutations per Bigtable write operation we're aiming for. The official max is
23-
# 100K; we're being conservative.
22+
# Default number of mutations per Bigtable write operation we're aiming for. The official
23+
# max is 100K; we're being conservative. Overridable via
24+
# ``BigtableOnlineStoreConfig.mutations_per_write``.
2425
MUTATIONS_PER_OP = 50_000
25-
# The Bigtable client library limits the connection pool size to 10. This imposes a
26-
# limitation to the concurrency we can get using a thread pool in each worker.
26+
# Default thread-pool size used to parallelize writes within a worker. The Bigtable client
27+
# library limits the connection pool size to 10, which bounds the useful concurrency.
28+
# Overridable via ``BigtableOnlineStoreConfig.write_concurrency``.
2729
BIGTABLE_CLIENT_CONNECTION_POOL_SIZE = 10
2830

2931

@@ -42,6 +44,16 @@ class BigtableOnlineStoreConfig(FeastConfigBaseModel):
4244
max_versions: int = 2
4345
"""The number of historical versions of data that will be kept around."""
4446

47+
mutations_per_write: PositiveInt = MUTATIONS_PER_OP
48+
"""(optional) Target number of cell mutations per Bigtable write request. Bigtable's
49+
hard limit is 100K; lower values produce smaller requests, which reduces the write
50+
burst a materialization places on the instance. Defaults to 50,000."""
51+
52+
write_concurrency: PositiveInt = BIGTABLE_CLIENT_CONNECTION_POOL_SIZE
53+
"""(optional) Maximum number of concurrent write requests (thread-pool size) issued
54+
per worker during ``online_write_batch``. Lower values reduce the write load placed
55+
on the Bigtable instance at the cost of longer materialization. Defaults to 10."""
56+
4557

4658
class BigtableOnlineStore(OnlineStore):
4759
_client: Optional[bigtable.Client] = None
@@ -133,10 +145,13 @@ def online_write_batch(
133145
# `columns_per_row` is used to calculate the number of rows we are allowed to
134146
# mutate in one request.
135147
columns_per_row = len(feature_view.features) + 1 # extra for event timestamp
136-
rows_per_write = MUTATIONS_PER_OP // columns_per_row
148+
# At least one row per request even when a feature view is very wide.
149+
rows_per_write = max(
150+
1, config.online_store.mutations_per_write // columns_per_row
151+
)
137152

138153
with futures.ThreadPoolExecutor(
139-
max_workers=BIGTABLE_CLIENT_CONNECTION_POOL_SIZE
154+
max_workers=config.online_store.write_concurrency
140155
) as executor:
141156
fs = []
142157
while data:

‎sdk/python/tests/unit/infra/online_store/test_bigtable_online_store.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
from dataclasses import dataclass, field
2+
from datetime import datetime, timezone
3+
from typing import List
4+
from unittest.mock import patch
5+
6+
import pytest
7+
8+
pytest.importorskip("google.cloud.bigtable")
9+
10+
from concurrent import futures # noqa: E402
11+
12+
import feast.infra.online_stores.bigtable as bigtable_module # noqa: E402
13+
from feast.infra.offline_stores.dask import DaskOfflineStoreConfig # noqa: E402
14+
from feast.infra.online_stores.bigtable import ( # noqa: E402
15+
BigtableOnlineStore,
16+
BigtableOnlineStoreConfig,
17+
)
18+
from feast.repo_config import RepoConfig # noqa: E402
19+
20+
21+
@dataclass
22+
class MockFeatureView:
23+
name: str
24+
features: List[object] = field(default_factory=list)
25+
26+
27+
def _make_repo_config(**online_store_kwargs) -> RepoConfig:
28+
return RepoConfig(
29+
registry="s3://test_registry/registry.db",
30+
project="test_bigtable",
31+
provider="local",
32+
online_store=BigtableOnlineStoreConfig(
33+
instance="test-instance",
34+
project_id="test-project",
35+
**online_store_kwargs,
36+
),
37+
offline_store=DaskOfflineStoreConfig(),
38+
entity_key_serialization_version=3,
39+
)
40+
41+
42+
def _rows(n: int):
43+
# Contents are irrelevant: _write_rows_to_bt is mocked, only the slicing
44+
# performed by online_write_batch is under test.
45+
return [(None, {}, datetime.now(timezone.utc), None) for _ in range(n)]
46+
47+
48+
def _run_write(store, config, feature_view, data, captured_chunk_sizes):
49+
def _capture(rows_to_write, **_):
50+
captured_chunk_sizes.append(len(rows_to_write))
51+
52+
with (
53+
patch.object(BigtableOnlineStore, "_get_client"),
54+
patch.object(BigtableOnlineStore, "_get_table_name", return_value="tbl"),
55+
patch.object(BigtableOnlineStore, "_write_rows_to_bt", side_effect=_capture),
56+
):
57+
store.online_write_batch(config, feature_view, data, None)
58+
59+
60+
def test_defaults_match_legacy_constants():
61+
config = BigtableOnlineStoreConfig(instance="test-instance")
62+
assert config.mutations_per_write == bigtable_module.MUTATIONS_PER_OP == 50_000
63+
assert (
64+
config.write_concurrency
65+
== bigtable_module.BIGTABLE_CLIENT_CONNECTION_POOL_SIZE
66+
== 10
67+
)
68+
69+
70+
@pytest.mark.parametrize("field_name", ["mutations_per_write", "write_concurrency"])
71+
@pytest.mark.parametrize("bad_value", [0, -1])
72+
def test_positive_int_validation(field_name, bad_value):
73+
from pydantic import ValidationError
74+
75+
with pytest.raises(ValidationError):
76+
BigtableOnlineStoreConfig(instance="test-instance", **{field_name: bad_value})
77+
78+
79+
def test_online_write_batch_respects_mutations_per_write():
80+
# 1 feature -> columns_per_row = 2 -> rows_per_write = max(1, 10 // 2) = 5
81+
store = BigtableOnlineStore()
82+
config = _make_repo_config(mutations_per_write=10)
83+
feature_view = MockFeatureView(name="fv", features=[object()])
84+
85+
captured: List[int] = []
86+
_run_write(store, config, feature_view, _rows(12), captured)
87+
88+
assert captured == [5, 5, 2]
89+
90+
91+
def test_rows_per_write_floors_to_one_for_wide_feature_view():
92+
# 10 features -> columns_per_row = 11 -> 2 // 11 = 0 -> floored to 1
93+
store = BigtableOnlineStore()
94+
config = _make_repo_config(mutations_per_write=2)
95+
feature_view = MockFeatureView(name="fv", features=[object() for _ in range(10)])
96+
97+
captured: List[int] = []
98+
_run_write(store, config, feature_view, _rows(3), captured)
99+
100+
assert captured == [1, 1, 1]
101+
102+
103+
def test_write_concurrency_sets_thread_pool_size():
104+
store = BigtableOnlineStore()
105+
config = _make_repo_config(write_concurrency=3)
106+
feature_view = MockFeatureView(name="fv", features=[object()])
107+
108+
with (
109+
patch.object(BigtableOnlineStore, "_get_client"),
110+
patch.object(BigtableOnlineStore, "_get_table_name", return_value="tbl"),
111+
patch.object(BigtableOnlineStore, "_write_rows_to_bt"),
112+
patch(
113+
"feast.infra.online_stores.bigtable.futures.ThreadPoolExecutor",
114+
wraps=futures.ThreadPoolExecutor,
115+
) as mock_pool,
116+
):
117+
store.online_write_batch(config, feature_view, _rows(4), None)
118+
119+
mock_pool.assert_called_once_with(max_workers=3)

0 commit comments

Comments
 (0)