Visitar URL original
Merge latest PR #6550 into merged-spark-e2e · aniketpalu/feast@8dedeea · GitHub
Skip to content

Commit 8dedeea

Browse files
aniketpalucursoragent
aniketpalu
and
cursoragent
committed
Merge latest PR feast-dev#6550 into merged-spark-e2e
Take supports_batch isolation; keep _force_local skip for remote async materialize on both batch and per-FV paths. Co-authored-by: Cursor <cursoragent@cursor.com>
2 parents 1a74a4a + a38a56d commit 8dedeea

3 files changed

Lines changed: 219 additions & 76 deletions

File tree

‎sdk/python/feast/feature_store.py‎

Lines changed: 205 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,58 @@ def _submit_and_process_materialization_jobs(
527527
if first_error:
528528
raise first_error
529529

530+
def _materialize_fvs_batch(
531+
self,
532+
provider,
533+
fv_with_dates: list,
534+
end_date: datetime,
535+
tqdm_builder,
536+
disable_event_timestamp: bool = False,
537+
_force_local: bool = False,
538+
) -> None:
539+
"""Batch path: collect all FVs, submit to engine in one call.
540+
541+
Only used when ``provider.batch_engine.supports_batch`` is True.
542+
"""
543+
from feast.infra.common.materialization_job import MaterializationTask
544+
545+
tasks: list = []
546+
regular_fvs: list = []
547+
previous_states: dict = {}
548+
fv_start_dates: dict = {"__end_date__": end_date}
549+
550+
for feature_view, fv_start in fv_with_dates:
551+
# Skip transition when _force_local: async handler already set MATERIALIZING.
552+
if _force_local:
553+
previous_states[feature_view.name] = getattr(
554+
feature_view, "state", None
555+
)
556+
else:
557+
self._transition_fv_to_materializing(
558+
feature_view, regular_fvs, previous_states
559+
)
560+
regular_fvs.append(feature_view)
561+
fv_start_dates[feature_view.name] = fv_start
562+
tasks.append(
563+
MaterializationTask(
564+
project=self.project,
565+
feature_view=feature_view,
566+
start_time=fv_start,
567+
end_time=end_date,
568+
tqdm_builder=tqdm_builder,
569+
disable_event_timestamp=disable_event_timestamp,
570+
)
571+
)
572+
573+
if tasks:
574+
self._submit_and_process_materialization_jobs(
575+
provider,
576+
tasks,
577+
regular_fvs,
578+
previous_states,
579+
fv_start_dates,
580+
)
581+
530582
@property
531583
def openlineage_emitter(self) -> Optional[Any]:
532584
"""Gets the OpenLineage emitter of this feature store."""
@@ -2623,20 +2675,14 @@ def materialize_incremental(
26232675

26242676
_mat_start = time.monotonic()
26252677
try:
2626-
from feast.infra.common.materialization_job import (
2627-
MaterializationTask,
2628-
)
2629-
26302678
provider = self._get_provider()
26312679
end_date_tz = utils.make_tzaware(end_date) or _utc_now()
26322680

26332681
def tqdm_builder(length):
26342682
return tqdm(total=length, ncols=100)
26352683

2636-
tasks: list = []
2637-
regular_fvs: list = []
2638-
previous_states: dict = {}
2639-
fv_start_dates: dict = {"__end_date__": end_date_tz}
2684+
# (feature_view, start_date) — start_date is always set before append
2685+
regular_fvs_with_dates: list[tuple[Any, datetime]] = []
26402686

26412687
for feature_view in feature_views_to_materialize:
26422688
if isinstance(feature_view, OnDemandFeatureView):
@@ -2666,61 +2712,105 @@ def tqdm_builder(length):
26662712
)
26672713
continue
26682714

2669-
fv_start_date = feature_view.most_recent_end_time
2670-
if fv_start_date is None:
2715+
start_date = feature_view.most_recent_end_time
2716+
if start_date is None:
26712717
if feature_view.ttl is None:
26722718
raise Exception(
26732719
f"No start time found for feature view {feature_view.name}. materialize_incremental() requires"
26742720
f" either a ttl to be set or for materialize() to have been run at least once."
26752721
)
26762722
elif feature_view.ttl.total_seconds() > 0:
2677-
fv_start_date = _utc_now() - feature_view.ttl
2723+
start_date = _utc_now() - feature_view.ttl
26782724
else:
26792725
# TODO(felixwang9817): Find the earliest timestamp for this specific feature
26802726
# view from the offline store, and set the start date to that timestamp.
26812727
print(
26822728
f"Since the ttl is 0 for feature view {Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}, "
26832729
"the start date will be set to 1 year before the current time."
26842730
)
2685-
fv_start_date = _utc_now() - timedelta(weeks=52)
2686-
2687-
fv_start_date = utils.make_tzaware(fv_start_date)
2688-
fv_start_dates[feature_view.name] = fv_start_date
2731+
start_date = _utc_now() - timedelta(weeks=52)
26892732

2733+
start_date = utils.make_tzaware(start_date)
26902734
print(
26912735
f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}"
2692-
f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(fv_start_date.replace(microsecond=0))}{Style.RESET_ALL}"
2736+
f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(start_date.replace(microsecond=0))}{Style.RESET_ALL}"
26932737
f" to {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(end_date.replace(microsecond=0))}{Style.RESET_ALL}:"
26942738
)
26952739

