Skip to content

Commit d855776

Browse files
oavdeevoavdeev
committed
support spark k8s operator launcher
Signed-off-by: Oleg Avdeev <oleg.v.avdeev@gmail.com>
1 parent 6d15876 commit d855776

11 files changed

Lines changed: 723 additions & 24 deletions

File tree

‎sdk/python/feast/constants.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,15 @@ class ConfigOptions(metaclass=ConfigMeta):
180180
#: No. of executor memory for Dataproc cluster
181181
DATAPROC_EXECUTOR_MEMORY = "2g"
182182

183+
# namespace to use for Spark jobs launched using k8s spark operator
184+
SPARK_K8S_NAMESPACE = "default"
185+
186+
# expect k8s spark operator to be running in the same cluster as Feast
187+
SPARK_K8S_USE_INCLUSTER_CONFIG = True
188+
189+
# SparkApplication resource template
190+
SPARK_K8S_JOB_TEMPLATE_PATH = None
191+
183192
#: File format of historical retrieval features
184193
HISTORICAL_FEATURE_OUTPUT_FORMAT: str = "parquet"
185194

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import hashlib
33
import json
44
import os
5+
from base64 import b64encode
56
from datetime import datetime
67
from enum import Enum
78
from typing import Dict, List, Optional
@@ -243,15 +244,18 @@ def get_main_file_path(self) -> str:
243244
)
244245

245246
def get_arguments(self) -> List[str]:
247+
def json_b64_encode(obj) -> str:
248+
return b64encode(json.dumps(obj).encode("utf8")).decode("ascii")
249+
246250
return [
247251
"--feature-tables",
248-
json.dumps(self._feature_tables),
252+
json_b64_encode(self._feature_tables),
249253
"--feature-tables-sources",
250-
json.dumps(self._feature_tables_sources),
254+
json_b64_encode(self._feature_tables_sources),
251255
"--entity-source",
252-
json.dumps(self._entity_source),
256+
json_b64_encode(self._entity_source),
253257
"--destination",
254-
json.dumps(self._destination),
258+
json_b64_encode(self._destination),
255259
]
256260

257261
def get_destination_path(self) -> str:

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import abc
22
import argparse
33
import json
4+
from base64 import b64decode
45
from datetime import timedelta
56
from typing import Any, Dict, List, NamedTuple, Optional
67

@@ -794,13 +795,17 @@ def _feature_table_from_dict(dct: Dict[str, Any]) -> FeatureTable:
794795
)
795796

796797

798+
def json_b64_decode(s: str) -> Any:
799+
return json.loads(b64decode(s.encode("ascii")))
800+
801+
797802
if __name__ == "__main__":
798803
spark = SparkSession.builder.getOrCreate()
799804
args = _get_args()
800-
feature_tables_conf = json.loads(args.feature_tables)
801-
feature_tables_sources_conf = json.loads(args.feature_tables_sources)
802-
entity_source_conf = json.loads(args.entity_source)
803-
destination_conf = json.loads(args.destination)
805+
feature_tables_conf = json_b64_decode(args.feature_tables)
806+
feature_tables_sources_conf = json_b64_decode(args.feature_tables_sources)
807+
entity_source_conf = json_b64_decode(args.entity_source)
808+
destination_conf = json_b64_decode(args.destination)
804809
start_job(
805810
spark,
806811
entity_source_conf,

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,26 @@ def _get_optional(option):
6565
)
6666

6767

68+
def _k8s_launcher(config: Config) -> JobLauncher:
69+
from feast.pyspark.launchers import k8s
70+
71+
def _get_optional(option):
72+
if config.exists(option):
73+
return config.get(option)
74+
75+
return k8s.KubernetesJobLauncher(
76+
namespace=config.get(opt.SPARK_K8S_NAMESPACE),
77+
resource_template_path=_get_optional(opt.SPARK_K8S_JOB_TEMPLATE_PATH),
78+
staging_location=config.get(opt.SPARK_STAGING_LOCATION),
79+
incluster=config.getboolean(opt.SPARK_K8S_USE_INCLUSTER_CONFIG),
80+
)
81+
82+
6883
_launchers = {
6984
"standalone": _standalone_launcher,
7085
"dataproc": _dataproc_launcher,
7186
"emr": _emr_launcher,
87+
"k8s": _k8s_launcher,
7288
}
7389

7490

