def create_dagrun(self, *args, **kwargs): run_args = "{}" # TODO extract the run Args from the tasks marquez_jobrun_id = None try: marquez_jobrun_id = self.report_jobrun(run_args, kwargs['execution_date']) log.info(f'Successfully recorded job run.', airflow_dag_id=self.dag_id, marquez_run_id=marquez_jobrun_id, marquez_namespace=self.marquez_namespace) except Exception as e: log.error(f'Failed to record job run: {e}', airflow_dag_id=self.dag_id, marquez_namespace=self.marquez_namespace) pass run = super(DAG, self).create_dagrun(*args, **kwargs) if marquez_jobrun_id: try: self._job_id_mapping.set( JobIdMapping.make_key(run.dag_id, run.run_id), marquez_jobrun_id) except Exception as e: log.error(f'Failed job run lookup: {e}', airflow_dag_id=self.dag_id, airflow_run_id=run.run_id, marquez_run_id=marquez_jobrun_id, marquez_namespace=self.marquez_namespace) pass return run
def report_jobrun_change(self, job_name, run_id, **kwargs): session = kwargs.get('session') marquez_job_run_ids = self._job_id_mapping.pop( JobIdMapping.make_key(job_name, run_id), session) if marquez_job_run_ids: log.info('Found job runs.', airflow_dag_id=self.dag_id, airflow_job_id=job_name, airflow_run_id=run_id, marquez_run_ids=marquez_job_run_ids, marquez_namespace=self.marquez_namespace) ids = json.loads(marquez_job_run_ids) if kwargs.get('success'): for marquez_job_run_id in ids: self.get_marquez_client().mark_job_run_as_completed( marquez_job_run_id) else: for marquez_job_run_id in ids: self.get_marquez_client().mark_job_run_as_failed( marquez_job_run_id) state = 'COMPLETED' if kwargs.get('success') else 'FAILED' log.info(f'Marked job run(s) as {state}.', airflow_dag_id=self.dag_id, airflow_job_id=job_name, airflow_run_id=run_id, marquez_run_id=marquez_job_run_ids, marquez_namespace=self.marquez_namespace)
def report_jobrun_change(self, dagrun, **kwargs): session = kwargs.get('session') marquez_job_run_id = self._job_id_mapping.pop( JobIdMapping.make_key(dagrun.dag_id, dagrun.run_id), session) if marquez_job_run_id: log.info(f'Found job run.', airflow_dag_id=dagrun.dag_id, airflow_run_id=dagrun.run_id, marquez_run_id=marquez_job_run_id, marquez_namespace=self.marquez_namespace) if kwargs.get('success'): self.get_marquez_client().mark_job_run_as_completed( marquez_job_run_id) else: self.get_marquez_client().mark_job_run_as_failed( marquez_job_run_id) state = 'COMPLETED' if kwargs.get('success') else 'FAILED' log.info(f'Marked job run as {state}.', airflow_dag_id=dagrun.dag_id, airflow_run_id=dagrun.run_id, marquez_run_id=marquez_job_run_id, marquez_namespace=self.marquez_namespace)
def report_task(self, dag_run_id, execution_date, run_args, task, extractor): report_job_start_ms = self._now_ms() marquez_client = self.get_marquez_client() if execution_date: start_time = self._to_iso_8601(execution_date) end_time = self.compute_endtime(execution_date) else: start_time = None end_time = None if end_time: end_time = self._to_iso_8601(end_time) task_location = None try: if hasattr(task, 'file_path') and task.file_path: task_location = get_location(task.file_path) else: task_location = get_location(task.dag.fileloc) except Exception: log.warn('Unable to fetch the location') steps_metadata = [] if extractor: try: log.info(f'Using extractor {extractor.__name__}', task_type=task.__class__.__name__, airflow_dag_id=self.dag_id, task_id=task.task_id, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) steps_metadata = extractor(task).extract() except Exception as e: log.error(f'Failed to extract metadata {e}', airflow_dag_id=self.dag_id, task_id=task.task_id, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) else: log.warn('Unable to find an extractor.', task_type=task.__class__.__name__, airflow_dag_id=self.dag_id, task_id=task.task_id, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) task_name = f'{self.dag_id}.{task.task_id}' # If no extractor found or failed to extract metadata, # report the task metadata if not steps_metadata: steps_metadata = [StepMetadata(task_name)] # store all the JobRuns associated with a task marquez_jobrun_ids = [] for step in steps_metadata: input_datasets = [] output_datasets = [] try: input_datasets = self.register_datasets(step.inputs) except Exception as e: log.error(f'Failed to register inputs: {e}', inputs=str(step.inputs), airflow_dag_id=self.dag_id, task_id=task.task_id, step=step.name, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) try: output_datasets = self.register_datasets(step.outputs) except Exception as e: log.error(f'Failed to register outputs: {e}', outputs=str(step.outputs), airflow_dag_id=self.dag_id, task_id=task.task_id, step=step.name, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) marquez_client.create_job( job_name=step.name, job_type='BATCH', # job type location=(step.location or task_location), input_dataset=input_datasets, output_dataset=output_datasets, context=step.context, description=self.description, namespace_name=self.marquez_namespace) log.info(f'Successfully recorded job: {step.name}', airflow_dag_id=self.dag_id, marquez_namespace=self.marquez_namespace) marquez_jobrun_id = marquez_client.create_job_run( step.name, run_args=run_args, nominal_start_time=start_time, nominal_end_time=end_time).get('runId') if marquez_jobrun_id: marquez_jobrun_ids.append(marquez_jobrun_id) marquez_client.mark_job_run_as_started(marquez_jobrun_id) else: log.error(f'Failed to get run id: {step.name}', airflow_dag_id=self.dag_id, airflow_run_id=dag_run_id, marquez_namespace=self.marquez_namespace) log.info(f'Successfully recorded job run: {step.name}', airflow_dag_id=self.dag_id, airflow_dag_execution_time=start_time, marquez_run_id=marquez_jobrun_id, marquez_namespace=self.marquez_namespace, duration_ms=(self._now_ms() - report_job_start_ms)) # Store the mapping for all the steps associated with a task try: self._job_id_mapping.set( JobIdMapping.make_key(task_name, dag_run_id), json.dumps(marquez_jobrun_ids)) except Exception as e: log.error(f'Failed to set id mapping : {e}', airflow_dag_id=self.dag_id, task_id=task.task_id, airflow_run_id=dag_run_id, marquez_run_id=marquez_jobrun_ids, marquez_namespace=self.marquez_namespace)