update dags
This commit is contained in:
@@ -5,102 +5,93 @@ from datetime import timedelta
|
||||
|
||||
WORLD_SIZE = 2
|
||||
|
||||
# Адрес мастера — headless DNS пода воркера
|
||||
MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu"
|
||||
MASTER_PORT = "29500"
|
||||
|
||||
default_args = {
|
||||
"owner": "airflow",
|
||||
"retries": 0,
|
||||
"execution_timeout": timedelta(hours=2),
|
||||
"execution_timeout": timedelta(hours=1),
|
||||
}
|
||||
|
||||
|
||||
# -------------------------
|
||||
# SIMPLE PREP (NO VARIABLES)
|
||||
# -------------------------
|
||||
def prepare_training_func():
|
||||
print("Starting distributed training (no sync state needed)")
|
||||
return True
|
||||
|
||||
|
||||
# -------------------------
|
||||
# CORE TRAINING
|
||||
# -------------------------
|
||||
def run_training_node_func(rank, world_size):
|
||||
def run_training_node_func(rank: int, world_size: int):
|
||||
import os
|
||||
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
|
||||
import socket
|
||||
|
||||
print("=" * 80)
|
||||
print(f"RANK {rank}/{world_size} START {datetime.now()}")
|
||||
print(f"HOST: {socket.gethostname()}")
|
||||
print(f"RANK {rank}/{world_size} HOST={socket.gethostname()} {datetime.now()}")
|
||||
print("=" * 80)
|
||||
|
||||
# -------------------------
|
||||
# FIXED BARRIER (important)
|
||||
# -------------------------
|
||||
print(f"[{rank}] sync barrier (static sleep)")
|
||||
# ── статический барьер: ждём пока оба воркера подтянутся ──────────────────
|
||||
time.sleep(10)
|
||||
|
||||
# -------------------------
|
||||
# STATIC CONFIG (NO AIRFLOW VARIABLES)
|
||||
# -------------------------
|
||||
os.environ["MASTER_ADDR"] = MASTER_ADDR
|
||||
os.environ["MASTER_PORT"] = MASTER_PORT
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
os.environ["RANK"] = str(rank)
|
||||
# ── переменные окружения для NCCL + InfiniBand ────────────────────────────
|
||||
os.environ.update(
|
||||
{
|
||||
"MASTER_ADDR": MASTER_ADDR,
|
||||
"MASTER_PORT": MASTER_PORT,
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"RANK": str(rank),
|
||||
# InfiniBand: указываем IPoIB-интерфейс вместо eth0
|
||||
# Разрешаем NCCL использовать RDMA (IB verbs)
|
||||
"NCCL_IB_DISABLE": "0",
|
||||
# GPUDirect RDMA — если драйвер поддерживает
|
||||
"NCCL_P2P_DISABLE": "0",
|
||||
# Явно форсируем IB transport (опционально, NCCL сам выберет,
|
||||
# но полезно для отладки)
|
||||
"NCCL_NET": "IB",
|
||||
# Отладка: INFO покажет какой транспорт выбран,
|
||||
# поменяй на TRACE для полного вывода
|
||||
"NCCL_DEBUG": "INFO",
|
||||
"NCCL_DEBUG_SUBSYS": "NET,INIT",
|
||||
"TORCH_NCCL_BLOCKING_WAIT": "1",
|
||||
}
|
||||
)
|
||||
|
||||
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
|
||||
os.environ["NCCL_DEBUG"] = "INFO"
|
||||
os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"
|
||||
print(f"[{rank}] MASTER={MASTER_ADDR}:{MASTER_PORT} IB_IF=ib0")
|
||||
|
||||
print(f"[{rank}] MASTER = {MASTER_ADDR}:{MASTER_PORT}")
|
||||
|
||||
# -------------------------
|
||||
# INIT PROCESS GROUP
|
||||
# -------------------------
|
||||
# ── init process group ────────────────────────────────────────────────────
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method="env://",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
timeout=timedelta(minutes=5),
|
||||
)
|
||||
print(f"[{rank}] dist.init_process_group OK")
|
||||
|
||||
print(f"[{rank}] DDP INIT OK")
|
||||
|
||||
# -------------------------
|
||||
# GPU BINDING (CRITICAL FIX)
|
||||
# -------------------------
|
||||
# ── GPU binding ───────────────────────────────────────────────────────────
|
||||
torch.cuda.set_device(0)
|
||||
device = torch.device("cuda:0")
|
||||
print(f"[{rank}] GPU={torch.cuda.get_device_name(0)}")
|
||||
|
||||
print(f"[{rank}] GPU = {torch.cuda.get_device_name(0)}")
|
||||
|
||||
# -------------------------
|
||||
# MODEL
|
||||
# -------------------------
|
||||
# ── минимальная модель ────────────────────────────────────────────────────
|
||||
model = nn.Sequential(
|
||||
nn.Linear(10, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 10),
|
||||
).to(device)
|
||||
|
||||
ddp_model = nn.parallel.DistributedDataParallel(
|
||||
model,
|
||||
device_ids=[0],
|
||||
)
|
||||
|
||||
ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[0])
|
||||
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
|
||||
loss_fn = nn.MSELoss()
|
||||
|
||||
# -------------------------
|
||||
# TRAIN LOOP
|
||||
# -------------------------
|
||||
# ── allreduce smoke-test перед обучением ──────────────────────────────────
|
||||
probe = torch.ones(1).to(device) * rank
|
||||
dist.all_reduce(probe, op=dist.ReduceOp.SUM)
|
||||
expected = sum(range(world_size))
|
||||
assert probe.item() == expected, f"allreduce mismatch: got {probe.item()}"
|
||||
print(f"[{rank}] allreduce smoke-test PASSED (sum={probe.item()})")
|
||||
|
||||
# ── train loop ────────────────────────────────────────────────────────────
|
||||
for epoch in range(5):
|
||||
x = torch.randn(32, 10).to(device)
|
||||
y = torch.randn(32, 10).to(device)
|
||||
@@ -114,42 +105,21 @@ def run_training_node_func(rank, world_size):
|
||||
print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")
|
||||
|
||||
dist.destroy_process_group()
|
||||
|
||||
print(f"[{rank}] DONE")
|
||||
return {"rank": rank, "status": "ok"}
|
||||
|
||||
|
||||
# -------------------------
|
||||
# CLEANUP
|
||||
# -------------------------
|
||||
def cleanup_func():
|
||||
print("cleanup done")
|
||||
return True
|
||||
|
||||
|
||||
def summary_func(**context):
|
||||
print("training done")
|
||||
return True
|
||||
|
||||
|
||||
# -------------------------
|
||||
# DAG
|
||||
# -------------------------
|
||||
# ── DAG ───────────────────────────────────────────────────────────────────────
|
||||
with DAG(
|
||||
dag_id="ddp_airflow_stable",
|
||||
dag_id="ddp_ib_test",
|
||||
start_date=pendulum.today("UTC").add(days=-1),
|
||||
schedule=None,
|
||||
catchup=False,
|
||||
max_active_runs=1,
|
||||
max_active_tasks=2,
|
||||
max_active_tasks=WORLD_SIZE,
|
||||
default_args=default_args,
|
||||
) as dag:
|
||||
|
||||
prep = PythonOperator(
|
||||
task_id="prep",
|
||||
python_callable=prepare_training_func,
|
||||
queue="gpu",
|
||||
)
|
||||
|
||||
tasks = []
|
||||
for r in range(WORLD_SIZE):
|
||||
t = PythonOperator(
|
||||
@@ -160,18 +130,5 @@ with DAG(
|
||||
)
|
||||
tasks.append(t)
|
||||
|
||||
cleanup = PythonOperator(
|
||||
task_id="cleanup",
|
||||
python_callable=cleanup_func,
|
||||
trigger_rule="all_done",
|
||||
queue="gpu",
|
||||
)
|
||||
|
||||
summary = PythonOperator(
|
||||
task_id="summary",
|
||||
python_callable=summary_func,
|
||||
trigger_rule="all_done",
|
||||
queue="gpu",
|
||||
)
|
||||
|
||||
prep >> tasks >> cleanup >> summary
|
||||
# запускаем параллельно — без зависимостей между рангами
|
||||
# (prep/cleanup убраны, это минимальный тест)
|
||||
|
||||
Reference in New Issue
Block a user