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 = { default_args = {
"owner": "airflow", "owner": "airflow",
"retries": 0, "retries": 0,
"execution_timeout": timedelta(minutes=30), "execution_timeout": timedelta(minutes=20),
} }
# ------------------------- # -------------------------
# RANK 0 # RANK 0
# ------------------------- # -------------------------
def run_rank_0(**context): def rank0(**context):
import os, socket, torch, torch.distributed as dist import os, socket, torch
import torch.distributed as dist
hostname = socket.gethostname() addr = socket.gethostname() + ".airflow-worker-gpu"
master_addr = f"{hostname}.airflow-worker-gpu"
# share master with rank1 context["ti"].xcom_push(key="master_addr", value=addr)
context["ti"].xcom_push(key="master_addr", value=master_addr)
os.environ.update({ os.environ.update({
"MASTER_ADDR": master_addr, "MASTER_ADDR": addr,
"MASTER_PORT": MASTER_PORT, "MASTER_PORT": MASTER_PORT,
"WORLD_SIZE": "2", "WORLD_SIZE": "2",
"RANK": "0", "RANK": "0",
# NCCL (minimal stable IB config)
"NCCL_DEBUG": "INFO", "NCCL_DEBUG": "INFO",
"NCCL_SOCKET_IFNAME": "eth0", "NCCL_SOCKET_IFNAME": "eth0",
# IB stability (IMPORTANT) # IB
"NCCL_IB_DISABLE": "0",
"NCCL_IB_GID_INDEX": "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) torch.cuda.set_device(0)
x = torch.ones(1, device="cuda") x = torch.ones(1, device="cuda") * 1
dist.all_reduce(x) dist.all_reduce(x)
print(f"[rank0] all_reduce result = {x.item()}") print(f"[rank0] result = {x.item()}")
dist.destroy_process_group() dist.destroy_process_group()
@@ -56,48 +52,38 @@ def run_rank_0(**context):
# ------------------------- # -------------------------
# RANK 1 # RANK 1
# ------------------------- # -------------------------
def run_rank_1(**context): def rank1(**context):
import os, time, torch, torch.distributed as dist import os, time, torch
import torch.distributed as dist
# wait for master addr = None
master_addr = None
for _ in range(30): for _ in range(30):
master_addr = context["ti"].xcom_pull( addr = context["ti"].xcom_pull(task_ids="rank0", key="master_addr")
task_ids="train_rank_0", if addr:
key="master_addr",
)
if master_addr:
break break
time.sleep(2) time.sleep(2)
if not master_addr:
raise RuntimeError("master_addr not found")
os.environ.update({ os.environ.update({
"MASTER_ADDR": master_addr, "MASTER_ADDR": addr,
"MASTER_PORT": MASTER_PORT, "MASTER_PORT": MASTER_PORT,
"WORLD_SIZE": "2", "WORLD_SIZE": "2",
"RANK": "1", "RANK": "1",
# NCCL (same as rank0)
"NCCL_DEBUG": "INFO", "NCCL_DEBUG": "INFO",
"NCCL_SOCKET_IFNAME": "eth0", "NCCL_SOCKET_IFNAME": "eth0",
# IB stability "NCCL_IB_DISABLE": "0",
"NCCL_IB_GID_INDEX": "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) torch.cuda.set_device(0)
x = torch.ones(1, device="cuda") x = torch.ones(1, device="cuda") * 2
dist.all_reduce(x) dist.all_reduce(x)
print(f"[rank1] all_reduce result = {x.item()}") print(f"[rank1] result = {x.item()}")
dist.destroy_process_group() dist.destroy_process_group()
@@ -106,25 +92,24 @@ def run_rank_1(**context):
# DAG # DAG
# ------------------------- # -------------------------
with DAG( with DAG(
dag_id="ib_ddp_final_test", dag_id="ib_simple_test",
start_date=pendulum.today("UTC").add(days=-1), start_date=pendulum.today("UTC").add(days=-1),
schedule=None, schedule=None,
catchup=False, catchup=False,
max_active_runs=1, max_active_runs=1,
max_active_tasks=2,
default_args=default_args, default_args=default_args,
) as dag: ) as dag:
rank0 = PythonOperator( r0 = PythonOperator(
task_id="train_rank_0", task_id="rank0",
python_callable=run_rank_0, python_callable=rank0,
queue="gpu", queue="gpu",
) )
rank1 = PythonOperator( r1 = PythonOperator(
task_id="train_rank_1", task_id="rank1",
python_callable=run_rank_1, python_callable=rank1,
queue="gpu", queue="gpu",
) )
rank0 >> rank1 r0 >> r1