airflow_branching_improved.py

import random
import datetime as dt
from airflow.decorators import task, dag

default_args = {"owner": "airflow", "start_date": dt.datetime(2024, 2, 1), "retries": 1}

@task.branch()
def choose_branch():
    return random.choice(["task_a", "task_b"])  # Randomly chooses a branch

@task
def task_a():
    print("Task A is running")

@task
def task_b():
    print("Task B is running")

@dag(default_args=default_args, catchup=False, schedule="@daily")
def my_dag():
    choose_branch() >> [task_a(), task_b()]

my_dag()