Skip to content

Commit 42c4480

Browse files
pyalexpyalex
committed
launcher
move common job properties to job classes Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com> universal test
1 parent fe202c9 commit 42c4480

9 files changed

Lines changed: 532 additions & 98 deletions

File tree

‎sdk/python/feast/client.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import uuid
1818
from itertools import groupby
1919
from typing import Any, Dict, List, Optional, Union
20+
from datetime import datetime
2021

2122
import grpc
2223
import pandas as pd
@@ -78,11 +79,13 @@
7879
from feast.pyspark.launcher import (
7980
start_historical_feature_retrieval_job,
8081
start_historical_feature_retrieval_spark_session,
82+
start_offline_to_online_ingestion,
8183
)
8284
from feast.serving.ServingService_pb2 import (
8385
GetFeastServingInfoRequest,
8486
GetOnlineFeaturesRequestV2,
8587
)
88+
from feast.pyspark.abc import SparkJob
8689
from feast.serving.ServingService_pb2_grpc import ServingServiceStub
8790

8891
_logger = logging.getLogger(__name__)
@@ -883,3 +886,11 @@ def _get_feature_tables_from_feature_refs(
883886
]
884887
feature_tables.append(feature_table)
885888
return feature_tables
889+
890+
def start_offline_to_online_ingestion(
891+
self,
892+
feature_table: Union[FeatureTable, str],
893+
start: Union[datetime, str],
894+
end: Union[datetime, str],
895+
) -> SparkJob:
896+
return start_offline_to_online_ingestion(feature_table, start, end, self)

‎sdk/python/feast/constants.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,10 @@ class AuthProvider(Enum):
6767
# Spark Job Config
6868
CONFIG_SPARK_LAUNCHER = "spark_launcher" # standalone, dataproc, emr
6969

70+
CONFIG_SPARK_INGESTION_JOB_JAR = "spark_ingestion_jar"
71+
7072
CONFIG_SPARK_STANDALONE_MASTER = "spark_standalone_master"
73+
CONFIG_SPARK_HOME = "spark_home"
7174

7275
CONFIG_SPARK_DATAPROC_CLUSTER_NAME = "dataproc_cluster_name"
7376
CONFIG_SPARK_DATAPROC_PROJECT = "dataproc_project"
@@ -109,4 +112,6 @@ class AuthProvider(Enum):
109112
CONFIG_MAX_WAIT_INTERVAL_KEY: "60",
110113
# Authentication Provider - Google OpenID/OAuth
111114
CONFIG_AUTH_PROVIDER: "google",
115+
CONFIG_SPARK_LAUNCHER: "dataproc",
116+
CONFIG_SPARK_INGESTION_JOB_JAR: "gs://feast-jobs/feast-ingestion-spark-0.8-SNAPSHOT.jar",
112117
}

‎sdk/python/feast/pyspark/abc.py‎

Lines changed: 135 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import abc
2-
from typing import Dict, List
2+
import os
3+
import json
4+
from datetime import datetime
5+
from typing import Dict, List, Optional
6+
from enum import Enum
37

48

59
class SparkJobFailure(Exception):
@@ -10,6 +14,12 @@ class SparkJobFailure(Exception):
1014
pass
1115

1216

17+
class SparkJobStatus(Enum):
18+
IN_PROGRESS = 1
19+
FAILED = 2
20+
COMPLETED = 3
21+
22+
1323
class SparkJob(abc.ABC):
1424
"""
1525
Base class for all spark jobs
@@ -25,12 +35,87 @@ def get_id(self) -> str:
2535
"""
2636
raise NotImplementedError
2737

38+
@abc.abstractmethod
39+
def get_status(self) -> SparkJobStatus:
40+
"""
41+
Job Status retrieval
42+
43+
:return: SparkJobStatus
44+
"""
45+
raise NotImplementedError
46+
47+
@abc.abstractmethod
48+
def get_name(self) -> str:
49+
"""
50+
Getter for job name
51+
:return: Job name
52+
"""
53+
raise NotImplementedError
54+
55+
@abc.abstractmethod
56+
def get_main_file_path(self) -> str:
57+
"""
58+
Getter for jar | python path
59+
:return: Full path to file
60+
"""
61+
raise NotImplementedError
62+
63+
def get_class_name(self) -> Optional[str]:
64+
"""
65+
Getter for main class name if it's applicable
66+
:return: java class path, e.g. feast.ingestion.IngestionJob
67+
"""
68+
return None
69+
70+
@abc.abstractmethod
71+
def get_arguments(self) -> List[str]:
72+
"""
73+
Getter for job arguments
74+
E.g., ["--source", '{"kafka":...}', ...]
75+
:return: List of arguments
76+
"""
77+
raise NotImplementedError
78+
2879

