init commit
This commit is contained in:
43
dags/connection-check-default.py
Normal file
43
dags/connection-check-default.py
Normal file
@@ -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()
|
||||
43
dags/connection-check.py
Normal file
43
dags/connection-check.py
Normal file
@@ -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()
|
||||
272
dags/test-train-pytorch.py
Normal file
272
dags/test-train-pytorch.py
Normal file
@@ -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
|
||||
373
dags/test-train-tensorflow.py
Normal file
373
dags/test-train-tensorflow.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user