Skip to content

Commit 42bf26a

Browse files
oavdeevoavdeev
authored
EMR support for offline-to-online job (#1026)
Signed-off-by: Oleg Avdeev <oleg.v.avdeev@gmail.com>
1 parent d607d53 commit 42bf26a

4 files changed

Lines changed: 341 additions & 4 deletions

File tree

‎sdk/python/feast/cli.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -351,5 +351,25 @@ def project_list():
351351
print(tabulate(table, headers=["NAME"], tablefmt="plain"))
352352

353353

354+
@cli.command()
355+
@click.option(
356+
"--feature-table",
357+
"-t",
358+
help="Feature table name to ingest data into",
359+
required=True,
360+
)
361+
@click.option("--start-time", "-s", help="Interval start", required=True)
362+
@click.option("--end-time", "-e", help="Interval end", required=True)
363+
def sync_offline_to_online(feature_table: str, start_time: str, end_time: str):
364+
"""
365+
Sync offline store to online.
366+
"""
367+
import feast.pyspark.aws.jobs
368+
369+
client = Client()
370+
table = client.get_feature_table(feature_table)
371+
feast.pyspark.aws.jobs.sync_offline_to_online(client, table, start_time, end_time)
372+
373+
354374
if __name__ == "__main__":
355375
cli()

‎sdk/python/feast/client.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -435,7 +435,7 @@ def list_entities(
435435
entities.append(entity)
436436
return entities
437437

438-
def get_entity(self, name: str, project: str = None) -> Union[Entity, None]:
438+
def get_entity(self, name: str, project: str = None) -> Entity:
439439
"""
440440
Retrieves an entity.
441441
@@ -548,9 +548,7 @@ def list_feature_tables(
548548
feature_tables.append(feature_table)
549549
return feature_tables
550550

551-
def get_feature_table(
552-
self, name: str, project: str = None
553-
) -> Union[FeatureTable, None]:
551+
def get_feature_table(self, name: str, project: str = None) -> FeatureTable:
554552
"""
555553
Retrieves a feature table.
556554

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

Whitespace-only changes.

‎sdk/python/feast/pyspark/aws/jobs.py‎

Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
import hashlib
2+
import json
3+
import logging
4+
import os
5+
import random
6+
import string
7+
from typing import Any, Dict, Tuple
8+
9+
import boto3
10+
import botocore
11+
import yaml
12+
13+
from feast.client import Client
14+
from feast.feature_table import FeatureTable
15+
from feast.value_type import ValueType
16+
17+
log = logging.getLogger("aws")
18+
19+
# Config example:
20+
#
21+
# aws:
22+
# logS3Prefix: "..a prefix for logs.."
23+
# artifactS3Prefix: "..a prefix for jars.."
24+
# existingClusterId: "..." # You need to set either existingClusterId
25+
# runJobFlowTemplate: # or runJobFlowTemplate
26+
# Name: "feast-ingestion-test"
27+
# ReleaseLabel: emr-6.0.0
28+
# Instances:
29+
# InstanceFleets:
30+
# - InstanceFleetType: MASTER
31+
# TargetOnDemandCapacity: 0
32+
# TargetSpotCapacity: 1
33+
# LaunchSpecifications:
34+
# SpotSpecification:
35+
# TimeoutDurationMinutes: 60
36+
# TimeoutAction: TERMINATE_CLUSTER
37+
# InstanceTypeConfigs:
38+
# - WeightedCapacity: 1
39+
# EbsConfiguration:
40+
# EbsBlockDeviceConfigs:
41+
# - VolumeSpecification:
42+
# SizeInGB: 32
43+
# VolumeType: gp2
44+
# VolumesPerInstance: 2
45+
# BidPriceAsPercentageOfOnDemandPrice: 100
46+
# InstanceType: m4.xlarge
47+
# - InstanceFleetType: CORE
48+
# TargetOnDemandCapacity: 0
49+
# TargetSpotCapacity: 2
50+
# LaunchSpecifications:
51+
# SpotSpecification:
52+
# TimeoutDurationMinutes: 60
53+
# TimeoutAction: TERMINATE_CLUSTER
54+
# InstanceTypeConfigs:
55+
# - WeightedCapacity: 1
56+
# EbsConfiguration:
57+
# EbsBlockDeviceConfigs:
58+
# - VolumeSpecification:
59+
# SizeInGB: 32
60+
# VolumeType: gp2
61+
# VolumesPerInstance: 2
62+
# BidPriceAsPercentageOfOnDemandPrice: 100
63+
# InstanceType: m4.xlarge
64+
# Ec2SubnetIds:
65+
# - "..a subnet id within a VPC with a route to redis..."
66+
# AdditionalMasterSecurityGroups:
67+
# - "..a security group that allows access to redis..."
68+
# AdditionalSlaveSecurityGroups:
69+
# - "..a security group that allows access to redis..."
70+
# KeepJobFlowAliveWhenNoSteps: false
71+
# BootstrapActions:
72+
# - Name: "s3://aws-bigdata-blog/artifacts/resize_storage/resize_storage.sh"
73+
# ScriptBootstrapAction:
74+
# Path: "s3://aws-bigdata-blog/artifacts/resize_storage/resize_storage.sh"
75+
# Args:
76+
# - "--scaling-factor"
77+
# - "1.5"
78+
# Applications:
79+
# - Name: Hadoop
80+
# - Name: Hive
81+
# - Name: Spark
82+
# - Name: Livy
83+
# JobFlowRole: my-spark-node
84+
# ServiceRole: my-worker-node
85+
# ScaleDownBehavior: TERMINATE_AT_TASK_COMPLETION
86+
# redisConfig:
87+
# host: my.redis.com
88+
# port: 6379
89+
# ssl: true
90+
91+
SUPPORTED_EMR_VERSION = "emr-6.0.0"
92+
93+
94+
def _sanity_check_config(config, config_path: str):
95+
"""
96+
Sanity check the config. We don't really have to do this here but if the spark job fails
97+
you'll only find out much later and this is annoying. Those are not exhaustive, just
98+
some checks to help debugging common configuration issues.
99+
"""
100+
aws_config = config.get("aws", {})
101+
102+
if ("runJobFlowTemplate" not in aws_config) and (
103+
"existingClusterId" not in aws_config
104+
):
105+
log.error("{config_path}: either clusterId or runJobFlowTemplate should be set")
106+
elif "runJobFlowTemplate" in aws_config:
107+
runJobFlowTemplate = aws_config["runJobFlowTemplate"]
108+
releaseLabel = runJobFlowTemplate.get("ReleaseLabel")
109+
if releaseLabel != SUPPORTED_EMR_VERSION:
110+
log.warn(
111+
f"{config_path}: ReleaseLabel is set to {releaseLabel}. Recommended: {SUPPORTED_EMR_VERSION}"
112+
)
113+
114+
if "redisConfig" not in config:
115+
log.error("{config_path}: redisConfig is not set")
116+
117+
118+
def _get_config_path() -> str:
119+
return os.environ["JOB_SERVICE_CONFIG_PATH"]
120+
121+
122+
def _load_job_service_config(config_path: str):
123+
with open(config_path) as f:
124+
config = yaml.load(f)
125+
_sanity_check_config(config, config_path)
126+
return config
127+
128+
129+
def _random_string(length) -> str:
130+
return "".join(random.choice(string.ascii_letters) for _ in range(length))
131+
132+
133+
def _batch_source_to_json(batch_source):
134+
return {
135+
"file": {
136+
"path": batch_source.file_options.file_url,
137+
"mapping": dict(batch_source.field_mapping),
138+
"timestampColumn": batch_source.timestamp_column,
139+
}
140+
}
141+
142+
143+
def _feature_table_to_json(client: Client, feature_table):
144+
return {
145+
"features": [
146+
{"name": f.name, "type": ValueType(f.dtype).name}
147+
for f in feature_table.features
148+
],
149+
"project": "default",
150+
"name": feature_table.name,
151+
"entities": [
152+
{"name": n, "type": client.get_entity(n).value_type}
153+
for n in feature_table.entities
154+
],
155+
}
156+
157+
158+
def _s3_split_path(path: str) -> Tuple[str, str]:
159+
""" Convert s3:// url to (bucket, key) """
160+
assert path.startswith("s3://")
161+
_, _, bucket, key = path.split("/", 3)
162+
return bucket, key
163+
164+
165+
def _hash_file(local_path: str) -> str:
166+
""" Compute sha256 hash of a file """
167+
h = hashlib.sha256()
168+
with open(local_path, "rb") as f:
169+
for block in iter(lambda: f.read(2 ** 20), b""):
170+
h.update(block)
171+
return h.hexdigest()
172+
173+
174+
def _s3_upload(local_path: str, remote_path: str) -> str:
175+
"""
176+
Upload a local file to S3. We store the file sha256 sum in S3 metadata and skip the upload
177+
if the file hasn't changed.
178+
"""
179+
bucket, key = _s3_split_path(remote_path)
180+
client = boto3.client("s3")
181+
182+
sha256sum = _hash_file(local_path)
183+
184+
try:
185+
head_response = client.head_object(Bucket=bucket, Key=key)
186+
if head_response["Metadata"]["sha256sum"] == sha256sum:
187+
# File already exists
188+
return remote_path
189+
else:
190+
log.info("Uploading {local_path} to {remote_path}")
191+
client.upload_file(
192+
local_path,
193+
bucket,
194+
key,
195+
ExtraArgs={"Metadata": {"sha256sum": sha256sum}},
196+
)
197+
return remote_path
198+
except botocore.exceptions.ClientError as e:
199+
if e.response["Error"]["Code"] == "404":
200+
log.info("Uploading {local_path} to {remote_path}")
201+
client.upload_file(
202+
local_path,
203+
bucket,
204+
key,
205+
ExtraArgs={"Metadata": {"sha256sum": sha256sum}},
206+
)
207+
return remote_path
208+
else:
209+
raise
210+
211+
212+
def _upload_jar(jar_s3_prefix: str, local_path: str) -> str:
213+
return _s3_upload(
214+
local_path, os.path.join(jar_s3_prefix, os.path.basename(local_path))
215+
)
216+
217+
218+
def _get_jar_s3_path(config) -> str:
219+
"""
220+
Extract job jar path from the configuration, upload it to S3 if necessary and return S3 path.
221+
"""
222+
jar_path = os.environ.get("INGESTION_JOB_JAR_PATH")
223+
if jar_path is None:
224+
raise ValueError("INGESTION_JOB_JAR_PATH not set")
225+
elif jar_path.startswith("s3://"):
226+
return jar_path
227+
else:
228+
artifactS3Prefix = config.get("aws").get("artifactS3Prefix")
229+
if artifactS3Prefix:
230+
return _upload_jar(artifactS3Prefix, jar_path)
231+
else:
232+
raise ValueError("artifactS3Prefix must be set")
233+
234+
235+
def _sync_offline_to_online_step(
236+
client: Client, config, feature_table, start_ts: str, end_ts: str
237+
) -> Dict[str, Any]:
238+
feature_table_json = _feature_table_to_json(client, feature_table)
239+
source_json = _batch_source_to_json(feature_table.batch_source)
240+
241+
return {
242+
"Name": "Feast Ingestion",
243+
"HadoopJarStep": {
244+
# TODO: generate those from proto
245+
"Properties": [
246+
{
247+
"Key": "feast.step_metadata.job_type",
248+
"Value": "OFFLINE_TO_ONLINE_JOB",
249+
},
250+
{
251+
"Key": "feast.step_metadata.offline_to_online.table_name",
252+
"Value": feature_table.name,
253+
},
254+
{
255+
"Key": "feast.step_metadata.offline_to_online.start_ts",
256+
"Value": start_ts,
257+
},
258+
{
259+
"Key": "feast.step_metadata.offline_to_online.end_ts",
260+
"Value": end_ts,
261+
},
262+
],
263+
"Args": [
264+
"spark-submit",
265+
"--class",
266+
"feast.ingestion.IngestionJob",
267+
"--packages",
268+
"com.google.cloud.spark:spark-bigquery-with-dependencies_2.12:0.17.2",
269+
_get_jar_s3_path(config),
270+
"--mode",
271+
"offline",
272+
"--feature-table",
273+
json.dumps(feature_table_json),
274+
"--source",
275+
json.dumps(source_json),
276+
"--redis",
277+
json.dumps(config["redisConfig"]),
278+
"--start",
279+
start_ts,
280+
"--end",
281+
end_ts,
282+
],
283+
"Jar": "command-runner.jar",
284+
},
285+
}
286+
287+
288+
def _submit_emr_job(step: Dict[str, Any], config: Dict[str, Any]):
289+
aws_config = config.get("aws", {})
290+
291+
emr = boto3.client("emr", region_name=aws_config.get("region"))
292+
293+
if "existingClusterId" in aws_config:
294+
step["ActionOnFailure"] = "CONTINUE"
295+
step_ids = emr.add_job_flow_steps(
296+
JobFlowId=aws_config["existingClusterId"], Steps=[step],
297+
)
298+
print(step_ids)
299+
else:
300+
jobTemplate = aws_config["runJobFlowTemplate"]
301+
step["ActionOnFailure"] = "TERMINATE_CLUSTER"
302+
303+
jobTemplate["Steps"] = [step]
304+
305+
if aws_config.get("logS3Prefix"):
306+
jobTemplate["LogUri"] = os.path.join(
307+
aws_config["logS3Prefix"], _random_string(5)
308+
)
309+
310+
job = emr.run_job_flow(**jobTemplate)
311+
print(job)
312+
313+
314+
def sync_offline_to_online(
315+
client: Client, feature_table: FeatureTable, start_ts: str, end_ts: str
316+
):
317+
config = _load_job_service_config(_get_config_path())
318+
step = _sync_offline_to_online_step(client, config, feature_table, start_ts, end_ts)
319+
_submit_emr_job(step, config)

0 commit comments

Comments
 (0)