diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 75f406c8e2c..032d1f7e18c 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -424,6 +424,113 @@ 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() + + 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.""" @@ -2249,7 +2356,21 @@ def materialize_incremental( _mat_start = time.monotonic() try: - # TODO paging large loads + 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: @@ -2278,15 +2399,15 @@ def materialize_incremental( ) continue - start_date = feature_view.most_recent_end_time - if start_date is None: + 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: - start_date = _utc_now() - feature_view.ttl + 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. @@ -2294,80 +2415,39 @@ def materialize_incremental( 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() + fv_start_date = _utc_now() - timedelta(weeks=52) + + 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(start_date.replace(microsecond=0))}{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}:" ) - def tqdm_builder(length): - return tqdm(total=length, ncols=100) - - start_date = utils.make_tzaware(start_date) - end_date = utils.make_tzaware(end_date) or _utc_now() - - # 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 - ): - 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 - ) - - 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, + self._transition_fv_to_materializing( + feature_view, regular_fvs, previous_states + ) + regular_fvs.append(feature_view) + tasks.append( + MaterializationTask( project=self.project, + feature_view=feature_view, + start_time=fv_start_date, + end_time=end_date_tz, tqdm_builder=tqdm_builder, ) - 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 - ) - raise - finally: - _tracker = _get_track_materialization() - if _tracker is not None: - _tracker( - feature_view.name, - fv_success, - time.monotonic() - fv_start, - ) + ) - if not isinstance(feature_view, OnDemandFeatureView): - self.registry.apply_materialization( - feature_view, - self.project, - start_date, - end_date, - ) + if tasks: + self._submit_and_process_materialization_jobs( + provider, + tasks, + regular_fvs, + previous_states, + fv_start_dates, + ) materialized_fv_names = [ fv.name @@ -2459,7 +2539,22 @@ def materialize( _mat_start = time.monotonic() try: - # TODO paging large loads + 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: @@ -2473,76 +2568,34 @@ def materialize( 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) - - 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 - ): - 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 - ) - - 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, + 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=feature_view, + start_time=start_date, + end_time=end_date, tqdm_builder=tqdm_builder, disable_event_timestamp=disable_event_timestamp, ) - 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 - ) - raise - finally: - _tracker = _get_track_materialization() - if _tracker is not None: - _tracker( - feature_view.name, - fv_success, - time.monotonic() - fv_start, - ) + ) - self.registry.apply_materialization( - feature_view, - self.project, - start_date, - end_date, + if tasks: + self._submit_and_process_materialization_jobs( + provider, + tasks, + regular_fvs, + previous_states, + fv_start_dates, ) materialized_fv_names = [ 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..066aac3ad86 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/Dockerfile @@ -0,0 +1,34 @@ +FROM apache/spark:4.0.1 + +USER root + +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 + +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 + +# 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]' 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 \ + -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 + +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..cf19001d395 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/compute.py @@ -0,0 +1,422 @@ +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, FeatureViewState +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 + + _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 in _FILE_BASED_ONLINE: + raise ValueError( + 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." + ) + + offline_type = getattr(repo_config.offline_store, "type", "") + if offline_type in _FILE_BASED_OFFLINE: + raise ValueError( + 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() + 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. + + 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] + + job_id = uuid.uuid4().hex[:8] + + try: + self._create_configmap(job_id, tasks) + except ApiException as e: + job = SparkApplicationMaterializationJob( + job_id, + self.config.namespace, + self.custom_api, + error=Exception(f"ConfigMap 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 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. + + 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"} + + 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_configmap(self, job_id: str, tasks: List[MaterializationTask]): + feast_config_yaml = yaml.dump( + self._build_driver_repo_config(), default_flow_style=False + ) + 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": "ConfigMap", + "metadata": { + "name": f"feast-sa-{job_id}", + "namespace": self.config.namespace, + "labels": {"feast-materializer": "configmap", **self.config.labels}, + }, + "data": { + "feature_store.yaml": feast_config_yaml, + "materialization_config.yaml": mat_config_yaml, + }, + } + 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_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", "") + value = entry.get("value", "") + if name and value: + driver_env_conf[f"spark.kubernetes.driverEnv.{name}"] = value + + 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 {}), + **driver_env_conf, + }, + "restartPolicy": { + "type": self.config.restart_policy, + "onFailureRetries": self.config.max_retries, + "onFailureRetryInterval": 30, + }, + "timeToLiveSeconds": self.config.ttl_seconds_after_finished, + "volumes": [ + { + "name": "feast-config", + "configMap": {"name": 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_config_map( + 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..7f766ce8db2 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/config.py @@ -0,0 +1,67 @@ +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] = [] + + @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..ef52189fc01 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/job.py @@ -0,0 +1,115 @@ +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..c48e12735d4 --- /dev/null +++ b/sdk/python/feast/infra/compute_engines/spark_application/main.py @@ -0,0 +1,182 @@ +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_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() + 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 + + +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), + ) + + 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 + + +def main(): + 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_CONFIGMAP_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 ff23907d8a2..de83a2db163 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..8a9fbccde13 --- /dev/null +++ b/sdk/python/tests/unit/infra/compute_engines/test_spark_application.py @@ -0,0 +1,307 @@ +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, +) +from feast.feature_view import FeatureViewState + + +def _make_repo_config( + online_store_type="redis", + 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 = registry_type + config.batch_engine = SparkApplicationComputeEngineConfig( + image="quay.io/test/feast-spark:latest", + 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": registry_type, "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 file-based online stores ── + +def test_rejects_sqlite_online_store(): + with pytest.raises(ValueError, match="sqlite"): + _make_engine(online_store_type="sqlite") + + +def test_rejects_faiss_online_store(): + with pytest.raises(ValueError, match="faiss"): + _make_engine(online_store_type="faiss") + + +# ── 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 3: EC-2 rejects file-based registries ── + +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 + + +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="spark") + d = engine._build_driver_repo_config() + assert d["batch_engine"]["type"] == "spark.engine" + 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() + assert d["batch_engine"] == {"type": "spark.engine"} + assert d["offline_store"]["spark_conf"]["spark.existing"] == "value" + + +# ── Test 7: 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 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_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 + 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 + ) + + 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") +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" + + +# ── 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()