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..4577d257dd0 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 @@ -26,7 +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), + config.get(opt.SPARK_STANDALONE_MASTER), + config.get(opt.SPARK_HOME), + _parse_additional_spark_options(config), ) @@ -41,6 +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), ) @@ -57,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), ) @@ -126,6 +131,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, None) + 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 +243,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 +292,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/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/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py index cfbfb828108..c2b52af5f69 100644 --- a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py +++ b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py @@ -208,6 +208,7 @@ def __init__( 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 @@ -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,6 +250,7 @@ def __init__( self.executor_instances = executor_instances self.executor_cores = executor_cores self.executor_memory = executor_memory + self.additional_options = additional_options or {} def _stage_file(self, file_path: str, job_id: str) -> str: if not os.path.isfile(file_path): @@ -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, }, } } 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/test_pyspark/__init__.py b/sdk/python/tests/test_pyspark/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/tests/test_pyspark/test_launchers.py b/sdk/python/tests/test_pyspark/test_launchers.py new file mode 100644 index 00000000000..102cc64fb83 --- /dev/null +++ b/sdk/python/tests/test_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) 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={} )