From 09b2c0041fd9ce535c4bd98fd78002d7eacd9eab Mon Sep 17 00:00:00 2001
From: George Stykalin <gstykalin@mail.ru>
Date: Wed, 15 Apr 2026 11:42:19 +0300
Subject: [PATCH] update dags

---
 dags/test-train-pytorch.py | 77 +++++++++++++++-----------------------
 1 file changed, 31 insertions(+), 46 deletions(-)

diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 34e7d2c..1263714 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -8,47 +8,43 @@ MASTER_PORT = "29500"
 default_args = {
     "owner": "airflow",
     "retries": 0,
-    "execution_timeout": timedelta(minutes=30),
+    "execution_timeout": timedelta(minutes=20),
 }
 
 
 # -------------------------
 # RANK 0
 # -------------------------
-def run_rank_0(**context):
-    import os, socket, torch, torch.distributed as dist
+def rank0(**context):
+    import os, socket, torch
+    import torch.distributed as dist
 
-    hostname = socket.gethostname()
-    master_addr = f"{hostname}.airflow-worker-gpu"
+    addr = socket.gethostname() + ".airflow-worker-gpu"
 
-    # share master with rank1
-    context["ti"].xcom_push(key="master_addr", value=master_addr)
+    context["ti"].xcom_push(key="master_addr", value=addr)
 
     os.environ.update({
-        "MASTER_ADDR": master_addr,
+        "MASTER_ADDR": addr,
         "MASTER_PORT": MASTER_PORT,
         "WORLD_SIZE": "2",
         "RANK": "0",
 
-        # NCCL (minimal stable IB config)
         "NCCL_DEBUG": "INFO",
         "NCCL_SOCKET_IFNAME": "eth0",
 
-        # IB stability (IMPORTANT)
+        # IB
+        "NCCL_IB_DISABLE": "0",
         "NCCL_IB_GID_INDEX": "0",
-        "NCCL_IB_USE_INLINE": "0",
-        "NCCL_IB_TIMEOUT": "22",
-        "NCCL_IB_RETRY_CNT": "7",
     })
 
-    dist.init_process_group("nccl", init_method="env://")
+    dist.init_process_group("nccl")
 
     torch.cuda.set_device(0)
 
-    x = torch.ones(1, device="cuda")
+    x = torch.ones(1, device="cuda") * 1
     dist.all_reduce(x)
 
-    print(f"[rank0] all_reduce result = {x.item()}")
+    print(f"[rank0] result = {x.item()}")
 
     dist.destroy_process_group()
 
@@ -56,48 +52,38 @@ def run_rank_0(**context):
 # -------------------------
 # RANK 1
 # -------------------------
-def run_rank_1(**context):
-    import os, time, torch, torch.distributed as dist
+def rank1(**context):
+    import os, time, torch
+    import torch.distributed as dist
 
-    # wait for master
-    master_addr = None
+    addr = None
     for _ in range(30):
-        master_addr = context["ti"].xcom_pull(
-            task_ids="train_rank_0",
-            key="master_addr",
-        )
-        if master_addr:
+        addr = context["ti"].xcom_pull(task_ids="rank0", key="master_addr")
+        if addr:
             break
         time.sleep(2)
 
-    if not master_addr:
-        raise RuntimeError("master_addr not found")
-
     os.environ.update({
-        "MASTER_ADDR": master_addr,
+        "MASTER_ADDR": addr,
         "MASTER_PORT": MASTER_PORT,
         "WORLD_SIZE": "2",
         "RANK": "1",
 
-        # NCCL (same as rank0)
         "NCCL_DEBUG": "INFO",
         "NCCL_SOCKET_IFNAME": "eth0",
 
-        # IB stability
+        "NCCL_IB_DISABLE": "0",
         "NCCL_IB_GID_INDEX": "0",
-        "NCCL_IB_USE_INLINE": "0",
-        "NCCL_IB_TIMEOUT": "22",
-        "NCCL_IB_RETRY_CNT": "7",
     })
 
-    dist.init_process_group("nccl", init_method="env://")
+    dist.init_process_group("nccl")
 
     torch.cuda.set_device(0)
 
-    x = torch.ones(1, device="cuda")
+    x = torch.ones(1, device="cuda") * 2
     dist.all_reduce(x)
 
-    print(f"[rank1] all_reduce result = {x.item()}")
+    print(f"[rank1] result = {x.item()}")
 
     dist.destroy_process_group()
 
@@ -106,25 +92,24 @@ def run_rank_1(**context):
 # DAG
 # -------------------------
 with DAG(
-    dag_id="ib_ddp_final_test",
+    dag_id="ib_simple_test",
     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,
+    r0 = PythonOperator(
+        task_id="rank0",
+        python_callable=rank0,
         queue="gpu",
     )
 
-    rank1 = PythonOperator(
-        task_id="train_rank_1",
-        python_callable=run_rank_1,
+    r1 = PythonOperator(
+        task_id="rank1",
+        python_callable=rank1,
         queue="gpu",
     )
 
-    rank0 >> rank1
+    r0 >> r1