change user to root for airflow images

dags/test-train-pytorch.py

@@ -3,7 +3,6 @@ from airflow.providers.standard.operators.python import PythonOperator
 import pendulum
 from datetime import timedelta
 
-WORLD_SIZE = 2
 MASTER_PORT = "29500"
 
 default_args = {
@@ -13,153 +12,101 @@ default_args = {
 }
 
 
-# ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
+# -------------------------
+# RANK 0
+# -------------------------
 def run_rank_0(**context):
-    import os, socket, torch
+    import os, socket, torch, torch.distributed as dist
-    import torch.distributed as dist
-    import torch.nn as nn
-    import torch.optim as optim
 
     hostname = socket.gethostname()
-    fqdn = f"{hostname}.airflow-worker-gpu"
+    master_addr = f"{hostname}.airflow-worker-gpu"
-    print(f"[rank0] hostname={fqdn}")
 
-    # Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться
+    # share master with rank1
-    context["ti"].xcom_push(key="master_addr", value=fqdn)
+    context["ti"].xcom_push(key="master_addr", value=master_addr)
 
     os.environ.update({
-        "MASTER_ADDR": fqdn,
+        "MASTER_ADDR": master_addr,
         "MASTER_PORT": MASTER_PORT,
         "WORLD_SIZE": "2",
-        "NCCL_IB_GID_INDEX": "0",
         "RANK": "0",
-        "NCCL_SOCKET_IFNAME": "eth0",   # bootstrap/rendezvous через eth
+
-        "NCCL_IB_DISABLE": "0",         # трафик через IB
+        # NCCL (minimal stable IB config)
-        "NCCL_NET": "IB",
-        "NCCL_IB_USE_INLINE": "1",      # без ibv_reg_mr pinned memory
-        "NCCL_BUFFSIZE": "1048576",
-        "NCCL_P2P_DISABLE": "1",        # разные ноды — p2p не нужен
-        "NCCL_SHM_DISABLE": "1",        # shm только для локальных рангов
         "NCCL_DEBUG": "INFO",
-        "NCCL_DEBUG_SUBSYS": "NET,INIT",
+        "NCCL_SOCKET_IFNAME": "eth0",
-        "TORCH_NCCL_BLOCKING_WAIT": "1",
+
+        # 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(
+    dist.init_process_group("nccl", init_method="env://")
-        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)
+    x = torch.ones(1, device="cuda")
+    dist.all_reduce(x)
+
+    print(f"[rank0] all_reduce result = {x.item()}")
 
     dist.destroy_process_group()
-    print("[rank0] DONE")
 
 
-# ── rank 1: ждёт XCom от rank0, потом коннектится ────────────────────────────
+# -------------------------
+# RANK 1
+# -------------------------
 def run_rank_1(**context):
-    import os, socket, time, torch
+    import os, time, torch, torch.distributed as dist
-    import torch.distributed as dist
-    import torch.nn as nn
-    import torch.optim as optim
 
-    # Polling XCom пока rank0 не запишет master_addr
+    # wait for master
     master_addr = None
-    for attempt in range(30):
+    for _ 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(2)
-        time.sleep(5)
 
     if not master_addr:
-        raise RuntimeError("rank1: timeout waiting for master_addr from rank0")
+        raise RuntimeError("master_addr not found")
-
-    print(f"[rank1] master_addr={master_addr}")
 
     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 (same as rank0)
-        "NCCL_NET": "IB",
-        "NCCL_IB_USE_INLINE": "1",      # без ibv_reg_mr pinned memory
-        "NCCL_BUFFSIZE": "1048576",
-        "NCCL_P2P_DISABLE": "1",
-        "NCCL_SHM_DISABLE": "1",
         "NCCL_DEBUG": "INFO",
-        "NCCL_DEBUG_SUBSYS": "NET,INIT",
+        "NCCL_SOCKET_IFNAME": "eth0",
-        "TORCH_NCCL_BLOCKING_WAIT": "1",
+
+        # 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(
+    dist.init_process_group("nccl", init_method="env://")
-        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)
+    x = torch.ones(1, device="cuda")
+    dist.all_reduce(x)
+
+    print(f"[rank1] all_reduce result = {x.item()}")
 
     dist.destroy_process_group()
-    print("[rank1] DONE")
 
 
-# ── общий train loop ──────────────────────────────────────────────────────────
+# -------------------------
-def _train(rank: int, device):
+# DAG
-    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",
+    dag_id="ib_ddp_final_test",
     start_date=pendulum.today("UTC").add(days=-1),
     schedule=None,
     catchup=False,
@@ -172,42 +119,12 @@ with DAG(
         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
-    [rank0, rank1]
-