178 lines
4.1 KiB
Python
178 lines
4.1 KiB
Python
from airflow import DAG
|
|
from airflow.operators.python import PythonOperator
|
|
import pendulum
|
|
from datetime import timedelta
|
|
|
|
WORLD_SIZE = 2
|
|
|
|
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
|
|
MASTER_PORT = "29500"
|
|
|
|
default_args = {
|
|
"owner": "airflow",
|
|
"retries": 0,
|
|
"execution_timeout": timedelta(hours=2),
|
|
}
|
|
|
|
|
|
# -------------------------
|
|
# SIMPLE PREP (NO VARIABLES)
|
|
# -------------------------
|
|
def prepare_training_func():
|
|
print("Starting distributed training (no sync state needed)")
|
|
return True
|
|
|
|
|
|
# -------------------------
|
|
# CORE TRAINING
|
|
# -------------------------
|
|
def run_training_node_func(rank, world_size):
|
|
import os
|
|
import time
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from datetime import datetime
|
|
import socket
|
|
|
|
print("=" * 80)
|
|
print(f"RANK {rank}/{world_size} START {datetime.now()}")
|
|
print(f"HOST: {socket.gethostname()}")
|
|
print("=" * 80)
|
|
|
|
# -------------------------
|
|
# FIXED BARRIER (important)
|
|
# -------------------------
|
|
print(f"[{rank}] sync barrier (static sleep)")
|
|
time.sleep(10)
|
|
|
|
# -------------------------
|
|
# STATIC CONFIG (NO AIRFLOW VARIABLES)
|
|
# -------------------------
|
|
os.environ["MASTER_ADDR"] = MASTER_ADDR
|
|
os.environ["MASTER_PORT"] = MASTER_PORT
|
|
os.environ["WORLD_SIZE"] = str(world_size)
|
|
os.environ["RANK"] = str(rank)
|
|
|
|
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
|
|
os.environ["NCCL_DEBUG"] = "INFO"
|
|
os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"
|
|
|
|
print(f"[{rank}] MASTER = {MASTER_ADDR}:{MASTER_PORT}")
|
|
|
|
# -------------------------
|
|
# INIT PROCESS GROUP
|
|
# -------------------------
|
|
dist.init_process_group(
|
|
backend="nccl",
|
|
init_method="env://",
|
|
rank=rank,
|
|
world_size=world_size,
|
|
)
|
|
|
|
print(f"[{rank}] DDP INIT OK")
|
|
|
|
# -------------------------
|
|
# GPU BINDING (CRITICAL FIX)
|
|
# -------------------------
|
|
torch.cuda.set_device(0)
|
|
device = torch.device("cuda:0")
|
|
|
|
print(f"[{rank}] GPU = {torch.cuda.get_device_name(0)}")
|
|
|
|
# -------------------------
|
|
# MODEL
|
|
# -------------------------
|
|
model = nn.Sequential(
|
|
nn.Linear(10, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, 10),
|
|
).to(device)
|
|
|
|
ddp_model = nn.parallel.DistributedDataParallel(
|
|
model,
|
|
device_ids=[0],
|
|
)
|
|
|
|
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
|
|
loss_fn = nn.MSELoss()
|
|
|
|
# -------------------------
|
|
# TRAIN LOOP
|
|
# -------------------------
|
|
for epoch in range(5):
|
|
x = torch.randn(32, 10).to(device)
|
|
y = torch.randn(32, 10).to(device)
|
|
|
|
optimizer.zero_grad()
|
|
out = ddp_model(x)
|
|
loss = loss_fn(out, y)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
return {"rank": rank, "status": "ok"}
|
|
|
|
|
|
# -------------------------
|
|
# CLEANUP
|
|
# -------------------------
|
|
def cleanup_func():
|
|
print("cleanup done")
|
|
return True
|
|
|
|
|
|
def summary_func(**context):
|
|
print("training done")
|
|
return True
|
|
|
|
|
|
# -------------------------
|
|
# DAG
|
|
# -------------------------
|
|
with DAG(
|
|
dag_id="ddp_airflow_stable",
|
|
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:
|
|
|
|
prep = PythonOperator(
|
|
task_id="prep",
|
|
python_callable=prepare_training_func,
|
|
queue="gpu",
|
|
)
|
|
|
|
tasks = []
|
|
for r in range(WORLD_SIZE):
|
|
t = PythonOperator(
|
|
task_id=f"train_rank_{r}",
|
|
python_callable=run_training_node_func,
|
|
op_kwargs={"rank": r, "world_size": WORLD_SIZE},
|
|
queue="gpu",
|
|
)
|
|
tasks.append(t)
|
|
|
|
cleanup = PythonOperator(
|
|
task_id="cleanup",
|
|
python_callable=cleanup_func,
|
|
trigger_rule="all_done",
|
|
queue="gpu",
|
|
)
|
|
|
|
summary = PythonOperator(
|
|
task_id="summary",
|
|
python_callable=summary_func,
|
|
trigger_rule="all_done",
|
|
queue="gpu",
|
|
)
|
|
|
|
prep >> tasks >> cleanup >> summary
|