update dags

dags/test-train-pytorch.py

@@ -1,79 +1,142 @@
 from airflow import DAG
-from airflow.operators.python import PythonOperator
+from airflow.providers.standard.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),
+    "execution_timeout": timedelta(minutes=30),
 }
 
 
-def run_training_node_func(rank: int, world_size: int):
+# ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
-    import os
+def run_rank_0(**context):
-    import time
+    import os, resource, socket, torch
-    import socket
+    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
-    from datetime import datetime
 
-    print("=" * 80)
+    # allreduce smoke-test
-    print(f"RANK {rank}/{world_size}  HOST={socket.gethostname()}  {datetime.now()}")
+    probe = torch.ones(1, device=device) * rank
-    print("=" * 80)
+    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")
 
-    # ── статический барьер: ждём пока оба воркера подтянутся ──────────────────
-    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(),
@@ -84,51 +147,66 @@ def run_training_node_func(rank: int, world_size: int):
     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)
+        x = torch.randn(32, 10, device=device)
-        y = torch.randn(32, 10).to(device)
+        y = torch.randn(32, 10, device=device)
-
         optimizer.zero_grad()
-        out = ddp_model(x)
+        loss = loss_fn(ddp_model(x), y)
-        loss = loss_fn(out, y)
         loss.backward()
         optimizer.step()
-
+        print(f"[rank{rank}] epoch={epoch} loss={loss.item():.4f}")
-        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",
+    dag_id="ddp_ib_xcom",
     start_date=pendulum.today("UTC").add(days=-1),
     schedule=None,
     catchup=False,
     max_active_runs=1,
-    max_active_tasks=WORLD_SIZE,
+    max_active_tasks=2,
     default_args=default_args,
 ) as dag:
 
-    tasks = []
+    rank0 = PythonOperator(
-    for r in range(WORLD_SIZE):
+        task_id="train_rank_0",
-        t = PythonOperator(
+        python_callable=run_rank_0,
-            task_id=f"train_rank_{r}",
-            python_callable=run_training_node_func,
-            op_kwargs={"rank": r, "world_size": WORLD_SIZE},
         queue="gpu",
+        executor_config={
+            "pod_override": {
+                "spec": {
+                    "containers": [{
+                        "name": "base",
+                        "securityContext": {
+                            "capabilities": {
+                                "add": ["IPC_LOCK", "SYS_RESOURCE"]
+                            }
+                        }
+                    }]
+                }
+            }
+        },
     )
-        tasks.append(t)
 
-    # запускаем параллельно — без зависимостей между рангами
+    rank1 = PythonOperator(
-    # (prep/cleanup убраны, это минимальный тест)
+        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]