131 lines
2.9 KiB
Python
131 lines
2.9 KiB
Python
from airflow import DAG
|
|
from airflow.providers.standard.operators.python import PythonOperator
|
|
import pendulum
|
|
from datetime import timedelta
|
|
|
|
MASTER_PORT = "29500"
|
|
|
|
default_args = {
|
|
"owner": "airflow",
|
|
"retries": 0,
|
|
"execution_timeout": timedelta(minutes=30),
|
|
}
|
|
|
|
|
|
# -------------------------
|
|
# RANK 0
|
|
# -------------------------
|
|
def run_rank_0(**context):
|
|
import os, socket, torch, torch.distributed as dist
|
|
|
|
hostname = socket.gethostname()
|
|
master_addr = f"{hostname}.airflow-worker-gpu"
|
|
|
|
# share master with rank1
|
|
context["ti"].xcom_push(key="master_addr", value=master_addr)
|
|
|
|
os.environ.update({
|
|
"MASTER_ADDR": master_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)
|
|
"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://")
|
|
|
|
torch.cuda.set_device(0)
|
|
|
|
x = torch.ones(1, device="cuda")
|
|
dist.all_reduce(x)
|
|
|
|
print(f"[rank0] all_reduce result = {x.item()}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
# -------------------------
|
|
# RANK 1
|
|
# -------------------------
|
|
def run_rank_1(**context):
|
|
import os, time, torch, torch.distributed as dist
|
|
|
|
# wait for master
|
|
master_addr = None
|
|
for _ in range(30):
|
|
master_addr = context["ti"].xcom_pull(
|
|
task_ids="train_rank_0",
|
|
key="master_addr",
|
|
)
|
|
if master_addr:
|
|
break
|
|
time.sleep(2)
|
|
|
|
if not master_addr:
|
|
raise RuntimeError("master_addr not found")
|
|
|
|
os.environ.update({
|
|
"MASTER_ADDR": master_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_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://")
|
|
|
|
torch.cuda.set_device(0)
|
|
|
|
x = torch.ones(1, device="cuda")
|
|
dist.all_reduce(x)
|
|
|
|
print(f"[rank1] all_reduce result = {x.item()}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
# -------------------------
|
|
# DAG
|
|
# -------------------------
|
|
with DAG(
|
|
dag_id="ib_ddp_final_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,
|
|
queue="gpu",
|
|
)
|
|
|
|
rank1 = PythonOperator(
|
|
task_id="train_rank_1",
|
|
python_callable=run_rank_1,
|
|
queue="gpu",
|
|
)
|
|
|
|
rank0 >> rank1
|