‎sdk/python/feast/pyspark/launchers/gcloud/dataproc.py‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import json
21
import os
32
import time
43
import uuid
@@ -265,7 +264,7 @@ def _stage_file(self, file_path: str, job_id: str) -> str:
265264
return blob_uri_str
266265

267266
def dataproc_submit(
268-
self, job_params: SparkJobParameters
267+
self, job_params: SparkJobParameters, extra_properties: Dict[str, str]
269268
) -> Tuple[Job, Callable[[], Job], Callable[[], None]]:
270269
local_job_id = str(uuid.uuid4())
271270
main_file_uri = self._stage_file(job_params.get_main_file_path(), local_job_id)
@@ -280,18 +279,22 @@ def dataproc_submit(
280279
job_config["labels"][self.JOB_HASH_LABEL_KEY] = job_params.get_job_hash()
281280

282281
if job_params.get_class_name():
282+
properties = {
283+
"spark.yarn.user.classpath.first": "true",
284+
"spark.executor.instances": self.executor_instances,
285+
"spark.executor.cores": self.executor_cores,
286+
"spark.executor.memory": self.executor_memory,
287+
}
288+
289+
properties.update(extra_properties)
290+
283291
job_config.update(
284292
{
285293
"spark_job": {
286294
"jar_file_uris": [main_file_uri] + self.EXTERNAL_JARS,
287295
"main_class": job_params.get_class_name(),
288296
"args": job_params.get_arguments(),
289-
"properties": {
290-
"spark.yarn.user.classpath.first": "true",
291-
"spark.executor.instances": self.executor_instances,
292-
"spark.executor.cores": self.executor_cores,
293-
"spark.executor.memory": self.executor_memory,
294-
},
297+
"properties": properties,
295298
}
296299
}
297300
)
@@ -302,6 +305,7 @@ def dataproc_submit(
302305
"main_python_file_uri": main_file_uri,
303306
"jar_file_uris": self.EXTERNAL_JARS,
304307
"args": job_params.get_arguments(),
308+
"properties": extra_properties if extra_properties else {},
305309
}
306310
}
307311
)
@@ -332,21 +336,23 @@ def dataproc_cancel(self, job_id):
332336
def historical_feature_retrieval(
333337
self, job_params: RetrievalJobParameters
334338
) -> RetrievalJob:
335-
job, refresh_fn, cancel_fn = self.dataproc_submit(job_params)
339+
job, refresh_fn, cancel_fn = self.dataproc_submit(
340+
job_params, {"dev.feast.outputuri": job_params.get_destination_path()}
341+
)
336342
return DataprocRetrievalJob(
337343
job, refresh_fn, cancel_fn, job_params.get_destination_path()
338344
)
339345

340346
def offline_to_online_ingestion(
341347
self, ingestion_job_params: BatchIngestionJobParameters
342348
) -> BatchIngestionJob:
343-
job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params)
349+
job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params, {})
344350
return DataprocBatchIngestionJob(job, refresh_fn, cancel_fn)
345351

346352
def start_stream_to_online_ingestion(
347353
self, ingestion_job_params: StreamIngestionJobParameters
348354
) -> StreamIngestionJob:
349-
job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params)
355+
job, refresh_fn, cancel_fn = self.dataproc_submit(ingestion_job_params, {})
350356
job_hash = ingestion_job_params.get_job_hash()
351357
return DataprocStreamingIngestionJob(job, refresh_fn, cancel_fn, job_hash)
352358

@@ -368,7 +374,7 @@ def _dataproc_job_to_spark_job(self, job: Job) -> SparkJob:
368374
cancel_fn = partial(self.dataproc_cancel, job_id)
369375

370376
if job_type == SparkJobType.HISTORICAL_RETRIEVAL.name.lower():
371-
output_path = json.loads(job.pyspark_job.args[-1])["path"]
377+
output_path = job.pyspark_job.properties.get("dev.feast.outputuri")
372378
return DataprocRetrievalJob(job, refresh_fn, cancel_fn, output_path)
373379

374380
if job_type == SparkJobType.BATCH_INGESTION.name.lower():

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
from .k8s import (
2+
KubernetesBatchIngestionJob,
3+
KubernetesJobLauncher,
4+
KubernetesRetrievalJob,
5+
KubernetesStreamIngestionJob,
6+
)
7+
8+
__all__ = [
9+
"KubernetesRetrievalJob",
10+
"KubernetesBatchIngestionJob",
11+
"KubernetesStreamIngestionJob",
12+
"KubernetesJobLauncher",
13+
]

0 commit comments

Comments
 (0)