from airflow import DAG
from airflow.providers.standard.operators.python import PythonOperator
import pendulum
from datetime import timedelta

MASTER_PORT = "29500"

default_args = {
    "owner": "airflow",
    "retries": 0,
    "execution_timeout": timedelta(minutes=20),
}


# -------------------------
# RANK 0
# -------------------------
def rank0(**context):
    import os, socket, torch
    import torch.distributed as dist

    addr = socket.gethostname() + ".airflow-worker-gpu"

    context["ti"].xcom_push(key="master_addr", value=addr)

    os.environ.update({
        "MASTER_ADDR": addr,
        "MASTER_PORT": MASTER_PORT,
        "WORLD_SIZE": "2",
        "RANK": "0",

        "NCCL_DEBUG": "INFO",
        "NCCL_SOCKET_IFNAME": "eth0",

        # IB
        "NCCL_IB_DISABLE": "0",
        "NCCL_IB_GID_INDEX": "0",
    })

    dist.init_process_group("nccl")

    torch.cuda.set_device(0)

    x = torch.ones(1, device="cuda") * 1
    dist.all_reduce(x)

    print(f"[rank0] result = {x.item()}")

    dist.destroy_process_group()


# -------------------------
# RANK 1
# -------------------------
def rank1(**context):
    import os, time, torch
    import torch.distributed as dist

    addr = None
    for _ in range(30):
        addr = context["ti"].xcom_pull(task_ids="rank0", key="master_addr")
        if addr:
            break
        time.sleep(2)

    os.environ.update({
        "MASTER_ADDR": addr,
        "MASTER_PORT": MASTER_PORT,
        "WORLD_SIZE": "2",
        "RANK": "1",

        "NCCL_DEBUG": "INFO",
        "NCCL_SOCKET_IFNAME": "eth0",

        "NCCL_IB_DISABLE": "0",
        "NCCL_IB_GID_INDEX": "0",
    })

    dist.init_process_group("nccl")

    torch.cuda.set_device(0)

    x = torch.ones(1, device="cuda") * 2
    dist.all_reduce(x)

    print(f"[rank1] result = {x.item()}")

    dist.destroy_process_group()


# -------------------------
# DAG
# -------------------------
with DAG(
    dag_id="ib_simple_test",
    start_date=pendulum.today("UTC").add(days=-1),
    schedule=None,
    catchup=False,
    max_active_runs=1,
    default_args=default_args,
) as dag:

    r0 = PythonOperator(
        task_id="rank0",
        python_callable=rank0,
        queue="gpu",
    )

    r1 = PythonOperator(
        task_id="rank1",
        python_callable=rank1,
        queue="gpu",
    )

    r0 >> r1