Skip to content

Commit 256d65d

Browse files
vkagamlykvkagamlyk
committed
fix(aerospike): add client init lock and batch chunking
Guard lazy client creation with a lock to avoid connection leaks under concurrent first use, and chunk batch reads/writes by batch_max_records so large materializations stay under Aerospike server batch limits. Signed-off-by: Valentyn Kahamlyk <valentin.kagamlyk@gmail.com>
1 parent b2b98bc commit 256d65d

3 files changed

Lines changed: 273 additions & 106 deletions

File tree

‎docs/reference/online-stores/aerospike.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ online_store:
6464
read_timeout_ms: 150 # hard deadline for a single-record get
6565
write_timeout_ms: 300 # hard deadline for a single-record put/operate
6666
batch_total_timeout_ms: 500 # hard deadline for online_read / online_write_batch
67+
batch_max_records: 1000 # chunk size for batch_write / batch_operate
6768
socket_timeout_ms: 50 # per-attempt deadline so max_retries can fire
6869
max_retries: 2
6970
```
@@ -76,6 +77,12 @@ online_store:
7677
> deadline; without it, `max_retries` effectively never fires because the
7778
> first attempt is allowed to consume the entire total deadline.
7879

80+
> **Batch chunking.** `online_read` and `online_write_batch` split large
81+
> requests into chunks of at most `batch_max_records` (default `1000`).
82+
> Aerospike enforces a per-node batch limit via the server `batch-max-requests`
83+
> setting (historically `5000`). Lower `batch_max_records` if your cluster cap
84+
> is tighter; raise it only when the server limit and client timeouts allow.
85+
7986
### Aerospike Enterprise with authentication
8087

8188
> Requires Aerospike Enterprise Edition. The Community Edition server has no built-in user/security model and will reject these config keys.

‎sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike.py‎

Lines changed: 161 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,21 @@
33
import asyncio
44
import functools
55
import importlib
6+
import threading
67
from datetime import datetime, timezone
78
from logging import getLogger
89
from typing import (
910
Any,
1011
Callable,
1112
Dict,
13+
Iterator,
1214
List,
1315
Literal,
1416
Optional,
1517
Sequence,
1618
Set,
1719
Tuple,
20+
TypeVar,
1821
Union,
1922
)
2023

@@ -88,6 +91,13 @@
8891
# ordering; subsequent puts keep it.
8992
_ORDERED_MAP_POLICY: Dict[str, Any] = {"map_order": aerospike.MAP_KEY_ORDERED}
9093

94+
# Aerospike server ``batch-max-requests`` defaults to 5000 on many clusters (0
95+
# means unlimited on newer releases). Stay well under that so materialization
96+
# and wide feature-server requests do not trip BatchMaxRequestError (code 151).
97+
_DEFAULT_BATCH_MAX_RECORDS: int = 1_000
98+
99+
_T = TypeVar("_T")
100+
91101

92102
def _datetime_to_epoch_ms(dt: datetime) -> int:
93103
"""Convert a datetime to int64 epoch milliseconds.
@@ -256,6 +266,15 @@ def hook(
256266
max_retries: int = 2
257267
"""Maximum number of automatic retries on transient errors."""
258268

269+
batch_max_records: int = _DEFAULT_BATCH_MAX_RECORDS
270+
"""Maximum records per ``batch_write`` / ``batch_operate`` call.
271+
272+
Aerospike enforces a per-node batch size via the server ``batch-max-requests``
273+
setting (historically 5000). Feast chunks read and write paths to this limit
274+
so large materializations and wide online-serving requests do not fail the
275+
whole batch when the server cap is exceeded.
276+
"""
277+
259278
client_kwargs: Dict[str, Any] = {}
260279
"""Escape hatch for any Aerospike client configuration not surfaced above.
261280
Merged into the client config passed to ``aerospike.client()``."""
@@ -290,64 +309,87 @@ def __init__(self) -> None:
290309
# store to a different hook) re-resolves on next call.
291310
self._prewriting_hook: Optional[PrewritingHook] = None
292311
self._prewriting_hook_spec: Optional[str] = None
312+
self._client_lock = threading.Lock()
293313

294314
# ------------------------------------------------------------------
295315
# Lifecycle / connection management
296316
# ------------------------------------------------------------------
317+
@staticmethod
318+
def _chunked(items: Sequence[_T], chunk_size: int) -> Iterator[Sequence[_T]]:
319+
"""Yield slices of ``items`` no larger than ``chunk_size``."""
320+
if chunk_size <= 0:
321+
raise ValueError(f"chunk_size must be positive, got {chunk_size}")
322+
for start in range(0, len(items), chunk_size):
323+
yield items[start : start + chunk_size]
324+
325+
def _batch_max_records(self, config: RepoConfig) -> int:
326+
store_cfg = config.online_store
327+
if not isinstance(store_cfg, AerospikeOnlineStoreConfig):
328+
raise RuntimeError(f"{config.online_store.type = }. It must be aerospike.")
329+
return store_cfg.batch_max_records
330+
297331
def _get_client(self, config: RepoConfig) -> aerospike.Client:
298332
"""Lazily create and cache an Aerospike client on first use.
299333
300334
The underlying C client maintains its own connection pool, so a single
301-
cached instance is safe to share across calls on this store.
335+
cached instance is safe to share across calls on this store. Creation
336+
is guarded by a lock so concurrent first callers in a threaded feature
337+
server do not leak extra connections.
302338
"""
303339
if self._client is not None:
304340
return self._client
305341

306-
if not isinstance(config.online_store, AerospikeOnlineStoreConfig):
307-
raise RuntimeError(f"{config.online_store.type = }. It must be aerospike.")
308-
store_cfg = config.online_store
342+
with self._client_lock:
343+
if self._client is not None:
344+
return self._client
309345

310-
read_policy: Dict[str, Any] = {
311-
"total_timeout": store_cfg.read_timeout_ms,
312-
"max_retries": store_cfg.max_retries,
313-
}
314-
write_policy: Dict[str, Any] = {
315-
"total_timeout": store_cfg.write_timeout_ms,
316-
"max_retries": store_cfg.max_retries,
317-
}
318-
batch_policy: Dict[str, Any] = {
319-
"total_timeout": store_cfg.batch_total_timeout_ms,
320-
"max_retries": store_cfg.max_retries,
321-
}
322-
if store_cfg.socket_timeout_ms is not None:
323-
# socket_timeout is the per-attempt deadline; without it,
324-
# total_timeout is the whole budget and retries never fire.
325-
read_policy["socket_timeout"] = store_cfg.socket_timeout_ms
326-
write_policy["socket_timeout"] = store_cfg.socket_timeout_ms
327-
batch_policy["socket_timeout"] = store_cfg.socket_timeout_ms
328-
329-
client_config: Dict[str, Any] = {
330-
"hosts": [tuple(h) for h in store_cfg.hosts],
331-
"policies": {
332-
"read": read_policy,
333-
"write": write_policy,
334-
"batch": batch_policy,
335-
},
336-
**store_cfg.client_kwargs,
337-
}
338-
if store_cfg.user:
339-
if store_cfg.password is None:
340-
raise ValueError(
341-
"AerospikeOnlineStoreConfig.user is set but password is not."
346+
if not isinstance(config.online_store, AerospikeOnlineStoreConfig):
347+
raise RuntimeError(
348+
f"{config.online_store.type = }. It must be aerospike."
342349
)
343-
client_config["user"] = store_cfg.user
344-
client_config["password"] = store_cfg.password.get_secret_value()
345-
client_config["auth_mode"] = _AUTH_MODE_TO_CONSTANT[store_cfg.auth_mode]
346-
if store_cfg.tls:
347-
client_config["tls"] = store_cfg.tls
350+
store_cfg = config.online_store
351+
352+
read_policy: Dict[str, Any] = {
353+
"total_timeout": store_cfg.read_timeout_ms,
354+
"max_retries": store_cfg.max_retries,
355+
}
356+
write_policy: Dict[str, Any] = {
357+
"total_timeout": store_cfg.write_timeout_ms,
358+
"max_retries": store_cfg.max_retries,
359+
}
360+
batch_policy: Dict[str, Any] = {
361+
"total_timeout": store_cfg.batch_total_timeout_ms,
362+
"max_retries": store_cfg.max_retries,
363+
}
364+
if store_cfg.socket_timeout_ms is not None:
365+
# socket_timeout is the per-attempt deadline; without it,
366+
# total_timeout is the whole budget and retries never fire.
367+
read_policy["socket_timeout"] = store_cfg.socket_timeout_ms
368+
write_policy["socket_timeout"] = store_cfg.socket_timeout_ms
369+
batch_policy["socket_timeout"] = store_cfg.socket_timeout_ms
370+
371+
client_config: Dict[str, Any] = {
372+
"hosts": [tuple(h) for h in store_cfg.hosts],
373+
"policies": {
374+
"read": read_policy,
375+
"write": write_policy,
376+
"batch": batch_policy,
377+
},
378+
**store_cfg.client_kwargs,
379+
}
380+
if store_cfg.user:
381+
if store_cfg.password is None:
382+
raise ValueError(
383+
"AerospikeOnlineStoreConfig.user is set but password is not."
384+
)
385+
client_config["user"] = store_cfg.user
386+
client_config["password"] = store_cfg.password.get_secret_value()
387+
client_config["auth_mode"] = _AUTH_MODE_TO_CONSTANT[store_cfg.auth_mode]
388+
if store_cfg.tls:
389+
client_config["tls"] = store_cfg.tls
348390

349-
self._client = aerospike.client(client_config).connect()
350-
return self._client
391+
self._client = aerospike.client(client_config).connect()
392+
return self._client
351393

352394
def _set_name(self, config: RepoConfig, fv_name: Optional[str] = None) -> str:
353395
"""Resolve the Aerospike set name for a feature view.
@@ -568,17 +610,23 @@ def online_write_batch(
568610
client = self._get_client(config)
569611
namespace = self._namespace_for_fv(config, table.name)
570612
set_name = self._set_name(config, table.name)
571-
batch = self._build_batch_writes(config, table, data, namespace, set_name)
572-
if batch.batch_records:
573-
client.batch_write(batch)
574-
# Per-record result codes must be inspected: client.batch_write
575-
# only raises if the whole request was rejected. A partial failure
576-
# (e.g. a single-partition timeout) is otherwise silent, which in
577-
# an online-serving path presents downstream as "model saw stale
578-
# features" weeks after the fact.
579-
self._raise_on_batch_errors(batch.batch_records, set_name, op="write")
580-
if progress:
581-
progress(len(data))
613+
chunk_size = self._batch_max_records(config)
614+
written = 0
615+
for chunk in self._chunked(data, chunk_size):
616+
batch = self._build_batch_writes(
617+
config, table, list(chunk), namespace, set_name
618+
)
619+
if batch.batch_records:
620+
client.batch_write(batch)
621+
# Per-record result codes must be inspected: client.batch_write
622+
# only raises if the whole request was rejected. A partial failure
623+
# (e.g. a single-partition timeout) is otherwise silent, which in
624+
# an online-serving path presents downstream as "model saw stale
625+
# features" weeks after the fact.
626+
self._raise_on_batch_errors(batch.batch_records, set_name, op="write")
627+
written += len(chunk)
628+
if progress:
629+
progress(written)
582630

583631
# ------------------------------------------------------------------
584632
# Read path
@@ -618,59 +666,64 @@ def online_read(
618666
client = self._get_client(config)
619667
ns = self._namespace_for_fv(config, table.name)
620668
set_name = self._set_name(config, table.name)
621-
622-
keys = [
623-
(
624-
ns,
625-
set_name,
626-
bytearray(
627-
serialize_entity_key(
628-
k,
629-
entity_key_serialization_version=config.entity_key_serialization_version,
630-
)
631-
),
632-
)
633-
for k in entity_keys
634-
]
635669
read_ops = self._build_read_ops(table.name, requested_features)
636-
637-
batch = client.batch_operate(keys, read_ops)
670+
chunk_size = self._batch_max_records(config)
638671

639672
# ``ids`` and ``docs`` use immutable ``bytes`` because ``bytearray`` is
640673
# unhashable and can't key a dict. Keys on the wire must stay
641674
# ``bytearray`` (see ``_aerospike_key``) — we only convert here for
642675
# lookup.
643-
ids = [bytes(user_key) for _, _, user_key in keys]
676+
ids: List[bytes] = []
644677
docs: Dict[bytes, Dict[str, Any]] = {}
645-
# batch_operate preserves request order. We pair each response with
646-
# the original user-key rather than ``br.key[2]``: the Aerospike
647-
# client may return the key in a different representation (e.g. only
648-
# the first byte as a str when the write didn't use POLICY_KEY_SEND
649-
# for reads).
650-
for user_key, br in zip(ids, batch.batch_records):
651-
if br.result == _AS_ERR_RECORD_NOT_FOUND:
652-
continue
653-
if br.result == _AS_ERR_OP_NOT_APPLICABLE:
654-
# The record exists but the nested feature-view slot doesn't;
655-
# treat as a miss to match the OnlineStore contract.
656-
continue
657-
if br.result != _AS_OK:
658-
raise RuntimeError(
659-
f"Aerospike batch_operate returned a non-OK status for "
660-
f"entity (ns={ns}, set={set_name}): result={br.result}"
678+
679+
for entity_chunk in self._chunked(entity_keys, chunk_size):
680+
keys = [
681+
(
682+
ns,
683+
set_name,
684+
bytearray(
685+
serialize_entity_key(
686+
k,
687+
entity_key_serialization_version=config.entity_key_serialization_version,
688+
)
689+
),
661690
)
662-
if br.record is None:
663-
continue
664-
_, _, bins = br.record
665-
raw_features = bins.get("features") if bins else None
666-
fv_event_ts_ms = bins.get("event_ts") if bins else None
667-
fv_features = self._normalize_projected_features(raw_features)
668-
docs[user_key] = {
669-
"features": {table.name: fv_features}
670-
if fv_features is not None
671-
else {},
672-
"event_timestamps": {table.name: _epoch_ms_to_datetime(fv_event_ts_ms)},
673-
}
691+
for k in entity_chunk
692+
]
693+
batch = client.batch_operate(keys, read_ops)
694+
chunk_ids = [bytes(user_key) for _, _, user_key in keys]
695+
ids.extend(chunk_ids)
696+
# batch_operate preserves request order. We pair each response with
697+
# the original user-key rather than ``br.key[2]``: the Aerospike
698+
# client may return the key in a different representation (e.g. only
699+
# the first byte as a str when the write didn't use POLICY_KEY_SEND
700+
# for reads).
701+
for user_key, br in zip(chunk_ids, batch.batch_records):
702+
if br.result == _AS_ERR_RECORD_NOT_FOUND:
703+
continue
704+
if br.result == _AS_ERR_OP_NOT_APPLICABLE:
705+
# The record exists but the nested feature-view slot doesn't;
706+
# treat as a miss to match the OnlineStore contract.
707+
continue
708+
if br.result != _AS_OK:
709+
raise RuntimeError(
710+
f"Aerospike batch_operate returned a non-OK status for "
711+
f"entity (ns={ns}, set={set_name}): result={br.result}"
712+
)
713+
if br.record is None:
714+
continue
715+
_, _, bins = br.record
716+
raw_features = bins.get("features") if bins else None
717+
fv_event_ts_ms = bins.get("event_ts") if bins else None
718+
fv_features = self._normalize_projected_features(raw_features)
719+
docs[user_key] = {
720+
"features": {table.name: fv_features}
721+
if fv_features is not None
722+
else {},
723+
"event_timestamps": {
724+
table.name: _epoch_ms_to_datetime(fv_event_ts_ms)
725+
},
726+
}
674727

675728
return self._convert_raw_docs_to_proto(ids, docs, table)
676729

@@ -867,10 +920,11 @@ async def initialize(self, config: RepoConfig) -> None:
867920

868921
async def close(self) -> None:
869922
"""Release the cached Aerospike client, if any."""
870-
if self._client is None:
871-
return
872-
client = self._client
873-
self._client = None
923+
with self._client_lock:
924+
if self._client is None:
925+
return
926+
client = self._client
927+
self._client = None
874928
loop = asyncio.get_running_loop()
875929
await loop.run_in_executor(None, client.close)
876930

@@ -987,6 +1041,7 @@ def teardown(
9871041
# meaning.
9881042
for ns, set_name in sorted(pairs):
9891043
client.truncate(ns, set_name, 0)
990-
if self._client is not None:
991-
self._client.close()
992-
self._client = None
1044+
with self._client_lock:
1045+
if self._client is not None:
1046+
self._client.close()
1047+
self._client = None

0 commit comments

Comments
 (0)