116 lines
2.3 KiB
Python
116 lines
2.3 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=20),
|
|
}
|
|
|
|
|
|
# -------------------------
|
|
# RANK 0
|
|
# -------------------------
|
|
def rank0(**context):
|
|
import os, socket, torch
|
|
import torch.distributed as dist
|
|
|
|
addr = socket.gethostname() + ".airflow-worker-gpu"
|
|
|
|
context["ti"].xcom_push(key="master_addr", value=addr)
|
|
|
|
os.environ.update({
|
|
"MASTER_ADDR": addr,
|
|
"MASTER_PORT": MASTER_PORT,
|
|
"WORLD_SIZE": "2",
|
|
"RANK": "0",
|
|
|
|
"NCCL_DEBUG": "INFO",
|
|
"NCCL_SOCKET_IFNAME": "eth0",
|
|
|
|
# IB
|
|
"NCCL_IB_DISABLE": "0",
|
|
"NCCL_IB_GID_INDEX": "0",
|
|
})
|
|
|
|
dist.init_process_group("nccl")
|
|
|
|
torch.cuda.set_device(0)
|
|
|
|
x = torch.ones(1, device="cuda") * 1
|
|
dist.all_reduce(x)
|
|
|
|
print(f"[rank0] result = {x.item()}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
# -------------------------
|
|
# RANK 1
|
|
# -------------------------
|
|
def rank1(**context):
|
|
import os, time, torch
|
|
import torch.distributed as dist
|
|
|
|
addr = None
|
|
for _ in range(30):
|
|
addr = context["ti"].xcom_pull(task_ids="rank0", key="master_addr")
|
|
if addr:
|
|
break
|
|
time.sleep(2)
|
|
|
|
os.environ.update({
|
|
"MASTER_ADDR": addr,
|
|
"MASTER_PORT": MASTER_PORT,
|
|
"WORLD_SIZE": "2",
|
|
"RANK": "1",
|
|
|
|
"NCCL_DEBUG": "INFO",
|
|
"NCCL_SOCKET_IFNAME": "eth0",
|
|
|
|
"NCCL_IB_DISABLE": "0",
|
|
"NCCL_IB_GID_INDEX": "0",
|
|
})
|
|
|
|
dist.init_process_group("nccl")
|
|
|
|
torch.cuda.set_device(0)
|
|
|
|
x = torch.ones(1, device="cuda") * 2
|
|
dist.all_reduce(x)
|
|
|
|
print(f"[rank1] result = {x.item()}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
# -------------------------
|
|
# DAG
|
|
# -------------------------
|
|
with DAG(
|
|
dag_id="ib_simple_test",
|
|
start_date=pendulum.today("UTC").add(days=-1),
|
|
schedule=None,
|
|
catchup=False,
|
|
max_active_runs=1,
|
|
default_args=default_args,
|
|
) as dag:
|
|
|
|
r0 = PythonOperator(
|
|
task_id="rank0",
|
|
python_callable=rank0,
|
|
queue="gpu",
|
|
)
|
|
|
|
r1 = PythonOperator(
|
|
task_id="rank1",
|
|
python_callable=rank1,
|
|
queue="gpu",
|
|
)
|
|
|
|
r0 >> r1
|