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

WORLD_SIZE = 2
MASTER_PORT = "29500"

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


# ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
def run_rank_0(**context):
    import os, resource, socket, torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim

    hostname = socket.gethostname()
    fqdn = f"{hostname}.airflow-worker-gpu"
    print(f"[rank0] hostname={fqdn}")

    # Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться
    context["ti"].xcom_push(key="master_addr", value=fqdn)

    # Снимаем лимит на locked memory (нужно для ibv_reg_mr)
    resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))

    os.environ.update({
        "MASTER_ADDR": fqdn,
        "MASTER_PORT": MASTER_PORT,
        "WORLD_SIZE": "2",
        "RANK": "0",
        "NCCL_SOCKET_IFNAME": "eth0",   # bootstrap/rendezvous через eth
        "NCCL_IB_DISABLE": "0",         # трафик через IB
        "NCCL_NET": "IB",
        "NCCL_P2P_DISABLE": "1",        # разные ноды — p2p не нужен
        "NCCL_SHM_DISABLE": "1",        # shm только для локальных рангов
        "NCCL_DEBUG": "INFO",
        "NCCL_DEBUG_SUBSYS": "NET,INIT",
        "TORCH_NCCL_BLOCKING_WAIT": "1",
    })

    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=0,
        world_size=2,
        timeout=timedelta(minutes=5),
    )
    print("[rank0] process group OK")

    torch.cuda.set_device(0)
    device = torch.device("cuda:0")
    print(f"[rank0] GPU={torch.cuda.get_device_name(0)}")

    _train(rank=0, device=device)

    dist.destroy_process_group()
    print("[rank0] DONE")


# ── rank 1: ждёт XCom от rank0, потом коннектится ────────────────────────────
def run_rank_1(**context):
    import os, resource, socket, time, torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim

    # Polling XCom пока rank0 не запишет master_addr
    master_addr = None
    for attempt in range(30):
        master_addr = context["ti"].xcom_pull(
            task_ids="train_rank_0",
            key="master_addr",
        )
        if master_addr:
            break
        print(f"[rank1] waiting for master_addr (attempt {attempt+1}/30)...")
        time.sleep(5)

    if not master_addr:
        raise RuntimeError("rank1: timeout waiting for master_addr from rank0")

    print(f"[rank1] master_addr={master_addr}")

    resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))

    os.environ.update({
        "MASTER_ADDR": master_addr,
        "MASTER_PORT": MASTER_PORT,
        "WORLD_SIZE": "2",
        "RANK": "1",
        "NCCL_SOCKET_IFNAME": "eth0",
        "NCCL_IB_DISABLE": "0",
        "NCCL_NET": "IB",
        "NCCL_P2P_DISABLE": "1",
        "NCCL_SHM_DISABLE": "1",
        "NCCL_DEBUG": "INFO",
        "NCCL_DEBUG_SUBSYS": "NET,INIT",
        "TORCH_NCCL_BLOCKING_WAIT": "1",
    })

    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=1,
        world_size=2,
        timeout=timedelta(minutes=5),
    )
    print("[rank1] process group OK")

    torch.cuda.set_device(0)
    device = torch.device("cuda:0")
    print(f"[rank1] GPU={torch.cuda.get_device_name(0)}")

    _train(rank=1, device=device)

    dist.destroy_process_group()
    print("[rank1] DONE")


# ── общий train loop ──────────────────────────────────────────────────────────
def _train(rank: int, device):
    import torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim

    # allreduce smoke-test
    probe = torch.ones(1, device=device) * rank
    dist.all_reduce(probe, op=dist.ReduceOp.SUM)
    assert probe.item() == 1.0, f"allreduce mismatch: {probe.item()}"
    print(f"[rank{rank}] allreduce smoke-test PASSED")

    model = nn.Sequential(
        nn.Linear(10, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    ).to(device)

    ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[0])
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    loss_fn = nn.MSELoss()

    for epoch in range(5):
        x = torch.randn(32, 10, device=device)
        y = torch.randn(32, 10, device=device)
        optimizer.zero_grad()
        loss = loss_fn(ddp_model(x), y)
        loss.backward()
        optimizer.step()
        print(f"[rank{rank}] epoch={epoch} loss={loss.item():.4f}")


# ── DAG ───────────────────────────────────────────────────────────────────────
with DAG(
    dag_id="ddp_ib_xcom",
    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",
        executor_config={
            "pod_override": {
                "spec": {
                    "containers": [{
                        "name": "base",
                        "securityContext": {
                            "capabilities": {
                                "add": ["IPC_LOCK", "SYS_RESOURCE"]
                            }
                        }
                    }]
                }
            }
        },
    )

    rank1 = PythonOperator(
        task_id="train_rank_1",
        python_callable=run_rank_1,
        queue="gpu",
        executor_config={
            "pod_override": {
                "spec": {
                    "containers": [{
                        "name": "base",
                        "securityContext": {
                            "capabilities": {
                                "add": ["IPC_LOCK", "SYS_RESOURCE"]
                            }
                        }
                    }]
                }
            }
        },
    )

    # rank1 стартует сразу, но ждёт master_addr через XCom polling
    [rank0, rank1]