From 526f201f416d08ed9cdb6e49c95d5c40509ad5e9 Mon Sep 17 00:00:00 2001 From: George Stykalin Date: Mon, 12 Jan 2026 16:25:19 +0300 Subject: [PATCH] init commit --- dags/connection-check-default.py | 43 ++++ dags/connection-check.py | 43 ++++ dags/test-train-pytorch.py | 272 ++++++++++++++++++++++ dags/test-train-tensorflow.py | 373 +++++++++++++++++++++++++++++++ 4 files changed, 731 insertions(+) create mode 100644 dags/connection-check-default.py create mode 100644 dags/connection-check.py create mode 100644 dags/test-train-pytorch.py create mode 100644 dags/test-train-tensorflow.py diff --git a/dags/connection-check-default.py b/dags/connection-check-default.py new file mode 100644 index 0000000..5c35861 --- /dev/null +++ b/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() diff --git a/dags/connection-check.py b/dags/connection-check.py new file mode 100644 index 0000000..f06d701 --- /dev/null +++ b/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() diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py new file mode 100644 index 0000000..2f185cf --- /dev/null +++ b/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 diff --git a/dags/test-train-tensorflow.py b/dags/test-train-tensorflow.py new file mode 100644 index 0000000..091b459 --- /dev/null +++ b/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