init commit

dags/connection-check-default.py

@@ -0,0 +1,43 @@
+from airflow import DAG
+from airflow.decorators import task
+import pendulum
+
+TARGET_HOST = "airflow-worker-default-0.airflow-worker-default"
+
+with DAG(
+    dag_id='debug_network_check',
+    # UPDATED: New argument name
+    schedule=None,
+    # UPDATED: Explicit pendulum date
+    start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
+    catchup=False
+) as dag:
+
+    @task(queue='default')
+    def check_network():
+        import subprocess
+        import socket
+        
+        print(f"I am running on hostname: {socket.gethostname()}")
+        print(f"Attempting to resolve and ping: {TARGET_HOST}")
+        
+        # 1. DNS Check
+        try:
+            ip = socket.gethostbyname(TARGET_HOST)
+            print(f"DNS Resolution Success: {TARGET_HOST} -> {ip}")
+        except Exception as e:
+            print(f"DNS Resolution Failed: {e}")
+            raise
+
+        # 2. Ping Check
+        cmd = ["ping", "-c", "4", TARGET_HOST]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        
+        if result.returncode == 0:
+            print("Ping Success!")
+            print(result.stdout)
+        else:
+            print("Ping Failed (might be firewall, check if TCP port 29500 works instead)")
+            print(result.stderr)
+
+    check_network()

dags/connection-check.py

@@ -0,0 +1,43 @@
+from airflow import DAG
+from airflow.decorators import task
+import pendulum
+
+TARGET_HOST = "airflow-k8s-0140-worker-gpu-0.airflow-k8s-0140-worker-gpu"
+
+with DAG(
+    dag_id='debug_gpu_network_check',
+    # UPDATED: New argument name
+    schedule=None,
+    # UPDATED: Explicit pendulum date
+    start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
+    catchup=False
+) as dag:
+
+    @task(queue='gpu')
+    def check_network():
+        import subprocess
+        import socket
+        
+        print(f"I am running on hostname: {socket.gethostname()}")
+        print(f"Attempting to resolve and ping: {TARGET_HOST}")
+        
+        # 1. DNS Check
+        try:
+            ip = socket.gethostbyname(TARGET_HOST)
+            print(f"DNS Resolution Success: {TARGET_HOST} -> {ip}")
+        except Exception as e:
+            print(f"DNS Resolution Failed: {e}")
+            raise
+
+        # 2. Ping Check
+        cmd = ["ping", "-c", "4", TARGET_HOST]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        
+        if result.returncode == 0:
+            print("Ping Success!")
+            print(result.stdout)
+        else:
+            print("Ping Failed (might be firewall, check if TCP port 29500 works instead)")
+            print(result.stderr)
+
+    check_network()

dags/test-train-pytorch.py

