Skip to content

Commit 664ab78

Browse files
pyalexpyalex
committed
split SparkJob into 2 classes
Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 42c4480 commit 664ab78

7 files changed

Lines changed: 142 additions & 173 deletions

File tree

‎sdk/python/feast/client.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
import multiprocessing
1616
import shutil
1717
import uuid
18+
from datetime import datetime
1819
from itertools import groupby
1920
from typing import Any, Dict, List, Optional, Union
20-
from datetime import datetime
2121

2222
import grpc
2323
import pandas as pd
@@ -75,7 +75,7 @@
7575
_write_partitioned_table_from_source,
7676
)
7777
from feast.online_response import OnlineResponse, _infer_online_entity_rows
78-
from feast.pyspark.abc import RetrievalJob
78+
from feast.pyspark.abc import RetrievalJob, SparkJob
7979
from feast.pyspark.launcher import (
8080
start_historical_feature_retrieval_job,
8181
start_historical_feature_retrieval_spark_session,
@@ -85,7 +85,6 @@
8585
GetFeastServingInfoRequest,
8686
GetOnlineFeaturesRequestV2,
8787
)
88-
from feast.pyspark.abc import SparkJob
8988
from feast.serving.ServingService_pb2_grpc import ServingServiceStub
9089

9190
_logger = logging.getLogger(__name__)
@@ -893,4 +892,4 @@ def start_offline_to_online_ingestion(
893892
start: Union[datetime, str],
894893
end: Union[datetime, str],
895894
) -> SparkJob:
896-
return start_offline_to_online_ingestion(feature_table, start, end, self)
895+
return start_offline_to_online_ingestion(feature_table, start, end, self) # type: ignore

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

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import abc
2-
import os
32
import json
3+
import os
44
from datetime import datetime
5-
from typing import Dict, List, Optional
65
from enum import Enum
6+
from typing import Dict, List, Optional
77

88

99
class SparkJobFailure(Exception):
@@ -44,6 +44,8 @@ def get_status(self) -> SparkJobStatus:
4444
"""
4545
raise NotImplementedError
4646

47+
48+
class SparkJobParameters(abc.ABC):
4749
@abc.abstractmethod
4850
def get_name(self) -> str:
4951
"""
@@ -77,11 +79,7 @@ def get_arguments(self) -> List[str]:
7779
raise NotImplementedError
7880

7981

