Skip to content

Commit cc5eeec

Browse files
oavdeevoavdeev
committed
CLI command to start/stop/list streaming ingestion job on emr
Signed-off-by: Oleg Avdeev <oleg.v.avdeev@gmail.com>
1 parent 2cd019c commit cc5eeec

2 files changed

Lines changed: 255 additions & 5 deletions

File tree

‎sdk/python/feast/cli.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -371,5 +371,58 @@ def sync_offline_to_online(feature_table: str, start_time: str, end_time: str):
371371
feast.pyspark.aws.jobs.sync_offline_to_online(client, table, start_time, end_time)
372372

373373

374+
@cli.command()
375+
@click.option(
376+
"--feature-table",
377+
"-t",
378+
help="Feature table name to ingest data into",
379+
required=True,
380+
)
381+
@click.option(
382+
"--jar", "-j", help="Feature table name to ingest data into", default="",
383+
)
384+
def start_stream_to_online(feature_table: str, jar: str):
385+
"""
386+
Start stream to online sync job.
387+
"""
388+
import feast.pyspark.aws.jobs
389+
390+
client = Client()
391+
table = client.get_feature_table(feature_table)
392+
feast.pyspark.aws.jobs.start_stream_to_online(client, table, [jar] if jar else [])
393+
394+
395+
@cli.command()
396+
@click.option(
397+
"--feature-table",
398+
"-t",
399+
help="Feature table name to ingest data into",
400+
required=True,
401+
)
402+
def stop_stream_to_online(feature_table: str):
403+
"""
404+
Start stream to online sync job.
405+
"""
406+
import feast.pyspark.aws.jobs
407+
408+
feast.pyspark.aws.jobs.stop_stream_to_online(feature_table)
409+
410+
411+
@cli.command()
412+
def list_emr_jobs():
413+
"""
414+
List jobs.
415+
"""
416+
from tabulate import tabulate
417+
418+
import feast.pyspark.aws.jobs
419+
420+
jobs = feast.pyspark.aws.jobs.list_jobs(None, None)
421+
422+
print(
423+
tabulate(jobs, headers=feast.pyspark.aws.jobs.JobInfo._fields, tablefmt="plain")
424+
)
425+
426+
374427
if __name__ == "__main__":
375428
cli()

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

Lines changed: 202 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@
44
import os
55
import random
66
import string
7-
from typing import Any, Dict, Tuple
7+
import time
8+
from typing import Any, Dict, List, NamedTuple, Optional, Tuple
89

910
import boto3
1011
import botocore
@@ -89,6 +90,13 @@
8990
# ssl: true
9091

9192
SUPPORTED_EMR_VERSION = "emr-6.0.0"
93+
STREAM_TO_ONLINE_JOB_TYPE = "STREAM_TO_ONLINE_JOB"
94+
OFFLINE_TO_ONLINE_JOB_TYPE = "OFFLINE_TO_ONLINE_JOB"
95+
96+
97+
# EMR Step states considered "active", i.e. not terminated
98+
ACTIVE_STEP_STATES = ["PENDING", "CANCEL_PENDING", "RUNNING"]
99+
TERMINAL_STEP_STATES = ["COMPLETED", "CANCELLED", "FAILED", "INTERRUPTED"]
92100

93101

94102
def _sanity_check_config(config, config_path: str):
@@ -121,7 +129,7 @@ def _get_config_path() -> str:
121129

122130
def _load_job_service_config(config_path: str):
123131
with open(config_path) as f:
124-
config = yaml.load(f)
132+
config = yaml.safe_load(f)
125133
_sanity_check_config(config, config_path)
126134
return config
127135

@@ -140,6 +148,18 @@ def _batch_source_to_json(batch_source):
140148
}
141149

142150