@@ -0,0 +1,272 @@
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from airflow.models import Variable
+import pendulum
+from datetime import timedelta
+import json
+
+# --- CONFIGURATION ---
+WORLD_SIZE = 2
+MASTER_ADDR = "airflow-k8s-0140-worker-gpu-0.airflow-k8s-0140-worker-gpu"
+MASTER_PORT = "29500"
+NCCL_TIMEOUT = 1800
+
+default_args = {
+  '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
+
+
+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
+
+  print(f"{'='*60}")
+  print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}")
+  print(f"Hostname: {socket.gethostname()}")
+  print(f"{'='*60}")
+
+  # STEP 1: Signal that this worker is ready
+  print(f"[{rank}] Signaling ready state...")
+  max_wait = 300  # 5 minutes
+  start_wait = time.time()
+
+  while time.time() - start_wait < max_wait:
+      try:
+          sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
+
+          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']}")
+
+          # 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}] Waiting for other workers... ({len(sync_state.get('ready_workers', []))}/{world_size} ready)")
+          time.sleep(2)
+
+      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!")
+
+  # Small delay to ensure all workers see the ready state
+  time.sleep(3)
+
+  # 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'
+
+  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 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
+
+  # STEP 5: Setup GPU
+  device = torch.device("cuda:0")
+  torch.cuda.set_device(device)
+
+  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")
+
+  # 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"""
+  try:
+      Variable.delete('ddp_sync_state')
+      print("Sync state cleaned up")
+  except:
+      pass
+  return True
+
+
+def training_summary_func(**context):
+  """Aggregate and display training results"""
+  ti = context['ti']
+
+  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}
+
+
+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,
+) as dag:
+
+  # Preparation task
+  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)
+
+  # Cleanup task
+  cleanup = PythonOperator(
+      task_id='cleanup_sync_state',
+      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

dags/test-train-tensorflow.py

@@ -0,0 +1,373 @@
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from airflow.models import Variable
+import pendulum
+from datetime import timedelta
+import json
+
+# --- CONFIGURATION ---
+WORLD_SIZE = 2
+WORKER_PORT = 12345
+# Cluster spec will be built dynamically based on actual worker hostnames
+
+default_args = {
+    '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 = {
+        'worker_addresses': {},  # {task_index: 'hostname:port'}
+        'ready_workers': [],
+        'completed_workers': [],  # Track workers that finished training
+        'training_started': False,
+        'start_time': None
+    }
+    Variable.set('tf_sync_state', json.dumps(sync_state))
+    print("Training preparation complete. Sync state initialized.")
+    return True
+
+
+def run_training_node_func(task_index):
+    """Execute TensorFlow distributed training"""
+    import os
+    import socket
+    from datetime import datetime
+    import time
+
+    hostname = socket.gethostname()
+    my_address = f"{hostname}.airflow-worker-gpu:{WORKER_PORT}"
+
+    print(f"{'='*60}")
+    print(f"Worker {task_index}/{WORLD_SIZE} - Starting at {datetime.now()}")
+    print(f"Hostname: {hostname}")
+    print(f"My address: {my_address}")
+    print(f"{'='*60}")
+
+    # STEP 1: Register worker address and wait for all workers
+    print(f"[Worker-{task_index}] Registering address and waiting for all workers...")
+    max_wait = 300  # 5 minutes
+    start_wait = time.time()
+    cluster_spec = None
+
+    while time.time() - start_wait < max_wait:
+        try:
+            sync_state = json.loads(Variable.get('tf_sync_state', default_var='{}'))
+            worker_addresses = sync_state.setdefault('worker_addresses', {})
+
+            # Register my address
+            if str(task_index) not in worker_addresses:
+                worker_addresses[str(task_index)] = my_address
+                sync_state['worker_addresses'] = worker_addresses
+                Variable.set('tf_sync_state', json.dumps(sync_state))
+                print(f"[Worker-{task_index}] Registered address: {my_address}")
+
+            # Mark as ready
+            if task_index not in sync_state.get('ready_workers', []):
+                sync_state.setdefault('ready_workers', []).append(task_index)
+                Variable.set('tf_sync_state', json.dumps(sync_state))
+
+            # STEP 2: Wait for all workers to register
+            if len(worker_addresses) == WORLD_SIZE:
+                print(f"[Worker-{task_index}] All {WORLD_SIZE} workers registered!")
+                # Build cluster spec from registered addresses
+                worker_list = [worker_addresses[str(i)] for i in range(WORLD_SIZE)]
+                cluster_spec = {'worker': worker_list}
+                print(f"[Worker-{task_index}] Cluster spec: {cluster_spec}")
+                break
+
+            print(f"[Worker-{task_index}] Waiting for workers... ({len(worker_addresses)}/{WORLD_SIZE} registered)")
+            time.sleep(2)
+
+        except Exception as e:
+            print(f"[Worker-{task_index}] Error during sync: {e}")
+            time.sleep(2)
+    else:
+        raise RuntimeError(f"[Worker-{task_index}] Timeout waiting for all workers to register!")
+
+    if cluster_spec is None:
+        raise RuntimeError(f"[Worker-{task_index}] Failed to build cluster spec!")
+
+    # Small delay to ensure all workers see the complete cluster
+    time.sleep(3)
+
+    # STEP 3: Configure TF_CONFIG BEFORE importing TensorFlow (CRITICAL!)
+    tf_config = {
+        'cluster': cluster_spec,
+        'task': {
+            'type': 'worker',
+            'index': task_index
+        }
+    }
+
+    os.environ['TF_CONFIG'] = json.dumps(tf_config)
+
+    print(f"[Worker-{task_index}] TF_CONFIG set:")
+    print(f"  Cluster: {cluster_spec}")
+    print(f"  Task type: worker")
+    print(f"  Task index: {task_index}")
+
+    # NOW import TensorFlow after TF_CONFIG is set
+    import tensorflow as tf
+    print(f"TensorFlow version: {tf.__version__}")
+
+    # STEP 4: Initialize distribution strategy IMMEDIATELY (CRITICAL!)
+    # Must happen before any other TensorFlow operations
+    print(f"[Worker-{task_index}] Initializing MultiWorkerMirroredStrategy...")
+
+    # Configure communication options for NCCL
+    communication_options = tf.distribute.experimental.CommunicationOptions(
+        implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
+    )
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy(
+        communication_options=communication_options
+    )
+
+    print(f"[Worker-{task_index}] ✓ Strategy initialized!")
+    print(f"  Number of devices in sync: {strategy.num_replicas_in_sync}")
+
+    # STEP 5: GPU info (after strategy is initialized)
+    gpus = tf.config.list_physical_devices('GPU')
+    print(f"[Worker-{task_index}] Available GPUs: {len(gpus)}")
+    if gpus:
+        print(f"[Worker-{task_index}] GPU: {gpus[0].name}")
+
+    # STEP 6: Define model and training within strategy scope
+    with strategy.scope():
+        # Create model
+        model = tf.keras.Sequential([
+            tf.keras.layers.Dense(128, activation='relu', input_shape=(10,)),
+            tf.keras.layers.Dense(128, activation='relu'),
+            tf.keras.layers.Dense(10)
+        ])
+
+        # Compile model
+        model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9),
+            loss=tf.keras.losses.MeanSquaredError(),
+            metrics=['mae']
+        )
+
+        print(f"[Worker-{task_index}] Model created with {model.count_params()} parameters")
+
+    # STEP 7: Create synthetic dataset with manual sharding
+    print(f"[Worker-{task_index}] Creating dataset...")
+
+    import numpy as np
+
+    batch_size = 32
+    total_samples = 160
+    steps_per_epoch = 5
+    num_workers = strategy.num_replicas_in_sync
+
+    # Create synthetic data - same for all workers initially
+    np.random.seed(42)
+    x_data = np.random.randn(total_samples, 10).astype(np.float32)
+    y_data = np.random.randn(total_samples, 10).astype(np.float32)
+
+    # Manually shard data for this worker
+    # Each worker gets a disjoint subset
+    samples_per_worker = total_samples // num_workers
+    start_idx = task_index * samples_per_worker
+    end_idx = start_idx + samples_per_worker
+
+    x_worker = x_data[start_idx:end_idx]
+    y_worker = y_data[start_idx:end_idx]
+
+    print(f"[Worker-{task_index}] Worker data: indices {start_idx}:{end_idx} ({len(x_worker)} samples)")
+
+    # Create dataset for this worker
+    dataset = tf.data.Dataset.from_tensor_slices((x_worker, y_worker))
+    dataset = dataset.shuffle(len(x_worker))
+    dataset = dataset.batch(batch_size, drop_remainder=True)
+    dataset = dataset.repeat()
+
+    # Disable auto-sharding since we manually sharded
+    options = tf.data.Options()
+    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
+    dataset = dataset.with_options(options)
+
+    print(f"[Worker-{task_index}] Dataset created (batch size: {batch_size})")
+
+    # STEP 8: Custom training loop (avoids PerReplica issues)
+    print(f"\n[Worker-{task_index}] {'='*60}")
+    print(f"[Worker-{task_index}] Starting Training Loop")
+    print(f"[Worker-{task_index}] {'='*60}")
+
+    num_epochs = 10
+    num_batches = 2  # Each worker will train on 2 batches per epoch
+
+    @tf.function
+    def train_step(x, y):
+        with tf.GradientTape() as tape:
+            predictions = model(x, training=True)
+            loss = model.compiled_loss(y, predictions)
+
+        gradients = tape.gradient(loss, model.trainable_variables)
+        model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+
+        return loss
+
+    history_loss = []
+
+    for epoch in range(num_epochs):
+        epoch_losses = []
+
+        # Get batches from dataset
+        batch_iter = iter(dataset)
+        for batch_idx in range(num_batches):
+            x_batch, y_batch = next(batch_iter)
+
+            # Distributed training step
+            per_replica_loss = strategy.run(train_step, args=(x_batch, y_batch))
+
+            # Reduce loss across replicas
+            loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, axis=None)
+            epoch_losses.append(float(loss.numpy()))
+
+        avg_loss = sum(epoch_losses) / len(epoch_losses)
+        history_loss.append(avg_loss)
+
+        if task_index == 0:
+            print(f"[Worker-{task_index}] Epoch {epoch+1}/{num_epochs} | Loss: {avg_loss:.6f}")
+
+    print(f"\n[Worker-{task_index}] {'='*60}")
+    print(f"[Worker-{task_index}] Training Complete!")
+    print(f"[Worker-{task_index}] {'='*60}")
+
+    # Get final metrics
+    final_loss = history_loss[-1]
+
+    if task_index == 0:
+        print(f"[Worker-{task_index}] Final Loss: {final_loss:.6f}")
+
+    # STEP 9: Synchronization barrier - wait for all workers to finish training
+    print(f"[Worker-{task_index}] Waiting at barrier for all workers to complete...")
+    import time
+    max_wait = 120
+    start_wait = time.time()
+
+    while time.time() - start_wait < max_wait:
+        try:
+            sync_state = json.loads(Variable.get('tf_sync_state', default_var='{}'))
+            completed_workers = sync_state.get('completed_workers', [])
+
+            if task_index not in completed_workers:
+                completed_workers.append(task_index)
+                sync_state['completed_workers'] = completed_workers
+                Variable.set('tf_sync_state', json.dumps(sync_state))
+                print(f"[Worker-{task_index}] Marked as completed ({len(completed_workers)}/{WORLD_SIZE})")
+
+            if len(completed_workers) == WORLD_SIZE:
+                print(f"[Worker-{task_index}] All workers completed! Proceeding to cleanup.")
+                break
+
+            time.sleep(2)
+        except Exception as e:
+            print(f"[Worker-{task_index}] Error during completion sync: {e}")
+            time.sleep(2)
+    else:
+        print(f"[Worker-{task_index}] Warning: Timeout at completion barrier")
+
+    time.sleep(2)  # Small delay to ensure all workers pass barrier
+
+    # STEP 10: Save model (only chief worker)
+    if task_index == 0:
+        model_path = '/tmp/tf_distributed_model.keras'
+        print(f"[Worker-{task_index}] Saving model to {model_path}...")
+        model.save(model_path)
+        print(f"[Worker-{task_index}] ✓ Model saved successfully!")
+
+    print(f"[Worker-{task_index}] Exiting gracefully...")
+
+    return {
+        'worker_index': task_index,
+        'final_loss': float(final_loss),
+        'epochs_completed': num_epochs,
+        'status': 'success'
+    }
+
+
+def cleanup_sync_state_func():
+    """Clean up synchronization state"""
+    try:
+        Variable.delete('tf_sync_state')
+        print("Sync state cleaned up")
+    except:
+        pass
+    return True
+
+
+def training_summary_func(**context):
+    """Aggregate and display training results"""
+    ti = context['ti']
+
+    print(f"\n{'='*60}")
+    print(f"TENSORFLOW DISTRIBUTED TRAINING SUMMARY")
+    print(f"{'='*60}")
+
+    for worker_idx in range(WORLD_SIZE):
+        try:
+            result = ti.xcom_pull(task_ids=f"train_worker_{worker_idx}")
+            if result and result.get('status') == 'success':
+                print(f"Worker {result['worker_index']}: ✓ 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 worker {worker_idx}: {e}")
+
+    print(f"\n✓ Distributed training completed across {WORLD_SIZE} workers!")
+    return {'status': 'success', 'workers': WORLD_SIZE}
+
+
+with DAG(
+    dag_id='tensorflow_distributed_training_multiworker',
+    default_args=default_args,
+    schedule=None,
+    start_date=pendulum.today('UTC').add(days=-1),
+    catchup=False,
+    tags=['gpu', 'ml', 'distributed', 'tensorflow'],
+    max_active_runs=1,
+    max_active_tasks=10,
+) as dag:
+
+    # Preparation task
+    prep = PythonOperator(
+        task_id='prepare_training',
+        python_callable=prepare_training_func,
+        queue='gpu'
+    )
+
+    # Training tasks - one per worker
+    training_tasks = []
+    for i in range(WORLD_SIZE):
+        task = PythonOperator(
+            task_id=f'train_worker_{i}',
+            python_callable=run_training_node_func,
+            op_kwargs={'task_index': i},
+            queue='gpu',
+        )
+        training_tasks.append(task)
+
+    # Cleanup task
+    cleanup = PythonOperator(
+        task_id='cleanup_sync_state',
+        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