2980
class RetrievalJob(SparkJob):
3081
"""
3182
Container for the historical feature retrieval job result
3283
"""
3384

85+
def __init__(
86+
self,
87+
feature_tables: List[Dict],
88+
feature_tables_sources: List[Dict],
89+
entity_source: Dict,
90+
destination: Dict,
91+
**kwargs,
92+
):
93+
super().__init__(**kwargs)
94+
self._feature_tables = feature_tables
95+
self._feature_tables_sources = feature_tables_sources
96+
self._entity_source = entity_source
97+
self._destination = destination
98+
99+
def get_name(self) -> str:
100+
return f"HistoryRetrieval-{self.get_id()}"
101+
102+
def get_main_file_path(self) -> str:
103+
return os.path.join(
104+
os.path.dirname(__file__), "historical_feature_retrieval_job.py"
105+
)
106+
107+
def get_arguments(self) -> List[str]:
108+
return [
109+
"--feature-tables",
110+
json.dumps(self._feature_tables),
111+
"--feature-tables-sources",
112+
json.dumps(self._feature_tables_sources),
113+
"--entity-source",
114+
json.dumps(self._entity_source),
115+
"--destination",
116+
json.dumps(self._destination),
117+
]
118+
34119
@abc.abstractmethod
35120
def get_output_file_uri(self, timeout_sec=None):
36121
"""
@@ -54,7 +139,44 @@ def get_output_file_uri(self, timeout_sec=None):
54139

55140

56141
class IngestionJob(SparkJob):
57-
pass
142+
def __init__(
143+
self,
144+
feature_table: Dict,
145+
source: Dict,
146+
start: datetime,
147+
end: datetime,
148+
jar: str,
149+
**kwargs,
150+
):
151+
super().__init__(**kwargs)
152+
self._feature_table = feature_table
153+
self._source = source
154+
self._start = start
155+
self._end = end
156+
self._jar = jar
157+
158+
def get_name(self) -> str:
159+
return f"BatchIngestion-{self.get_id()}"
160+
161+
def get_main_file_path(self) -> str:
162+
return self._jar
163+
164+
def get_class_name(self) -> Optional[str]:
165+
return "feast.ingestion.IngestionJob"
166+
167+
def get_arguments(self) -> List[str]:
168+
return [
169+
"--mode",
170+
"offline",
171+
"--feature-table",
172+
json.dumps(self._feature_table),
173+
"--source",
174+
json.dumps(self._source),
175+
"--start",
176+
self._start.strftime("%Y-%m-%dT%H:%M:%S"),
177+
"--end",
178+
self._end.strftime("%Y-%m-%dT%H:%M:%S"),
179+
]
58180

59181

60182
class JobLauncher(abc.ABC):
@@ -65,20 +187,16 @@ class JobLauncher(abc.ABC):
65187
@abc.abstractmethod
66188
def historical_feature_retrieval(
67189
self,
68-
pyspark_script: str,
69190
entity_source_conf: Dict,
70191
feature_tables_sources_conf: List[Dict],
71192
feature_tables_conf: List[Dict],
72193
destination_conf: Dict,
73-
job_id: str,
74194
**kwargs,
75195
) -> RetrievalJob:
76196
"""
77197
Submits a historical feature retrieval job to a Spark cluster.
78198
79199
Args:
80-
pyspark_script (str): Local file path to the pyspark script for historical feature
81-
retrieval.
82200
entity_source_conf (Dict): Entity data source configuration.
83201
feature_tables_sources_conf (List[Dict]): List of feature tables data sources configurations.
84202
feature_tables_conf (List[Dict]): List of feature table specification.
@@ -191,3 +309,14 @@ def historical_feature_retrieval(
191309
str: file uri to the result file.
192310
"""
193311
raise NotImplementedError
312+
313+
@abc.abstractmethod
314+
def offline_to_online_ingestion(
315+
self,
316+
jar_path: str,
317+
source_conf: Dict,
318+
feature_table_conf: Dict,
319+
start: datetime,
320+
end: datetime,
321+
) -> IngestionJob:
322+
raise NotImplementedError

‎sdk/python/feast/pyspark/launcher.py‎

Lines changed: 46 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,17 @@
1-
import pathlib
21
from typing import TYPE_CHECKING, List, Union
32

3+
from datetime import datetime
4+
from urllib.parse import urlparse
5+
import tempfile
6+
import shutil
7+
48
from feast.config import Config
5-
from feast.constants import (
6-
CONFIG_SPARK_DATAPROC_CLUSTER_NAME,
7-
CONFIG_SPARK_DATAPROC_PROJECT,
8-
CONFIG_SPARK_DATAPROC_REGION,
9-
CONFIG_SPARK_DATAPROC_STAGING_LOCATION,
10-
CONFIG_SPARK_LAUNCHER,
11-
CONFIG_SPARK_STANDALONE_MASTER,
12-
)
9+
from feast.constants import *
1310
from feast.data_source import BigQuerySource, DataSource, FileSource
1411
from feast.feature_table import FeatureTable
15-
from feast.pyspark.abc import JobLauncher, RetrievalJob
12+
from feast.pyspark.abc import JobLauncher, RetrievalJob, IngestionJob
1613
from feast.value_type import ValueType
14+
from feast.staging.storage_client import get_staging_client
1715

1816
if TYPE_CHECKING:
1917
from feast.client import Client
@@ -23,7 +21,7 @@ def _standalone_launcher(config: Config) -> JobLauncher:
2321
from feast.pyspark.launchers import standalone
2422

2523
return standalone.StandaloneClusterLauncher(
26-
config.get(CONFIG_SPARK_STANDALONE_MASTER)
24+
config.get(CONFIG_SPARK_STANDALONE_MASTER), config.get(CONFIG_SPARK_HOME)
2725
)
2826

2927

@@ -51,7 +49,7 @@ def resolve_launcher(config: Config) -> JobLauncher:
5149
}
5250

5351

54-
def source_to_argument(source: DataSource):
52+
def _source_to_argument(source: DataSource):
5553
common_properties = {
5654
"field_mapping": dict(source.field_mapping),
5755
"event_timestamp_column": source.event_timestamp_column,
@@ -72,7 +70,7 @@ def source_to_argument(source: DataSource):
7270
return {kind: properties}
7371

7472

75-
def feature_table_to_argument(client: "Client", feature_table: FeatureTable):
73+
def _feature_table_to_argument(client: "Client", feature_table: FeatureTable):
7674
return {
7775
"features": [
7876
{"name": f.name, "type": ValueType(f.dtype).name}
@@ -102,13 +100,13 @@ def start_historical_feature_retrieval_spark_session(
102100
spark_session = SparkSession.builder.getOrCreate()
103101
return retrieve_historical_features(
104102
spark=spark_session,
105-
entity_source_conf=source_to_argument(entity_source),
103+
entity_source_conf=_source_to_argument(entity_source),
106104
feature_tables_sources_conf=[
107-
source_to_argument(feature_table.batch_source)
105+
_source_to_argument(feature_table.batch_source)
108106
for feature_table in feature_tables
109107
],
110108
feature_tables_conf=[
111-
feature_table_to_argument(client, feature_table)
109+
_feature_table_to_argument(client, feature_table)
112110
for feature_table in feature_tables
113111
],
114112
)
@@ -123,22 +121,45 @@ def start_historical_feature_retrieval_job(
123121
job_id: str,
124122
) -> RetrievalJob:
125123
launcher = resolve_launcher(client._config)
126-
retrieval_job_pyspark_script = str(
127-
pathlib.Path(__file__).parent.absolute()
128-
/ "pyspark"
129-
/ "historical_feature_retrieval_job.py"
130-
)
131124
return launcher.historical_feature_retrieval(
132-
pyspark_script=retrieval_job_pyspark_script,
133-
entity_source_conf=source_to_argument(entity_source),
125+
entity_source_conf=_source_to_argument(entity_source),
134126
feature_tables_sources_conf=[
135-
source_to_argument(feature_table.batch_source)
127+
_source_to_argument(feature_table.batch_source)
136128
for feature_table in feature_tables
137129
],
138130
feature_tables_conf=[
139-
feature_table_to_argument(client, feature_table)
131+
_feature_table_to_argument(client, feature_table)
140132
for feature_table in feature_tables
141133
],
142134
destination_conf={"format": output_format, "path": output_path},
143135
job_id=job_id,
144136
)
137+
138+
139+
def _download_jar(remote_jar: str) -> str:
140+
remote_jar_parts = urlparse(remote_jar)
141+
142+
f = tempfile.NamedTemporaryFile(suffix=".jar", delete=False)
143+
with f:
144+
shutil.copyfileobj(
145+
get_staging_client(remote_jar_parts.scheme).download_file(remote_jar_parts),
146+
f,
147+
)
148+
149+
return f.name
150+
151+
152+
def start_offline_to_online_ingestion(
153+
feature_table: FeatureTable, start: datetime, end: datetime, client: Client
154+
) -> IngestionJob:
155+
156+
launcher = resolve_launcher(client._config)
157+
local_jar_path = _download_jar(client._config.get(CONFIG_SPARK_INGESTION_JOB_JAR))
158+
159+
return launcher.offline_to_online_ingestion(
160+
jar_path=local_jar_path,
161+
source_conf=_source_to_argument(feature_table.batch_source),
162+
feature_table_conf=_feature_table_to_argument(client, feature_table),
163+
start=start,
164+
end=end,
165+
)

0 commit comments

Comments
 (0)