151+
def _stream_source_to_json(stream_source):
152+
return {
153+
"kafka": {
154+
"bootstrapServers": stream_source.kafka_options.bootstrap_servers,
155+
"mapping": dict(stream_source.field_mapping),
156+
"topic": stream_source.kafka_options.topic,
157+
"timestampColumn": stream_source.timestamp_column,
158+
"classpath": stream_source.kafka_options.class_path,
159+
}
160+
}
161+
162+
143163
def _feature_table_to_json(client: Client, feature_table):
144164
return {
145165
"features": [
@@ -215,7 +235,7 @@ def _upload_jar(jar_s3_prefix: str, local_path: str) -> str:
215235
)
216236

217237

218-
def _get_jar_s3_path(config) -> str:
238+
def _get_ingestion_jar_s3_path(config) -> str:
219239
"""
220240
Extract job jar path from the configuration, upload it to S3 if necessary and return S3 path.
221241
"""
@@ -245,7 +265,7 @@ def _sync_offline_to_online_step(
245265
"Properties": [
246266
{
247267
"Key": "feast.step_metadata.job_type",
248-
"Value": "OFFLINE_TO_ONLINE_JOB",
268+
"Value": OFFLINE_TO_ONLINE_JOB_TYPE,
249269
},
250270
{
251271
"Key": "feast.step_metadata.offline_to_online.table_name",
@@ -266,7 +286,7 @@ def _sync_offline_to_online_step(
266286
"feast.ingestion.IngestionJob",
267287
"--packages",
268288
"com.google.cloud.spark:spark-bigquery-with-dependencies_2.12:0.17.2",
269-
_get_jar_s3_path(config),
289+
_get_ingestion_jar_s3_path(config),
270290
"--mode",
271291
"offline",
272292
"--feature-table",
@@ -317,3 +337,180 @@ def sync_offline_to_online(
317337
config = _load_job_service_config(_get_config_path())
318338
step = _sync_offline_to_online_step(client, config, feature_table, start_ts, end_ts)
319339
_submit_emr_job(step, config)
340+
341+
342+
def _stream_ingestion_step(
343+
client: Client, config, feature_table, jars: List[str]
344+
) -> Dict[str, Any]:
345+
feature_table_json = _feature_table_to_json(client, feature_table)
346+
source_json = _stream_source_to_json(feature_table.stream_source)
347+
348+
if jars:
349+
jars_args = ["--jars", ",".join(jars)]
350+
else:
351+
jars_args = []
352+
353+
return {
354+
"Name": "Feast Streaming Ingestion",
355+
"HadoopJarStep": {
356+
"Properties": [
357+
{
358+
"Key": "feast.step_metadata.job_type",
359+
"Value": STREAM_TO_ONLINE_JOB_TYPE,
360+
},
361+
{
362+
"Key": "feast.step_metadata.stream_to_online.table_name",
363+
"Value": feature_table.name,
364+
},
365+
],
366+
"Args": ["spark-submit", "--class", "feast.ingestion.IngestionJob"]
367+
+ jars_args
368+
+ [
369+
"--packages",
370+
"com.google.cloud.spark:spark-bigquery-with-dependencies_2.12:0.17.2",
371+
_get_ingestion_jar_s3_path(config),
372+
"--mode",
373+
"online",
374+
"--feature-table",
375+
json.dumps(feature_table_json),
376+
"--source",
377+
json.dumps(source_json),
378+
"--redis",
379+
json.dumps(config["redisConfig"]),
380+
],
381+
"Jar": "command-runner.jar",
382+
},
383+
}
384+
385+
386+
def start_stream_to_online(
387+
client: Client, feature_table: FeatureTable, jars: List[str]
388+
):
389+
if _get_stream_to_online_job(client, feature_table):
390+
raise Exception("Job already running")
391+
392+
config = _load_job_service_config(_get_config_path())
393+
step = _stream_ingestion_step(client, config, feature_table, jars)
394+
_submit_emr_job(step, config)
395+
396+
397+
class JobInfo(NamedTuple):
398+
job_type: str
399+
cluster_id: str
400+
step_id: str
401+
table_name: str
402+
state: str
403+
404+
405+
def list_jobs(
406+
job_type: Optional[str], table_name: Optional[str], active_only=True
407+
) -> List[JobInfo]:
408+
"""
409+
List Feast EMR jobs.
410+
411+
Args:
412+
job_type: optional filter by job type
413+
table_name: optional filter by table name
414+
active_only: filter only for "active" jobs, that is the ones that are running or pending, not terminated
415+
416+
Returns:
417+
A list of jobs.
418+
"""
419+
config = _load_job_service_config(_get_config_path())
420+
aws_config = config.get("aws", {})
421+
emr = boto3.client("emr", region_name=aws_config.get("region"))
422+
paginator = emr.get_paginator("list_clusters")
423+
res: List[JobInfo] = []
424+
for page in paginator.paginate(
425+
ClusterStates=["STARTING", "BOOTSTRAPPING", "RUNNING", "WAITING", "TERMINATING"]
426+
):
427+
for cluster in page["Clusters"]:
428+
cluster_id = cluster["Id"]
429+
step_paginator = emr.get_paginator("list_steps")
430+
431+
list_steps_params = dict(ClusterId=cluster_id)
432+
if active_only:
433+
list_steps_params["StepStates"] = ACTIVE_STEP_STATES
434+
435+
for step_page in step_paginator.paginate(**list_steps_params):
436+
for step in step_page["Steps"]:
437+
props = step["Config"]["Properties"]
438+
if "feast.step_metadata.job_type" not in props:
439+
continue
440+
441+
step_table_name = props.get(
442+
"feast.step_metadata.stream_to_online.table_name"
443+
) or props.get("feast.step_metadata.offline_to_online.table_name")
444+
step_job_type = props["feast.step_metadata.job_type"]
445+
446+
if table_name and step_table_name != table_name:
447+
continue
448+
449+
if job_type and step_job_type != job_type:
450+
continue
451+
452+
res.append(
453+
JobInfo(
454+
job_type=step_job_type,
455+
cluster_id=cluster_id,
456+
step_id=step["Id"],
457+
state=step["Status"]["State"],
458+
table_name=step_table_name,
459+
)
460+
)
461+
return res
462+
463+
464+
def _get_stream_to_online_job(
465+
client: Client, feature_table: FeatureTable
466+
) -> List[JobInfo]:
467+
return list_jobs(
468+
job_type=STREAM_TO_ONLINE_JOB_TYPE,
469+
table_name=feature_table.name,
470+
active_only=True,
471+
)
472+
473+
474+
def _wait_for_job_state(
475+
emr_client, job: JobInfo, desired_states: List[str], timeout_seconds=90
476+
):
477+
"""
478+
Wait up to timeout seconds for job to go into one of the desired states.
479+
"""
480+
start_time = time.time()
481+
while time.time() - start_time < timeout_seconds:
482+
response = emr_client.describe_step(
483+
ClusterId=job.cluster_id, StepId=job.step_id
484+
)
485+
state = response["Step"]["Status"]["State"]
486+
if state in desired_states:
487+
return
488+
else:
489+
time.sleep(0.5)
490+
else:
491+
raise TimeoutError(
492+
f'Timeout waiting for job state to become {"|".join(desired_states)}'
493+
)
494+
495+
496+
def _cancel_job(job_type, table_name):
497+
"""
498+
Cancel a EMR job.
499+
"""
500+
jobs = list_jobs(job_type=job_type, table_name=table_name, active_only=True)
501+
config = _load_job_service_config(_get_config_path())
502+
aws_config = config.get("aws", {})
503+
504+
emr = boto3.client("emr", region_name=aws_config.get("region"))
505+
for job in jobs:
506+
emr.cancel_steps(ClusterId=job.cluster_id, StepIds=[job.step_id])
507+
508+
for job in jobs:
509+
_wait_for_job_state(emr, job, TERMINAL_STEP_STATES)
510+
511+
512+
def stop_stream_to_online(table_name: str):
513+
"""
514+
Stop offline-to-online ingestion job for the table.
515+
"""
516+
_cancel_job(STREAM_TO_ONLINE_JOB_TYPE, table_name)

0 commit comments

Comments
 (0)