From b98e6cdadbceecda0af1b7aa81cd4021239d116c Mon Sep 17 00:00:00 2001 From: mike0sv Date: Thu, 26 Nov 2020 21:59:50 +0300 Subject: [PATCH 1/6] Option for additional spark configuration Signed-off-by: mike0sv --- sdk/python/feast/constants.py | 3 ++ sdk/python/feast/pyspark/launcher.py | 25 ++++++++++++++-- .../pyspark/launchers/standalone/local.py | 14 ++++++++- sdk/python/tests/pyspark/__init__.py | 0 sdk/python/tests/pyspark/test_launchers.py | 29 +++++++++++++++++++ 5 files changed, 67 insertions(+), 4 deletions(-) create mode 100644 sdk/python/tests/pyspark/__init__.py create mode 100644 sdk/python/tests/pyspark/test_launchers.py diff --git a/sdk/python/feast/constants.py b/sdk/python/feast/constants.py index 8b1db7b76bd..ced4a0f1078 100644 --- a/sdk/python/feast/constants.py +++ b/sdk/python/feast/constants.py @@ -154,6 +154,9 @@ class ConfigOptions(metaclass=ConfigMeta): #: Directory where Spark is installed SPARK_HOME: Optional[str] = None + #: Addtional config options for Spark + SPARK_ADDITIONAL_OPTS: Optional[str] = None + #: Dataproc cluster to run Feast Spark Jobs in DATAPROC_CLUSTER_NAME: Optional[str] = None diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index d928a4f129e..3d9910680db 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -1,5 +1,6 @@ +import re from datetime import datetime -from typing import TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING, Dict, List, Union from feast.config import Config from feast.constants import ConfigOptions as opt @@ -126,6 +127,26 @@ def _feature_table_to_argument( } +def _quoted_split(string, delimiter): + for token in re.findall(f'(?:".*?"|[^{delimiter}])+', string): + if token.startswith('"') and token.endswith('"'): + token = token[1:-1] + yield token + + +def _parse_additional_spark_options(config: Config) -> Dict[str, str]: + options_string = config.get(opt.SPARK_ADDITIONAL_OPTS) + if options_string is None: + return {} + try: + return dict( + _quoted_split(opt_val, "=") + for opt_val in _quoted_split(options_string, ";") + ) + except ValueError: + raise ValueError(f"Cannot parse {opt.SPARK_ADDITIONAL_OPTS}: {options_string}") + + def start_historical_feature_retrieval_spark_session( client: "Client", project: str, @@ -218,7 +239,6 @@ def start_offline_to_online_ingestion( start: datetime, end: datetime, ) -> BatchIngestionJob: - launcher = resolve_launcher(client._config) return launcher.offline_to_online_ingestion( @@ -268,7 +288,6 @@ def get_stream_to_online_ingestion_params( def start_stream_to_online_ingestion( client: "Client", project: str, feature_table: FeatureTable, extra_jars: List[str] ) -> StreamIngestionJob: - launcher = resolve_launcher(client._config) return launcher.start_stream_to_online_ingestion( diff --git a/sdk/python/feast/pyspark/launchers/standalone/local.py b/sdk/python/feast/pyspark/launchers/standalone/local.py index 9783cbe52e7..49f3eca064c 100644 --- a/sdk/python/feast/pyspark/launchers/standalone/local.py +++ b/sdk/python/feast/pyspark/launchers/standalone/local.py @@ -224,7 +224,12 @@ class StandaloneClusterLauncher(JobLauncher): BQ_CONNECTOR_VERSION = "2.12:0.17.3" - def __init__(self, master_url: str, spark_home: str = None): + def __init__( + self, + master_url: str, + spark_home: str = None, + additional_options: Dict[str, str] = None, + ): """ This launcher executes the spark-submit script in a subprocess. The subprocess will run until the Pyspark driver exits. @@ -235,9 +240,12 @@ def __init__(self, master_url: str, spark_home: str = None): spark_home (str): Local file path to Spark installation directory. If not provided, the environmental variable `SPARK_HOME` will be used instead. + additional_options (Dict[str, str]): + Additional configuration options for Spark job """ self.master_url = master_url self.spark_home = spark_home if spark_home else os.getenv("SPARK_HOME") + self.additional_options = additional_options @property def spark_submit_script_path(self): @@ -285,6 +293,10 @@ def spark_submit( ] ) + if self.additional_options is not None: + for option, value in self.additional_options.items(): + submission_cmd.extend(["--conf", f'"{option}"="{value}"']) + submission_cmd.append(job_params.get_main_file_path()) submission_cmd.extend(job_params.get_arguments()) diff --git a/sdk/python/tests/pyspark/__init__.py b/sdk/python/tests/pyspark/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/tests/pyspark/test_launchers.py b/sdk/python/tests/pyspark/test_launchers.py new file mode 100644 index 00000000000..102cc64fb83 --- /dev/null +++ b/sdk/python/tests/pyspark/test_launchers.py @@ -0,0 +1,29 @@ +import pytest + +from feast.config import Config +from feast.constants import ConfigOptions as opt +from feast.pyspark.launcher import _parse_additional_spark_options + + +class TestSparkAdditionalOpts: + def parse(self, options_string): + return _parse_additional_spark_options( + Config(options={opt.SPARK_ADDITIONAL_OPTS: options_string}) + ) + + def test_normal_options(self): + options_string = "option1=aaaa;option2=bbb" + assert self.parse(options_string) == {"option1": "aaaa", "option2": "bbb"} + + def test_value_with_delimiter(self): + options_string = 'option1=aaaa;option2="b;b"' + assert self.parse(options_string) == {"option1": "aaaa", "option2": "b;b"} + + def test_value_with_another_delimiter(self): + options_string = 'option1=aaaa;option2="b=b"' + assert self.parse(options_string) == {"option1": "aaaa", "option2": "b=b"} + + def test_error_on_wrong_format(self): + options_string = "option1=aaaa;option2" + with pytest.raises(ValueError): + self.parse(options_string) From eb3709ced577c7051f9344f6d4c1c88e7eeebe7a Mon Sep 17 00:00:00 2001 From: mike0sv Date: Thu, 26 Nov 2020 22:15:05 +0300 Subject: [PATCH 2/6] standalone and dataproc config Signed-off-by: mike0sv --- sdk/python/feast/pyspark/launcher.py | 1 + .../pyspark/launchers/gcloud/dataproc.py | 73 ++++++++++--------- 2 files changed, 40 insertions(+), 34 deletions(-) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index 3d9910680db..c39589f2d0e 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -28,6 +28,7 @@ def _standalone_launcher(config: Config) -> JobLauncher: return standalone.StandaloneClusterLauncher( config.get(opt.SPARK_STANDALONE_MASTER), config.get(opt.SPARK_HOME), + _parse_additional_spark_options(config) ) diff --git a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py index cfbfb828108..e701996a3d2 100644 --- a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py +++ b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py @@ -27,7 +27,7 @@ class DataprocJobMixin: def __init__( - self, job: Job, refresh_fn: Callable[[], Job], cancel_fn: Callable[[], None] + self, job: Job, refresh_fn: Callable[[], Job], cancel_fn: Callable[[], None] ): """ Implementation of common methods for different types of SparkJob running on Dataproc cluster. @@ -60,18 +60,18 @@ def get_status(self) -> SparkJobStatus: self._job = self._refresh_fn() status = self._job.status if status.state in ( - JobStatus.State.ERROR, - JobStatus.State.CANCEL_PENDING, - JobStatus.State.CANCEL_STARTED, - JobStatus.State.CANCELLED, + JobStatus.State.ERROR, + JobStatus.State.CANCEL_PENDING, + JobStatus.State.CANCEL_STARTED, + JobStatus.State.CANCELLED, ): return SparkJobStatus.FAILED elif status.state == JobStatus.State.RUNNING: return SparkJobStatus.IN_PROGRESS elif status.state in ( - JobStatus.State.PENDING, - JobStatus.State.SETUP_DONE, - JobStatus.State.STATE_UNSPECIFIED, + JobStatus.State.PENDING, + JobStatus.State.SETUP_DONE, + JobStatus.State.STATE_UNSPECIFIED, ): return SparkJobStatus.STARTING @@ -95,9 +95,9 @@ def get_error_message(self) -> Optional[str]: if status.state == JobStatus.State.ERROR: return status.details elif status.state in ( - JobStatus.State.CANCEL_PENDING, - JobStatus.State.CANCEL_STARTED, - JobStatus.State.CANCELLED, + JobStatus.State.CANCEL_PENDING, + JobStatus.State.CANCEL_STARTED, + JobStatus.State.CANCELLED, ): return "Job was cancelled." return None @@ -138,11 +138,11 @@ class DataprocRetrievalJob(DataprocJobMixin, RetrievalJob): """ def __init__( - self, - job: Job, - refresh_fn: Callable[[], Job], - cancel_fn: Callable[[], None], - output_file_uri: str, + self, + job: Job, + refresh_fn: Callable[[], Job], + cancel_fn: Callable[[], None], + output_file_uri: str, ): """ This is the returned historical feature retrieval job result for DataprocClusterLauncher. @@ -175,11 +175,11 @@ class DataprocStreamingIngestionJob(DataprocJobMixin, StreamIngestionJob): """ def __init__( - self, - job: Job, - refresh_fn: Callable[[], Job], - cancel_fn: Callable[[], None], - job_hash: str, + self, + job: Job, + refresh_fn: Callable[[], Job], + cancel_fn: Callable[[], None], + job_hash: str, ) -> None: super().__init__(job, refresh_fn, cancel_fn) self._job_hash = job_hash @@ -200,14 +200,15 @@ class DataprocClusterLauncher(JobLauncher): JOB_HASH_LABEL_KEY = "feast_job_hash" def __init__( - self, - cluster_name: str, - staging_location: str, - region: str, - project_id: str, - executor_instances: str, - executor_cores: str, - executor_memory: str, + self, + cluster_name: str, + staging_location: str, + region: str, + project_id: str, + executor_instances: str, + executor_cores: str, + executor_memory: str, + additional_options: Dict[str, str] ): """ Initialize a dataproc job controller client, used internally for job submission and result @@ -228,6 +229,8 @@ def __init__( Number of cores for dataproc job. executor_memory (str): Amount of memory for dataproc job. + additional_options (Dict[str, str]): + Additional configuration options for Spark job """ self.cluster_name = cluster_name @@ -247,19 +250,20 @@ def __init__( self.executor_instances = executor_instances self.executor_cores = executor_cores self.executor_memory = executor_memory + self.additional_options = additional_options def _stage_file(self, file_path: str, job_id: str) -> str: if not os.path.isfile(file_path): return file_path staging_client = get_staging_client("gs") - blob_path = os.path.join(self.remote_path, job_id, os.path.basename(file_path),) + blob_path = os.path.join(self.remote_path, job_id, os.path.basename(file_path), ) staging_client.upload_file(file_path, self.staging_bucket, blob_path) return f"gs://{self.staging_bucket}/{blob_path}" def dataproc_submit( - self, job_params: SparkJobParameters + self, job_params: SparkJobParameters ) -> Tuple[Job, Callable[[], Job], Callable[[], None]]: local_job_id = str(uuid.uuid4()) main_file_uri = self._stage_file(job_params.get_main_file_path(), local_job_id) @@ -285,6 +289,7 @@ def dataproc_submit( "spark.executor.instances": self.executor_instances, "spark.executor.cores": self.executor_cores, "spark.executor.memory": self.executor_memory, + **self.additional_options }, } } @@ -324,7 +329,7 @@ def dataproc_cancel(self, job_id): ) def historical_feature_retrieval( - self, job_params: RetrievalJobParameters + self, job_params: RetrievalJobParameters ) -> RetrievalJob: job, refresh_fn, cancel_fn = self.dataproc_submit(job_params) return DataprocRetrievalJob( @@ -332,13 +337,13 @@ def historical_feature_retrieval( ) def offline_to_online_ingestion( - self, ingestion_job_params: BatchIngestionJobParameters + self, ingestion_job_params: BatchIngestionJobParameters ) -> BatchIngestionJob: job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params) return DataprocBatchIngestionJob(job, refresh_fn, cancel_fn) def start_stream_to_online_ingestion( - self, ingestion_job_params: StreamIngestionJobParameters + self, ingestion_job_params: StreamIngestionJobParameters ) -> StreamIngestionJob: job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params) job_hash = ingestion_job_params.get_job_hash() From 6ec62e4d3a85ecdd36dc98bca58835d59738f446 Mon Sep 17 00:00:00 2001 From: mike0sv Date: Thu, 26 Nov 2020 22:16:49 +0300 Subject: [PATCH 3/6] dataproc config Signed-off-by: mike0sv --- sdk/python/feast/pyspark/launcher.py | 1 + sdk/python/feast/pyspark/launchers/gcloud/dataproc.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index c39589f2d0e..fd415335dad 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -43,6 +43,7 @@ def _dataproc_launcher(config: Config) -> JobLauncher: executor_instances=config.get(opt.DATAPROC_EXECUTOR_INSTANCES), executor_cores=config.get(opt.DATAPROC_EXECUTOR_CORES), executor_memory=config.get(opt.DATAPROC_EXECUTOR_MEMORY), + additional_options=_parse_additional_spark_options(config) ) diff --git a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py index e701996a3d2..b4333facca6 100644 --- a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py +++ b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py @@ -208,7 +208,7 @@ def __init__( executor_instances: str, executor_cores: str, executor_memory: str, - additional_options: Dict[str, str] + additional_options: Dict[str, str]=None ): """ Initialize a dataproc job controller client, used internally for job submission and result @@ -250,7 +250,7 @@ def __init__( self.executor_instances = executor_instances self.executor_cores = executor_cores self.executor_memory = executor_memory - self.additional_options = additional_options + self.additional_options = additional_options or {} def _stage_file(self, file_path: str, job_id: str) -> str: if not os.path.isfile(file_path): From ad05891566c65bf33d86cdbf978817a794a95795 Mon Sep 17 00:00:00 2001 From: mike0sv Date: Thu, 26 Nov 2020 22:17:17 +0300 Subject: [PATCH 4/6] linter Signed-off-by: mike0sv --- sdk/python/feast/pyspark/launcher.py | 7 +- .../pyspark/launchers/gcloud/dataproc.py | 72 +++++++++---------- 2 files changed, 40 insertions(+), 39 deletions(-) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index fd415335dad..0b3eaee59cd 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -27,8 +27,9 @@ def _standalone_launcher(config: Config) -> JobLauncher: from feast.pyspark.launchers import standalone return standalone.StandaloneClusterLauncher( - config.get(opt.SPARK_STANDALONE_MASTER), config.get(opt.SPARK_HOME), - _parse_additional_spark_options(config) + config.get(opt.SPARK_STANDALONE_MASTER), + config.get(opt.SPARK_HOME), + _parse_additional_spark_options(config), ) @@ -43,7 +44,7 @@ def _dataproc_launcher(config: Config) -> JobLauncher: executor_instances=config.get(opt.DATAPROC_EXECUTOR_INSTANCES), executor_cores=config.get(opt.DATAPROC_EXECUTOR_CORES), executor_memory=config.get(opt.DATAPROC_EXECUTOR_MEMORY), - additional_options=_parse_additional_spark_options(config) + additional_options=_parse_additional_spark_options(config), ) diff --git a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py index b4333facca6..c2b52af5f69 100644 --- a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py +++ b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py @@ -27,7 +27,7 @@ class DataprocJobMixin: def __init__( - self, job: Job, refresh_fn: Callable[[], Job], cancel_fn: Callable[[], None] + self, job: Job, refresh_fn: Callable[[], Job], cancel_fn: Callable[[], None] ): """ Implementation of common methods for different types of SparkJob running on Dataproc cluster. @@ -60,18 +60,18 @@ def get_status(self) -> SparkJobStatus: self._job = self._refresh_fn() status = self._job.status if status.state in ( - JobStatus.State.ERROR, - JobStatus.State.CANCEL_PENDING, - JobStatus.State.CANCEL_STARTED, - JobStatus.State.CANCELLED, + JobStatus.State.ERROR, + JobStatus.State.CANCEL_PENDING, + JobStatus.State.CANCEL_STARTED, + JobStatus.State.CANCELLED, ): return SparkJobStatus.FAILED elif status.state == JobStatus.State.RUNNING: return SparkJobStatus.IN_PROGRESS elif status.state in ( - JobStatus.State.PENDING, - JobStatus.State.SETUP_DONE, - JobStatus.State.STATE_UNSPECIFIED, + JobStatus.State.PENDING, + JobStatus.State.SETUP_DONE, + JobStatus.State.STATE_UNSPECIFIED, ): return SparkJobStatus.STARTING @@ -95,9 +95,9 @@ def get_error_message(self) -> Optional[str]: if status.state == JobStatus.State.ERROR: return status.details elif status.state in ( - JobStatus.State.CANCEL_PENDING, - JobStatus.State.CANCEL_STARTED, - JobStatus.State.CANCELLED, + JobStatus.State.CANCEL_PENDING, + JobStatus.State.CANCEL_STARTED, + JobStatus.State.CANCELLED, ): return "Job was cancelled." return None @@ -138,11 +138,11 @@ class DataprocRetrievalJob(DataprocJobMixin, RetrievalJob): """ def __init__( - self, - job: Job, - refresh_fn: Callable[[], Job], - cancel_fn: Callable[[], None], - output_file_uri: str, + self, + job: Job, + refresh_fn: Callable[[], Job], + cancel_fn: Callable[[], None], + output_file_uri: str, ): """ This is the returned historical feature retrieval job result for DataprocClusterLauncher. @@ -175,11 +175,11 @@ class DataprocStreamingIngestionJob(DataprocJobMixin, StreamIngestionJob): """ def __init__( - self, - job: Job, - refresh_fn: Callable[[], Job], - cancel_fn: Callable[[], None], - job_hash: str, + self, + job: Job, + refresh_fn: Callable[[], Job], + cancel_fn: Callable[[], None], + job_hash: str, ) -> None: super().__init__(job, refresh_fn, cancel_fn) self._job_hash = job_hash @@ -200,15 +200,15 @@ class DataprocClusterLauncher(JobLauncher): JOB_HASH_LABEL_KEY = "feast_job_hash" def __init__( - self, - cluster_name: str, - staging_location: str, - region: str, - project_id: str, - executor_instances: str, - executor_cores: str, - executor_memory: str, - additional_options: Dict[str, str]=None + self, + cluster_name: str, + staging_location: str, + region: str, + project_id: str, + executor_instances: str, + executor_cores: str, + executor_memory: str, + additional_options: Dict[str, str] = None, ): """ Initialize a dataproc job controller client, used internally for job submission and result @@ -257,13 +257,13 @@ def _stage_file(self, file_path: str, job_id: str) -> str: return file_path staging_client = get_staging_client("gs") - blob_path = os.path.join(self.remote_path, job_id, os.path.basename(file_path), ) + blob_path = os.path.join(self.remote_path, job_id, os.path.basename(file_path),) staging_client.upload_file(file_path, self.staging_bucket, blob_path) return f"gs://{self.staging_bucket}/{blob_path}" def dataproc_submit( - self, job_params: SparkJobParameters + self, job_params: SparkJobParameters ) -> Tuple[Job, Callable[[], Job], Callable[[], None]]: local_job_id = str(uuid.uuid4()) main_file_uri = self._stage_file(job_params.get_main_file_path(), local_job_id) @@ -289,7 +289,7 @@ def dataproc_submit( "spark.executor.instances": self.executor_instances, "spark.executor.cores": self.executor_cores, "spark.executor.memory": self.executor_memory, - **self.additional_options + **self.additional_options, }, } } @@ -329,7 +329,7 @@ def dataproc_cancel(self, job_id): ) def historical_feature_retrieval( - self, job_params: RetrievalJobParameters + self, job_params: RetrievalJobParameters ) -> RetrievalJob: job, refresh_fn, cancel_fn = self.dataproc_submit(job_params) return DataprocRetrievalJob( @@ -337,13 +337,13 @@ def historical_feature_retrieval( ) def offline_to_online_ingestion( - self, ingestion_job_params: BatchIngestionJobParameters + self, ingestion_job_params: BatchIngestionJobParameters ) -> BatchIngestionJob: job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params) return DataprocBatchIngestionJob(job, refresh_fn, cancel_fn) def start_stream_to_online_ingestion( - self, ingestion_job_params: StreamIngestionJobParameters + self, ingestion_job_params: StreamIngestionJobParameters ) -> StreamIngestionJob: job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params) job_hash = ingestion_job_params.get_job_hash() From 04858411a688ca8bed7dc2c5e58a91798a1eb613 Mon Sep 17 00:00:00 2001 From: mike0sv Date: Fri, 27 Nov 2020 03:01:09 +0300 Subject: [PATCH 5/6] default value Signed-off-by: mike0sv --- sdk/python/feast/pyspark/launcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index 0b3eaee59cd..b5cef0f029f 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -138,7 +138,7 @@ def _quoted_split(string, delimiter): def _parse_additional_spark_options(config: Config) -> Dict[str, str]: - options_string = config.get(opt.SPARK_ADDITIONAL_OPTS) + options_string = config.get(opt.SPARK_ADDITIONAL_OPTS, None) if options_string is None: return {} try: From 61a607930b68b00e40a68803592e30ecff88cf6e Mon Sep 17 00:00:00 2001 From: mike0sv Date: Fri, 27 Nov 2020 03:36:31 +0300 Subject: [PATCH 6/6] emr job and fix import error by rename Signed-off-by: mike0sv --- sdk/python/feast/pyspark/launcher.py | 1 + sdk/python/feast/pyspark/launchers/aws/emr.py | 7 +++++++ .../feast/pyspark/launchers/aws/emr_utils.py | 19 ++++++++++++++++--- .../{pyspark => test_pyspark}/__init__.py | 0 .../test_launchers.py | 0 tests/integration/fixtures/launchers.py | 1 + 6 files changed, 25 insertions(+), 3 deletions(-) rename sdk/python/tests/{pyspark => test_pyspark}/__init__.py (100%) rename sdk/python/tests/{pyspark => test_pyspark}/test_launchers.py (100%) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index b5cef0f029f..4577d257dd0 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -61,6 +61,7 @@ def _get_optional(option): new_cluster_template_path=_get_optional(opt.EMR_CLUSTER_TEMPLATE_PATH), staging_location=config.get(opt.SPARK_STAGING_LOCATION), emr_log_location=config.get(opt.EMR_LOG_LOCATION), + additional_options=_parse_additional_spark_options(config), ) diff --git a/sdk/python/feast/pyspark/launchers/aws/emr.py b/sdk/python/feast/pyspark/launchers/aws/emr.py index 42b5348c298..0c68750a21e 100644 --- a/sdk/python/feast/pyspark/launchers/aws/emr.py +++ b/sdk/python/feast/pyspark/launchers/aws/emr.py @@ -145,6 +145,7 @@ def __init__( new_cluster_template_path: Optional[str], staging_location: str, emr_log_location: str, + additional_options: Dict[str, str], ): """ Initialize a dataproc job controller client, used internally for job submission and result @@ -162,6 +163,8 @@ def __init__( An S3 staging location for artifacts. emr_log_location: S3 location for EMR logs. + additional_options: + Additional configuration options for Spark job """ assert existing_cluster_id or new_cluster_template_path @@ -177,6 +180,7 @@ def __init__( self._staging_location = staging_location self._emr_log_location = emr_log_location self._region = region + self._additional_options = additional_options def _emr_client(self): @@ -230,6 +234,7 @@ def historical_feature_retrieval( step = _historical_retrieval_step( pyspark_script_path, + conf=self._additional_options, args=job_params.get_arguments(), output_file_uri=job_params.get_destination_path(), ) @@ -260,6 +265,7 @@ def offline_to_online_ingestion( step = _sync_offline_to_online_step( jar_s3_path, ingestion_job_params.get_feature_table_name(), + self._additional_options, args=ingestion_job_params.get_arguments(), ) @@ -293,6 +299,7 @@ def start_stream_to_online_ingestion( jar_s3_path, extra_jar_paths, ingestion_job_params.get_feature_table_name(), + self._additional_options, args=ingestion_job_params.get_arguments(), job_hash=job_hash, ) diff --git a/sdk/python/feast/pyspark/launchers/aws/emr_utils.py b/sdk/python/feast/pyspark/launchers/aws/emr_utils.py index 72df3f92d8c..e8e1731e85e 100644 --- a/sdk/python/feast/pyspark/launchers/aws/emr_utils.py +++ b/sdk/python/feast/pyspark/launchers/aws/emr_utils.py @@ -135,7 +135,7 @@ def _upload_jar(jar_s3_prefix: str, local_path: str) -> str: def _sync_offline_to_online_step( - jar_path: str, feature_table_name: str, args: List[str], + jar_path: str, feature_table_name: str, conf: Dict[str, str], args: List[str], ) -> Dict[str, Any]: return { @@ -155,6 +155,7 @@ def _sync_offline_to_online_step( "spark-submit", "--class", "feast.ingestion.IngestionJob", + *_prepare_conf_args(conf), "--packages", "com.google.cloud.spark:spark-bigquery-with-dependencies_2.12:0.17.2", jar_path, @@ -344,8 +345,17 @@ def _upload_dataframe(s3prefix: str, df: pandas.DataFrame) -> str: ) +def _prepare_conf_args(conf: Dict[str, str]): + return [ + _ for name, value in conf.items() for _ in ["--conf", f'"{name}"="{value}"'] + ] + + def _historical_retrieval_step( - pyspark_script_path: str, args: List[str], output_file_uri: str, + pyspark_script_path: str, + conf: Dict[str, str], + args: List[str], + output_file_uri: str, ) -> Dict[str, Any]: return { @@ -361,7 +371,8 @@ def _historical_retrieval_step( "Value": output_file_uri, }, ], - "Args": ["spark-submit", pyspark_script_path] + args, + "Args": ["spark-submit", *_prepare_conf_args(conf), pyspark_script_path] + + args, "Jar": "command-runner.jar", }, } @@ -371,6 +382,7 @@ def _stream_ingestion_step( jar_path: str, extra_jar_paths: List[str], feature_table_name: str, + conf: Dict[str, str], args: List[str], job_hash: str, ) -> Dict[str, Any]: @@ -395,6 +407,7 @@ def _stream_ingestion_step( {"Key": "feast.step_metadata.job_hash", "Value": job_hash}, ], "Args": ["spark-submit", "--class", "feast.ingestion.IngestionJob"] + + _prepare_conf_args(conf) + jars_args + [ "--packages", diff --git a/sdk/python/tests/pyspark/__init__.py b/sdk/python/tests/test_pyspark/__init__.py similarity index 100% rename from sdk/python/tests/pyspark/__init__.py rename to sdk/python/tests/test_pyspark/__init__.py diff --git a/sdk/python/tests/pyspark/test_launchers.py b/sdk/python/tests/test_pyspark/test_launchers.py similarity index 100% rename from sdk/python/tests/pyspark/test_launchers.py rename to sdk/python/tests/test_pyspark/test_launchers.py diff --git a/tests/integration/fixtures/launchers.py b/tests/integration/fixtures/launchers.py index d289d974ac4..9820cda918d 100644 --- a/tests/integration/fixtures/launchers.py +++ b/tests/integration/fixtures/launchers.py @@ -20,4 +20,5 @@ def dataproc_launcher(pytestconfig) -> DataprocClusterLauncher: executor_instances=executor_instances, executor_cores=executor_cores, executor_memory=executor_memory, + additional_options={} )