From 3474f89e3fee6d7a79a3dab2136b4a111ad672ef Mon Sep 17 00:00:00 2001 From: Oleg Avdeev Date: Tue, 20 Oct 2020 11:44:44 +0300 Subject: [PATCH] implement EMR job cancelling Signed-off-by: Oleg Avdeev --- sdk/python/feast/pyspark/launchers/aws/emr.py | 3 +- .../feast/pyspark/launchers/aws/emr_utils.py | 32 +++++++------------ 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/sdk/python/feast/pyspark/launchers/aws/emr.py b/sdk/python/feast/pyspark/launchers/aws/emr.py index c23d07fb59a..4dd5773e889 100644 --- a/sdk/python/feast/pyspark/launchers/aws/emr.py +++ b/sdk/python/feast/pyspark/launchers/aws/emr.py @@ -25,6 +25,7 @@ SUCCEEDED_STEP_STATES, TERMINAL_STEP_STATES, EmrJobRef, + _cancel_job, _get_job_state, _historical_retrieval_step, _load_new_cluster_template, @@ -63,7 +64,7 @@ def get_status(self) -> SparkJobStatus: raise Exception("Invalid EMR state") def cancel(self): - raise NotImplementedError + _cancel_job(self._emr_client, self._job_ref) class EmrRetrievalJob(EmrJobMixin, RetrievalJob): diff --git a/sdk/python/feast/pyspark/launchers/aws/emr_utils.py b/sdk/python/feast/pyspark/launchers/aws/emr_utils.py index 75c68e63181..f421390c4f9 100644 --- a/sdk/python/feast/pyspark/launchers/aws/emr_utils.py +++ b/sdk/python/feast/pyspark/launchers/aws/emr_utils.py @@ -240,6 +240,9 @@ def _get_stream_to_online_job(emr_client, table_name: str) -> List[JobInfo]: class EmrJobRef(NamedTuple): + """ EMR job reference. step_id can be None when using on-demand clusters, in that case each + cluster has only one step """ + cluster_id: str step_id: Optional[str] @@ -304,28 +307,17 @@ def _wait_for_step_state( ) -def _cancel_job(emr_client, job_type: str, table_name: str): - """ - Cancel a EMR job. - """ - jobs = list_jobs( - emr_client, job_type=job_type, table_name=table_name, active_only=True - ) - - for job in jobs: - emr_client.cancel_steps(ClusterId=job.cluster_id, StepIds=[job.step_id]) - - for job in jobs: - _wait_for_job_state( - emr_client, EmrJobRef(job.cluster_id, job.step_id), TERMINAL_STEP_STATES, 90 - ) +def _cancel_job(emr_client, job: EmrJobRef): + if job.step_id is None: + step_id = _get_first_step_id(emr_client, job.cluster_id) + else: + step_id = job.step_id + emr_client.cancel_steps(ClusterId=job.cluster_id, StepIds=[step_id]) -def stop_stream_to_online(emr_client, table_name: str): - """ - Stop offline-to-online ingestion job for the table. - """ - _cancel_job(emr_client, STREAM_TO_ONLINE_JOB_TYPE, table_name) + _wait_for_job_state( + emr_client, EmrJobRef(job.cluster_id, step_id), TERMINAL_STEP_STATES, 180 + ) def _upload_dataframe(s3prefix: str, df: pandas.DataFrame) -> str: