From 560d5c6280085d357c26b685ce99e561a9cec2e5 Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Wed, 24 Jun 2026 14:25:05 +0530 Subject: [PATCH 1/4] feat: Add SparkApplicationComputeEngine for Kubernetes-native batch materialization Adds a new batch compute engine that submits materialization jobs as SparkApplication CRDs via the Kubeflow Spark Operator. One 'feast materialize' call creates one SparkApplication pod that processes all feature views using distributed Spark, rather than running in-process on the Feast server. Key changes: - Refactor materialize()/materialize_incremental() to pass all tasks to the engine in a single batch call instead of looping per feature view. Existing engines are unaffected (base class loops tasks internally via _materialize_one). - Add public get_provider() method on FeatureStore. - New spark_application engine: config, compute, job, driver script, Dockerfile. - 12 unit tests covering config, validation, CR structure, state mapping, timeout, cleanup, and job naming. --- sdk/python/feast/feature_store.py | 329 +++++++++------- .../spark_application/.dockerignore | 8 + .../spark_application/Dockerfile | 21 ++ .../spark_application/__init__.py | 0 .../spark_application/compute.py | 356 ++++++++++++++++++ .../spark_application/config.py | 68 ++++ .../compute_engines/spark_application/job.py | 109 ++++++ .../compute_engines/spark_application/main.py | 179 +++++++++ sdk/python/feast/repo_config.py | 1 + .../compute_engines/test_spark_application.py | 217 +++++++++++ 10 files changed, 1162 insertions(+), 126 deletions(-) create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/.dockerignore create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/Dockerfile create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/__init__.py create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/compute.py create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/config.py create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/job.py create mode 100644 sdk/python/feast/infra/compute_engines/spark_application/main.py create mode 100644 sdk/python/tests/unit/infra/compute_engines/test_spark_application.py diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 3dcd4189d3c..68fffdd0e15 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -424,6 +424,10 @@ def _get_provider(self) -> Provider: # TODO: Bake self.repo_path into self.config so that we dont only have one interface to paths return self.provider + def get_provider(self) -> Provider: + """Public accessor for the provider instance.""" + return self._get_provider() + @property def openlineage_emitter(self) -> Optional[Any]: """Gets the OpenLineage emitter of this feature store.""" @@ -2136,7 +2140,6 @@ def materialize_incremental( _mat_start = time.monotonic() try: - # TODO paging large loads for feature_view in feature_views_to_materialize: if isinstance(feature_view, OnDemandFeatureView): if feature_view.write_to_online_store: @@ -2163,98 +2166,138 @@ def materialize_incremental( end_date, full_feature_names=full_feature_names, ) - continue - start_date = feature_view.most_recent_end_time - if start_date is None: - if feature_view.ttl is None: - raise Exception( - f"No start time found for feature view {feature_view.name}. materialize_incremental() requires" - f" either a ttl to be set or for materialize() to have been run at least once." - ) - elif feature_view.ttl.total_seconds() > 0: - start_date = _utc_now() - feature_view.ttl - else: - # TODO(felixwang9817): Find the earliest timestamp for this specific feature - # view from the offline store, and set the start date to that timestamp. - print( - f"Since the ttl is 0 for feature view {Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}, " - "the start date will be set to 1 year before the current time." - ) - start_date = _utc_now() - timedelta(weeks=52) - provider = self._get_provider() - print( - f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}" - f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(start_date.replace(microsecond=0))}{Style.RESET_ALL}" - f" to {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(end_date.replace(microsecond=0))}{Style.RESET_ALL}:" + regular_fvs = [ + fv + for fv in feature_views_to_materialize + if not isinstance(fv, OnDemandFeatureView) + ] + + if regular_fvs: + # Deferred import: top-level causes circular import via + # feast.__init__ -> feature_store -> materialization_job -> feast + from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + MaterializationTask, ) - def tqdm_builder(length): - return tqdm(total=length, ncols=100) + provider = self._get_provider() + end_date_tz = utils.make_tzaware(end_date) or _utc_now() + + # Compute per-FV start dates and transition all to MATERIALIZING + # before submitting work to the engine. + previous_states: dict = {} + fv_start_dates: dict = {} + + for fv in regular_fvs: + fv_start_date = fv.most_recent_end_time + if fv_start_date is None: + if fv.ttl is None: + raise Exception( + f"No start time found for feature view {fv.name}. materialize_incremental() requires" + f" either a ttl to be set or for materialize() to have been run at least once." + ) + elif fv.ttl.total_seconds() > 0: + fv_start_date = _utc_now() - fv.ttl + else: + print( + f"Since the ttl is 0 for feature view {Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}, " + "the start date will be set to 1 year before the current time." + ) + fv_start_date = _utc_now() - timedelta(weeks=52) - start_date = utils.make_tzaware(start_date) - end_date = utils.make_tzaware(end_date) or _utc_now() + fv_start_dates[fv.name] = utils.make_tzaware(fv_start_date) - # Transition state to MATERIALIZING before starting. - # Only enforce when the state machine is active (not STATE_UNSPECIFIED). - previous_state = getattr(feature_view, "state", None) - if ( - hasattr(feature_view, "state") - and feature_view.state != FeatureViewState.STATE_UNSPECIFIED - ): - if not feature_view.state.can_transition_to( - FeatureViewState.MATERIALIZING + print( + f"{Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}" + f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(fv_start_date.replace(microsecond=0))}{Style.RESET_ALL}" + f" to {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(end_date.replace(microsecond=0))}{Style.RESET_ALL}:" + ) + + previous_states[fv.name] = getattr(fv, "state", None) + if ( + hasattr(fv, "state") + and fv.state != FeatureViewState.STATE_UNSPECIFIED ): - raise ValueError( - f"FeatureView {feature_view.name} cannot transition " - f"from {feature_view.state.name} to MATERIALIZING." + if not fv.state.can_transition_to( + FeatureViewState.MATERIALIZING + ): + raise ValueError( + f"FeatureView {fv.name} cannot transition " + f"from {fv.state.name} to MATERIALIZING." + ) + fv.state = FeatureViewState.MATERIALIZING + self.registry.apply_feature_view( + fv, self.project, commit=True ) - feature_view.state = FeatureViewState.MATERIALIZING - self.registry.apply_feature_view( - feature_view, self.project, commit=True - ) - fv_start = time.monotonic() - fv_success = True - try: - provider.materialize_single_feature_view( - config=self.config, - feature_view=feature_view, - start_date=start_date, - end_date=end_date, - registry=self.registry, + def tqdm_builder(length): + return tqdm(total=length, ncols=100) + + tasks = [ + MaterializationTask( project=self.project, + feature_view=fv, + start_time=fv_start_dates[fv.name], + end_time=end_date_tz, tqdm_builder=tqdm_builder, ) + for fv in regular_fvs + ] + + fv_batch_start = time.monotonic() + try: + jobs = provider.batch_engine.materialize( + self.registry, tasks + ) except Exception: - fv_success = False - # Roll back state to previous value on failure. - if ( - hasattr(feature_view, "state") - and previous_state is not None - and previous_state != FeatureViewState.STATE_UNSPECIFIED - ): - feature_view.state = previous_state - self.registry.apply_feature_view( - feature_view, self.project, commit=True - ) + for fv in regular_fvs: + prev = previous_states[fv.name] + if ( + prev is not None + and prev != FeatureViewState.STATE_UNSPECIFIED + ): + fv.state = prev + self.registry.apply_feature_view( + fv, self.project, commit=True + ) raise - finally: + + errors = [] + for fv, job in zip(regular_fvs, jobs): + fv_status = job.status() + fv_success = fv_status != MaterializationJobStatus.ERROR + + if fv_status == MaterializationJobStatus.SUCCEEDED: + self.registry.apply_materialization( + fv, + self.project, + fv_start_dates[fv.name], + end_date_tz, + ) + elif fv_status == MaterializationJobStatus.ERROR: + prev = previous_states[fv.name] + if ( + prev is not None + and prev != FeatureViewState.STATE_UNSPECIFIED + ): + fv.state = prev + self.registry.apply_feature_view( + fv, self.project, commit=True + ) + if job.error(): + errors.append(job.error()) + _tracker = _get_track_materialization() if _tracker is not None: _tracker( - feature_view.name, + fv.name, fv_success, - time.monotonic() - fv_start, + time.monotonic() - fv_batch_start, ) - if not isinstance(feature_view, OnDemandFeatureView): - self.registry.apply_materialization( - feature_view, - self.project, - start_date, - end_date, - ) + if errors: + raise errors[0] materialized_fv_names = [ fv.name @@ -2346,7 +2389,6 @@ def materialize( _mat_start = time.monotonic() try: - # TODO paging large loads for feature_view in feature_views_to_materialize: if isinstance(feature_view, OnDemandFeatureView): if feature_view.write_to_online_store: @@ -2359,78 +2401,113 @@ def materialize( end_date, full_feature_names=full_feature_names, ) - continue - provider = self._get_provider() - print( - f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}:" - ) - def tqdm_builder(length): - return tqdm(total=length, ncols=100) + regular_fvs = [ + fv + for fv in feature_views_to_materialize + if not isinstance(fv, OnDemandFeatureView) + ] + if regular_fvs: + # Deferred import: top-level causes circular import via + # feast.__init__ -> feature_store -> materialization_job -> feast + from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + MaterializationTask, + ) + + provider = self._get_provider() start_date = utils.make_tzaware(start_date) end_date = utils.make_tzaware(end_date) - # Transition state to MATERIALIZING before starting. - # Only enforce when the state machine is active (not STATE_UNSPECIFIED). - previous_state = getattr(feature_view, "state", None) - if ( - hasattr(feature_view, "state") - and feature_view.state != FeatureViewState.STATE_UNSPECIFIED - ): - if not feature_view.state.can_transition_to( - FeatureViewState.MATERIALIZING + # Transition all FVs to MATERIALIZING before submitting work. + previous_states: dict = {} + for fv in regular_fvs: + print( + f"{Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}:" + ) + previous_states[fv.name] = getattr(fv, "state", None) + if ( + hasattr(fv, "state") + and fv.state != FeatureViewState.STATE_UNSPECIFIED ): - raise ValueError( - f"FeatureView {feature_view.name} cannot transition " - f"from {feature_view.state.name} to MATERIALIZING." + if not fv.state.can_transition_to( + FeatureViewState.MATERIALIZING + ): + raise ValueError( + f"FeatureView {fv.name} cannot transition " + f"from {fv.state.name} to MATERIALIZING." + ) + fv.state = FeatureViewState.MATERIALIZING + self.registry.apply_feature_view( + fv, self.project, commit=True ) - feature_view.state = FeatureViewState.MATERIALIZING - self.registry.apply_feature_view( - feature_view, self.project, commit=True - ) - fv_start = time.monotonic() - fv_success = True - try: - provider.materialize_single_feature_view( - config=self.config, - feature_view=feature_view, - start_date=start_date, - end_date=end_date, - registry=self.registry, + def tqdm_builder(length): + return tqdm(total=length, ncols=100) + + tasks = [ + MaterializationTask( project=self.project, + feature_view=fv, + start_time=start_date, + end_time=end_date, tqdm_builder=tqdm_builder, disable_event_timestamp=disable_event_timestamp, ) + for fv in regular_fvs + ] + + fv_start = time.monotonic() + try: + jobs = provider.batch_engine.materialize( + self.registry, tasks + ) except Exception: - fv_success = False - # Roll back state to previous value on failure. - if ( - hasattr(feature_view, "state") - and previous_state is not None - and previous_state != FeatureViewState.STATE_UNSPECIFIED - ): - feature_view.state = previous_state - self.registry.apply_feature_view( - feature_view, self.project, commit=True - ) + for fv in regular_fvs: + prev = previous_states[fv.name] + if ( + prev is not None + and prev != FeatureViewState.STATE_UNSPECIFIED + ): + fv.state = prev + self.registry.apply_feature_view( + fv, self.project, commit=True + ) raise - finally: + + errors = [] + for fv, job in zip(regular_fvs, jobs): + fv_status = job.status() + fv_success = fv_status != MaterializationJobStatus.ERROR + + if fv_status == MaterializationJobStatus.SUCCEEDED: + self.registry.apply_materialization( + fv, self.project, start_date, end_date, + ) + elif fv_status == MaterializationJobStatus.ERROR: + prev = previous_states[fv.name] + if ( + prev is not None + and prev != FeatureViewState.STATE_UNSPECIFIED + ): + fv.state = prev + self.registry.apply_feature_view( + fv, self.project, commit=True + ) + if job.error(): + errors.append(job.error()) + _tracker = _get_track_materialization() if _tracker is not None: _tracker( - feature_view.name, + fv.name, fv_success, time.monotonic() - fv_start, ) - self.registry.apply_materialization( - feature_view, - self.project, - start_date, - end_date, - ) + if errors: + raise errors[0] materialized_fv_names = [ fv.name diff --git a/sdk/python/feast/infra/compute_engines/spark_application/.dockerignore b/sdk/python/feast/infra/compute_engines/spark_application/.dockerignore new file mode 100644 index 00000000000..0a4119f63bc --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/.dockerignore @@ -0,0 +1,8 @@ +feast/.git +feast/__pycache__ +feast/**/__pycache__ +feast/.mypy_cache +feast/tests +feast/.pixi +**/*.pyc +**/.pytest_cache diff --git a/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile new file mode 100644 index 00000000000..db3a4cea26b --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile @@ -0,0 +1,21 @@ +FROM apache/spark:4.0.1 + +USER root + +RUN pip install --no-cache-dir \ + "feast[redis]==0.64.0" \ + pyyaml \ + kubernetes \ + tqdm + +# Hadoop AWS JARs for S3A filesystem support + AWS SDK v2 bundle +RUN curl -sL https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.4.1/hadoop-aws-3.4.1.jar \ + -o /opt/spark/jars/hadoop-aws-3.4.1.jar && \ + curl -sL https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/1.12.782/aws-java-sdk-bundle-1.12.782.jar \ + -o /opt/spark/jars/aws-java-sdk-bundle-1.12.782.jar && \ + curl -sL https://repo1.maven.org/maven2/software/amazon/awssdk/bundle/2.28.4/bundle-2.28.4.jar \ + -o /opt/spark/jars/bundle-2.28.4.jar + +COPY feast/infra/compute_engines/spark_application/main.py /opt/feast/main.py + +USER spark diff --git a/sdk/python/feast/infra/compute_engines/spark_application/__init__.py b/sdk/python/feast/infra/compute_engines/spark_application/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/feast/infra/compute_engines/spark_application/compute.py b/sdk/python/feast/infra/compute_engines/spark_application/compute.py new file mode 100644 index 00000000000..6c0aecaf738 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/compute.py @@ -0,0 +1,356 @@ +import logging +import time +import uuid +from typing import List, Optional, Sequence, Union + +import pyarrow as pa +import yaml +from kubernetes import client +from kubernetes import config as k8s_config +from kubernetes.client.exceptions import ApiException + +from feast import RepoConfig +from feast.batch_feature_view import BatchFeatureView +from feast.entity import Entity +from feast.feature_view import FeatureView +from feast.infra.common.materialization_job import ( + MaterializationJob, + MaterializationJobStatus, + MaterializationTask, +) +from feast.infra.common.retrieval_task import HistoricalRetrievalTask +from feast.infra.compute_engines.base import ComputeEngine +from feast.infra.offline_stores.offline_store import OfflineStore +from feast.infra.online_stores.online_store import OnlineStore +from feast.infra.registry.base_registry import BaseRegistry +from feast.on_demand_feature_view import OnDemandFeatureView +from feast.stream_feature_view import StreamFeatureView + +from .config import SparkApplicationComputeEngineConfig # noqa: F401 — required for Feast config resolution +from .job import SparkApplicationMaterializationJob + +logger = logging.getLogger(__name__) + + +class SparkApplicationComputeEngine(ComputeEngine): + def __init__( + self, + *, + repo_config: RepoConfig, + offline_store: OfflineStore, + online_store: OnlineStore, + **kwargs, + ): + super().__init__( + repo_config=repo_config, + offline_store=offline_store, + online_store=online_store, + **kwargs, + ) + self.config = repo_config.batch_engine + + # EC-3: SQLite online store — data written inside pod is lost on termination + online_type = getattr(repo_config.online_store, "type", "") + if online_type == "sqlite": + raise ValueError( + "spark_application engine cannot use SQLite online store. " + "SQLite is file-based — data written inside the " + "SparkApplication pod is lost when the pod terminates. " + "Use a network-accessible store: redis, postgres, etc." + ) + + # EC-2: Registry access — pod must be able to reach registry + if not self.config.registry_address: + registry = repo_config.registry + if hasattr(registry, "path") and registry.path: + path = registry.path + is_remote = any( + path.startswith(s) + for s in ("s3://", "gs://", "hdfs://", "http://", "https://", "postgresql", "mysql") + ) + if not is_remote: + raise ValueError( + f"Registry path '{path}' is a local file. " + f"SparkApplication pods cannot access the Feast " + f"server's filesystem. Either:\n" + f" 1. Set registry_address to the Feast registry " + f"gRPC endpoint (recommended), or\n" + f" 2. Use a remote registry path (s3://, gs://), or\n" + f" 3. Switch to registry_type: 'sql' or 'remote'." + ) + + k8s_config.load_config() + self.k8s_client = client.ApiClient() + self.core_v1 = client.CoreV1Api(self.k8s_client) + self.custom_api = client.CustomObjectsApi(self.k8s_client) + self._server_id = uuid.uuid4().hex[:8] + + def update( + self, + project: str, + views_to_delete: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]], + views_to_keep: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView]], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + ): + pass + + def teardown_infra( + self, + project: str, + fvs: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]], + entities: Sequence[Entity], + ): + pass + + def get_historical_features( + self, registry: BaseRegistry, task: HistoricalRetrievalTask + ) -> pa.Table: + raise NotImplementedError( + "SparkApplicationComputeEngine does not yet support get_historical_features(). " + "This is planned for Phase 2." + ) + + def materialize( + self, + registry: BaseRegistry, + tasks: Union[MaterializationTask, List[MaterializationTask]], + **kwargs, + ) -> List[MaterializationJob]: + """Batch all materialization tasks into a single SparkApplication.""" + if isinstance(tasks, MaterializationTask): + tasks = [tasks] + + job_id = uuid.uuid4().hex[:8] + + try: + self._create_secret(job_id, tasks) + except ApiException as e: + job = SparkApplicationMaterializationJob( + job_id, self.config.namespace, self.custom_api, + error=Exception(f"Secret creation failed: {e.reason}"), + ) + return [job for _ in tasks] + + try: + cr = self._build_spark_application_cr(job_id) + self.custom_api.create_namespaced_custom_object( + group="sparkoperator.k8s.io", + version="v1beta2", + namespace=self.config.namespace, + plural="sparkapplications", + body=cr, + ) + except ApiException as e: + self._cleanup(job_id) + job = SparkApplicationMaterializationJob( + job_id, self.config.namespace, self.custom_api, + error=Exception(f"SparkApplication creation failed: {e.reason}"), + ) + return [job for _ in tasks] + + job = SparkApplicationMaterializationJob(job_id, self.config.namespace, self.custom_api) + self._wait_for_completion(job) + return [job for _ in tasks] + + def _build_driver_repo_config(self) -> dict: + """Build feature_store.yaml for the SparkApplication driver pod. + + Two rewrites: + 1. batch_engine → spark.engine: Pod uses SparkComputeEngine with the + active SparkSession (from spark-submit). Enables distributed reads + via SparkReadNode and distributed writes via mapInArrow across executors. + This is NOT recursive — SparkComputeEngine uses the local session, + it does not create CRDs. + 2. registry → remote (if registry_address set): Pod can't access + server's local filesystem. Routes registry ops via gRPC. + + offline_store is NOT rewritten — respects user's configured data sources. + User should configure offline_store: spark for full distributed performance. + """ + config_dict = self.repo_config.model_dump(by_alias=True) + + config_dict["batch_engine"] = {"type": "spark.engine"} + if self.config.spark_conf: + config_dict["batch_engine"]["spark_conf"] = self.config.spark_conf + + if self.config.registry_address: + config_dict["registry"] = { + "registry_type": "remote", + "path": self.config.registry_address, + } + + return config_dict + + def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): + feast_config_yaml = yaml.dump(self._build_driver_repo_config()) + mat_config = { + "operation": "materialize", + "tasks": [ + { + "feature_view": task.feature_view.name, + "start_time": task.start_time.isoformat(), + "end_time": task.end_time.isoformat(), + } + for task in tasks + ], + } + if self.config.concurrency > 1: + mat_config["concurrency"] = self.config.concurrency + mat_config_yaml = yaml.dump(mat_config) + manifest = { + "apiVersion": "v1", + "kind": "Secret", + "metadata": { + "name": f"feast-sa-{job_id}", + "namespace": self.config.namespace, + "labels": {"feast-materializer": "secret", **self.config.labels}, + }, + "stringData": { + "feature_store.yaml": feast_config_yaml, + "materialization_config.yaml": mat_config_yaml, + }, + } + self.core_v1.create_namespaced_secret( + namespace=self.config.namespace, body=manifest + ) + + def _build_spark_application_cr(self, job_id: str) -> dict: + spec = { + "type": "Python", + "mode": "cluster", + "pythonVersion": "3", + "image": self.config.image, + "imagePullPolicy": "IfNotPresent", + "mainApplicationFile": "local:///opt/feast/main.py", + "sparkVersion": self.config.spark_version, + "sparkConf": { + "spark.scheduler.mode": "FAIR", + **(self.config.spark_conf or {}), + "spark.kubernetes.driverEnv.FEAST_SECRET_NAME": f"feast-sa-{job_id}", + "spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE": self.config.namespace, + }, + "restartPolicy": { + "type": self.config.restart_policy, + "onFailureRetries": self.config.max_retries, + "onFailureRetryInterval": 30, + }, + "timeToLiveSeconds": self.config.ttl_seconds_after_finished, + "volumes": [ + {"name": "feast-config", "secret": {"secretName": f"feast-sa-{job_id}"}}, + *self.config.volumes, + ], + "driver": { + "cores": self.config.driver_cores, + "memory": self.config.driver_memory, + "serviceAccount": self.config.service_account, + "volumeMounts": [ + {"name": "feast-config", "mountPath": "/var/feast/"}, + *self.config.volume_mounts, + ], + }, + "executor": { + "instances": max(self.config.executor_instances, 1), + "cores": self.config.executor_cores, + "memory": self.config.executor_memory, + }, + } + + if self.config.image_pull_secrets: + spec["imagePullSecrets"] = self.config.image_pull_secrets + if self.config.hadoop_conf: + spec["hadoopConf"] = self.config.hadoop_conf + if self.config.py_files: + spec["deps"] = {"pyFiles": self.config.py_files} + if self.config.env: + spec["driver"]["env"] = self.config.env + spec["executor"]["env"] = self.config.env + if self.config.env_from: + spec["driver"]["envFrom"] = self.config.env_from + spec["executor"]["envFrom"] = self.config.env_from + if self.config.node_selector: + spec["driver"]["nodeSelector"] = self.config.node_selector + spec["executor"]["nodeSelector"] = self.config.node_selector + if self.config.tolerations: + spec["driver"]["tolerations"] = self.config.tolerations + spec["executor"]["tolerations"] = self.config.tolerations + if self.config.volume_mounts: + spec["executor"]["volumeMounts"] = self.config.volume_mounts + + return { + "apiVersion": "sparkoperator.k8s.io/v1beta2", + "kind": "SparkApplication", + "metadata": { + "name": f"feast-sa-{job_id}", + "namespace": self.config.namespace, + "labels": { + "feast-materializer": "sparkapplication", + "feast-job-id": job_id, + "feast-server-id": self._server_id, + **self._kueue_labels(), + **self.config.labels, + }, + }, + "spec": spec, + } + + def _kueue_labels(self) -> dict: + if self.config.queue_name: + return {"kueue.x-k8s.io/queue-name": self.config.queue_name} + return {} + + def _wait_for_completion(self, job: SparkApplicationMaterializationJob): + start = time.monotonic() + deadline = start + self.config.job_timeout_seconds + while time.monotonic() < deadline: + status = job.status() + elapsed = time.monotonic() - start + logger.info( + f"SparkApplication {job.job_id()} status={status.name} elapsed={elapsed:.0f}s" + ) + if status == MaterializationJobStatus.ERROR: + logs = self._get_driver_logs(job._job_id) + if logs: + logger.error(f"Driver logs (last 50 lines):\n{logs}") + return + if status == MaterializationJobStatus.SUCCEEDED: + return + time.sleep(self.config.poll_interval_seconds) + self._cleanup(job._job_id) + job._error = Exception( + f"SparkApplication {job.job_id()} did not complete " + f"within {self.config.job_timeout_seconds}s" + ) + + def _get_driver_logs(self, job_id: str, tail_lines: int = 50) -> Optional[str]: + """Fetch last N lines of driver pod logs for error diagnostics.""" + try: + pods = self.core_v1.list_namespaced_pod( + namespace=self.config.namespace, + label_selector=f"spark-role=driver,sparkoperator.k8s.io/app-name=feast-sa-{job_id}", + ) + if pods.items: + return self.core_v1.read_namespaced_pod_log( + name=pods.items[0].metadata.name, + namespace=self.config.namespace, + tail_lines=tail_lines, + ) + except ApiException: + logger.warning(f"Could not retrieve driver logs for feast-sa-{job_id}") + return None + + def _cleanup(self, job_id: str): + for fn in [ + lambda: self.custom_api.delete_namespaced_custom_object( + "sparkoperator.k8s.io", "v1beta2", self.config.namespace, + "sparkapplications", f"feast-sa-{job_id}", + ), + lambda: self.core_v1.delete_namespaced_secret( + f"feast-sa-{job_id}", self.config.namespace, + ), + ]: + try: + fn() + except ApiException as e: + if e.status != 404: + logger.warning(f"Cleanup failed: {e.reason}") diff --git a/sdk/python/feast/infra/compute_engines/spark_application/config.py b/sdk/python/feast/infra/compute_engines/spark_application/config.py new file mode 100644 index 00000000000..101c5d864c0 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/config.py @@ -0,0 +1,68 @@ +import warnings +from typing import Dict, List, Literal, Optional + +from pydantic import StrictStr, model_validator + +from feast.repo_config import FeastConfigBaseModel + + +class SparkApplicationComputeEngineConfig(FeastConfigBaseModel): + """Batch Compute Engine config for SparkApplication CRDs via Kubeflow Spark Operator.""" + + type: Literal["spark_application"] = "spark_application" + + image: StrictStr + image_pull_secrets: List[str] = [] + + namespace: StrictStr = "default" + service_account: StrictStr = "" + + driver_cores: int = 1 + driver_memory: StrictStr = "1g" + executor_instances: int = 1 + executor_cores: int = 1 + executor_memory: StrictStr = "1g" + + spark_conf: Optional[Dict[str, str]] = None + hadoop_conf: Optional[Dict[str, str]] = None + spark_version: StrictStr = "4.0.1" + staging_location: Optional[str] = None + + env: List[dict] = [] + env_from: List[dict] = [] + + queue_name: Optional[str] = None + + job_timeout_seconds: int = 3600 + poll_interval_seconds: int = 10 + ttl_seconds_after_finished: int = 3600 + restart_policy: StrictStr = "Never" + max_retries: int = 3 + + concurrency: int = 1 + + labels: Dict[str, str] = {} + + volumes: List[dict] = [] + volume_mounts: List[dict] = [] + py_files: List[str] = [] + node_selector: Optional[Dict[str, str]] = None + tolerations: List[dict] = [] + registry_address: Optional[str] = None + + @model_validator(mode="after") + def _validate_config(self) -> "SparkApplicationComputeEngineConfig": + if self.staging_location: + warnings.warn( + "staging_location is configured but only used for " + "get_historical_features (not yet supported by this engine). " + "It will be ignored for materialize operations.", + stacklevel=2, + ) + for i, entry in enumerate(self.env): + if "name" not in entry: + raise ValueError( + f"env[{i}] is missing required 'name' key. " + f"Each env entry must have at least 'name' and 'value'. Got: {entry}" + ) + return self diff --git a/sdk/python/feast/infra/compute_engines/spark_application/job.py b/sdk/python/feast/infra/compute_engines/spark_application/job.py new file mode 100644 index 00000000000..c93dea7d1ae --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/job.py @@ -0,0 +1,109 @@ +import logging +import time +from typing import Optional + +from kubernetes import client +from kubernetes.client.exceptions import ApiException + +from feast.infra.common.materialization_job import ( + MaterializationJob, + MaterializationJobStatus, +) + +logger = logging.getLogger(__name__) + +_MAX_POLL_RETRIES = 3 +_RETRY_BACKOFF_BASE = 2 + +_STATE_MAP = { + "": MaterializationJobStatus.WAITING, + "SUBMITTED": MaterializationJobStatus.WAITING, + "RUNNING": MaterializationJobStatus.RUNNING, + "COMPLETED": MaterializationJobStatus.SUCCEEDED, + "FAILED": MaterializationJobStatus.ERROR, + "SUBMISSION_FAILED": MaterializationJobStatus.ERROR, + "PENDING_RERUN": MaterializationJobStatus.WAITING, + "INVALIDATING": MaterializationJobStatus.WAITING, + "SUCCEEDING": MaterializationJobStatus.RUNNING, + "FAILING": MaterializationJobStatus.RUNNING, + "SUSPENDING": MaterializationJobStatus.WAITING, + "SUSPENDED": MaterializationJobStatus.WAITING, + "RESUMING": MaterializationJobStatus.WAITING, + "UNKNOWN": MaterializationJobStatus.RUNNING, +} +assert len(_STATE_MAP) == 14 + + +class SparkApplicationMaterializationJob(MaterializationJob): + def __init__( + self, + job_id: str, + namespace: str, + custom_api: client.CustomObjectsApi, + error: Optional[BaseException] = None, + ): + super().__init__() + self._job_id = job_id + self.namespace = namespace + self.custom_api = custom_api + self._error: Optional[BaseException] = error + + def status(self) -> MaterializationJobStatus: + if self._error is not None: + return MaterializationJobStatus.ERROR + + obj = self._get_cr_with_retry() + if obj is None: + return MaterializationJobStatus.ERROR if self._error else MaterializationJobStatus.RUNNING + + state = obj.get("status", {}).get("applicationState", {}).get("state", "") + result = _STATE_MAP.get(state, MaterializationJobStatus.WAITING) + if result == MaterializationJobStatus.ERROR: + msg = obj.get("status", {}).get("applicationState", {}).get( + "errorMessage", f"SparkApplication failed: {state}" + ) + self._error = Exception(msg) + return result + + def _get_cr_with_retry(self) -> Optional[dict]: + """Fetch SparkApplication CR with exponential backoff on transient errors.""" + last_exc = None + for attempt in range(_MAX_POLL_RETRIES): + try: + return self.custom_api.get_namespaced_custom_object( + group="sparkoperator.k8s.io", + version="v1beta2", + namespace=self.namespace, + plural="sparkapplications", + name=f"feast-sa-{self._job_id}", + ) + except ApiException as e: + if e.status == 404: + self._error = Exception( + f"SparkApplication feast-sa-{self._job_id} not found" + ) + return None + last_exc = e + wait = _RETRY_BACKOFF_BASE**attempt + logger.warning( + f"API poll attempt {attempt + 1}/{_MAX_POLL_RETRIES} failed " + f"(HTTP {e.status}), retrying in {wait}s" + ) + time.sleep(wait) + self._error = Exception( + f"Failed to poll SparkApplication after {_MAX_POLL_RETRIES} attempts: " + f"{last_exc.reason if last_exc else 'unknown'}" + ) + return None + + def error(self) -> Optional[BaseException]: + return self._error + + def should_be_retried(self) -> bool: + return False + + def job_id(self) -> str: + return f"feast-sa-{self._job_id}" + + def url(self) -> Optional[str]: + return None diff --git a/sdk/python/feast/infra/compute_engines/spark_application/main.py b/sdk/python/feast/infra/compute_engines/spark_application/main.py new file mode 100644 index 00000000000..73a9db3e885 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/main.py @@ -0,0 +1,179 @@ +import base64 +import logging +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import yaml + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("feast.spark_application.driver") + + +def _load_config_from_secret(): + """Load feast and materialization config from a Kubernetes Secret.""" + from kubernetes import client, config as k8s_config + + k8s_config.load_incluster_config() + v1 = client.CoreV1Api() + secret = v1.read_namespaced_secret( + name=os.environ["FEAST_SECRET_NAME"], + namespace=os.environ["FEAST_SECRET_NAMESPACE"], + ) + feast_config = yaml.safe_load(base64.b64decode(secret.data["feature_store.yaml"]).decode()) + mat_config = yaml.safe_load(base64.b64decode(secret.data["materialization_config.yaml"]).decode()) + return feast_config, mat_config + + +def _load_config_from_files(): + """Load feast and materialization config from mounted files.""" + with open("/var/feast/feature_store.yaml") as f: + feast_config = yaml.safe_load(f) + with open("/var/feast/materialization_config.yaml") as f: + mat_config = yaml.safe_load(f) + return feast_config, mat_config + + +def _materialize_one_fv(spark_session, config, task_info): + """Materialize a single feature view in a worker thread. + + Each thread gets its own FeatureStore instance to avoid race conditions + in Feast's usage.py call_stack (not thread-safe). + SparkSession visibility is handled via the getActiveSession patch in main(). + """ + from datetime import datetime, timezone + from feast import FeatureStore, RepoConfig + from tqdm import tqdm + + fv_name = task_info["feature_view"] + logger.info(f"Thread started: {fv_name}") + + # Per-thread FeatureStore avoids race in feast/usage.py call_stack + thread_config = RepoConfig(**config) + thread_store = FeatureStore(config=thread_config) + fv = thread_store.get_feature_view(fv_name) + provider = thread_store.get_provider() + + start = datetime.fromisoformat(task_info["start_time"]) + end = datetime.fromisoformat(task_info["end_time"]) + if start.tzinfo is None: + start = start.replace(tzinfo=timezone.utc) + if end.tzinfo is None: + end = end.replace(tzinfo=timezone.utc) + + t0 = time.time() + provider.materialize_single_feature_view( + config=thread_config, + feature_view=fv, + start_date=start, + end_date=end, + registry=thread_store.registry, + project=thread_store.project, + tqdm_builder=lambda length: tqdm(total=length, ncols=100), + ) + elapsed = time.time() - t0 + return fv_name, elapsed + + +def main(): + if os.environ.get("FEAST_SECRET_NAME"): + logger.info("Loading config from Secret via K8s API") + feast_config, mat_config = _load_config_from_secret() + elif os.path.exists("/var/feast/feature_store.yaml"): + logger.info("Loading config from mounted files") + feast_config, mat_config = _load_config_from_files() + else: + raise RuntimeError( + "No config source found. Set FEAST_SECRET_NAME env var " + "or mount config at /var/feast/" + ) + + from feast import RepoConfig + from pyspark.sql import SparkSession + + RepoConfig(**feast_config) # validate config eagerly before any Spark work + operation = mat_config["operation"] + + if operation == "materialize": + tasks = mat_config.get("tasks", []) + if not tasks: + tasks = [{ + "feature_view": mat_config["feature_view"], + "start_time": mat_config["start_time"], + "end_time": mat_config["end_time"], + }] + + concurrency = mat_config.get("concurrency", 1) + total = len(tasks) + total_start = time.time() + + logger.info( + f"Starting materialization: {total} feature views, " + f"concurrency={concurrency}" + ) + + # Get the active SparkSession (created by spark-submit) + spark = SparkSession.getActiveSession() + if spark is None: + spark = SparkSession.builder.getOrCreate() + logger.info(f"SparkSession: {spark.sparkContext.applicationId}") + + if concurrency > 1: + _original_get_active = SparkSession.getActiveSession + SparkSession.getActiveSession = staticmethod(lambda: spark) + logger.info("Patched SparkSession.getActiveSession for thread safety") + + try: + if concurrency <= 1: + # Sequential mode + succeeded, failed = 0, 0 + for i, task in enumerate(tasks, 1): + fv_name = task["feature_view"] + logger.info(f"[{i}/{total}] Materializing: {fv_name}") + try: + name, elapsed = _materialize_one_fv(spark, feast_config, task) + succeeded += 1 + logger.info(f"[{i}/{total}] Completed: {name} ({elapsed:.1f}s)") + except Exception: + failed += 1 + logger.exception(f"[{i}/{total}] Failed: {fv_name}") + else: + succeeded, failed = 0, 0 + with ThreadPoolExecutor(max_workers=concurrency) as executor: + future_to_fv = { + executor.submit(_materialize_one_fv, spark, feast_config, task): task["feature_view"] + for task in tasks + } + for future in as_completed(future_to_fv): + fv_name = future_to_fv[future] + try: + name, elapsed = future.result() + succeeded += 1 + logger.info(f"Completed: {name} ({elapsed:.1f}s)") + except Exception: + failed += 1 + logger.exception(f"Failed: {fv_name}") + finally: + if concurrency > 1: + SparkSession.getActiveSession = _original_get_active + logger.info("Restored SparkSession.getActiveSession") + + total_elapsed = time.time() - total_start + logger.info( + f"Materialization batch complete: " + f"{succeeded} succeeded, {failed} failed, {total} total, " + f"elapsed={total_elapsed:.1f}s" + ) + if failed > 0 and succeeded == 0: + sys.exit(1) + else: + raise ValueError(f"Unknown operation: {operation}") + + +if __name__ == "__main__": + try: + main() + except Exception: + logger.exception("Driver failed") + sys.exit(1) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 06529cea0f2..23f460ccf2c 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -51,6 +51,7 @@ "spark.engine": "feast.infra.compute_engines.spark.compute.SparkComputeEngine", "ray.engine": "feast.infra.compute_engines.ray.compute.RayComputeEngine", "flink.engine": "feast.infra.compute_engines.flink.compute.FlinkComputeEngine", + "spark_application": "feast.infra.compute_engines.spark_application.compute.SparkApplicationComputeEngine", } LEGACY_ONLINE_STORE_CLASS_FOR_TYPE = { diff --git a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py new file mode 100644 index 00000000000..5f4eb6fc9d3 --- /dev/null +++ b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py @@ -0,0 +1,217 @@ +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest + +from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + MaterializationTask, +) +from feast.infra.compute_engines.spark_application.config import ( + SparkApplicationComputeEngineConfig, +) +from feast.infra.compute_engines.spark_application.job import ( + SparkApplicationMaterializationJob, + _STATE_MAP, +) + + +def _make_repo_config( + online_store_type="redis", + registry_path="s3://bucket/registry.db", + registry_address=None, + offline_store_type="dask", + spark_conf=None, +): + """Build a mock RepoConfig for testing.""" + config = MagicMock() + config.online_store = MagicMock() + config.online_store.type = online_store_type + config.registry = MagicMock() + config.registry.path = registry_path + config.registry.registry_type = "file" + config.batch_engine = SparkApplicationComputeEngineConfig( + image="quay.io/test/feast-spark:latest", + registry_address=registry_address, + spark_conf=spark_conf, + ) + config.model_dump = MagicMock(return_value={ + "project": "test", + "provider": "local", + "batch_engine": {"type": "spark_application"}, + "offline_store": {"type": offline_store_type, "spark_conf": {"spark.existing": "value"}}, + "online_store": {"type": online_store_type}, + "registry": {"registry_type": "file", "path": registry_path}, + }) + return config + + +@patch("feast.infra.compute_engines.spark_application.compute.k8s_config") +@patch("feast.infra.compute_engines.spark_application.compute.client") +def _make_engine(mock_client, mock_k8s_config, **kwargs): + """Create engine with mocked K8s client.""" + from feast.infra.compute_engines.spark_application.compute import ( + SparkApplicationComputeEngine, + ) + repo_config = _make_repo_config(**kwargs) + return SparkApplicationComputeEngine( + repo_config=repo_config, offline_store=None, online_store=None + ) + + +# ── Test 1: Config defaults + required field ── + +def test_config_defaults_and_required_image(): + c = SparkApplicationComputeEngineConfig(image="quay.io/test:v1") + assert c.type == "spark_application" + assert c.namespace == "default" + assert c.executor_instances == 1 + assert c.restart_policy == "Never" + assert c.max_retries == 3 + + with pytest.raises(Exception): + SparkApplicationComputeEngineConfig() # image is required + + +# ── Test 2: EC-3 rejects SQLite ── + +def test_rejects_sqlite_online_store(): + with pytest.raises(ValueError, match="SQLite"): + _make_engine(online_store_type="sqlite") + + +# ── Test 3: EC-2 rejects local registry without registry_address ── + +def test_rejects_local_registry_without_address(): + with pytest.raises(ValueError, match="local file"): + _make_engine(registry_path="/local/registry.db") + + +# ── Test 4: EC-2 accepts local registry WITH registry_address ── + +def test_accepts_local_registry_with_address(): + engine = _make_engine(registry_path="/local/registry.db", registry_address="feast:6570") + assert engine is not None + + +# ── Test 5: _build_driver_repo_config — two rewrites (batch_engine + registry) ── + +def test_build_driver_repo_config_rewrites(): + engine = _make_engine( + offline_store_type="dask", + registry_address="feast-server:6566", + ) + d = engine._build_driver_repo_config() + assert d["batch_engine"]["type"] == "spark.engine" + assert d["offline_store"]["type"] == "dask" # NOT rewritten — respects user intent + assert d["registry"] == {"registry_type": "remote", "path": "feast-server:6566"} + + +# ── Test 6: _build_driver_repo_config — spark_conf placed in batch_engine ── + +def test_build_driver_repo_config_spark_conf(): + engine = _make_engine(spark_conf={"spark.new": "from_engine", "spark.existing": "from_engine"}) + d = engine._build_driver_repo_config() + assert d["batch_engine"]["spark_conf"]["spark.new"] == "from_engine" + assert d["batch_engine"]["spark_conf"]["spark.existing"] == "from_engine" + # offline_store spark_conf left untouched + assert d["offline_store"]["spark_conf"]["spark.existing"] == "value" + + +# ── Test 7: _build_driver_repo_config — no registry rewrite without address ── + +def test_build_driver_repo_config_no_registry_rewrite(): + engine = _make_engine(registry_address=None) + d = engine._build_driver_repo_config() + assert d["registry"] == {"registry_type": "file", "path": "s3://bucket/registry.db"} + + +# ── Test 8: CR structure ── + +def test_cr_structure(): + engine = _make_engine() + cr = engine._build_spark_application_cr("abcd1234") + assert cr["apiVersion"] == "sparkoperator.k8s.io/v1beta2" + assert cr["kind"] == "SparkApplication" + assert cr["spec"]["type"] == "Python" + assert cr["spec"]["mode"] == "cluster" + assert cr["spec"]["mainApplicationFile"] == "local:///opt/feast/main.py" + assert "driver" in cr["spec"] + assert "executor" in cr["spec"] + assert cr["metadata"]["name"] == "feast-sa-abcd1234" + + +# ── Test 9: Status mapping covers all 14 states ── + +def test_state_map_coverage(): + assert len(_STATE_MAP) == 14 + assert _STATE_MAP["COMPLETED"] == MaterializationJobStatus.SUCCEEDED + assert _STATE_MAP["FAILED"] == MaterializationJobStatus.ERROR + assert _STATE_MAP["SUBMISSION_FAILED"] == MaterializationJobStatus.ERROR + assert _STATE_MAP["RUNNING"] == MaterializationJobStatus.RUNNING + assert _STATE_MAP[""] == MaterializationJobStatus.WAITING + + +# ── Test 10: Cleanup swallows 404 ── + +@patch("feast.infra.compute_engines.spark_application.compute.k8s_config") +@patch("feast.infra.compute_engines.spark_application.compute.client") +def test_cleanup_swallows_404(mock_client, mock_k8s_config): + from kubernetes.client.exceptions import ApiException + from feast.infra.compute_engines.spark_application.compute import ( + SparkApplicationComputeEngine, + ) + + repo_config = _make_repo_config() + engine = SparkApplicationComputeEngine( + repo_config=repo_config, offline_store=None, online_store=None + ) + + # Mock both delete calls to raise 404 + engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException(status=404) + engine.core_v1.delete_namespaced_secret.side_effect = ApiException(status=404) + + # Should not raise + engine._cleanup("test-id") + + +# ── Test 11: Timeout sets error on job (does not raise) ── + +@patch("feast.infra.compute_engines.spark_application.compute.k8s_config") +@patch("feast.infra.compute_engines.spark_application.compute.client") +@patch("feast.infra.compute_engines.spark_application.compute.time") +def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): + from feast.infra.compute_engines.spark_application.compute import ( + SparkApplicationComputeEngine, + ) + + repo_config = _make_repo_config() + repo_config.batch_engine = SparkApplicationComputeEngineConfig( + image="test", job_timeout_seconds=1, poll_interval_seconds=1, + ) + engine = SparkApplicationComputeEngine( + repo_config=repo_config, offline_store=None, online_store=None + ) + + # Calls: start(0), while-check(0), elapsed(0), sleep, while-check(2 > deadline=1) → exit + mock_time.monotonic.side_effect = [0, 0, 0, 2] + mock_time.sleep = MagicMock() + + mock_job = MagicMock() + mock_job.status.return_value = MaterializationJobStatus.RUNNING + mock_job._job_id = "test123" + mock_job._error = None + mock_job.job_id.return_value = "feast-sa-test123" + + engine._wait_for_completion(mock_job) + assert mock_job._error is not None + assert "did not complete" in str(mock_job._error) + + +# ── Test 12: Job naming < 63 chars ── + +def test_job_naming_under_63_chars(): + mock_api = MagicMock() + job = SparkApplicationMaterializationJob("abcdef12", "default", mock_api) + assert len(job.job_id()) <= 63 + assert job.job_id() == "feast-sa-abcdef12" From 5804df5991999f26c4f1637e0520fd364c55b25d Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Fri, 3 Jul 2026 13:16:27 +0530 Subject: [PATCH 2/4] feat: Add per-FV result reporting and clean up Dockerfile - Pod calls apply_materialization via gRPC after each successful FV, setting state to AVAILABLE_ONLINE. Server reads FV state post-completion to determine per-FV success/failure in batched SparkApplication runs. - registry_address is now mandatory (simplified from complex path heuristic). - Dockerfile rewritten to install feast from source (matches K8s engine pattern). - Unit tests updated: 15/15 pass (3 new tests for _build_per_fv_jobs). - E2E validated: 5 FVs x 9600 rows, 5 executors, all AVAILABLE_ONLINE. Signed-off-by: Aniket Paluskar --- sdk/python/feast/feature_store.py | 412 +++++++++--------- .../spark_application/Dockerfile | 24 +- .../spark_application/compute.py | 88 ++-- .../compute_engines/spark_application/main.py | 4 + .../compute_engines/test_spark_application.py | 119 ++++- 5 files changed, 373 insertions(+), 274 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 68fffdd0e15..faf44ff327c 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -428,6 +428,115 @@ def get_provider(self) -> Provider: """Public accessor for the provider instance.""" return self._get_provider() + def _rollback_fv_states( + self, + feature_views: list, + previous_states: dict, + ) -> None: + """Restore feature views to their pre-materialization states.""" + for fv in feature_views: + prev = previous_states.get(fv.name) + if ( + hasattr(fv, "state") + and prev is not None + and prev != FeatureViewState.STATE_UNSPECIFIED + ): + fv.state = prev + self.registry.apply_feature_view( + fv, self.project, commit=True + ) + + def _transition_fv_to_materializing( + self, + feature_view, + already_transitioned: list, + previous_states: dict, + ) -> None: + """ + Transition a feature view to MATERIALIZING state. + + Rolls back all already-transitioned FVs if this one can't transition. + """ + previous_state = getattr(feature_view, "state", None) + if ( + hasattr(feature_view, "state") + and feature_view.state != FeatureViewState.STATE_UNSPECIFIED + ): + if not feature_view.state.can_transition_to( + FeatureViewState.MATERIALIZING + ): + self._rollback_fv_states(already_transitioned, previous_states) + raise ValueError( + f"FeatureView {feature_view.name} cannot transition " + f"from {feature_view.state.name} to MATERIALIZING." + ) + feature_view.state = FeatureViewState.MATERIALIZING + self.registry.apply_feature_view( + feature_view, self.project, commit=True + ) + previous_states[feature_view.name] = previous_state + + def _submit_and_process_materialization_jobs( + self, + provider, + tasks: list, + regular_fvs: list, + previous_states: dict, + fv_start_dates: dict, + ) -> None: + """ + Submit all tasks to the engine in one call and process the results. + + For each returned job: record watermark on success, roll back state on + error. If the engine itself raises, all states are rolled back. + """ + from feast.infra.common.materialization_job import ( + MaterializationJobStatus, + ) + + batch_start = time.monotonic() + try: + jobs = provider.batch_engine.materialize(self.registry, tasks) + except Exception: + self._rollback_fv_states(regular_fvs, previous_states) + raise + + first_error = None + succeeded_fvs = [] + failed_fvs = [] + + for fv, job in zip(regular_fvs, jobs): + fv_status = job.status() + + if fv_status == MaterializationJobStatus.ERROR: + failed_fvs.append(fv) + if first_error is None and job.error(): + first_error = job.error() + else: + succeeded_fvs.append(fv) + + if failed_fvs: + self._rollback_fv_states(failed_fvs, previous_states) + + for fv in succeeded_fvs: + self.registry.apply_materialization( + fv, + self.project, + fv_start_dates[fv.name], + fv_start_dates["__end_date__"], + ) + + _tracker = _get_track_materialization() + if _tracker is not None: + elapsed = time.monotonic() - batch_start + for fv in succeeded_fvs: + _tracker(fv.name, True, elapsed) + for fv in failed_fvs: + _tracker(fv.name, False, elapsed) + + if first_error: + raise first_error + @property def openlineage_emitter(self) -> Optional[Any]: """Gets the OpenLineage emitter of this feature store.""" @@ -2140,6 +2249,21 @@ def materialize_incremental( _mat_start = time.monotonic() try: + from feast.infra.common.materialization_job import ( + MaterializationTask, + ) + + provider = self._get_provider() + end_date_tz = utils.make_tzaware(end_date) or _utc_now() + + def tqdm_builder(length): + return tqdm(total=length, ncols=100) + + tasks: list = [] + regular_fvs: list = [] + previous_states: dict = {} + fv_start_dates: dict = {"__end_date__": end_date_tz} + for feature_view in feature_views_to_materialize: if isinstance(feature_view, OnDemandFeatureView): if feature_view.write_to_online_store: @@ -2166,138 +2290,54 @@ def materialize_incremental( end_date, full_feature_names=full_feature_names, ) + continue - regular_fvs = [ - fv - for fv in feature_views_to_materialize - if not isinstance(fv, OnDemandFeatureView) - ] - - if regular_fvs: - # Deferred import: top-level causes circular import via - # feast.__init__ -> feature_store -> materialization_job -> feast - from feast.infra.common.materialization_job import ( - MaterializationJobStatus, - MaterializationTask, - ) - - provider = self._get_provider() - end_date_tz = utils.make_tzaware(end_date) or _utc_now() - - # Compute per-FV start dates and transition all to MATERIALIZING - # before submitting work to the engine. - previous_states: dict = {} - fv_start_dates: dict = {} - - for fv in regular_fvs: - fv_start_date = fv.most_recent_end_time - if fv_start_date is None: - if fv.ttl is None: - raise Exception( - f"No start time found for feature view {fv.name}. materialize_incremental() requires" - f" either a ttl to be set or for materialize() to have been run at least once." - ) - elif fv.ttl.total_seconds() > 0: - fv_start_date = _utc_now() - fv.ttl - else: - print( - f"Since the ttl is 0 for feature view {Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}, " - "the start date will be set to 1 year before the current time." - ) - fv_start_date = _utc_now() - timedelta(weeks=52) - - fv_start_dates[fv.name] = utils.make_tzaware(fv_start_date) - - print( - f"{Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}" - f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(fv_start_date.replace(microsecond=0))}{Style.RESET_ALL}" - f" to {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(end_date.replace(microsecond=0))}{Style.RESET_ALL}:" - ) - - previous_states[fv.name] = getattr(fv, "state", None) - if ( - hasattr(fv, "state") - and fv.state != FeatureViewState.STATE_UNSPECIFIED - ): - if not fv.state.can_transition_to( - FeatureViewState.MATERIALIZING - ): - raise ValueError( - f"FeatureView {fv.name} cannot transition " - f"from {fv.state.name} to MATERIALIZING." - ) - fv.state = FeatureViewState.MATERIALIZING - self.registry.apply_feature_view( - fv, self.project, commit=True + fv_start_date = feature_view.most_recent_end_time + if fv_start_date is None: + if feature_view.ttl is None: + raise Exception( + f"No start time found for feature view {feature_view.name}. materialize_incremental() requires" + f" either a ttl to be set or for materialize() to have been run at least once." + ) + elif feature_view.ttl.total_seconds() > 0: + fv_start_date = _utc_now() - feature_view.ttl + else: + # TODO(felixwang9817): Find the earliest timestamp for this specific feature + # view from the offline store, and set the start date to that timestamp. + print( + f"Since the ttl is 0 for feature view {Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}, " + "the start date will be set to 1 year before the current time." ) + fv_start_date = _utc_now() - timedelta(weeks=52) - def tqdm_builder(length): - return tqdm(total=length, ncols=100) + fv_start_date = utils.make_tzaware(fv_start_date) + fv_start_dates[feature_view.name] = fv_start_date + + print( + f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}" + f" from {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(fv_start_date.replace(microsecond=0))}{Style.RESET_ALL}" + f" to {Style.BRIGHT + Fore.GREEN}{utils.make_tzaware(end_date.replace(microsecond=0))}{Style.RESET_ALL}:" + ) - tasks = [ + self._transition_fv_to_materializing( + feature_view, regular_fvs, previous_states + ) + regular_fvs.append(feature_view) + tasks.append( MaterializationTask( project=self.project, - feature_view=fv, - start_time=fv_start_dates[fv.name], + feature_view=feature_view, + start_time=fv_start_date, end_time=end_date_tz, tqdm_builder=tqdm_builder, ) - for fv in regular_fvs - ] - - fv_batch_start = time.monotonic() - try: - jobs = provider.batch_engine.materialize( - self.registry, tasks - ) - except Exception: - for fv in regular_fvs: - prev = previous_states[fv.name] - if ( - prev is not None - and prev != FeatureViewState.STATE_UNSPECIFIED - ): - fv.state = prev - self.registry.apply_feature_view( - fv, self.project, commit=True - ) - raise - - errors = [] - for fv, job in zip(regular_fvs, jobs): - fv_status = job.status() - fv_success = fv_status != MaterializationJobStatus.ERROR - - if fv_status == MaterializationJobStatus.SUCCEEDED: - self.registry.apply_materialization( - fv, - self.project, - fv_start_dates[fv.name], - end_date_tz, - ) - elif fv_status == MaterializationJobStatus.ERROR: - prev = previous_states[fv.name] - if ( - prev is not None - and prev != FeatureViewState.STATE_UNSPECIFIED - ): - fv.state = prev - self.registry.apply_feature_view( - fv, self.project, commit=True - ) - if job.error(): - errors.append(job.error()) - - _tracker = _get_track_materialization() - if _tracker is not None: - _tracker( - fv.name, - fv_success, - time.monotonic() - fv_batch_start, - ) + ) - if errors: - raise errors[0] + if tasks: + self._submit_and_process_materialization_jobs( + provider, tasks, regular_fvs, + previous_states, fv_start_dates, + ) materialized_fv_names = [ fv.name @@ -2389,6 +2429,22 @@ def materialize( _mat_start = time.monotonic() try: + from feast.infra.common.materialization_job import ( + MaterializationTask, + ) + + provider = self._get_provider() + start_date = utils.make_tzaware(start_date) + end_date = utils.make_tzaware(end_date) + + def tqdm_builder(length): + return tqdm(total=length, ncols=100) + + tasks: list = [] + regular_fvs: list = [] + previous_states: dict = {} + fv_start_dates: dict = {"__end_date__": end_date} + for feature_view in feature_views_to_materialize: if isinstance(feature_view, OnDemandFeatureView): if feature_view.write_to_online_store: @@ -2401,113 +2457,33 @@ def materialize( end_date, full_feature_names=full_feature_names, ) + continue - regular_fvs = [ - fv - for fv in feature_views_to_materialize - if not isinstance(fv, OnDemandFeatureView) - ] - - if regular_fvs: - # Deferred import: top-level causes circular import via - # feast.__init__ -> feature_store -> materialization_job -> feast - from feast.infra.common.materialization_job import ( - MaterializationJobStatus, - MaterializationTask, + print( + f"{Style.BRIGHT + Fore.GREEN}{feature_view.name}{Style.RESET_ALL}:" ) - provider = self._get_provider() - start_date = utils.make_tzaware(start_date) - end_date = utils.make_tzaware(end_date) - - # Transition all FVs to MATERIALIZING before submitting work. - previous_states: dict = {} - for fv in regular_fvs: - print( - f"{Style.BRIGHT + Fore.GREEN}{fv.name}{Style.RESET_ALL}:" - ) - previous_states[fv.name] = getattr(fv, "state", None) - if ( - hasattr(fv, "state") - and fv.state != FeatureViewState.STATE_UNSPECIFIED - ): - if not fv.state.can_transition_to( - FeatureViewState.MATERIALIZING - ): - raise ValueError( - f"FeatureView {fv.name} cannot transition " - f"from {fv.state.name} to MATERIALIZING." - ) - fv.state = FeatureViewState.MATERIALIZING - self.registry.apply_feature_view( - fv, self.project, commit=True - ) - - def tqdm_builder(length): - return tqdm(total=length, ncols=100) - - tasks = [ + self._transition_fv_to_materializing( + feature_view, regular_fvs, previous_states + ) + regular_fvs.append(feature_view) + fv_start_dates[feature_view.name] = start_date + tasks.append( MaterializationTask( project=self.project, - feature_view=fv, + feature_view=feature_view, start_time=start_date, end_time=end_date, tqdm_builder=tqdm_builder, disable_event_timestamp=disable_event_timestamp, ) - for fv in regular_fvs - ] - - fv_start = time.monotonic() - try: - jobs = provider.batch_engine.materialize( - self.registry, tasks - ) - except Exception: - for fv in regular_fvs: - prev = previous_states[fv.name] - if ( - prev is not None - and prev != FeatureViewState.STATE_UNSPECIFIED - ): - fv.state = prev - self.registry.apply_feature_view( - fv, self.project, commit=True - ) - raise - - errors = [] - for fv, job in zip(regular_fvs, jobs): - fv_status = job.status() - fv_success = fv_status != MaterializationJobStatus.ERROR - - if fv_status == MaterializationJobStatus.SUCCEEDED: - self.registry.apply_materialization( - fv, self.project, start_date, end_date, - ) - elif fv_status == MaterializationJobStatus.ERROR: - prev = previous_states[fv.name] - if ( - prev is not None - and prev != FeatureViewState.STATE_UNSPECIFIED - ): - fv.state = prev - self.registry.apply_feature_view( - fv, self.project, commit=True - ) - if job.error(): - errors.append(job.error()) - - _tracker = _get_track_materialization() - if _tracker is not None: - _tracker( - fv.name, - fv_success, - time.monotonic() - fv_start, - ) + ) - if errors: - raise errors[0] + if tasks: + self._submit_and_process_materialization_jobs( + provider, tasks, regular_fvs, + previous_states, fv_start_dates, + ) materialized_fv_names = [ fv.name diff --git a/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile index db3a4cea26b..268e377b4bc 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile +++ b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile @@ -2,11 +2,23 @@ FROM apache/spark:4.0.1 USER root -RUN pip install --no-cache-dir \ - "feast[redis]==0.64.0" \ - pyyaml \ - kubernetes \ - tqdm +RUN apt-get update && apt-get install --no-install-suggests --no-install-recommends --yes git && rm -rf /var/lib/apt/lists/* + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv + +WORKDIR /app + +COPY sdk/python/feast/infra/compute_engines/spark_application/main.py /opt/feast/main.py + +# Copy necessary parts of the Feast codebase +COPY sdk/python sdk/python +COPY protos protos +COPY pyproject.toml pyproject.toml +COPY README.md README.md + +# setuptools_scm needs .git to infer the version. +# https://github.com/pypa/setuptools_scm#usage-from-docker +RUN --mount=source=.git,target=.git,type=bind uv pip install --system --no-cache-dir '.[redis,grpcio,k8s]' # Hadoop AWS JARs for S3A filesystem support + AWS SDK v2 bundle RUN curl -sL https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.4.1/hadoop-aws-3.4.1.jar \ @@ -16,6 +28,4 @@ RUN curl -sL https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.4.1/h curl -sL https://repo1.maven.org/maven2/software/amazon/awssdk/bundle/2.28.4/bundle-2.28.4.jar \ -o /opt/spark/jars/bundle-2.28.4.jar -COPY feast/infra/compute_engines/spark_application/main.py /opt/feast/main.py - USER spark diff --git a/sdk/python/feast/infra/compute_engines/spark_application/compute.py b/sdk/python/feast/infra/compute_engines/spark_application/compute.py index 6c0aecaf738..c25c4eb1999 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/compute.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/compute.py @@ -23,6 +23,7 @@ from feast.infra.offline_stores.offline_store import OfflineStore from feast.infra.online_stores.online_store import OnlineStore from feast.infra.registry.base_registry import BaseRegistry +from feast.feature_view import FeatureViewState from feast.on_demand_feature_view import OnDemandFeatureView from feast.stream_feature_view import StreamFeatureView @@ -49,7 +50,6 @@ def __init__( ) self.config = repo_config.batch_engine - # EC-3: SQLite online store — data written inside pod is lost on termination online_type = getattr(repo_config.online_store, "type", "") if online_type == "sqlite": raise ValueError( @@ -59,25 +59,16 @@ def __init__( "Use a network-accessible store: redis, postgres, etc." ) - # EC-2: Registry access — pod must be able to reach registry + # EC-2: registry_address required — pod reports materialization results + # back to the Feast server via gRPC. Without it, the server cannot + # determine per-FV success/failure after a batched SparkApplication run. if not self.config.registry_address: - registry = repo_config.registry - if hasattr(registry, "path") and registry.path: - path = registry.path - is_remote = any( - path.startswith(s) - for s in ("s3://", "gs://", "hdfs://", "http://", "https://", "postgresql", "mysql") - ) - if not is_remote: - raise ValueError( - f"Registry path '{path}' is a local file. " - f"SparkApplication pods cannot access the Feast " - f"server's filesystem. Either:\n" - f" 1. Set registry_address to the Feast registry " - f"gRPC endpoint (recommended), or\n" - f" 2. Use a remote registry path (s3://, gs://), or\n" - f" 3. Switch to registry_type: 'sql' or 'remote'." - ) + raise ValueError( + "registry_address is required for spark_application engine. " + "Set it to the Feast server's gRPC endpoint " + "(e.g., feast-server.namespace.svc.cluster.local:6566). " + "The Feast Operator sets this automatically." + ) k8s_config.load_config() self.k8s_client = client.ApiClient() @@ -117,7 +108,13 @@ def materialize( tasks: Union[MaterializationTask, List[MaterializationTask]], **kwargs, ) -> List[MaterializationJob]: - """Batch all materialization tasks into a single SparkApplication.""" + """Batch all materialization tasks into a single SparkApplication. + + The pod calls apply_materialization (via gRPC → Feast server → registry) + for each FV it successfully materializes. After the pod finishes, we read + each FV's state from the registry: AVAILABLE_ONLINE = succeeded, + still MATERIALIZING = failed. + """ if isinstance(tasks, MaterializationTask): tasks = [tasks] @@ -151,7 +148,7 @@ def materialize( job = SparkApplicationMaterializationJob(job_id, self.config.namespace, self.custom_api) self._wait_for_completion(job) - return [job for _ in tasks] + return self._build_per_fv_jobs(registry, tasks, job_id, job) def _build_driver_repo_config(self) -> dict: """Build feature_store.yaml for the SparkApplication driver pod. @@ -168,11 +165,9 @@ def _build_driver_repo_config(self) -> dict: offline_store is NOT rewritten — respects user's configured data sources. User should configure offline_store: spark for full distributed performance. """ - config_dict = self.repo_config.model_dump(by_alias=True) + config_dict = self.repo_config.model_dump(by_alias=True, mode="json") config_dict["batch_engine"] = {"type": "spark.engine"} - if self.config.spark_conf: - config_dict["batch_engine"]["spark_conf"] = self.config.spark_conf if self.config.registry_address: config_dict["registry"] = { @@ -182,8 +177,38 @@ def _build_driver_repo_config(self) -> dict: return config_dict + def _build_per_fv_jobs( + self, + registry: BaseRegistry, + tasks: List[MaterializationTask], + job_id: str, + job: SparkApplicationMaterializationJob, + ) -> List[MaterializationJob]: + """Read each FV's state from registry to determine per-FV success/failure. + + The pod calls apply_materialization for each succeeded FV, which sets + state to AVAILABLE_ONLINE. FVs still in MATERIALIZING were not processed. + """ + if len(tasks) <= 1: + return [job for _ in tasks] + + jobs: List[MaterializationJob] = [] + for task in tasks: + fv = registry.get_feature_view(task.feature_view.name, task.project) + if getattr(fv, "state", None) == FeatureViewState.AVAILABLE_ONLINE: + jobs.append(job) + else: + jobs.append(SparkApplicationMaterializationJob( + job_id, self.config.namespace, self.custom_api, + error=Exception( + f"Feature view '{task.feature_view.name}' was not " + f"materialized by the SparkApplication pod" + ), + )) + return jobs + def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): - feast_config_yaml = yaml.dump(self._build_driver_repo_config()) + feast_config_yaml = yaml.dump(self._build_driver_repo_config(), default_flow_style=False) mat_config = { "operation": "materialize", "tasks": [ @@ -216,6 +241,16 @@ def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): ) def _build_spark_application_cr(self, job_id: str) -> dict: + driver_env_conf = { + "spark.kubernetes.driverEnv.FEAST_SECRET_NAME": f"feast-sa-{job_id}", + "spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE": self.config.namespace, + } + for entry in self.config.env: + name = entry.get("name", "") + value = entry.get("value", "") + if name and value: + driver_env_conf[f"spark.kubernetes.driverEnv.{name}"] = value + spec = { "type": "Python", "mode": "cluster", @@ -227,8 +262,7 @@ def _build_spark_application_cr(self, job_id: str) -> dict: "sparkConf": { "spark.scheduler.mode": "FAIR", **(self.config.spark_conf or {}), - "spark.kubernetes.driverEnv.FEAST_SECRET_NAME": f"feast-sa-{job_id}", - "spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE": self.config.namespace, + **driver_env_conf, }, "restartPolicy": { "type": self.config.restart_policy, diff --git a/sdk/python/feast/infra/compute_engines/spark_application/main.py b/sdk/python/feast/infra/compute_engines/spark_application/main.py index 73a9db3e885..d5946807d20 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/main.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/main.py @@ -72,6 +72,10 @@ def _materialize_one_fv(spark_session, config, task_info): project=thread_store.project, tqdm_builder=lambda length: tqdm(total=length, ncols=100), ) + + thread_store.registry.apply_materialization(fv, thread_store.project, start, end) + logger.info(f"Applied materialization metadata for {fv_name}") + elapsed = time.time() - t0 return fv_name, elapsed diff --git a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py index 5f4eb6fc9d3..27d7e9451f8 100644 --- a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py +++ b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py @@ -14,12 +14,13 @@ SparkApplicationMaterializationJob, _STATE_MAP, ) +from feast.feature_view import FeatureViewState def _make_repo_config( online_store_type="redis", registry_path="s3://bucket/registry.db", - registry_address=None, + registry_address="feast-server:6566", offline_store_type="dask", spark_conf=None, ): @@ -80,11 +81,11 @@ def test_rejects_sqlite_online_store(): _make_engine(online_store_type="sqlite") -# ── Test 3: EC-2 rejects local registry without registry_address ── +# ── Test 3: registry_address is mandatory ── -def test_rejects_local_registry_without_address(): - with pytest.raises(ValueError, match="local file"): - _make_engine(registry_path="/local/registry.db") +def test_rejects_missing_registry_address(): + with pytest.raises(ValueError, match="registry_address is required"): + _make_engine(registry_address=None) # ── Test 4: EC-2 accepts local registry WITH registry_address ── @@ -107,26 +108,16 @@ def test_build_driver_repo_config_rewrites(): assert d["registry"] == {"registry_type": "remote", "path": "feast-server:6566"} -# ── Test 6: _build_driver_repo_config — spark_conf placed in batch_engine ── +# ── Test 6: _build_driver_repo_config — batch_engine is just type (no spark_conf copy) ── -def test_build_driver_repo_config_spark_conf(): - engine = _make_engine(spark_conf={"spark.new": "from_engine", "spark.existing": "from_engine"}) +def test_build_driver_repo_config_batch_engine_minimal(): + engine = _make_engine(spark_conf={"spark.new": "from_engine"}) d = engine._build_driver_repo_config() - assert d["batch_engine"]["spark_conf"]["spark.new"] == "from_engine" - assert d["batch_engine"]["spark_conf"]["spark.existing"] == "from_engine" - # offline_store spark_conf left untouched + assert d["batch_engine"] == {"type": "spark.engine"} assert d["offline_store"]["spark_conf"]["spark.existing"] == "value" -# ── Test 7: _build_driver_repo_config — no registry rewrite without address ── - -def test_build_driver_repo_config_no_registry_rewrite(): - engine = _make_engine(registry_address=None) - d = engine._build_driver_repo_config() - assert d["registry"] == {"registry_type": "file", "path": "s3://bucket/registry.db"} - - -# ── Test 8: CR structure ── +# ── Test 7: CR structure ── def test_cr_structure(): engine = _make_engine() @@ -141,6 +132,16 @@ def test_cr_structure(): assert cr["metadata"]["name"] == "feast-sa-abcd1234" +# ── Test 8: CR sparkConf includes driver env passthrough ── + +def test_cr_driver_env_passthrough(): + engine = _make_engine() + cr = engine._build_spark_application_cr("abcd1234") + spark_conf = cr["spec"]["sparkConf"] + assert spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAME"] == "feast-sa-abcd1234" + assert spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE"] == "default" + + # ── Test 9: Status mapping covers all 14 states ── def test_state_map_coverage(): @@ -167,11 +168,9 @@ def test_cleanup_swallows_404(mock_client, mock_k8s_config): repo_config=repo_config, offline_store=None, online_store=None ) - # Mock both delete calls to raise 404 engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException(status=404) engine.core_v1.delete_namespaced_secret.side_effect = ApiException(status=404) - # Should not raise engine._cleanup("test-id") @@ -188,6 +187,7 @@ def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): repo_config = _make_repo_config() repo_config.batch_engine = SparkApplicationComputeEngineConfig( image="test", job_timeout_seconds=1, poll_interval_seconds=1, + registry_address="feast-server:6566", ) engine = SparkApplicationComputeEngine( repo_config=repo_config, offline_store=None, online_store=None @@ -215,3 +215,78 @@ def test_job_naming_under_63_chars(): job = SparkApplicationMaterializationJob("abcdef12", "default", mock_api) assert len(job.job_id()) <= 63 assert job.job_id() == "feast-sa-abcdef12" + + +# ── Test 13: _build_per_fv_jobs — all succeeded ── + +def test_build_per_fv_jobs_all_succeeded(): + engine = _make_engine() + mock_registry = MagicMock() + + fv1 = MagicMock() + fv1.name = "fv_1" + fv1.state = FeatureViewState.AVAILABLE_ONLINE + fv2 = MagicMock() + fv2.name = "fv_2" + fv2.state = FeatureViewState.AVAILABLE_ONLINE + mock_registry.get_feature_view.side_effect = [fv1, fv2] + + task1 = MagicMock() + task1.feature_view.name = "fv_1" + task1.project = "test" + task2 = MagicMock() + task2.feature_view.name = "fv_2" + task2.project = "test" + + parent_job = SparkApplicationMaterializationJob("job1", "default", MagicMock()) + jobs = engine._build_per_fv_jobs(mock_registry, [task1, task2], "job1", parent_job) + + assert len(jobs) == 2 + assert all(j.status() != MaterializationJobStatus.ERROR for j in jobs) + + +# ── Test 14: _build_per_fv_jobs — partial failure ── + +def test_build_per_fv_jobs_partial_failure(): + engine = _make_engine() + mock_registry = MagicMock() + + fv_ok = MagicMock() + fv_ok.name = "fv_ok" + fv_ok.state = FeatureViewState.AVAILABLE_ONLINE + fv_fail = MagicMock() + fv_fail.name = "fv_fail" + fv_fail.state = FeatureViewState.MATERIALIZING + mock_registry.get_feature_view.side_effect = [fv_ok, fv_fail] + + task_ok = MagicMock() + task_ok.feature_view.name = "fv_ok" + task_ok.project = "test" + task_fail = MagicMock() + task_fail.feature_view.name = "fv_fail" + task_fail.project = "test" + + parent_job = SparkApplicationMaterializationJob("job1", "default", MagicMock()) + jobs = engine._build_per_fv_jobs(mock_registry, [task_ok, task_fail], "job1", parent_job) + + assert len(jobs) == 2 + assert jobs[0].status() != MaterializationJobStatus.ERROR + assert jobs[1].status() == MaterializationJobStatus.ERROR + assert "fv_fail" in str(jobs[1].error()) + + +# ── Test 15: _build_per_fv_jobs — single task returns parent job directly ── + +def test_build_per_fv_jobs_single_task(): + engine = _make_engine() + mock_registry = MagicMock() + task = MagicMock() + task.feature_view.name = "fv_1" + task.project = "test" + + parent_job = SparkApplicationMaterializationJob("job1", "default", MagicMock()) + jobs = engine._build_per_fv_jobs(mock_registry, [task], "job1", parent_job) + + assert len(jobs) == 1 + assert jobs[0] is parent_job + mock_registry.get_feature_view.assert_not_called() From c8d8397647f895424dda877c95cc4369a8efa4e4 Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Fri, 3 Jul 2026 13:23:03 +0530 Subject: [PATCH 3/4] Minor lint & formatting change Signed-off-by: Aniket Paluskar --- sdk/python/feast/feature_store.py | 26 ++++---- .../spark_application/compute.py | 64 ++++++++++++------ .../compute_engines/spark_application/job.py | 12 +++- .../compute_engines/spark_application/main.py | 34 ++++++---- .../compute_engines/test_spark_application.py | 65 ++++++++++++++----- 5 files changed, 137 insertions(+), 64 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index faf44ff327c..664e398cf0a 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -442,9 +442,7 @@ def _rollback_fv_states( and prev != FeatureViewState.STATE_UNSPECIFIED ): fv.state = prev - self.registry.apply_feature_view( - fv, self.project, commit=True - ) + self.registry.apply_feature_view(fv, self.project, commit=True) def _transition_fv_to_materializing( self, @@ -462,18 +460,14 @@ def _transition_fv_to_materializing( hasattr(feature_view, "state") and feature_view.state != FeatureViewState.STATE_UNSPECIFIED ): - if not feature_view.state.can_transition_to( - FeatureViewState.MATERIALIZING - ): + if not feature_view.state.can_transition_to(FeatureViewState.MATERIALIZING): self._rollback_fv_states(already_transitioned, previous_states) raise ValueError( f"FeatureView {feature_view.name} cannot transition " f"from {feature_view.state.name} to MATERIALIZING." ) feature_view.state = FeatureViewState.MATERIALIZING - self.registry.apply_feature_view( - feature_view, self.project, commit=True - ) + self.registry.apply_feature_view(feature_view, self.project, commit=True) previous_states[feature_view.name] = previous_state def _submit_and_process_materialization_jobs( @@ -2335,8 +2329,11 @@ def tqdm_builder(length): if tasks: self._submit_and_process_materialization_jobs( - provider, tasks, regular_fvs, - previous_states, fv_start_dates, + provider, + tasks, + regular_fvs, + previous_states, + fv_start_dates, ) materialized_fv_names = [ @@ -2481,8 +2478,11 @@ def tqdm_builder(length): if tasks: self._submit_and_process_materialization_jobs( - provider, tasks, regular_fvs, - previous_states, fv_start_dates, + provider, + tasks, + regular_fvs, + previous_states, + fv_start_dates, ) materialized_fv_names = [ diff --git a/sdk/python/feast/infra/compute_engines/spark_application/compute.py b/sdk/python/feast/infra/compute_engines/spark_application/compute.py index c25c4eb1999..4d95c56d925 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/compute.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/compute.py @@ -12,7 +12,7 @@ from feast import RepoConfig from feast.batch_feature_view import BatchFeatureView from feast.entity import Entity -from feast.feature_view import FeatureView +from feast.feature_view import FeatureView, FeatureViewState from feast.infra.common.materialization_job import ( MaterializationJob, MaterializationJobStatus, @@ -23,11 +23,12 @@ from feast.infra.offline_stores.offline_store import OfflineStore from feast.infra.online_stores.online_store import OnlineStore from feast.infra.registry.base_registry import BaseRegistry -from feast.feature_view import FeatureViewState from feast.on_demand_feature_view import OnDemandFeatureView from feast.stream_feature_view import StreamFeatureView -from .config import SparkApplicationComputeEngineConfig # noqa: F401 — required for Feast config resolution +from .config import ( + SparkApplicationComputeEngineConfig, # noqa: F401 — required for Feast config resolution +) from .job import SparkApplicationMaterializationJob logger = logging.getLogger(__name__) @@ -79,8 +80,12 @@ def __init__( def update( self, project: str, - views_to_delete: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]], - views_to_keep: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView]], + views_to_delete: Sequence[ + Union[BatchFeatureView, StreamFeatureView, FeatureView] + ], + views_to_keep: Sequence[ + Union[BatchFeatureView, StreamFeatureView, FeatureView, OnDemandFeatureView] + ], entities_to_delete: Sequence[Entity], entities_to_keep: Sequence[Entity], ): @@ -124,7 +129,9 @@ def materialize( self._create_secret(job_id, tasks) except ApiException as e: job = SparkApplicationMaterializationJob( - job_id, self.config.namespace, self.custom_api, + job_id, + self.config.namespace, + self.custom_api, error=Exception(f"Secret creation failed: {e.reason}"), ) return [job for _ in tasks] @@ -141,12 +148,16 @@ def materialize( except ApiException as e: self._cleanup(job_id) job = SparkApplicationMaterializationJob( - job_id, self.config.namespace, self.custom_api, + job_id, + self.config.namespace, + self.custom_api, error=Exception(f"SparkApplication creation failed: {e.reason}"), ) return [job for _ in tasks] - job = SparkApplicationMaterializationJob(job_id, self.config.namespace, self.custom_api) + job = SparkApplicationMaterializationJob( + job_id, self.config.namespace, self.custom_api + ) self._wait_for_completion(job) return self._build_per_fv_jobs(registry, tasks, job_id, job) @@ -198,17 +209,23 @@ def _build_per_fv_jobs( if getattr(fv, "state", None) == FeatureViewState.AVAILABLE_ONLINE: jobs.append(job) else: - jobs.append(SparkApplicationMaterializationJob( - job_id, self.config.namespace, self.custom_api, - error=Exception( - f"Feature view '{task.feature_view.name}' was not " - f"materialized by the SparkApplication pod" - ), - )) + jobs.append( + SparkApplicationMaterializationJob( + job_id, + self.config.namespace, + self.custom_api, + error=Exception( + f"Feature view '{task.feature_view.name}' was not " + f"materialized by the SparkApplication pod" + ), + ) + ) return jobs def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): - feast_config_yaml = yaml.dump(self._build_driver_repo_config(), default_flow_style=False) + feast_config_yaml = yaml.dump( + self._build_driver_repo_config(), default_flow_style=False + ) mat_config = { "operation": "materialize", "tasks": [ @@ -271,7 +288,10 @@ def _build_spark_application_cr(self, job_id: str) -> dict: }, "timeToLiveSeconds": self.config.ttl_seconds_after_finished, "volumes": [ - {"name": "feast-config", "secret": {"secretName": f"feast-sa-{job_id}"}}, + { + "name": "feast-config", + "secret": {"secretName": f"feast-sa-{job_id}"}, + }, *self.config.volumes, ], "driver": { @@ -376,11 +396,15 @@ def _get_driver_logs(self, job_id: str, tail_lines: int = 50) -> Optional[str]: def _cleanup(self, job_id: str): for fn in [ lambda: self.custom_api.delete_namespaced_custom_object( - "sparkoperator.k8s.io", "v1beta2", self.config.namespace, - "sparkapplications", f"feast-sa-{job_id}", + "sparkoperator.k8s.io", + "v1beta2", + self.config.namespace, + "sparkapplications", + f"feast-sa-{job_id}", ), lambda: self.core_v1.delete_namespaced_secret( - f"feast-sa-{job_id}", self.config.namespace, + f"feast-sa-{job_id}", + self.config.namespace, ), ]: try: diff --git a/sdk/python/feast/infra/compute_engines/spark_application/job.py b/sdk/python/feast/infra/compute_engines/spark_application/job.py index c93dea7d1ae..ef52189fc01 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/job.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/job.py @@ -54,13 +54,19 @@ def status(self) -> MaterializationJobStatus: obj = self._get_cr_with_retry() if obj is None: - return MaterializationJobStatus.ERROR if self._error else MaterializationJobStatus.RUNNING + return ( + MaterializationJobStatus.ERROR + if self._error + else MaterializationJobStatus.RUNNING + ) state = obj.get("status", {}).get("applicationState", {}).get("state", "") result = _STATE_MAP.get(state, MaterializationJobStatus.WAITING) if result == MaterializationJobStatus.ERROR: - msg = obj.get("status", {}).get("applicationState", {}).get( - "errorMessage", f"SparkApplication failed: {state}" + msg = ( + obj.get("status", {}) + .get("applicationState", {}) + .get("errorMessage", f"SparkApplication failed: {state}") ) self._error = Exception(msg) return result diff --git a/sdk/python/feast/infra/compute_engines/spark_application/main.py b/sdk/python/feast/infra/compute_engines/spark_application/main.py index d5946807d20..105c3f60ffe 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/main.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/main.py @@ -13,7 +13,8 @@ def _load_config_from_secret(): """Load feast and materialization config from a Kubernetes Secret.""" - from kubernetes import client, config as k8s_config + from kubernetes import client + from kubernetes import config as k8s_config k8s_config.load_incluster_config() v1 = client.CoreV1Api() @@ -21,8 +22,12 @@ def _load_config_from_secret(): name=os.environ["FEAST_SECRET_NAME"], namespace=os.environ["FEAST_SECRET_NAMESPACE"], ) - feast_config = yaml.safe_load(base64.b64decode(secret.data["feature_store.yaml"]).decode()) - mat_config = yaml.safe_load(base64.b64decode(secret.data["materialization_config.yaml"]).decode()) + feast_config = yaml.safe_load( + base64.b64decode(secret.data["feature_store.yaml"]).decode() + ) + mat_config = yaml.safe_load( + base64.b64decode(secret.data["materialization_config.yaml"]).decode() + ) return feast_config, mat_config @@ -43,9 +48,11 @@ def _materialize_one_fv(spark_session, config, task_info): SparkSession visibility is handled via the getActiveSession patch in main(). """ from datetime import datetime, timezone - from feast import FeatureStore, RepoConfig + from tqdm import tqdm + from feast import FeatureStore, RepoConfig + fv_name = task_info["feature_view"] logger.info(f"Thread started: {fv_name}") @@ -93,20 +100,23 @@ def main(): "or mount config at /var/feast/" ) - from feast import RepoConfig from pyspark.sql import SparkSession + from feast import RepoConfig + RepoConfig(**feast_config) # validate config eagerly before any Spark work operation = mat_config["operation"] if operation == "materialize": tasks = mat_config.get("tasks", []) if not tasks: - tasks = [{ - "feature_view": mat_config["feature_view"], - "start_time": mat_config["start_time"], - "end_time": mat_config["end_time"], - }] + tasks = [ + { + "feature_view": mat_config["feature_view"], + "start_time": mat_config["start_time"], + "end_time": mat_config["end_time"], + } + ] concurrency = mat_config.get("concurrency", 1) total = len(tasks) @@ -146,7 +156,9 @@ def main(): succeeded, failed = 0, 0 with ThreadPoolExecutor(max_workers=concurrency) as executor: future_to_fv = { - executor.submit(_materialize_one_fv, spark, feast_config, task): task["feature_view"] + executor.submit( + _materialize_one_fv, spark, feast_config, task + ): task["feature_view"] for task in tasks } for future in as_completed(future_to_fv): diff --git a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py index 27d7e9451f8..66c365f6ec7 100644 --- a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py +++ b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py @@ -1,20 +1,18 @@ -from datetime import datetime from unittest.mock import MagicMock, patch import pytest +from feast.feature_view import FeatureViewState from feast.infra.common.materialization_job import ( MaterializationJobStatus, - MaterializationTask, ) from feast.infra.compute_engines.spark_application.config import ( SparkApplicationComputeEngineConfig, ) from feast.infra.compute_engines.spark_application.job import ( - SparkApplicationMaterializationJob, _STATE_MAP, + SparkApplicationMaterializationJob, ) -from feast.feature_view import FeatureViewState def _make_repo_config( @@ -36,14 +34,19 @@ def _make_repo_config( registry_address=registry_address, spark_conf=spark_conf, ) - config.model_dump = MagicMock(return_value={ - "project": "test", - "provider": "local", - "batch_engine": {"type": "spark_application"}, - "offline_store": {"type": offline_store_type, "spark_conf": {"spark.existing": "value"}}, - "online_store": {"type": online_store_type}, - "registry": {"registry_type": "file", "path": registry_path}, - }) + config.model_dump = MagicMock( + return_value={ + "project": "test", + "provider": "local", + "batch_engine": {"type": "spark_application"}, + "offline_store": { + "type": offline_store_type, + "spark_conf": {"spark.existing": "value"}, + }, + "online_store": {"type": online_store_type}, + "registry": {"registry_type": "file", "path": registry_path}, + } + ) return config @@ -54,6 +57,7 @@ def _make_engine(mock_client, mock_k8s_config, **kwargs): from feast.infra.compute_engines.spark_application.compute import ( SparkApplicationComputeEngine, ) + repo_config = _make_repo_config(**kwargs) return SparkApplicationComputeEngine( repo_config=repo_config, offline_store=None, online_store=None @@ -62,6 +66,7 @@ def _make_engine(mock_client, mock_k8s_config, **kwargs): # ── Test 1: Config defaults + required field ── + def test_config_defaults_and_required_image(): c = SparkApplicationComputeEngineConfig(image="quay.io/test:v1") assert c.type == "spark_application" @@ -76,6 +81,7 @@ def test_config_defaults_and_required_image(): # ── Test 2: EC-3 rejects SQLite ── + def test_rejects_sqlite_online_store(): with pytest.raises(ValueError, match="SQLite"): _make_engine(online_store_type="sqlite") @@ -83,6 +89,7 @@ def test_rejects_sqlite_online_store(): # ── Test 3: registry_address is mandatory ── + def test_rejects_missing_registry_address(): with pytest.raises(ValueError, match="registry_address is required"): _make_engine(registry_address=None) @@ -90,13 +97,17 @@ def test_rejects_missing_registry_address(): # ── Test 4: EC-2 accepts local registry WITH registry_address ── + def test_accepts_local_registry_with_address(): - engine = _make_engine(registry_path="/local/registry.db", registry_address="feast:6570") + engine = _make_engine( + registry_path="/local/registry.db", registry_address="feast:6570" + ) assert engine is not None # ── Test 5: _build_driver_repo_config — two rewrites (batch_engine + registry) ── + def test_build_driver_repo_config_rewrites(): engine = _make_engine( offline_store_type="dask", @@ -110,6 +121,7 @@ def test_build_driver_repo_config_rewrites(): # ── Test 6: _build_driver_repo_config — batch_engine is just type (no spark_conf copy) ── + def test_build_driver_repo_config_batch_engine_minimal(): engine = _make_engine(spark_conf={"spark.new": "from_engine"}) d = engine._build_driver_repo_config() @@ -119,6 +131,7 @@ def test_build_driver_repo_config_batch_engine_minimal(): # ── Test 7: CR structure ── + def test_cr_structure(): engine = _make_engine() cr = engine._build_spark_application_cr("abcd1234") @@ -134,16 +147,21 @@ def test_cr_structure(): # ── Test 8: CR sparkConf includes driver env passthrough ── + def test_cr_driver_env_passthrough(): engine = _make_engine() cr = engine._build_spark_application_cr("abcd1234") spark_conf = cr["spec"]["sparkConf"] - assert spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAME"] == "feast-sa-abcd1234" + assert ( + spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAME"] + == "feast-sa-abcd1234" + ) assert spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE"] == "default" # ── Test 9: Status mapping covers all 14 states ── + def test_state_map_coverage(): assert len(_STATE_MAP) == 14 assert _STATE_MAP["COMPLETED"] == MaterializationJobStatus.SUCCEEDED @@ -155,10 +173,12 @@ def test_state_map_coverage(): # ── Test 10: Cleanup swallows 404 ── + @patch("feast.infra.compute_engines.spark_application.compute.k8s_config") @patch("feast.infra.compute_engines.spark_application.compute.client") def test_cleanup_swallows_404(mock_client, mock_k8s_config): from kubernetes.client.exceptions import ApiException + from feast.infra.compute_engines.spark_application.compute import ( SparkApplicationComputeEngine, ) @@ -168,7 +188,9 @@ def test_cleanup_swallows_404(mock_client, mock_k8s_config): repo_config=repo_config, offline_store=None, online_store=None ) - engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException(status=404) + engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException( + status=404 + ) engine.core_v1.delete_namespaced_secret.side_effect = ApiException(status=404) engine._cleanup("test-id") @@ -176,6 +198,7 @@ def test_cleanup_swallows_404(mock_client, mock_k8s_config): # ── Test 11: Timeout sets error on job (does not raise) ── + @patch("feast.infra.compute_engines.spark_application.compute.k8s_config") @patch("feast.infra.compute_engines.spark_application.compute.client") @patch("feast.infra.compute_engines.spark_application.compute.time") @@ -186,7 +209,9 @@ def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): repo_config = _make_repo_config() repo_config.batch_engine = SparkApplicationComputeEngineConfig( - image="test", job_timeout_seconds=1, poll_interval_seconds=1, + image="test", + job_timeout_seconds=1, + poll_interval_seconds=1, registry_address="feast-server:6566", ) engine = SparkApplicationComputeEngine( @@ -210,6 +235,7 @@ def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): # ── Test 12: Job naming < 63 chars ── + def test_job_naming_under_63_chars(): mock_api = MagicMock() job = SparkApplicationMaterializationJob("abcdef12", "default", mock_api) @@ -219,6 +245,7 @@ def test_job_naming_under_63_chars(): # ── Test 13: _build_per_fv_jobs — all succeeded ── + def test_build_per_fv_jobs_all_succeeded(): engine = _make_engine() mock_registry = MagicMock() @@ -247,6 +274,7 @@ def test_build_per_fv_jobs_all_succeeded(): # ── Test 14: _build_per_fv_jobs — partial failure ── + def test_build_per_fv_jobs_partial_failure(): engine = _make_engine() mock_registry = MagicMock() @@ -267,7 +295,9 @@ def test_build_per_fv_jobs_partial_failure(): task_fail.project = "test" parent_job = SparkApplicationMaterializationJob("job1", "default", MagicMock()) - jobs = engine._build_per_fv_jobs(mock_registry, [task_ok, task_fail], "job1", parent_job) + jobs = engine._build_per_fv_jobs( + mock_registry, [task_ok, task_fail], "job1", parent_job + ) assert len(jobs) == 2 assert jobs[0].status() != MaterializationJobStatus.ERROR @@ -277,6 +307,7 @@ def test_build_per_fv_jobs_partial_failure(): # ── Test 15: _build_per_fv_jobs — single task returns parent job directly ── + def test_build_per_fv_jobs_single_task(): engine = _make_engine() mock_registry = MagicMock() From 1e3fb3649d930523b39d59196faf7b4d41ea9524 Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Fri, 10 Jul 2026 03:13:17 +0530 Subject: [PATCH 4/4] =?UTF-8?q?feat:=20Switch=20to=20ConfigMap,=20remove?= =?UTF-8?q?=20registry=5Faddress,=20reject=20file-based=20stores=20-=20Con?= =?UTF-8?q?fig=20delivery:=20Secret=20=E2=86=92=20ConfigMap.=20Operator's?= =?UTF-8?q?=20ClusterRole=20already=20has=20=20=20full=20ConfigMap=20CRUD?= =?UTF-8?q?=20=E2=80=94=20avoids=20widening=20RBAC=20for=20Secrets=20in=20?= =?UTF-8?q?ODH.=20=20=20Matches=20KubernetesComputeEngine=20pattern.=20-?= =?UTF-8?q?=20Removed=20registry=5Faddress=20config=20field.=20Pod=20inher?= =?UTF-8?q?its=20server's=20registry=20=20=20config=20(SQL,=20Snowflake)?= =?UTF-8?q?=20directly=20and=20writes=20apply=5Fmaterialization()=20=20=20?= =?UTF-8?q?to=20the=20same=20database.=20Eliminates=20TLS=20certificate=20?= =?UTF-8?q?mounting=20complexity.=20-=20Reject=20file-based=20offline=20st?= =?UTF-8?q?ores=20(dask,=20file,=20duckdb)=20and=20registries=20=20=20(fil?= =?UTF-8?q?e)=20at=20=5F=5Finit=5F=5F(),=20same=20as=20existing=20sqlite/f?= =?UTF-8?q?aiss=20online=20store=20=20=20rejection.=20SparkApplication=20p?= =?UTF-8?q?od=20has=20ephemeral=20filesystem.=20-=20Dockerfile:=20added=20?= =?UTF-8?q?PYTHONPATH/SPARK=5FHOME=20for=20PySpark,=20added=20pymysql.=20-?= =?UTF-8?q?=2020/20=20unit=20tests=20pass.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Aniket Paluskar --- .../spark_application/Dockerfile | 5 +- .../spark_application/compute.py | 90 +++++++------ .../spark_application/config.py | 1 - .../compute_engines/spark_application/main.py | 53 +++----- .../compute_engines/test_spark_application.py | 120 ++++++++---------- 5 files changed, 125 insertions(+), 144 deletions(-) diff --git a/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile index 268e377b4bc..066aac3ad86 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile +++ b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile @@ -6,6 +6,9 @@ RUN apt-get update && apt-get install --no-install-suggests --no-install-recomme COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv +ENV PYTHONPATH="/opt/spark/python:/opt/spark/python/lib/py4j-0.10.9.9-src.zip" +ENV SPARK_HOME="/opt/spark" + WORKDIR /app COPY sdk/python/feast/infra/compute_engines/spark_application/main.py /opt/feast/main.py @@ -18,7 +21,7 @@ COPY README.md README.md # setuptools_scm needs .git to infer the version. # https://github.com/pypa/setuptools_scm#usage-from-docker -RUN --mount=source=.git,target=.git,type=bind uv pip install --system --no-cache-dir '.[redis,grpcio,k8s]' +RUN --mount=source=.git,target=.git,type=bind uv pip install --system --no-cache-dir '.[redis,grpcio,k8s]' pymysql # Hadoop AWS JARs for S3A filesystem support + AWS SDK v2 bundle RUN curl -sL https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/3.4.1/hadoop-aws-3.4.1.jar \ diff --git a/sdk/python/feast/infra/compute_engines/spark_application/compute.py b/sdk/python/feast/infra/compute_engines/spark_application/compute.py index 4d95c56d925..cf19001d395 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/compute.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/compute.py @@ -51,24 +51,39 @@ def __init__( ) self.config = repo_config.batch_engine + _FILE_BASED_ONLINE = {"sqlite", "faiss"} + _FILE_BASED_OFFLINE = {"dask", "file", "duckdb"} + _FILE_BASED_REGISTRY = {"file"} + online_type = getattr(repo_config.online_store, "type", "") - if online_type == "sqlite": + if online_type in _FILE_BASED_ONLINE: raise ValueError( - "spark_application engine cannot use SQLite online store. " - "SQLite is file-based — data written inside the " - "SparkApplication pod is lost when the pod terminates. " - "Use a network-accessible store: redis, postgres, etc." + f"spark_application engine cannot use '{online_type}' online store. " + f"File-based stores ({', '.join(sorted(_FILE_BASED_ONLINE))}) write " + "data inside the SparkApplication pod, which is lost when the pod " + "terminates. Use a network-accessible store: redis, postgres, etc." ) - # EC-2: registry_address required — pod reports materialization results - # back to the Feast server via gRPC. Without it, the server cannot - # determine per-FV success/failure after a batched SparkApplication run. - if not self.config.registry_address: + offline_type = getattr(repo_config.offline_store, "type", "") + if offline_type in _FILE_BASED_OFFLINE: raise ValueError( - "registry_address is required for spark_application engine. " - "Set it to the Feast server's gRPC endpoint " - "(e.g., feast-server.namespace.svc.cluster.local:6566). " - "The Feast Operator sets this automatically." + f"spark_application engine cannot use '{offline_type}' offline store. " + f"File-based stores ({', '.join(sorted(_FILE_BASED_OFFLINE))}) read " + "from the local filesystem, which is inaccessible from the " + "SparkApplication pod. Use a network-accessible store: spark, " + "bigquery, snowflake, redshift, etc." + ) + + registry_type = getattr( + repo_config.registry, "registry_type", "" + ) + if registry_type in _FILE_BASED_REGISTRY: + raise ValueError( + f"spark_application engine cannot use '{registry_type}' registry. " + f"File-based registries ({', '.join(sorted(_FILE_BASED_REGISTRY))}) " + "store data on the local filesystem, which is inaccessible from the " + "SparkApplication pod. Use a network-accessible registry: sql, " + "snowflake.registry, etc." ) k8s_config.load_config() @@ -126,13 +141,13 @@ def materialize( job_id = uuid.uuid4().hex[:8] try: - self._create_secret(job_id, tasks) + self._create_configmap(job_id, tasks) except ApiException as e: job = SparkApplicationMaterializationJob( job_id, self.config.namespace, self.custom_api, - error=Exception(f"Secret creation failed: {e.reason}"), + error=Exception(f"ConfigMap creation failed: {e.reason}"), ) return [job for _ in tasks] @@ -164,28 +179,21 @@ def materialize( def _build_driver_repo_config(self) -> dict: """Build feature_store.yaml for the SparkApplication driver pod. - Two rewrites: - 1. batch_engine → spark.engine: Pod uses SparkComputeEngine with the - active SparkSession (from spark-submit). Enables distributed reads - via SparkReadNode and distributed writes via mapInArrow across executors. - This is NOT recursive — SparkComputeEngine uses the local session, - it does not create CRDs. - 2. registry → remote (if registry_address set): Pod can't access - server's local filesystem. Routes registry ops via gRPC. - - offline_store is NOT rewritten — respects user's configured data sources. - User should configure offline_store: spark for full distributed performance. + One rewrite: batch_engine → spark.engine. Pod uses SparkComputeEngine + with the active SparkSession (from spark-submit). Enables distributed + reads via SparkReadNode and distributed writes via mapInArrow across + executors. This is NOT recursive — SparkComputeEngine uses the local + session, it does not create CRDs. + + offline_store and registry are NOT rewritten — the pod inherits the + server's config. File-based registries are rejected at __init__(), so + the registry is always network-accessible (SQL, Snowflake, etc.) and + the pod can write apply_materialization() directly. """ config_dict = self.repo_config.model_dump(by_alias=True, mode="json") config_dict["batch_engine"] = {"type": "spark.engine"} - if self.config.registry_address: - config_dict["registry"] = { - "registry_type": "remote", - "path": self.config.registry_address, - } - return config_dict def _build_per_fv_jobs( @@ -222,7 +230,7 @@ def _build_per_fv_jobs( ) return jobs - def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): + def _create_configmap(self, job_id: str, tasks: List[MaterializationTask]): feast_config_yaml = yaml.dump( self._build_driver_repo_config(), default_flow_style=False ) @@ -242,25 +250,25 @@ def _create_secret(self, job_id: str, tasks: List[MaterializationTask]): mat_config_yaml = yaml.dump(mat_config) manifest = { "apiVersion": "v1", - "kind": "Secret", + "kind": "ConfigMap", "metadata": { "name": f"feast-sa-{job_id}", "namespace": self.config.namespace, - "labels": {"feast-materializer": "secret", **self.config.labels}, + "labels": {"feast-materializer": "configmap", **self.config.labels}, }, - "stringData": { + "data": { "feature_store.yaml": feast_config_yaml, "materialization_config.yaml": mat_config_yaml, }, } - self.core_v1.create_namespaced_secret( + self.core_v1.create_namespaced_config_map( namespace=self.config.namespace, body=manifest ) def _build_spark_application_cr(self, job_id: str) -> dict: driver_env_conf = { - "spark.kubernetes.driverEnv.FEAST_SECRET_NAME": f"feast-sa-{job_id}", - "spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE": self.config.namespace, + "spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAME": f"feast-sa-{job_id}", + "spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAMESPACE": self.config.namespace, } for entry in self.config.env: name = entry.get("name", "") @@ -290,7 +298,7 @@ def _build_spark_application_cr(self, job_id: str) -> dict: "volumes": [ { "name": "feast-config", - "secret": {"secretName": f"feast-sa-{job_id}"}, + "configMap": {"name": f"feast-sa-{job_id}"}, }, *self.config.volumes, ], @@ -402,7 +410,7 @@ def _cleanup(self, job_id: str): "sparkapplications", f"feast-sa-{job_id}", ), - lambda: self.core_v1.delete_namespaced_secret( + lambda: self.core_v1.delete_namespaced_config_map( f"feast-sa-{job_id}", self.config.namespace, ), diff --git a/sdk/python/feast/infra/compute_engines/spark_application/config.py b/sdk/python/feast/infra/compute_engines/spark_application/config.py index 101c5d864c0..7f766ce8db2 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/config.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/config.py @@ -48,7 +48,6 @@ class SparkApplicationComputeEngineConfig(FeastConfigBaseModel): py_files: List[str] = [] node_selector: Optional[Dict[str, str]] = None tolerations: List[dict] = [] - registry_address: Optional[str] = None @model_validator(mode="after") def _validate_config(self) -> "SparkApplicationComputeEngineConfig": diff --git a/sdk/python/feast/infra/compute_engines/spark_application/main.py b/sdk/python/feast/infra/compute_engines/spark_application/main.py index 105c3f60ffe..c48e12735d4 100644 --- a/sdk/python/feast/infra/compute_engines/spark_application/main.py +++ b/sdk/python/feast/infra/compute_engines/spark_application/main.py @@ -1,4 +1,3 @@ -import base64 import logging import os import sys @@ -11,23 +10,18 @@ logger = logging.getLogger("feast.spark_application.driver") -def _load_config_from_secret(): - """Load feast and materialization config from a Kubernetes Secret.""" - from kubernetes import client - from kubernetes import config as k8s_config +def _load_config_from_configmap(): + """Load feast and materialization config from a Kubernetes ConfigMap.""" + from kubernetes import client, config as k8s_config k8s_config.load_incluster_config() v1 = client.CoreV1Api() - secret = v1.read_namespaced_secret( - name=os.environ["FEAST_SECRET_NAME"], - namespace=os.environ["FEAST_SECRET_NAMESPACE"], - ) - feast_config = yaml.safe_load( - base64.b64decode(secret.data["feature_store.yaml"]).decode() - ) - mat_config = yaml.safe_load( - base64.b64decode(secret.data["materialization_config.yaml"]).decode() + cm = v1.read_namespaced_config_map( + name=os.environ["FEAST_CONFIGMAP_NAME"], + namespace=os.environ["FEAST_CONFIGMAP_NAMESPACE"], ) + feast_config = yaml.safe_load(cm.data["feature_store.yaml"]) + mat_config = yaml.safe_load(cm.data["materialization_config.yaml"]) return feast_config, mat_config @@ -48,10 +42,8 @@ def _materialize_one_fv(spark_session, config, task_info): SparkSession visibility is handled via the getActiveSession patch in main(). """ from datetime import datetime, timezone - - from tqdm import tqdm - from feast import FeatureStore, RepoConfig + from tqdm import tqdm fv_name = task_info["feature_view"] logger.info(f"Thread started: {fv_name}") @@ -88,21 +80,20 @@ def _materialize_one_fv(spark_session, config, task_info): def main(): - if os.environ.get("FEAST_SECRET_NAME"): - logger.info("Loading config from Secret via K8s API") - feast_config, mat_config = _load_config_from_secret() + if os.environ.get("FEAST_CONFIGMAP_NAME"): + logger.info("Loading config from ConfigMap via K8s API") + feast_config, mat_config = _load_config_from_configmap() elif os.path.exists("/var/feast/feature_store.yaml"): logger.info("Loading config from mounted files") feast_config, mat_config = _load_config_from_files() else: raise RuntimeError( - "No config source found. Set FEAST_SECRET_NAME env var " + "No config source found. Set FEAST_CONFIGMAP_NAME env var " "or mount config at /var/feast/" ) - from pyspark.sql import SparkSession - from feast import RepoConfig + from pyspark.sql import SparkSession RepoConfig(**feast_config) # validate config eagerly before any Spark work operation = mat_config["operation"] @@ -110,13 +101,11 @@ def main(): if operation == "materialize": tasks = mat_config.get("tasks", []) if not tasks: - tasks = [ - { - "feature_view": mat_config["feature_view"], - "start_time": mat_config["start_time"], - "end_time": mat_config["end_time"], - } - ] + tasks = [{ + "feature_view": mat_config["feature_view"], + "start_time": mat_config["start_time"], + "end_time": mat_config["end_time"], + }] concurrency = mat_config.get("concurrency", 1) total = len(tasks) @@ -156,9 +145,7 @@ def main(): succeeded, failed = 0, 0 with ThreadPoolExecutor(max_workers=concurrency) as executor: future_to_fv = { - executor.submit( - _materialize_one_fv, spark, feast_config, task - ): task["feature_view"] + executor.submit(_materialize_one_fv, spark, feast_config, task): task["feature_view"] for task in tasks } for future in as_completed(future_to_fv): diff --git a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py index 66c365f6ec7..8a9fbccde13 100644 --- a/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py +++ b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py @@ -1,52 +1,50 @@ +from datetime import datetime from unittest.mock import MagicMock, patch import pytest -from feast.feature_view import FeatureViewState from feast.infra.common.materialization_job import ( MaterializationJobStatus, + MaterializationTask, ) from feast.infra.compute_engines.spark_application.config import ( SparkApplicationComputeEngineConfig, ) from feast.infra.compute_engines.spark_application.job import ( - _STATE_MAP, SparkApplicationMaterializationJob, + _STATE_MAP, ) +from feast.feature_view import FeatureViewState def _make_repo_config( online_store_type="redis", - registry_path="s3://bucket/registry.db", - registry_address="feast-server:6566", - offline_store_type="dask", + registry_type="sql", + registry_path="postgresql://user:pass@host:5432/feast", + offline_store_type="spark", spark_conf=None, ): """Build a mock RepoConfig for testing.""" config = MagicMock() config.online_store = MagicMock() config.online_store.type = online_store_type + config.offline_store = MagicMock() + config.offline_store.type = offline_store_type config.registry = MagicMock() config.registry.path = registry_path - config.registry.registry_type = "file" + config.registry.registry_type = registry_type config.batch_engine = SparkApplicationComputeEngineConfig( image="quay.io/test/feast-spark:latest", - registry_address=registry_address, spark_conf=spark_conf, ) - config.model_dump = MagicMock( - return_value={ - "project": "test", - "provider": "local", - "batch_engine": {"type": "spark_application"}, - "offline_store": { - "type": offline_store_type, - "spark_conf": {"spark.existing": "value"}, - }, - "online_store": {"type": online_store_type}, - "registry": {"registry_type": "file", "path": registry_path}, - } - ) + config.model_dump = MagicMock(return_value={ + "project": "test", + "provider": "local", + "batch_engine": {"type": "spark_application"}, + "offline_store": {"type": offline_store_type, "spark_conf": {"spark.existing": "value"}}, + "online_store": {"type": online_store_type}, + "registry": {"registry_type": registry_type, "path": registry_path}, + }) return config @@ -57,7 +55,6 @@ def _make_engine(mock_client, mock_k8s_config, **kwargs): from feast.infra.compute_engines.spark_application.compute import ( SparkApplicationComputeEngine, ) - repo_config = _make_repo_config(**kwargs) return SparkApplicationComputeEngine( repo_config=repo_config, offline_store=None, online_store=None @@ -66,7 +63,6 @@ def _make_engine(mock_client, mock_k8s_config, **kwargs): # ── Test 1: Config defaults + required field ── - def test_config_defaults_and_required_image(): c = SparkApplicationComputeEngineConfig(image="quay.io/test:v1") assert c.type == "spark_application" @@ -79,49 +75,57 @@ def test_config_defaults_and_required_image(): SparkApplicationComputeEngineConfig() # image is required -# ── Test 2: EC-3 rejects SQLite ── - +# ── Test 2: EC-3 rejects file-based online stores ── def test_rejects_sqlite_online_store(): - with pytest.raises(ValueError, match="SQLite"): + with pytest.raises(ValueError, match="sqlite"): _make_engine(online_store_type="sqlite") -# ── Test 3: registry_address is mandatory ── +def test_rejects_faiss_online_store(): + with pytest.raises(ValueError, match="faiss"): + _make_engine(online_store_type="faiss") -def test_rejects_missing_registry_address(): - with pytest.raises(ValueError, match="registry_address is required"): - _make_engine(registry_address=None) +# ── Test 2b: rejects file-based offline stores ── +@pytest.mark.parametrize("store_type", ["dask", "file", "duckdb"]) +def test_rejects_file_based_offline_store(store_type): + with pytest.raises(ValueError, match=store_type): + _make_engine(offline_store_type=store_type) -# ── Test 4: EC-2 accepts local registry WITH registry_address ── +# ── Test 3: EC-2 rejects file-based registries ── -def test_accepts_local_registry_with_address(): - engine = _make_engine( - registry_path="/local/registry.db", registry_address="feast:6570" - ) +def test_rejects_file_registry(): + with pytest.raises(ValueError, match="file.*registry"): + _make_engine(registry_type="file") + + +# ── Test 4: accepts network-accessible registries ── + +def test_accepts_sql_registry(): + engine = _make_engine(registry_type="sql") assert engine is not None -# ── Test 5: _build_driver_repo_config — two rewrites (batch_engine + registry) ── +def test_accepts_snowflake_registry(): + engine = _make_engine(registry_type="snowflake.registry") + assert engine is not None + +# ── Test 5: _build_driver_repo_config — one rewrite (batch_engine only) ── def test_build_driver_repo_config_rewrites(): - engine = _make_engine( - offline_store_type="dask", - registry_address="feast-server:6566", - ) + engine = _make_engine(offline_store_type="spark") d = engine._build_driver_repo_config() assert d["batch_engine"]["type"] == "spark.engine" - assert d["offline_store"]["type"] == "dask" # NOT rewritten — respects user intent - assert d["registry"] == {"registry_type": "remote", "path": "feast-server:6566"} + assert d["offline_store"]["type"] == "spark" # NOT rewritten — respects user intent + assert d["registry"]["registry_type"] == "sql" # NOT rewritten — pod uses SQL directly # ── Test 6: _build_driver_repo_config — batch_engine is just type (no spark_conf copy) ── - def test_build_driver_repo_config_batch_engine_minimal(): engine = _make_engine(spark_conf={"spark.new": "from_engine"}) d = engine._build_driver_repo_config() @@ -131,7 +135,6 @@ def test_build_driver_repo_config_batch_engine_minimal(): # ── Test 7: CR structure ── - def test_cr_structure(): engine = _make_engine() cr = engine._build_spark_application_cr("abcd1234") @@ -147,21 +150,16 @@ def test_cr_structure(): # ── Test 8: CR sparkConf includes driver env passthrough ── - def test_cr_driver_env_passthrough(): engine = _make_engine() cr = engine._build_spark_application_cr("abcd1234") spark_conf = cr["spec"]["sparkConf"] - assert ( - spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAME"] - == "feast-sa-abcd1234" - ) - assert spark_conf["spark.kubernetes.driverEnv.FEAST_SECRET_NAMESPACE"] == "default" + assert spark_conf["spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAME"] == "feast-sa-abcd1234" + assert spark_conf["spark.kubernetes.driverEnv.FEAST_CONFIGMAP_NAMESPACE"] == "default" # ── Test 9: Status mapping covers all 14 states ── - def test_state_map_coverage(): assert len(_STATE_MAP) == 14 assert _STATE_MAP["COMPLETED"] == MaterializationJobStatus.SUCCEEDED @@ -173,12 +171,10 @@ def test_state_map_coverage(): # ── Test 10: Cleanup swallows 404 ── - @patch("feast.infra.compute_engines.spark_application.compute.k8s_config") @patch("feast.infra.compute_engines.spark_application.compute.client") def test_cleanup_swallows_404(mock_client, mock_k8s_config): from kubernetes.client.exceptions import ApiException - from feast.infra.compute_engines.spark_application.compute import ( SparkApplicationComputeEngine, ) @@ -188,17 +184,14 @@ def test_cleanup_swallows_404(mock_client, mock_k8s_config): repo_config=repo_config, offline_store=None, online_store=None ) - engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException( - status=404 - ) - engine.core_v1.delete_namespaced_secret.side_effect = ApiException(status=404) + engine.custom_api.delete_namespaced_custom_object.side_effect = ApiException(status=404) + engine.core_v1.delete_namespaced_config_map.side_effect = ApiException(status=404) engine._cleanup("test-id") # ── Test 11: Timeout sets error on job (does not raise) ── - @patch("feast.infra.compute_engines.spark_application.compute.k8s_config") @patch("feast.infra.compute_engines.spark_application.compute.client") @patch("feast.infra.compute_engines.spark_application.compute.time") @@ -209,10 +202,7 @@ def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): repo_config = _make_repo_config() repo_config.batch_engine = SparkApplicationComputeEngineConfig( - image="test", - job_timeout_seconds=1, - poll_interval_seconds=1, - registry_address="feast-server:6566", + image="test", job_timeout_seconds=1, poll_interval_seconds=1, ) engine = SparkApplicationComputeEngine( repo_config=repo_config, offline_store=None, online_store=None @@ -235,7 +225,6 @@ def test_timeout_sets_error(mock_time, mock_client, mock_k8s_config): # ── Test 12: Job naming < 63 chars ── - def test_job_naming_under_63_chars(): mock_api = MagicMock() job = SparkApplicationMaterializationJob("abcdef12", "default", mock_api) @@ -245,7 +234,6 @@ def test_job_naming_under_63_chars(): # ── Test 13: _build_per_fv_jobs — all succeeded ── - def test_build_per_fv_jobs_all_succeeded(): engine = _make_engine() mock_registry = MagicMock() @@ -274,7 +262,6 @@ def test_build_per_fv_jobs_all_succeeded(): # ── Test 14: _build_per_fv_jobs — partial failure ── - def test_build_per_fv_jobs_partial_failure(): engine = _make_engine() mock_registry = MagicMock() @@ -295,9 +282,7 @@ def test_build_per_fv_jobs_partial_failure(): task_fail.project = "test" parent_job = SparkApplicationMaterializationJob("job1", "default", MagicMock()) - jobs = engine._build_per_fv_jobs( - mock_registry, [task_ok, task_fail], "job1", parent_job - ) + jobs = engine._build_per_fv_jobs(mock_registry, [task_ok, task_fail], "job1", parent_job) assert len(jobs) == 2 assert jobs[0].status() != MaterializationJobStatus.ERROR @@ -307,7 +292,6 @@ def test_build_per_fv_jobs_partial_failure(): # ── Test 15: _build_per_fv_jobs — single task returns parent job directly ── - def test_build_per_fv_jobs_single_task(): engine = _make_engine() mock_registry = MagicMock()