diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 939d95e..94fa594 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -1,9 +1,9 @@ from airflow import DAG from airflow.operators.python import PythonOperator -from airflow.models import Variable import pendulum from datetime import timedelta -import json +import os +import time # --- CONFIGURATION --- WORLD_SIZE = 2 @@ -11,263 +11,210 @@ MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu" MASTER_PORT = "29500" NCCL_TIMEOUT = 1800 +BARRIER_FILE = f"/tmp/ddp_barrier_{WORLD_SIZE}" + default_args = { - 'owner': 'airflow', - 'retries': 1, - 'retry_delay': timedelta(minutes=2), - 'execution_timeout': timedelta(hours=2), + 'owner': 'airflow', + 'retries': 1, + 'retry_delay': timedelta(minutes=2), + 'execution_timeout': timedelta(hours=2), } def prepare_training_func(): - """Initialize shared state for synchronization""" - sync_state = { - '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 + # reset barrier + try: + if os.path.exists(BARRIER_FILE): + os.remove(BARRIER_FILE) + except: + pass + + print("Barrier initialized") + return True def run_training_node_func(rank, world_size): - """Execute distributed training with synchronization""" - import os - import socket - import torch - import torch.distributed as dist - import torch.nn as nn - import torch.optim as optim - from datetime import datetime - import time + import socket + import torch + import torch.distributed as dist + import torch.nn as nn + import torch.optim as optim + from datetime import datetime - print(f"{'='*60}") - print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}") - print(f"Hostname: {socket.gethostname()}") - print(f"{'='*60}") + print("=" * 60) + print(f"Rank {rank}/{world_size} START {datetime.now()}") + print(f"Hostname: {socket.gethostname()}") + print("=" * 60) - # STEP 1: Signal that this worker is ready - print(f"[{rank}] Signaling ready state...") - max_wait = 300 # 5 minutes - start_wait = time.time() + # ========================= + # STEP 1: SAFE BARRIER + # ========================= + print(f"[{rank}] entering barrier sync...") - while time.time() - start_wait < max_wait: - try: - sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}')) + with open(BARRIER_FILE, "a+") as f: + f.write(f"{rank}\n") - if rank not in sync_state.get('ready_workers', []): - sync_state.setdefault('ready_workers', []).append(rank) - Variable.set('ddp_sync_state', json.dumps(sync_state)) - print(f"[{rank}] Marked as ready. Ready workers: {sync_state['ready_workers']}") + # wait for all ranks + while True: + try: + with open(BARRIER_FILE, "r") as f: + ready = set(f.read().strip().splitlines()) - # STEP 2: Wait for all workers to be ready - 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}] ready workers: {ready}") - print(f"[{rank}] Waiting for other workers... ({len(sync_state.get('ready_workers', []))}/{world_size} ready)") - time.sleep(2) + if len(ready) == world_size: + print(f"[{rank}] ALL WORKERS READY") + break - except Exception as e: - print(f"[{rank}] Error during sync: {e}") - time.sleep(2) - else: - raise RuntimeError(f"[{rank}] Timeout waiting for all workers to be ready!") + except FileNotFoundError: + pass - # Small delay to ensure all workers see the ready state - time.sleep(3) + time.sleep(1) - # STEP 3: Configure distributed environment - 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['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) - os.environ['NCCL_BLOCKING_WAIT'] = '1' + time.sleep(2) # small stabilization delay - print(f"[{rank}] Environment configured:") - print(f" MASTER_ADDR: {MASTER_ADDR}") - print(f" MASTER_PORT: {MASTER_PORT}") - print(f" RANK: {rank}") - print(f" WORLD_SIZE: {world_size}") + # ========================= + # STEP 2: NCCL ENV + # ========================= + os.environ['MASTER_ADDR'] = MASTER_ADDR + os.environ['MASTER_PORT'] = MASTER_PORT + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['RANK'] = str(rank) - # STEP 4: Initialize process group - 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 + # IMPORTANT FIX (IB stability toggle) + os.environ["NCCL_IB_DISABLE"] = "1" # <- FIX #2 (can turn OFF later) + os.environ["NCCL_SOCKET_IFNAME"] = "eth0" - # STEP 5: Setup GPU - device = torch.device("cuda:0") - torch.cuda.set_device(device) + os.environ['NCCL_DEBUG'] = 'INFO' + os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) + os.environ['NCCL_BLOCKING_WAIT'] = '1' - print(f"[{rank}] GPU Configuration:") - print(f" Device: {torch.cuda.get_device_name(0)}") - print(f" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB") + print(f"[{rank}] env ready") - # 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) + # ========================= + # STEP 3: INIT PROCESS GROUP + # ========================= + dist.init_process_group( + backend="nccl", + init_method="env://", + timeout=timedelta(minutes=10), + rank=rank, + world_size=world_size + ) - ddp_model = nn.parallel.DistributedDataParallel( - model, - device_ids=[0], - output_device=0 - ) + print(f"[{rank}] process group OK") - criterion = nn.MSELoss() - optimizer = optim.SGD(ddp_model.parameters(), lr=0.001, momentum=0.9) + # ========================= + # STEP 4: GPU + # ========================= + device = torch.device("cuda:0") + torch.cuda.set_device(device) - print(f"[{rank}] Model initialized with {sum(p.numel() for p in model.parameters())} parameters") + model = nn.Sequential( + nn.Linear(10, 128), + nn.ReLU(), + nn.Linear(128, 10) + ).to(device) - # STEP 7: Training loop - print(f"\n[{rank}] {'='*60}") - print(f"[{rank}] Starting Training Loop") - print(f"[{rank}] {'='*60}") + ddp_model = nn.parallel.DistributedDataParallel( + model, + device_ids=[0], + output_device=0 + ) - num_epochs = 10 - batch_size = 32 - num_batches = 5 + criterion = nn.MSELoss() + optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) - for epoch in range(num_epochs): - ddp_model.train() - epoch_loss = 0.0 + # ========================= + # STEP 5: TRAIN LOOP + # ========================= + for epoch in range(5): + ddp_model.train() + loss_sum = 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) + for _ in range(5): + x = torch.randn(32, 10).to(device) + y = torch.randn(32, 10).to(device) - optimizer.zero_grad() - outputs = ddp_model(inputs) - loss = criterion(outputs, labels) - loss.backward() - optimizer.step() + optimizer.zero_grad() + out = ddp_model(x) + loss = criterion(out, y) + loss.backward() + optimizer.step() - epoch_loss += loss.item() + loss_sum += loss.item() - avg_loss = epoch_loss / num_batches + print(f"[{rank}] epoch {epoch} loss {loss_sum/5:.4f}") - # 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() + dist.destroy_process_group() - if rank == 0: - print(f"[{rank}] Epoch {epoch+1}/{num_epochs} | Global Avg Loss: {global_avg_loss:.6f}") + print(f"[{rank}] DONE") - 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' - } + return {"rank": rank, "status": "success"} def cleanup_sync_state_func(): - """Clean up synchronization state""" - try: - Variable.delete('ddp_sync_state') - print("Sync state cleaned up") - except: - pass - return True + try: + if os.path.exists(BARRIER_FILE): + os.remove(BARRIER_FILE) + except: + pass + print("Barrier cleaned") + return True def training_summary_func(**context): - """Aggregate and display training results""" - ti = context['ti'] + ti = context['ti'] - print(f"\n{'='*60}") - print(f"DISTRIBUTED TRAINING SUMMARY") - print(f"{'='*60}") + print("\n=== SUMMARY ===") - 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}") + for rank in range(WORLD_SIZE): + res = ti.xcom_pull(task_ids=f"train_rank_{rank}") + print(f"Rank {rank}: {res}") - print(f"\nāœ“ Training completed!") - return {'status': 'success', 'workers': WORLD_SIZE} + return {"status": "done"} +# ========================= +# DAG +# ========================= with DAG( - dag_id='pytorch_distributed_training_ddp_production', - default_args=default_args, - schedule=None, - start_date=pendulum.today('UTC').add(days=-1), - catchup=False, - tags=['gpu', 'ml', 'distributed'], - max_active_runs=1, - max_active_tasks=10, + dag_id='pytorch_ddp_airflow_fixed', + default_args=default_args, + schedule=None, + start_date=pendulum.today('UTC').add(days=-1), + catchup=False, + max_active_runs=1, + tags=['gpu', 'ddp', 'fixed'], ) as dag: - # Preparation task - prep = PythonOperator( - task_id='prepare_training', - python_callable=prepare_training_func, - queue='gpu' - ) + prep = PythonOperator( + task_id='prepare_training', + python_callable=prepare_training_func, + queue='gpu', + ) - # Training tasks - one per rank - training_tasks = [] - for i in range(WORLD_SIZE): - task = PythonOperator( - task_id=f'train_rank_{i}', - python_callable=run_training_node_func, - op_kwargs={'rank': i, 'world_size': WORLD_SIZE}, - queue='gpu', - ) - training_tasks.append(task) + training_tasks = [] + for i in range(WORLD_SIZE): + t = PythonOperator( + task_id=f'train_rank_{i}', + python_callable=run_training_node_func, + op_kwargs={'rank': i, 'world_size': WORLD_SIZE}, + queue='gpu', + ) + training_tasks.append(t) - # Cleanup task - cleanup = PythonOperator( - task_id='cleanup_sync_state', - python_callable=cleanup_sync_state_func, - queue='gpu', - trigger_rule='all_done' - ) + cleanup = PythonOperator( + task_id='cleanup', + python_callable=cleanup_sync_state_func, + queue='gpu', + trigger_rule='all_done', + ) - # Summary task - summary = PythonOperator( - task_id='training_summary', - python_callable=training_summary_func, - queue='gpu', - trigger_rule='all_done' - ) - - # Set dependencies - prep >> training_tasks >> cleanup >> summary + summary = PythonOperator( + task_id='summary', + python_callable=training_summary_func, + trigger_rule='all_done', + ) + prep >> training_tasks >> cleanup >> summary