update dags
This commit is contained in:
@@ -1,295 +1,177 @@
|
|||||||
from airflow import DAG
|
from airflow import DAG
|
||||||
from airflow.operators.python import PythonOperator
|
from airflow.operators.python import PythonOperator
|
||||||
from airflow.models import Variable
|
|
||||||
import pendulum
|
import pendulum
|
||||||
from datetime import timedelta
|
from datetime import timedelta
|
||||||
import json
|
|
||||||
|
|
||||||
# --- CONFIGURATION ---
|
|
||||||
WORLD_SIZE = 2
|
WORLD_SIZE = 2
|
||||||
|
|
||||||
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
|
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
|
||||||
MASTER_PORT = "29500"
|
MASTER_PORT = "29500"
|
||||||
NCCL_TIMEOUT = 1800
|
|
||||||
|
|
||||||
default_args = {
|
default_args = {
|
||||||
'owner': 'airflow',
|
"owner": "airflow",
|
||||||
'retries': 1,
|
"retries": 0,
|
||||||
'retry_delay': timedelta(minutes=2),
|
"execution_timeout": timedelta(hours=2),
|
||||||
'execution_timeout': timedelta(hours=2),
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# -------------------------
|
||||||
|
# SIMPLE PREP (NO VARIABLES)
|
||||||
|
# -------------------------
|
||||||
def prepare_training_func():
|
def prepare_training_func():
|
||||||
"""Initialize shared state for synchronization"""
|
print("Starting distributed training (no sync state needed)")
|
||||||
sync_state = {
|
return True
|
||||||
'ready_workers': [],
|
|
||||||
'training_started': False,
|
|
||||||
'start_time': None
|
|
||||||
}
|
|
||||||
Variable.set('ddp_sync_state', json.dumps(sync_state))
|
|
||||||
print("Training preparation complete. Sync state initialized.")
|
|
||||||
return True
|
|
||||||
|
|
||||||
|
|
||||||
|
# -------------------------
|
||||||
|
# CORE TRAINING
|
||||||
|
# -------------------------
|
||||||
def run_training_node_func(rank, world_size):
|
def run_training_node_func(rank, world_size):
|
||||||
"""Execute distributed training with synchronization"""
|
import os
|
||||||
import os
|
import time
|
||||||
import socket
|
import torch
|
||||||
import torch
|
import torch.distributed as dist
|
||||||
import torch.distributed as dist
|
import torch.nn as nn
|
||||||
import torch.nn as nn
|
import torch.optim as optim
|
||||||
import torch.optim as optim
|
from datetime import datetime
|
||||||
from datetime import datetime
|
import socket
|
||||||
import time
|
|
||||||
|
|
||||||
print(f"{'='*60}")
|
print("=" * 80)
|
||||||
print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}")
|
print(f"RANK {rank}/{world_size} START {datetime.now()}")
|
||||||
print(f"Hostname: {socket.gethostname()}")
|
print(f"HOST: {socket.gethostname()}")
|
||||||
print(f"{'='*60}")
|
print("=" * 80)
|
||||||
|
|
||||||
# STEP 1: Signal that this worker is ready
|
# -------------------------
|
||||||
print(f"[{rank}] Signaling ready state...")
|
# FIXED BARRIER (important)
|
||||||
max_wait = 300 # 5 minutes
|
# -------------------------
|
||||||
start_wait = time.time()
|
print(f"[{rank}] sync barrier (static sleep)")
|
||||||
|
time.sleep(10)
|
||||||
|
|
||||||
while time.time() - start_wait < max_wait:
|
# -------------------------
|
||||||
try:
|
# STATIC CONFIG (NO AIRFLOW VARIABLES)
|
||||||
sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
|
# -------------------------
|
||||||
|
os.environ["MASTER_ADDR"] = MASTER_ADDR
|
||||||
|
os.environ["MASTER_PORT"] = MASTER_PORT
|
||||||
|
os.environ["WORLD_SIZE"] = str(world_size)
|
||||||
|
os.environ["RANK"] = str(rank)
|
||||||
|
|
||||||
if rank not in sync_state.get('ready_workers', []):
|
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
|
||||||
sync_state.setdefault('ready_workers', []).append(rank)
|
os.environ["NCCL_DEBUG"] = "INFO"
|
||||||
Variable.set('ddp_sync_state', json.dumps(sync_state))
|
os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"
|
||||||
print(f"[{rank}] Marked as ready. Ready workers: {sync_state['ready_workers']}")
|
|
||||||
|
|
||||||
# STEP 2: Wait for all workers to be ready
|
print(f"[{rank}] MASTER = {MASTER_ADDR}:{MASTER_PORT}")
|
||||||
if len(sync_state.get('ready_workers', [])) == world_size:
|
|
||||||
print(f"[{rank}] All {world_size} workers are ready! Proceeding to training...")
|
|
||||||
break
|
|
||||||
|
|
||||||
print(f"[{rank}] Waiting for other workers... ({len(sync_state.get('ready_workers', []))}/{world_size} ready)")
|
# -------------------------
|
||||||
time.sleep(2)
|
# INIT PROCESS GROUP
|
||||||
|
# -------------------------
|
||||||
|
dist.init_process_group(
|
||||||
|
backend="nccl",
|
||||||
|
init_method="env://",
|
||||||
|
rank=rank,
|
||||||
|
world_size=world_size,
|
||||||
|
)
|
||||||
|
|
||||||
except Exception as e:
|
print(f"[{rank}] DDP INIT OK")
|
||||||
print(f"[{rank}] Error during sync: {e}")
|
|
||||||
time.sleep(2)
|
|
||||||
else:
|
|
||||||
raise RuntimeError(f"[{rank}] Timeout waiting for all workers to be ready!")
|
|
||||||
|
|
||||||
# Small delay to ensure all workers see the ready state
|
# -------------------------
|
||||||
time.sleep(3)
|
# GPU BINDING (CRITICAL FIX)
|
||||||
|
# -------------------------
|
||||||
|
torch.cuda.set_device(0)
|
||||||
|
device = torch.device("cuda:0")
|
||||||
|
|
||||||
# STEP 3: Configure distributed Environment
|
print(f"[{rank}] GPU = {torch.cuda.get_device_name(0)}")
|
||||||
|
|
||||||
import socket
|
# -------------------------
|
||||||
|
# MODEL
|
||||||
|
# -------------------------
|
||||||
|
model = nn.Sequential(
|
||||||
|
nn.Linear(10, 128),
|
||||||
|
nn.ReLU(),
|
||||||
|
nn.Linear(128, 10),
|
||||||
|
).to(device)
|
||||||
|
|
||||||
if rank == 0:
|
ddp_model = nn.parallel.DistributedDataParallel(
|
||||||
# Rank 0 becomes master
|
model,
|
||||||
master_addr = socket.gethostbyname(socket.gethostname())
|
device_ids=[0],
|
||||||
Variable.set("MASTER_ADDR_DYNAMIC", master_addr)
|
)
|
||||||
print(f"[{rank}] Acting as MASTER at {master_addr}")
|
|
||||||
else:
|
|
||||||
# Other ranks wait for master
|
|
||||||
print(f"[{rank}] Waiting for MASTER_ADDR...")
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
master_addr = Variable.get("MASTER_ADDR_DYNAMIC", default_var=None)
|
|
||||||
if master_addr:
|
|
||||||
break
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
time.sleep(1)
|
|
||||||
|
|
||||||
print(f"[{rank}] Found MASTER at {master_addr}")
|
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
|
||||||
|
loss_fn = nn.MSELoss()
|
||||||
|
|
||||||
os.environ['MASTER_ADDR'] = MASTER_ADDR
|
# -------------------------
|
||||||
os.environ['MASTER_PORT'] = MASTER_PORT
|
# TRAIN LOOP
|
||||||
os.environ['WORLD_SIZE'] = str(world_size)
|
# -------------------------
|
||||||
os.environ['RANK'] = str(rank)
|
for epoch in range(5):
|
||||||
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
|
x = torch.randn(32, 10).to(device)
|
||||||
os.environ['NCCL_PORT_RANGE'] = '30000-30100'
|
y = torch.randn(32, 10).to(device)
|
||||||
os.environ['NCCL_DEBUG'] = 'INFO'
|
|
||||||
os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
|
|
||||||
os.environ['NCCL_BLOCKING_WAIT'] = '1'
|
|
||||||
|
|
||||||
print(f"[{rank}] Environment configured:")
|
optimizer.zero_grad()
|
||||||
print(f" MASTER_ADDR: {MASTER_ADDR}")
|
out = ddp_model(x)
|
||||||
print(f" MASTER_PORT: {MASTER_PORT}")
|
loss = loss_fn(out, y)
|
||||||
print(f" RANK: {rank}")
|
loss.backward()
|
||||||
print(f" WORLD_SIZE: {world_size}")
|
optimizer.step()
|
||||||
|
|
||||||
# STEP 4: Initialize process group
|
print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")
|
||||||
print(f"[{rank}] Initializing process group (backend=nccl)...")
|
|
||||||
try:
|
|
||||||
dist.init_process_group(
|
|
||||||
backend="nccl",
|
|
||||||
init_method="env://",
|
|
||||||
timeout=timedelta(minutes=10),
|
|
||||||
rank=rank,
|
|
||||||
world_size=world_size
|
|
||||||
)
|
|
||||||
print(f"[{rank}] ✓ Successfully joined distributed group!")
|
|
||||||
print(f" Process group size: {dist.get_world_size()}")
|
|
||||||
print(f" My rank: {dist.get_rank()}")
|
|
||||||
except Exception as e:
|
|
||||||
print(f"[{rank}] ✗ Failed to initialize process group!")
|
|
||||||
print(f" Error: {str(e)}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
# STEP 5: Setup GPU
|
dist.destroy_process_group()
|
||||||
device = torch.device("cuda:0")
|
|
||||||
torch.cuda.set_device(device)
|
|
||||||
|
|
||||||
print(f"[{rank}] GPU Configuration:")
|
return {"rank": rank, "status": "ok"}
|
||||||
print(f" Device: {torch.cuda.get_device_name(0)}")
|
|
||||||
print(f" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
|
|
||||||
|
|
||||||
# STEP 6: Define model
|
|
||||||
model = nn.Sequential(
|
|
||||||
nn.Linear(10, 128),
|
|
||||||
nn.ReLU(),
|
|
||||||
nn.Linear(128, 128),
|
|
||||||
nn.ReLU(),
|
|
||||||
nn.Linear(128, 10)
|
|
||||||
).to(device)
|
|
||||||
|
|
||||||
ddp_model = nn.parallel.DistributedDataParallel(
|
|
||||||
model,
|
|
||||||
device_ids=[0],
|
|
||||||
output_device=0
|
|
||||||
)
|
|
||||||
|
|
||||||
criterion = nn.MSELoss()
|
|
||||||
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001, momentum=0.9)
|
|
||||||
|
|
||||||
print(f"[{rank}] Model initialized with {sum(p.numel() for p in model.parameters())} parameters")
|
|
||||||
|
|
||||||
# STEP 7: Training loop
|
|
||||||
print(f"\n[{rank}] {'='*60}")
|
|
||||||
print(f"[{rank}] Starting Training Loop")
|
|
||||||
print(f"[{rank}] {'='*60}")
|
|
||||||
|
|
||||||
num_epochs = 10
|
|
||||||
batch_size = 32
|
|
||||||
num_batches = 5
|
|
||||||
|
|
||||||
for epoch in range(num_epochs):
|
|
||||||
ddp_model.train()
|
|
||||||
epoch_loss = 0.0
|
|
||||||
|
|
||||||
for batch_idx in range(num_batches):
|
|
||||||
torch.manual_seed(epoch * num_batches + batch_idx)
|
|
||||||
inputs = torch.randn(batch_size, 10).to(device)
|
|
||||||
labels = torch.randn(batch_size, 10).to(device)
|
|
||||||
|
|
||||||
optimizer.zero_grad()
|
|
||||||
outputs = ddp_model(inputs)
|
|
||||||
loss = criterion(outputs, labels)
|
|
||||||
loss.backward()
|
|
||||||
optimizer.step()
|
|
||||||
|
|
||||||
epoch_loss += loss.item()
|
|
||||||
|
|
||||||
avg_loss = epoch_loss / num_batches
|
|
||||||
|
|
||||||
# Synchronize loss across ranks
|
|
||||||
loss_tensor = torch.tensor([avg_loss]).to(device)
|
|
||||||
dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
|
|
||||||
global_avg_loss = loss_tensor.item()
|
|
||||||
|
|
||||||
if rank == 0:
|
|
||||||
print(f"[{rank}] Epoch {epoch+1}/{num_epochs} | Global Avg Loss: {global_avg_loss:.6f}")
|
|
||||||
|
|
||||||
print(f"\n[{rank}] {'='*60}")
|
|
||||||
print(f"[{rank}] Training Complete!")
|
|
||||||
print(f"[{rank}] {'='*60}")
|
|
||||||
|
|
||||||
# STEP 8: Cleanup
|
|
||||||
dist.destroy_process_group()
|
|
||||||
print(f"[{rank}] Process group destroyed. Finished at {datetime.now()}")
|
|
||||||
|
|
||||||
return {
|
|
||||||
'rank': rank,
|
|
||||||
'final_loss': global_avg_loss if rank == 0 else avg_loss,
|
|
||||||
'epochs_completed': num_epochs,
|
|
||||||
'status': 'success'
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def cleanup_sync_state_func():
|
# -------------------------
|
||||||
"""Clean up synchronization state"""
|
# CLEANUP
|
||||||
try:
|
# -------------------------
|
||||||
Variable.delete('ddp_sync_state')
|
def cleanup_func():
|
||||||
print("Sync state cleaned up")
|
print("cleanup done")
|
||||||
except:
|
return True
|
||||||
pass
|
|
||||||
return True
|
|
||||||
|
|
||||||
|
|
||||||
def training_summary_func(**context):
|
def summary_func(**context):
|
||||||
"""Aggregate and display training results"""
|
print("training done")
|
||||||
ti = context['ti']
|
return True
|
||||||
|
|
||||||
print(f"\n{'='*60}")
|
|
||||||
print(f"DISTRIBUTED TRAINING SUMMARY")
|
|
||||||
print(f"{'='*60}")
|
|
||||||
|
|
||||||
for rank in range(WORLD_SIZE):
|
|
||||||
try:
|
|
||||||
result = ti.xcom_pull(task_ids=f"train_rank_{rank}")
|
|
||||||
if result and result.get('status') == 'success':
|
|
||||||
print(f"Rank {result['rank']}: ✓ Completed {result['epochs_completed']} epochs")
|
|
||||||
print(f" Final loss: {result['final_loss']:.6f}")
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Warning: Could not get result for rank {rank}: {e}")
|
|
||||||
|
|
||||||
print(f"\n✓ Training completed!")
|
|
||||||
return {'status': 'success', 'workers': WORLD_SIZE}
|
|
||||||
|
|
||||||
|
|
||||||
|
# -------------------------
|
||||||
|
# DAG
|
||||||
|
# -------------------------
|
||||||
with DAG(
|
with DAG(
|
||||||
dag_id='pytorch_distributed_training_ddp_production',
|
dag_id="ddp_airflow_stable",
|
||||||
default_args=default_args,
|
start_date=pendulum.today("UTC").add(days=-1),
|
||||||
schedule=None,
|
schedule=None,
|
||||||
start_date=pendulum.today('UTC').add(days=-1),
|
catchup=False,
|
||||||
catchup=False,
|
max_active_runs=1,
|
||||||
tags=['gpu', 'ml', 'distributed'],
|
max_active_tasks=2,
|
||||||
max_active_runs=1,
|
default_args=default_args,
|
||||||
max_active_tasks=10,
|
|
||||||
) as dag:
|
) as dag:
|
||||||
|
|
||||||
# Preparation task
|
prep = PythonOperator(
|
||||||
prep = PythonOperator(
|
task_id="prep",
|
||||||
task_id='prepare_training',
|
python_callable=prepare_training_func,
|
||||||
python_callable=prepare_training_func,
|
queue="gpu",
|
||||||
queue='gpu'
|
)
|
||||||
)
|
|
||||||
|
|
||||||
# Training tasks - one per rank
|
tasks = []
|
||||||
training_tasks = []
|
for r in range(WORLD_SIZE):
|
||||||
for i in range(WORLD_SIZE):
|
t = PythonOperator(
|
||||||
task = PythonOperator(
|
task_id=f"train_rank_{r}",
|
||||||
task_id=f'train_rank_{i}',
|
python_callable=run_training_node_func,
|
||||||
python_callable=run_training_node_func,
|
op_kwargs={"rank": r, "world_size": WORLD_SIZE},
|
||||||
op_kwargs={'rank': i, 'world_size': WORLD_SIZE},
|
queue="gpu",
|
||||||
queue='gpu',
|
)
|
||||||
)
|
tasks.append(t)
|
||||||
training_tasks.append(task)
|
|
||||||
|
|
||||||
# Cleanup task
|
cleanup = PythonOperator(
|
||||||
cleanup = PythonOperator(
|
task_id="cleanup",
|
||||||
task_id='cleanup_sync_state',
|
python_callable=cleanup_func,
|
||||||
python_callable=cleanup_sync_state_func,
|
trigger_rule="all_done",
|
||||||
queue='gpu',
|
queue="gpu",
|
||||||
trigger_rule='all_done'
|
)
|
||||||
)
|
|
||||||
|
|
||||||
# Summary task
|
summary = PythonOperator(
|
||||||
summary = PythonOperator(
|
task_id="summary",
|
||||||
task_id='training_summary',
|
python_callable=summary_func,
|
||||||
python_callable=training_summary_func,
|
trigger_rule="all_done",
|
||||||
queue='gpu',
|
queue="gpu",
|
||||||
trigger_rule='all_done'
|
)
|
||||||
)
|
|
||||||
|
|
||||||
# Set dependencies
|
prep >> tasks >> cleanup >> summary
|
||||||
prep >> training_tasks >> cleanup >> summary
|
|
||||||
|
|||||||
Reference in New Issue
Block a user