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

WORLD_SIZE = 2

# Адрес мастера — headless DNS пода воркера
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
MASTER_PORT = "29500"

default_args = {
    "owner": "airflow",
    "retries": 0,
    "execution_timeout": timedelta(hours=1),
}


def run_training_node_func(rank: int, world_size: int):
    import os
    import time
    import socket
    import torch
    import torch.distributed as dist
    import torch.nn as nn
    import torch.optim as optim
    from datetime import datetime

    print("=" * 80)
    print(f"RANK {rank}/{world_size}  HOST={socket.gethostname()}  {datetime.now()}")
    print("=" * 80)

    # ── статический барьер: ждём пока оба воркера подтянутся ──────────────────
    time.sleep(10)

    # ── переменные окружения для NCCL + InfiniBand ────────────────────────────
    os.environ.update(
        {
            "MASTER_ADDR": MASTER_ADDR,
            "MASTER_PORT": MASTER_PORT,
            "WORLD_SIZE": str(world_size),
            "RANK": str(rank),
            # InfiniBand: указываем IPoIB-интерфейс вместо eth0
            # Разрешаем NCCL использовать RDMA (IB verbs)
            "NCCL_IB_DISABLE": "0",
            # GPUDirect RDMA — если драйвер поддерживает
            "NCCL_P2P_DISABLE": "0",
            # Явно форсируем IB transport (опционально, NCCL сам выберет,
            # но полезно для отладки)
            "NCCL_NET": "IB",
            # Отладка: INFO покажет какой транспорт выбран,
            # поменяй на TRACE для полного вывода
            "NCCL_DEBUG": "INFO",
            "NCCL_DEBUG_SUBSYS": "NET,INIT",
            "TORCH_NCCL_BLOCKING_WAIT": "1",
        }
    )

    print(f"[{rank}] MASTER={MASTER_ADDR}:{MASTER_PORT}  IB_IF=ib0")

    # ── init process group ────────────────────────────────────────────────────
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=rank,
        world_size=world_size,
        timeout=timedelta(minutes=5),
    )
    print(f"[{rank}] dist.init_process_group OK")

    # ── GPU binding ───────────────────────────────────────────────────────────
    torch.cuda.set_device(0)
    device = torch.device("cuda:0")
    print(f"[{rank}] GPU={torch.cuda.get_device_name(0)}")

    # ── минимальная модель ────────────────────────────────────────────────────
    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()

    # ── allreduce smoke-test перед обучением ──────────────────────────────────
    probe = torch.ones(1).to(device) * rank
    dist.all_reduce(probe, op=dist.ReduceOp.SUM)
    expected = sum(range(world_size))
    assert probe.item() == expected, f"allreduce mismatch: got {probe.item()}"
    print(f"[{rank}] allreduce smoke-test PASSED (sum={probe.item()})")

    # ── train loop ────────────────────────────────────────────────────────────
    for epoch in range(5):
        x = torch.randn(32, 10).to(device)
        y = torch.randn(32, 10).to(device)

        optimizer.zero_grad()
        out = ddp_model(x)
        loss = loss_fn(out, y)
        loss.backward()
        optimizer.step()

        print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")

    dist.destroy_process_group()
    print(f"[{rank}] DONE")
    return {"rank": rank, "status": "ok"}


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

    tasks = []
    for r in range(WORLD_SIZE):
        t = PythonOperator(
            task_id=f"train_rank_{r}",
            python_callable=run_training_node_func,
            op_kwargs={"rank": r, "world_size": WORLD_SIZE},
            queue="gpu",
        )
        tasks.append(t)

    # запускаем параллельно — без зависимостей между рангами
    # (prep/cleanup убраны, это минимальный тест)