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=30),
}


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

    hostname = socket.gethostname()
    master_addr = f"{hostname}.airflow-worker-gpu"

    # share master with rank1
    context["ti"].xcom_push(key="master_addr", value=master_addr)

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

        # NCCL (minimal stable IB config)
        "NCCL_DEBUG": "INFO",
        "NCCL_SOCKET_IFNAME": "eth0",

        # IB stability (IMPORTANT)
        "NCCL_IB_GID_INDEX": "0",
        "NCCL_IB_USE_INLINE": "0",
        "NCCL_IB_TIMEOUT": "22",
        "NCCL_IB_RETRY_CNT": "7",
    })

    dist.init_process_group("nccl", init_method="env://")

    torch.cuda.set_device(0)

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

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

    dist.destroy_process_group()


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

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

    if not master_addr:
        raise RuntimeError("master_addr not found")

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

        # NCCL (same as rank0)
        "NCCL_DEBUG": "INFO",
        "NCCL_SOCKET_IFNAME": "eth0",

        # IB stability
        "NCCL_IB_GID_INDEX": "0",
        "NCCL_IB_USE_INLINE": "0",
        "NCCL_IB_TIMEOUT": "22",
        "NCCL_IB_RETRY_CNT": "7",
    })

    dist.init_process_group("nccl", init_method="env://")

    torch.cuda.set_device(0)

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

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

    dist.destroy_process_group()


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

    rank0 = PythonOperator(
        task_id="train_rank_0",
        python_callable=run_rank_0,
        queue="gpu",
    )

    rank1 = PythonOperator(
        task_id="train_rank_1",
        python_callable=run_rank_1,
        queue="gpu",
    )

    rank0 >> rank1