def test_get_cluster_id_by_name(self): """ Test that we can resolve cluster id by cluster name. """ hook = EmrHook(aws_conn_id='aws_default', emr_conn_id='emr_default') job_flow = hook.create_job_flow({'Name': 'test_cluster', 'Instances': {'KeepJobFlowAliveWhenNoSteps': True}}) job_flow_id = job_flow['JobFlowId'] matching_cluster = hook.get_cluster_id_by_name('test_cluster', ['RUNNING', 'WAITING']) self.assertEqual(matching_cluster, job_flow_id) no_match = hook.get_cluster_id_by_name('foo', ['RUNNING', 'WAITING', 'BOOTSTRAPPING']) self.assertIsNone(no_match)
def execute(self, context): attempt = context['ti'].try_number logging.info('attempt: {}'.format(attempt)) emr = EmrHook(aws_conn_id=self.aws_conn_id).get_conn() job_flow_id = self.job_flow_id if not job_flow_id: job_flow_id = emr.get_cluster_id_by_name(self.job_flow_name, self.cluster_states) if self.do_xcom_push: context['ti'].xcom_push(key='job_flow_id', value=job_flow_id) step_name = self.step_name if attempt == 1 else "{} (attempt {})".format( self.step_name, attempt) action_on_failure = self.action_on_failure if attempt % 3 == 0: action_on_failure = 'TERMINATE_JOB_FLOW' spark_conf = self.get_spark_params_config(self.spark_params, self.spark_conf) steps = self.generate_spark_step(step_name, self.main_class, self.app_name, spark_conf, self.application_args, self.jar_path, action_on_failure) logging.info("spark_params: " + str(steps)) self.log.info('Adding steps to %s', job_flow_id) response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps) logging.info('Running Spark Job {} with JobFlow ID {}'.format( self.task_id, self.job_flow_id)) while True: step_id = response['StepIds'][0] logging.info('step id - {}'.format(step_id)) result = self.describe_step(emr, response) step_status = result['Step']['Status']['State'] logging.info('step status - {}'.format(step_status)) # step state can be 'PENDING'|'CANCEL_PENDING'|'RUNNING'|'COMPLETED'|'CANCELLED'|'FAILED'|'INTERRUPTED' if step_status == 'COMPLETED': break elif step_status != 'COMPLETED' and step_status != 'PENDING' and step_status != 'RUNNING': raise AirflowException('Spark job {} has failed'.format( self.task_id)) logging.info("Spark Job '{}' status is {}".format( self.task_id, step_status))
def execute(self, context): emr = EmrHook(aws_conn_id=self.aws_conn_id).get_conn() job_flow_id = self.job_flow_id if not job_flow_id: job_flow_id = emr.get_cluster_id_by_name(self.job_flow_name, self.cluster_states) if self.do_xcom_push: context['ti'].xcom_push(key='job_flow_id', value=job_flow_id) self.log.info('Adding steps to %s', job_flow_id) response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=self.steps) if not response['ResponseMetadata']['HTTPStatusCode'] == 200: raise AirflowException('Adding steps failed: %s' % response) else: self.log.info('Steps %s added to JobFlow', response['StepIds']) return response['StepIds']
def execute(self, context: Dict[str, Any]) -> List[str]: emr_hook = EmrHook(aws_conn_id=self.aws_conn_id) emr = emr_hook.get_conn() job_flow_id = self.job_flow_id or emr_hook.get_cluster_id_by_name( str(self.job_flow_name), self.cluster_states ) if not job_flow_id: raise AirflowException(f"No cluster found for name: {self.job_flow_name}") if self.do_xcom_push: context["ti"].xcom_push(key="job_flow_id", value=job_flow_id) self.log.info("Adding steps to %s", job_flow_id) # steps may arrive as a string representing a list # e.g. if we used XCom or a file then: steps="[{ step1 }, { step2 }]" steps = self.steps if isinstance(steps, str): steps = ast.literal_eval(steps) response = emr.add_job_flow_steps(JobFlowId=job_flow_id, Steps=steps) if not response["ResponseMetadata"]["HTTPStatusCode"] == 200: raise AirflowException("Adding steps failed: %s" % response) else: # Assumption : ONly a single step is submitted each time. step_ids = response["StepIds"] step_id = step_ids[0] if self.wait_for_completion: self.check_status( job_flow_id, step_id, self.describe_step, self.check_interval, ) self.log.info("Steps %s added to JobFlow", response["StepIds"]) return response["StepIds"]