80-
class RetrievalJob(SparkJob):
81-
"""
82-
Container for the historical feature retrieval job result
83-
"""
84-
82+
class RetrievalJobParameters(SparkJobParameters):
8583
def __init__(
8684
self,
8785
feature_tables: List[Dict],
@@ -90,14 +88,14 @@ def __init__(
9088
destination: Dict,
9189
**kwargs,
9290
):
93-
super().__init__(**kwargs)
9491
self._feature_tables = feature_tables
9592
self._feature_tables_sources = feature_tables_sources
9693
self._entity_source = entity_source
9794
self._destination = destination
9895

9996
def get_name(self) -> str:
100-
return f"HistoryRetrieval-{self.get_id()}"
97+
all_feature_tables_names = [ft["name"] for ft in self._feature_tables]
98+
return f"HistoryRetrieval-{all_feature_tables_names}"
10199

102100
def get_main_file_path(self) -> str:
103101
return os.path.join(
@@ -116,6 +114,12 @@ def get_arguments(self) -> List[str]:
116114
json.dumps(self._destination),
117115
]
118116

117+
118+
class RetrievalJob(SparkJob):
119+
"""
120+
Container for the historical feature retrieval job result
121+
"""
122+
119123
@abc.abstractmethod
120124
def get_output_file_uri(self, timeout_sec=None):
121125
"""
@@ -138,7 +142,7 @@ def get_output_file_uri(self, timeout_sec=None):
138142
raise NotImplementedError
139143

140144

141-
class IngestionJob(SparkJob):
145+
class IngestionJobParameters(SparkJobParameters):
142146
def __init__(
143147
self,
144148
feature_table: Dict,
@@ -148,15 +152,17 @@ def __init__(
148152
jar: str,
149153
**kwargs,
150154
):
151-
super().__init__(**kwargs)
152155
self._feature_table = feature_table
153156
self._source = source
154157
self._start = start
155158
self._end = end
156159
self._jar = jar
157160

158161
def get_name(self) -> str:
159-
return f"BatchIngestion-{self.get_id()}"
162+
return (
163+
f"BatchIngestion-{self._feature_table['name']}-"
164+
f"{self._start.strftime('%Y-%m-%d')}-{self._end.strftime('%Y-%m-%d')}"
165+
)
160166

161167
def get_main_file_path(self) -> str:
162168
return self._jar
@@ -179,6 +185,12 @@ def get_arguments(self) -> List[str]:
179185
]
180186

181187

188+
class IngestionJob(SparkJob):
189+
"""
190+
Container for the ingestion job result
191+
"""
192+
193+
182194
class JobLauncher(abc.ABC):
183195
"""
184196
Submits spark jobs to a spark cluster. Currently supports only historical feature retrieval jobs.

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

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,25 @@
1-
from typing import TYPE_CHECKING, List, Union
2-
1+
import shutil
2+
import tempfile
33
from datetime import datetime
4+
from typing import TYPE_CHECKING, List, Union
45
from urllib.parse import urlparse
5-
import tempfile
6-
import shutil
76

87
from feast.config import Config
9-
from feast.constants import *
8+
from feast.constants import (
9+
CONFIG_SPARK_DATAPROC_CLUSTER_NAME,
10+
CONFIG_SPARK_DATAPROC_PROJECT,
11+
CONFIG_SPARK_DATAPROC_REGION,
12+
CONFIG_SPARK_DATAPROC_STAGING_LOCATION,
13+
CONFIG_SPARK_HOME,
14+
CONFIG_SPARK_INGESTION_JOB_JAR,
15+
CONFIG_SPARK_LAUNCHER,
16+
CONFIG_SPARK_STANDALONE_MASTER,
17+
)
1018
from feast.data_source import BigQuerySource, DataSource, FileSource
1119
from feast.feature_table import FeatureTable
12-
from feast.pyspark.abc import JobLauncher, RetrievalJob, IngestionJob
13-
from feast.value_type import ValueType
20+
from feast.pyspark.abc import IngestionJob, JobLauncher, RetrievalJob
1421
from feast.staging.storage_client import get_staging_client
22+
from feast.value_type import ValueType
1523

1624
if TYPE_CHECKING:
1725
from feast.client import Client

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

Lines changed: 32 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,28 @@
11
import os
22
import uuid
33
from datetime import datetime
4+
from typing import Dict, List, cast
5+
from urllib.parse import urlparse
6+
7+
from google.api_core.operation import Operation
8+
from google.cloud import dataproc_v1, storage
9+
from google.cloud.dataproc_v1 import Job as DataprocJob
10+
from google.cloud.dataproc_v1 import JobStatus
411

512
from feast.pyspark.abc import (
13+
IngestionJob,
14+
IngestionJobParameters,
15+
JobLauncher,
616
RetrievalJob,
17+
RetrievalJobParameters,
718
SparkJobFailure,
8-
JobLauncher,
19+
SparkJobParameters,
920
SparkJobStatus,
10-
SparkJob,
11-
IngestionJob,
1221
)
13-
from google.api_core.operation import Operation
14-
from google.cloud import dataproc_v1
15-
from google.cloud import storage
16-
17-
from typing import Dict, List, cast
18-
from urllib.parse import urlparse
1922

20-
from google.cloud.dataproc_v1 import Job as DataprocJob, JobStatus
2123

22-
23-
class DataprocSparkJob(SparkJob):
24-
def __init__(self, **kwargs):
25-
super().__init__(**kwargs)
26-
self._operation = None # type: Operation
27-
28-
def set_operation(self, operation: Operation):
24+
class DataprocJobMixin:
25+
def __init__(self, operation: Operation):
2926
"""
3027
:param operation: (google.api.core.operation.Operation): A Future for the spark job result,
3128
returned by the dataproc client.
@@ -47,19 +44,19 @@ def get_status(self) -> SparkJobStatus:
4744
return SparkJobStatus.FAILED
4845

4946

50-
class DataprocRetrievalJob(RetrievalJob, DataprocSparkJob):
47+
class DataprocRetrievalJob(DataprocJobMixin, RetrievalJob):
5148
"""
5249
Historical feature retrieval job result for a Dataproc cluster
5350
"""
5451

55-
def __init__(self, output_file_uri, **kwargs):
52+
def __init__(self, operation: Operation, output_file_uri: str):
5653
"""
5754
This is the returned historical feature retrieval job result for DataprocClusterLauncher.
5855
5956
Args:
6057
output_file_uri (str): Uri to the historical feature retrieval job output file.
6158
"""
62-
super().__init__(**kwargs)
59+
super().__init__(operation)
6360
self._output_file_uri = output_file_uri
6461

6562
def get_output_file_uri(self, timeout_sec=None):
@@ -70,18 +67,10 @@ def get_output_file_uri(self, timeout_sec=None):
7067
return self._output_file_uri
7168

7269

73-
class DataprocIngestionJob(IngestionJob, DataprocSparkJob):
74-
def __init__(
75-
self,
76-
feature_table: Dict,
77-
source: Dict,
78-
start: datetime,
79-
end: datetime,
80-
jar: str,
81-
):
82-
super().__init__(
83-
feature_table=feature_table, source=source, start=start, end=end, jar=jar
84-
)
70+
class DataprocIngestionJob(DataprocJobMixin, IngestionJob):
71+
"""
72+
Ingestion job result for a Dataproc cluster
73+
"""
8574

8675

8776
class DataprocClusterLauncher(JobLauncher):
@@ -135,26 +124,24 @@ def _stage_files(self, pyspark_script: str, job_id: str) -> str:
135124

136125
return f"gs://{self.staging_bucket}/{blob_path}"
137126

138-
def dataproc_submit(self, job: DataprocSparkJob) -> Operation:
127+
def dataproc_submit(self, job_params: SparkJobParameters) -> Operation:
139128
local_job_id = str(uuid.uuid4())
140-
pyspark_gcs = self._stage_files(job.get_main_file_path(), local_job_id)
129+
pyspark_gcs = self._stage_files(job_params.get_main_file_path(), local_job_id)
141130
job_config = {
142131
"reference": {"job_id": local_job_id},
143132
"placement": {"cluster_name": self.cluster_name},
144133
"pyspark_job": {
145134
"main_python_file_uri": pyspark_gcs,
146-
"args": job.get_arguments(),
135+
"args": job_params.get_arguments(),
147136
},
148137
}
149-
operation = self.job_client.submit_job_as_operation(
138+
return self.job_client.submit_job_as_operation(
150139
request={
151140
"project_id": self.project_id,
152141
"region": self.region,
153142
"job": job_config,
154143
}
155144
)
156-
job.set_operation(operation)
157-
return operation
158145

159146
def historical_feature_retrieval(
160147
self,
@@ -164,15 +151,16 @@ def historical_feature_retrieval(
164151
destination_conf: Dict,
165152
**kwargs,
166153
) -> RetrievalJob:
167-
job = DataprocRetrievalJob(
168-
output_file_uri=destination_conf["path"],
154+
job_params = RetrievalJobParameters(
169155
feature_tables=feature_tables_conf,
170156
feature_tables_sources=feature_tables_sources_conf,
171157
entity_source=entity_source_conf,
172158
destination=destination_conf,
173159
)
174-
self.dataproc_submit(job)
175-
return job
160+
161+
return DataprocRetrievalJob(
162+
self.dataproc_submit(job_params), destination_conf["path"]
163+
)
176164

177165
def offline_to_online_ingestion(
178166
self,
@@ -182,12 +170,11 @@ def offline_to_online_ingestion(
182170
start: datetime,
183171
end: datetime,
184172
) -> IngestionJob:
185-
job = DataprocIngestionJob(
173+
job_params = IngestionJobParameters(
186174
feature_table=feature_table_conf,
187175
source=source_conf,
188176
start=start,
189177
end=end,
190178
jar=jar_path,
191179
)
192-
self.dataproc_submit(job)
193-
return job
180+
return DataprocIngestionJob(self.dataproc_submit(job_params))

0 commit comments

Comments
 (0)