gstykalin/test-dags
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update dags

Browse Source
This commit is contained in:
George Stykalin
2026-04-15 11:42:19 +03:00
parent 6c0105d85e
commit 09b2c0041f
1 changed files with 31 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
dags/test-train-pytorch.py
View File

@@ -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