2696-
# Skip transition when _force_local: async handler already set MATERIALIZING.
2697-
if _force_local:
2698-
previous_states[feature_view.name] = getattr(
2699-
feature_view, "state", None
2700-
)
2701-
else:
2702-
self._transition_fv_to_materializing(
2703-
feature_view, regular_fvs, previous_states
2704-
)
2705-
regular_fvs.append(feature_view)
2706-
tasks.append(
2707-
MaterializationTask(
2708-
project=self.project,
2709-
feature_view=feature_view,
2710-
start_time=fv_start_date,
2711-
end_time=end_date_tz,
2712-
tqdm_builder=tqdm_builder,
2713-
)
2714-
)
2740+
regular_fvs_with_dates.append((feature_view, start_date))
27152741

2716-
if tasks:
2717-
self._submit_and_process_materialization_jobs(
2742+
# batch_engine is on PassthroughProvider (concrete); same access as
2743+
# _submit_and_process_materialization_jobs via untyped provider.
2744+
if getattr(provider, "batch_engine").supports_batch:
2745+
self._materialize_fvs_batch(
27182746
provider,
2719-
tasks,
2720-
regular_fvs,
2721-
previous_states,
2722-
fv_start_dates,
2747+
regular_fvs_with_dates,
2748+
end_date_tz,
2749+
tqdm_builder,
2750+
_force_local=_force_local,
27232751
)
2752+
else:
2753+
for feature_view, start_date in regular_fvs_with_dates:
2754+
# Transition state to MATERIALIZING before starting.
2755+
# Skip when _force_local: async handler already set MATERIALIZING.
2756+
previous_state = getattr(feature_view, "state", None)
2757+
if (
2758+
not _force_local
2759+
and hasattr(feature_view, "state")
2760+
and feature_view.state != FeatureViewState.STATE_UNSPECIFIED
2761+
):
2762+
if not feature_view.state.can_transition_to(
2763+
FeatureViewState.MATERIALIZING
2764+
):
2765+
raise ValueError(
2766+
f"FeatureView {feature_view.name} cannot transition "
2767+
f"from {feature_view.state.name} to MATERIALIZING."
2768+
)
2769+
feature_view.state = FeatureViewState.MATERIALIZING
2770+
self.registry.apply_feature_view(
2771+
feature_view, self.project, commit=True
2772+
)
2773+
2774+
fv_start = time.monotonic()
2775+
fv_success = True
2776+
try:
2777+
provider.materialize_single_feature_view(
2778+
config=self.config,
2779+
feature_view=feature_view,
2780+
start_date=start_date,
2781+
end_date=end_date_tz,
2782+
registry=self.registry,
2783+
project=self.project,
2784+
tqdm_builder=tqdm_builder,
2785+
)
2786+
except Exception:
2787+
fv_success = False
2788+
# Roll back state to previous value on failure.
2789+
if (
2790+
hasattr(feature_view, "state")
2791+
and previous_state is not None
2792+
and previous_state != FeatureViewState.STATE_UNSPECIFIED
2793+
):
2794+
feature_view.state = previous_state
2795+
self.registry.apply_feature_view(
2796+
feature_view, self.project, commit=True
2797+
)
2798+
raise
2799+
finally:
2800+
_tracker = _get_track_materialization()
2801+
if _tracker is not None:
2802+
_tracker(
2803+
feature_view.name,
2804+
fv_success,
2805+
time.monotonic() - fv_start,
2806+
)
2807+
2808+
self.registry.apply_materialization(
2809+
feature_view,
2810+
self.project,
2811+
start_date,
2812+
end_date_tz,
2813+
)
27242814

27252815
materialized_fv_names = [
27262816
fv.name
@@ -2837,21 +2927,14 @@ def materialize(
28372927

28382928
_mat_start = time.monotonic()
28392929
try:
2840-
from feast.infra.common.materialization_job import (
2841-
MaterializationTask,
2842-
)
2843-
28442930
provider = self._get_provider()
28452931
start_date = utils.make_tzaware(start_date)
28462932
end_date = utils.make_tzaware(end_date)
28472933

28482934
def tqdm_builder(length):
28492935
return tqdm(total=length, ncols=100)
28502936

2851-
tasks: list = []
2852-
regular_fvs: list = []
2853-
previous_states: dict = {}
2854-
fv_start_dates: dict = {"__end_date__": end_date}
2937+
regular_fvs_with_dates: list[tuple[Any, datetime]] = []
28552938

28562939
for feature_view in feature_views_to_materialize:
28572940
if isinstance(feature_view, OnDemandFeatureView):
@@ -2871,36 +2954,82 @@ def tqdm_builder(length):
28712954
f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}:"
28722955
)
28732956

2874-
# Skip transition when _force_local: async handler already set MATERIALIZING.
2875-
if _force_local:
2876-
previous_states[feature_view.name] = getattr(
2877-
feature_view, "state", None
2878-
)
2879-
else:
2880-
self._transition_fv_to_materializing(
2881-
feature_view, regular_fvs, previous_states
2882-
)
2883-
regular_fvs.append(feature_view)
2884-
fv_start_dates[feature_view.name] = start_date
2885-
tasks.append(
2886-
MaterializationTask(
2887-
project=self.project,
2888-
feature_view=feature_view,
2889-
start_time=start_date,
2890-
end_time=end_date,
2891-
tqdm_builder=tqdm_builder,
2892-
disable_event_timestamp=disable_event_timestamp,
2893-
)
2894-
)
2957+
regular_fvs_with_dates.append((feature_view, start_date))
28952958

2896-
if tasks:
2897-
self._submit_and_process_materialization_jobs(
2959+
# batch_engine is on PassthroughProvider (concrete); same access as
2960+
# _submit_and_process_materialization_jobs via untyped provider.
2961+
if getattr(provider, "batch_engine").supports_batch:
2962+
self._materialize_fvs_batch(
28982963
provider,
2899-
tasks,
2900-
regular_fvs,
2901-
previous_states,
2902-
fv_start_dates,
2964+
regular_fvs_with_dates,
2965+
end_date,
2966+
tqdm_builder,
2967+
disable_event_timestamp=disable_event_timestamp,
2968+
_force_local=_force_local,
29032969
)
2970+
else:
2971+
for feature_view, fv_start in regular_fvs_with_dates:
2972+
# Transition state to MATERIALIZING before starting.
2973+
# Skip when _force_local: async handler already set MATERIALIZING.
2974+
previous_state = getattr(feature_view, "state", None)
2975+
if (
2976+
not _force_local
2977+
and hasattr(feature_view, "state")
2978+
and feature_view.state != FeatureViewState.STATE_UNSPECIFIED
2979+
):
2980+
if not feature_view.state.can_transition_to(
2981+
FeatureViewState.MATERIALIZING
2982+
):
2983+
raise ValueError(
2984+
f"FeatureView {feature_view.name} cannot transition "
2985+
f"from {feature_view.state.name} to MATERIALIZING."
2986+
)
2987+
feature_view.state = FeatureViewState.MATERIALIZING
2988+
self.registry.apply_feature_view(
2989+
feature_view, self.project, commit=True
2990+
)
2991+
2992+
fv_start_time = time.monotonic()
2993+
fv_success = True
2994+
try:
2995+
provider.materialize_single_feature_view(
2996+
config=self.config,
2997+
feature_view=feature_view,
2998+
start_date=fv_start,
2999+
end_date=end_date,
3000+
registry=self.registry,
3001+
project=self.project,
3002+
tqdm_builder=tqdm_builder,
3003+
disable_event_timestamp=disable_event_timestamp,
3004+
)
3005+
except Exception:
3006+
fv_success = False
3007+
# Roll back state to previous value on failure.
3008+
if (
3009+
hasattr(feature_view, "state")
3010+
and previous_state is not None
3011+
and previous_state != FeatureViewState.STATE_UNSPECIFIED
3012+
):
3013+
feature_view.state = previous_state
3014+
self.registry.apply_feature_view(
3015+
feature_view, self.project, commit=True
3016+
)
3017+
raise
3018+
finally:
3019+
_tracker = _get_track_materialization()
3020+
if _tracker is not None:
3021+
_tracker(
3022+
feature_view.name,
3023+
fv_success,
3024+
time.monotonic() - fv_start_time,
3025+
)
3026+
3027+
self.registry.apply_materialization(
3028+
feature_view,
3029+
self.project,
3030+
fv_start,
3031+
end_date,
3032+
)
29043033

29053034
materialized_fv_names = [
29063035
fv.name

‎sdk/python/feast/infra/compute_engines/base.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,16 @@ def teardown_infra(
8585
"""
8686
pass
8787

88+
@property
89+
def supports_batch(self) -> bool:
90+
"""Whether this engine can accept all tasks in a single materialize() call.
91+
92+
When True, feature_store.py collects all FV tasks upfront and submits
93+
them in one batch. When False (default), the standard per-FV loop
94+
through ``provider.materialize_single_feature_view()`` is used.
95+
"""
96+
return False
97+
8898
def materialize(
8999
self,
90100
registry: BaseRegistry,

‎sdk/python/feast/infra/compute_engines/spark_application/compute.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@ def __init__(
9797
self.custom_api = client.CustomObjectsApi(self.k8s_client)
9898
self._server_id = uuid.uuid4().hex[:8]
9999

100+
@property
101+
def supports_batch(self) -> bool:
102+
return True
103+
100104
def update(
101105
self,
102106
project: str,

0 commit comments

Comments
 (0)