update dags

dags/test-train-pytorch.py

@@ -5,102 +5,93 @@ 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=2),
+    "execution_timeout": timedelta(hours=1),
 }
 
 
-# -------------------------
+def run_training_node_func(rank: int, world_size: int):
-# SIMPLE PREP (NO VARIABLES)
-# -------------------------
-def prepare_training_func():
-    print("Starting distributed training (no sync state needed)")
-    return True
-
-
-# -------------------------
-# CORE TRAINING
-# -------------------------
-def run_training_node_func(rank, world_size):
     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
-    import socket
 
     print("=" * 80)
-    print(f"RANK {rank}/{world_size} START {datetime.now()}")
+    print(f"RANK {rank}/{world_size}  HOST={socket.gethostname()}  {datetime.now()}")
-    print(f"HOST: {socket.gethostname()}")
     print("=" * 80)
 
-    # -------------------------
+    # ── статический барьер: ждём пока оба воркера подтянутся ──────────────────
-    # FIXED BARRIER (important)
-    # -------------------------
-    print(f"[{rank}] sync barrier (static sleep)")
     time.sleep(10)
 
-    # -------------------------
+    # ── переменные окружения для NCCL + InfiniBand ────────────────────────────
-    # STATIC CONFIG (NO AIRFLOW VARIABLES)
+    os.environ.update(
-    # -------------------------
+        {
-    os.environ["MASTER_ADDR"] = MASTER_ADDR
+            "MASTER_ADDR": MASTER_ADDR,
-    os.environ["MASTER_PORT"] = MASTER_PORT
+            "MASTER_PORT": MASTER_PORT,
-    os.environ["WORLD_SIZE"] = str(world_size)
+            "WORLD_SIZE": str(world_size),
-    os.environ["RANK"] = str(rank)
+            "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",
+        }
+    )
 
-    os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
+    print(f"[{rank}] MASTER={MASTER_ADDR}:{MASTER_PORT}  IB_IF=ib0")
-    os.environ["NCCL_DEBUG"] = "INFO"
-    os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"
 
-    print(f"[{rank}] MASTER = {MASTER_ADDR}:{MASTER_PORT}")
+    # ── init process group ────────────────────────────────────────────────────
-
-    # -------------------------
-    # 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")
 
-    print(f"[{rank}] DDP INIT OK")
+    # ── GPU binding ───────────────────────────────────────────────────────────
-
-    # -------------------------
-    # GPU BINDING (CRITICAL FIX)
-    # -------------------------
     torch.cuda.set_device(0)
     device = torch.device("cuda:0")
-
     print(f"[{rank}] GPU={torch.cuda.get_device_name(0)}")
 
-    # -------------------------
+    # ── минимальная модель ────────────────────────────────────────────────────
-    # MODEL
-    # -------------------------
     model = nn.Sequential(
         nn.Linear(10, 128),
         nn.ReLU(),
         nn.Linear(128, 10),
     ).to(device)
 
-    ddp_model = nn.parallel.DistributedDataParallel(
+    ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[0])
-        model,
-        device_ids=[0],
-    )
-
     optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
     loss_fn = nn.MSELoss()
 
-    # -------------------------
+    # ── allreduce smoke-test перед обучением ──────────────────────────────────
-    # TRAIN LOOP
+    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)
@@ -114,42 +105,21 @@ def run_training_node_func(rank, world_size):
         print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")
 
     dist.destroy_process_group()
-
+    print(f"[{rank}] DONE")
     return {"rank": rank, "status": "ok"}
 
 
-# -------------------------
+# ── DAG ───────────────────────────────────────────────────────────────────────
-# CLEANUP
-# -------------------------
-def cleanup_func():
-    print("cleanup done")
-    return True
-
-
-def summary_func(**context):
-    print("training done")
-    return True
-
-
-# -------------------------
-# DAG
-# -------------------------
 with DAG(
-    dag_id="ddp_airflow_stable",
+    dag_id="ddp_ib_test",
     start_date=pendulum.today("UTC").add(days=-1),
     schedule=None,
     catchup=False,
     max_active_runs=1,
-    max_active_tasks=2,
+    max_active_tasks=WORLD_SIZE,
     default_args=default_args,
 ) as dag:
 
-    prep = PythonOperator(
-        task_id="prep",
-        python_callable=prepare_training_func,
-        queue="gpu",
-    )
-
     tasks = []
     for r in range(WORLD_SIZE):
         t = PythonOperator(
@@ -160,18 +130,5 @@ with DAG(
         )
         tasks.append(t)
 
-    cleanup = PythonOperator(
+    # запускаем параллельно — без зависимостей между рангами
-        task_id="cleanup",
+    # (prep/cleanup убраны, это минимальный тест)
-        python_callable=cleanup_func,
-        trigger_rule="all_done",
-        queue="gpu",
-    )
-
-    summary = PythonOperator(
-        task_id="summary",
-        python_callable=summary_func,
-        trigger_rule="all_done",
-        queue="gpu",
-    )
-
-    prep >> tasks >> cleanup >> summary