Skip to content

Commit 112e94d

Browse files
oavdeevoavdeev
authored
emr streaming job launcher (#1065)
Signed-off-by: Oleg Avdeev <oleg.v.avdeev@gmail.com>
1 parent c28941e commit 112e94d

9 files changed

Lines changed: 249 additions & 48 deletions

File tree

‎sdk/python/feast/cli.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -387,11 +387,10 @@ def start_stream_to_online(feature_table: str, jar: str):
387387
"""
388388
Start stream to online sync job.
389389
"""
390-
import feast.pyspark.aws.jobs
391390

392391
client = Client()
393392
table = client.get_feature_table(feature_table)
394-
feast.pyspark.aws.jobs.start_stream_to_online(client, table, [jar] if jar else [])
393+
client.start_stream_to_online_ingestion(table, [jar] if jar else [])
395394

396395

397396
@cli.command()

‎sdk/python/feast/client.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@
8181
start_historical_feature_retrieval_job,
8282
start_historical_feature_retrieval_spark_session,
8383
start_offline_to_online_ingestion,
84+
start_stream_to_online_ingestion,
8485
)
8586
from feast.serving.ServingService_pb2 import (
8687
GetFeastServingInfoRequest,
@@ -886,10 +887,15 @@ def _get_feature_tables_from_feature_refs(
886887
return feature_tables
887888

888889
def start_offline_to_online_ingestion(
889-
self, feature_table: Union[FeatureTable, str], start: datetime, end: datetime,
890+
self, feature_table: FeatureTable, start: datetime, end: datetime,
890891
) -> SparkJob:
891892
return start_offline_to_online_ingestion(feature_table, start, end, self) # type: ignore
892893

894+
def start_stream_to_online_ingestion(
895+
self, feature_table: FeatureTable, extra_jars: Optional[List[str]] = None,
896+
) -> SparkJob:
897+
return start_stream_to_online_ingestion(feature_table, extra_jars or [], self)
898+
893899
def stage_dataframe(
894900
self,
895901
df: pd.DataFrame,

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

Lines changed: 77 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -254,7 +254,7 @@ def get_output_file_uri(self, timeout_sec=None):
254254
raise NotImplementedError
255255

256256

257-
class IngestionJobParameters(SparkJobParameters):
257+
class BatchIngestionJobParameters(SparkJobParameters):
258258
def __init__(
259259
self,
260260
feature_table: Dict,
@@ -310,12 +310,68 @@ def get_arguments(self) -> List[str]:
310310
]
311311

312312

313-
class IngestionJob(SparkJob):
313+
class StreamIngestionJobParameters(SparkJobParameters):
314+
def __init__(
315+
self,
316+
feature_table: Dict,
317+
source: Dict,
318+
jar: str,
319+
extra_jars: List[str],
320+
redis_host: str,
321+
redis_port: int,
322+
redis_ssl: bool,
323+
):
324+
self._feature_table = feature_table
325+
self._source = source
326+
self._jar = jar
327+
self._extra_jars = extra_jars
328+
self._redis_host = redis_host
329+
self._redis_port = redis_port
330+
self._redis_ssl = redis_ssl
331+
332+
def get_name(self) -> str:
333+
return f"StreamIngestion-{self.get_feature_table_name()}"
334+
335+
def _get_redis_config(self):
336+
return dict(host=self._redis_host, port=self._redis_port, ssl=self._redis_ssl)
337+
338+
def get_feature_table_name(self) -> str:
339+
return self._feature_table["name"]
340+
341+
def get_main_file_path(self) -> str:
342+
return self._jar
343+
344+
def get_extra_jar_paths(self) -> List[str]:
345+
return self._extra_jars
346+
347+
def get_class_name(self) -> Optional[str]:
348+
return "feast.ingestion.IngestionJob"
349+
350+
def get_arguments(self) -> List[str]:
351+
return [
352+
"--mode",
353+
"online",
354+
"--feature-table",
355+
json.dumps(self._feature_table),
356+
"--source",
357+
json.dumps(self._source),
358+
"--redis",
359+
json.dumps(self._get_redis_config()),
360+
]
361+
362+
363+
class BatchIngestionJob(SparkJob):
314364
"""
315365
Container for the ingestion job result
316366
"""
317367

318368

369+
class StreamIngestionJob(SparkJob):
370+
"""
371+
Container for the streaming ingestion job result
372+
"""
373+
374+
319375
class JobLauncher(abc.ABC):
320376
"""
321377
Submits spark jobs to a spark cluster. Currently supports only historical feature retrieval jobs.
@@ -339,8 +395,8 @@ def historical_feature_retrieval(
339395

340396
@abc.abstractmethod
341397
def offline_to_online_ingestion(
342-
self, ingestion_job_params: IngestionJobParameters
343-
) -> IngestionJob:
398+
self, ingestion_job_params: BatchIngestionJobParameters
399+
) -> BatchIngestionJob:
344400
"""
345401
Submits a batch ingestion job to a Spark cluster.
346402
@@ -349,7 +405,23 @@ def offline_to_online_ingestion(
349405
during execution, or timeout.
350406
351407
Returns:
352-
IngestionJob: wrapper around remote job that can be used to check when job completed.
408+
BatchIngestionJob: wrapper around remote job that can be used to check when job completed.
409+
"""
410+
raise NotImplementedError
411+
412+
@abc.abstractmethod
413+
def start_stream_to_online_ingestion(
414+
self, ingestion_job_params: StreamIngestionJobParameters
415+
) -> StreamIngestionJob:
416+
"""
417+
Starts a stream ingestion job to a Spark cluster.
418+
419+
Raises:
420+
SparkJobFailure: The spark job submission failed, encountered error
421+
during execution, or timeout.
422+
423+
Returns:
424+
StreamIngestionJob: wrapper around remote job.
353425
"""
354426
raise NotImplementedError
355427

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

Lines changed: 40 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import shutil
22
import tempfile
33
from datetime import datetime
4-
from typing import TYPE_CHECKING, List, Union, cast
4+
from typing import TYPE_CHECKING, List, Union
55
from urllib.parse import urlparse
66

77
from feast.config import Config
@@ -23,14 +23,16 @@
2323
CONFIG_SPARK_LAUNCHER,
2424
CONFIG_SPARK_STANDALONE_MASTER,
2525
)
26-
from feast.data_source import BigQuerySource, DataSource, FileSource
26+
from feast.data_source import BigQuerySource, DataSource, FileSource, KafkaSource
2727
from feast.feature_table import FeatureTable
2828
from feast.pyspark.abc import (
29-
IngestionJob,
30-
IngestionJobParameters,
29+
BatchIngestionJob,
30+
BatchIngestionJobParameters,
3131
JobLauncher,
3232
RetrievalJob,
3333
RetrievalJobParameters,
34+
StreamIngestionJob,
35+
StreamIngestionJobParameters,
3436
)
3537
from feast.staging.storage_client import get_staging_client
3638
from feast.value_type import ValueType
@@ -85,10 +87,7 @@ def resolve_launcher(config: Config) -> JobLauncher:
8587
return _launchers[config.get(CONFIG_SPARK_LAUNCHER)](config)
8688

8789

88-
_SOURCES = {
89-
FileSource: ("file", "file_options"),
90-
BigQuerySource: ("bq", "bigquery_options"),
91-
}
90+
_SOURCES = {FileSource: "file", BigQuerySource: "bq", KafkaSource: "kafka"}
9291

9392

9493
def _source_to_argument(source: DataSource):
@@ -99,16 +98,19 @@ def _source_to_argument(source: DataSource):
9998
"date_partition_column": source.date_partition_column,
10099
}
101100

102-
kind, option_field = _SOURCES[type(source)]
101+
kind = _SOURCES[type(source)]
103102
properties = {**common_properties}
104-
if type(source) == FileSource:
105-
file_source = cast(FileSource, source)
106-
properties["path"] = file_source.file_options.file_url
107-
properties["format"] = str(file_source.file_options.file_format)
103+
if isinstance(source, FileSource):
104+
properties["path"] = source.file_options.file_url
105+
properties["format"] = str(source.file_options.file_format)
106+
return {kind: properties}
107+
if isinstance(source, BigQuerySource):
108+
properties["table_ref"] = source.bigquery_options.table_ref
108109
return {kind: properties}
109-
if type(source) == BigQuerySource:
110-
bq_source = cast(BigQuerySource, source)
111-
properties["table_ref"] = bq_source.bigquery_options.table_ref
110+
if isinstance(source, KafkaSource):
111+
properties["topic"] = source.kafka_options.topic
112+
properties["classpath"] = source.kafka_options.class_path
113+
properties["bootstrap_servers"] = source.kafka_options.bootstrap_servers
112114
return {kind: properties}
113115
raise NotImplementedError(f"Unsupported Datasource: {type(source)}")
114116

@@ -194,13 +196,13 @@ def _download_jar(remote_jar: str) -> str:
194196

195197
def start_offline_to_online_ingestion(
196198
feature_table: FeatureTable, start: datetime, end: datetime, client: "Client"
197-
) -> IngestionJob:
199+
) -> BatchIngestionJob:
198200

199201
launcher = resolve_launcher(client._config)
200202
local_jar_path = _download_jar(client._config.get(CONFIG_SPARK_INGESTION_JOB_JAR))
201203

202204
return launcher.offline_to_online_ingestion(
203-
IngestionJobParameters(
205+
BatchIngestionJobParameters(
204206
jar=local_jar_path,
205207
source=_source_to_argument(feature_table.batch_source),
206208
feature_table=_feature_table_to_argument(client, feature_table),
@@ -213,6 +215,26 @@ def start_offline_to_online_ingestion(
213215
)
214216

215217

218+
def start_stream_to_online_ingestion(
219+
feature_table: FeatureTable, extra_jars: List[str], client: "Client"
220+
) -> StreamIngestionJob:
221+
222+
launcher = resolve_launcher(client._config)
223+
local_jar_path = _download_jar(client._config.get(CONFIG_SPARK_INGESTION_JOB_JAR))
224+
225+
return launcher.start_stream_to_online_ingestion(
226+
StreamIngestionJobParameters(
227+
jar=local_jar_path,
228+
extra_jars=extra_jars,
229+
source=_source_to_argument(feature_table.stream_source),
230+
feature_table=_feature_table_to_argument(client, feature_table),
231+
redis_host=client._config.get(CONFIG_REDIS_HOST),
232+
redis_port=client._config.getint(CONFIG_REDIS_PORT),
233+
redis_ssl=client._config.getboolean(CONFIG_REDIS_SSL),
234+
)
235+
)
236+
237+
216238
def stage_dataframe(
217239
df, event_timestamp_column: str, created_timestamp_column: str, client: "Client"
218240
) -> FileSource:

‎sdk/python/feast/pyspark/launchers/aws/__init__.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,13 @@
1-
from .emr import EmrClusterLauncher, EmrIngestionJob, EmrRetrievalJob
1+
from .emr import (
2+
EmrBatchIngestionJob,
3+
EmrClusterLauncher,
4+
EmrRetrievalJob,
5+
EmrStreamIngestionJob,
6+
)
27

3-
__all__ = ["EmrRetrievalJob", "EmrIngestionJob", "EmrClusterLauncher"]
8+
__all__ = [
9+
"EmrRetrievalJob",
10+
"EmrBatchIngestionJob",
11+
"EmrStreamIngestionJob",
12+
"EmrClusterLauncher",
13+
]

‎sdk/python/feast/pyspark/launchers/aws/emr.py‎

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,22 @@
11
import os
22
import tempfile
33
from io import BytesIO
4-
from typing import Any, Dict, Optional
4+
from typing import Any, Dict, List, Optional
55

66
import boto3
77
import pandas
88

99
from feast.data_format import ParquetFormat
1010
from feast.data_source import FileSource
1111
from feast.pyspark.abc import (
12-
IngestionJob,
13-
IngestionJobParameters,
12+
BatchIngestionJob,
13+
BatchIngestionJobParameters,
1414
JobLauncher,
1515
RetrievalJob,
1616
RetrievalJobParameters,
1717
SparkJobStatus,
18+
StreamIngestionJob,
19+
StreamIngestionJobParameters,
1820
)
1921

2022
from .emr_utils import (
@@ -28,6 +30,7 @@
2830
_load_new_cluster_template,
2931
_random_string,
3032
_s3_upload,
33+
_stream_ingestion_step,
3134
_sync_offline_to_online_step,
3235
_upload_jar,
3336
_wait_for_job_state,
@@ -82,7 +85,7 @@ def get_output_file_uri(self, timeout_sec=None):
8285
return self._output_file_uri
8386

8487

85-
class EmrIngestionJob(EmrJobMixin, IngestionJob):
88+
class EmrBatchIngestionJob(EmrJobMixin, BatchIngestionJob):
8689
"""
8790
Ingestion job result for a EMR cluster
8891
"""
@@ -91,6 +94,15 @@ def __init__(self, emr_client, job_ref: EmrJobRef):
9194
super().__init__(emr_client, job_ref)
9295

9396

97+
class EmrStreamIngestionJob(EmrJobMixin, StreamIngestionJob):
98+
"""
99+
Ingestion streaming job for a EMR cluster
100+
"""
101+
102+
def __init__(self, emr_client, job_ref: EmrJobRef):
103+
super().__init__(emr_client, job_ref)
104+
105+
94106
class EmrClusterLauncher(JobLauncher):
95107
"""
96108
Submits jobs to an existing or new EMR cluster. Requires boto3 as an additional dependency.
@@ -203,8 +215,8 @@ def historical_feature_retrieval(
203215
)
204216

205217
def offline_to_online_ingestion(
206-
self, ingestion_job_params: IngestionJobParameters
207-
) -> IngestionJob:
218+
self, ingestion_job_params: BatchIngestionJobParameters
219+
) -> BatchIngestionJob:
208220
"""
209221
Submits a batch ingestion job to a Spark cluster.
210222
@@ -213,7 +225,7 @@ def offline_to_online_ingestion(
213225
during execution, or timeout.
214226
215227
Returns:
216-
IngestionJob: wrapper around remote job that can be used to check when job completed.
228+
BatchIngestionJob: wrapper around remote job that can be used to check when job completed.
217229
"""
218230

219231
jar_s3_path = _upload_jar(
@@ -227,7 +239,38 @@ def offline_to_online_ingestion(
227239

228240
job_ref = self._submit_emr_job(step)
229241

230-
return EmrIngestionJob(self._emr_client(), job_ref)
242+
return EmrBatchIngestionJob(self._emr_client(), job_ref)
243+
244+
def start_stream_to_online_ingestion(
245+
self, ingestion_job_params: StreamIngestionJobParameters
246+
) -> StreamIngestionJob:
247+
"""
248+
Starts a stream ingestion job on a Spark cluster.
249+
250+
Returns:
251+
StreamIngestionJob: wrapper around remote job that can be used to check on the job.
252+
"""
253+
jar_s3_path = _upload_jar(
254+
self._staging_location, ingestion_job_params.get_main_file_path()
255+
)
256+
257+
extra_jar_paths: List[str] = []
258+
for extra_jar in ingestion_job_params.get_extra_jar_paths():
259+
if extra_jar.startswith("s3://"):
260+
extra_jar_paths.append(extra_jar)
261+
else:
262+
extra_jar_paths.append(_upload_jar(self._staging_location, extra_jar))
263+
264+
step = _stream_ingestion_step(
265+
jar_s3_path,
266+
extra_jar_paths,
267+
ingestion_job_params.get_feature_table_name(),
268+
args=ingestion_job_params.get_arguments(),
269+
)
270+
271+
job_ref = self._submit_emr_job(step)
272+
273+
return EmrStreamIngestionJob(self._emr_client(), job_ref)
231274

232275
def stage_dataframe(
233276
self, df: pandas.DataFrame, event_timestamp: str, created_timestamp_column: str

0 commit comments

Comments
 (0)