Airflow: Zoom in Sub Dag button in customized SubDagOperator

Viewed 212

I've designed a customized SubDagOperator. Everything works fine except that the "Zoom in Sub Dag" button doesn't appear. It seems that if the airflow UI doesn't recognise the task as a subdagoperator itself the button is not shown. I've tried to override the task_type property, as it was mentioned in and old issue, but it doesn't work for me. Do you know if it's possible to see the button with customized SubDagOperators?

Airflow version: 1.10.12

Here is my try:

class EmrSubdagOperator(SubDagOperator):
  template_fields = ()
  template_ext = ()
  


  @apply_defaults
  def __init__(self, *args, **kwargs):

    dag = kwargs.get('dag')
    task_id = kwargs.get('task_id')
    spark_steps = kwargs.get('spark_steps')
    job_flow_overrides = kwargs.get('job_flow_overrides')

    subdag = DAG(
        '{}.{}'.format(dag.dag_id, task_id),
        schedule_interval=dag.schedule_interval,
        start_date=dag.start_date
        )

    cluster_creator = EmrCreateJobFlowOperator(
        dag=subdag,
        task_id='create_job_flow',
        job_flow_overrides=job_flow_overrides
    )

    step_adder = EmrAddStepsOperator(
        dag=subdag,
        task_id='add_steps',
        job_flow_id="{{ task_instance.xcom_pull(task_ids='create_job_flow', key='return_value') }}",
        aws_conn_id='aws_default',
        steps=spark_steps,
    )
    
    #Check if the last step is completed, skipped or terminated
    last_step = len(spark_steps) - 1 
    step_checker = EmrStepSensor(
        dag=subdag,
        task_id='watch_step',
        job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
        step_id="{{ task_instance.xcom_pull(task_ids='add_steps', key='return_value')[" + str(last_step) +"]}}",
        aws_conn_id='aws_default',
    )

    cluster_checker = EmrJobFlowSensor(
        dag=subdag,
        task_id='watch_cluster',
        job_flow_id="{{ task_instance.xcom_pull('create_job_flow', key='return_value') }}",
        aws_conn_id='aws_default',
    )
    
    
    super(EmrSubdagOperator, self).__init__(subdag=subdag, *args, **kwargs)
    self.spark_steps = spark_steps
    self.job_flow_overrides = job_flow_overrides

  @property
  def task_type(self):
      return 'SubDagOperator'
2 Answers

I had the same problem. I resolved that, doing:

class DynamicTasksOperator(SubDagOperator):
 def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self._task_type = 'SubDagOperator'

 @property
 def task_type(self):
    return self._task_type

I have made the following changes in the airflow.cfg and its working for me now.

store_serialized_dags = False
store_dag_code = False
Related