update dags
This commit is contained in:
@@ -1,79 +1,142 @@
|
||||
from airflow import DAG
|
||||
from airflow.operators.python import PythonOperator
|
||||
from airflow.providers.standard.operators.python import PythonOperator
|
||||
import pendulum
|
||||
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=1),
|
||||
"execution_timeout": timedelta(minutes=30),
|
||||
}
|
||||
|
||||
|
||||
def run_training_node_func(rank: int, world_size: int):
|
||||
import os
|
||||
import time
|
||||
import socket
|
||||
# ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
|
||||
def run_rank_0(**context):
|
||||
import os, resource, socket, torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
hostname = socket.gethostname()
|
||||
fqdn = f"{hostname}.airflow-worker-gpu"
|
||||
print(f"[rank0] hostname={fqdn}")
|
||||
|
||||
# Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться
|
||||
context["ti"].xcom_push(key="master_addr", value=fqdn)
|
||||
|
||||
# Снимаем лимит на locked memory (нужно для ibv_reg_mr)
|
||||
resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
|
||||
|
||||
os.environ.update({
|
||||
"MASTER_ADDR": fqdn,
|
||||
"MASTER_PORT": MASTER_PORT,
|
||||
"WORLD_SIZE": "2",
|
||||
"RANK": "0",
|
||||
"NCCL_SOCKET_IFNAME": "eth0", # bootstrap/rendezvous через eth
|
||||
"NCCL_IB_DISABLE": "0", # трафик через IB
|
||||
"NCCL_NET": "IB",
|
||||
"NCCL_P2P_DISABLE": "1", # разные ноды — p2p не нужен
|
||||
"NCCL_SHM_DISABLE": "1", # shm только для локальных рангов
|
||||
"NCCL_DEBUG": "INFO",
|
||||
"NCCL_DEBUG_SUBSYS": "NET,INIT",
|
||||
"TORCH_NCCL_BLOCKING_WAIT": "1",
|
||||
})
|
||||
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method="env://",
|
||||
rank=0,
|
||||
world_size=2,
|
||||
timeout=timedelta(minutes=5),
|
||||
)
|
||||
print("[rank0] process group OK")
|
||||
|
||||
torch.cuda.set_device(0)
|
||||
device = torch.device("cuda:0")
|
||||
print(f"[rank0] GPU={torch.cuda.get_device_name(0)}")
|
||||
|
||||
_train(rank=0, device=device)
|
||||
|
||||
dist.destroy_process_group()
|
||||
print("[rank0] DONE")
|
||||
|
||||
|
||||
# ── rank 1: ждёт XCom от rank0, потом коннектится ────────────────────────────
|
||||
def run_rank_1(**context):
|
||||
import os, resource, socket, time, torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
# Polling XCom пока rank0 не запишет master_addr
|
||||
master_addr = None
|
||||
for attempt in range(30):
|
||||
master_addr = context["ti"].xcom_pull(
|
||||
task_ids="train_rank_0",
|
||||
key="master_addr",
|
||||
)
|
||||
if master_addr:
|
||||
break
|
||||
print(f"[rank1] waiting for master_addr (attempt {attempt+1}/30)...")
|
||||
time.sleep(5)
|
||||
|
||||
if not master_addr:
|
||||
raise RuntimeError("rank1: timeout waiting for master_addr from rank0")
|
||||
|
||||
print(f"[rank1] master_addr={master_addr}")
|
||||
|
||||
resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
|
||||
|
||||
os.environ.update({
|
||||
"MASTER_ADDR": master_addr,
|
||||
"MASTER_PORT": MASTER_PORT,
|
||||
"WORLD_SIZE": "2",
|
||||
"RANK": "1",
|
||||
"NCCL_SOCKET_IFNAME": "eth0",
|
||||
"NCCL_IB_DISABLE": "0",
|
||||
"NCCL_NET": "IB",
|
||||
"NCCL_P2P_DISABLE": "1",
|
||||
"NCCL_SHM_DISABLE": "1",
|
||||
"NCCL_DEBUG": "INFO",
|
||||
"NCCL_DEBUG_SUBSYS": "NET,INIT",
|
||||
"TORCH_NCCL_BLOCKING_WAIT": "1",
|
||||
})
|
||||
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method="env://",
|
||||
rank=1,
|
||||
world_size=2,
|
||||
timeout=timedelta(minutes=5),
|
||||
)
|
||||
print("[rank1] process group OK")
|
||||
|
||||
torch.cuda.set_device(0)
|
||||
device = torch.device("cuda:0")
|
||||
print(f"[rank1] GPU={torch.cuda.get_device_name(0)}")
|
||||
|
||||
_train(rank=1, device=device)
|
||||
|
||||
dist.destroy_process_group()
|
||||
print("[rank1] DONE")
|
||||
|
||||
|
||||
# ── общий train loop ──────────────────────────────────────────────────────────
|
||||
def _train(rank: int, device):
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from datetime import datetime
|
||||
|
||||
print("=" * 80)
|
||||
print(f"RANK {rank}/{world_size} HOST={socket.gethostname()} {datetime.now()}")
|
||||
print("=" * 80)
|
||||
# allreduce smoke-test
|
||||
probe = torch.ones(1, device=device) * rank
|
||||
dist.all_reduce(probe, op=dist.ReduceOp.SUM)
|
||||
assert probe.item() == 1.0, f"allreduce mismatch: {probe.item()}"
|
||||
print(f"[rank{rank}] allreduce smoke-test PASSED")
|
||||
|
||||
# ── статический барьер: ждём пока оба воркера подтянутся ──────────────────
|
||||
time.sleep(10)
|
||||
|
||||
# ── переменные окружения для 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",
|
||||
}
|
||||
)
|
||||
|
||||
print(f"[{rank}] MASTER={MASTER_ADDR}:{MASTER_PORT} IB_IF=ib0")
|
||||
|
||||
# ── 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")
|
||||
|
||||
# ── GPU binding ───────────────────────────────────────────────────────────
|
||||
torch.cuda.set_device(0)
|
||||
device = torch.device("cuda:0")
|
||||
print(f"[{rank}] GPU={torch.cuda.get_device_name(0)}")
|
||||
|
||||
# ── минимальная модель ────────────────────────────────────────────────────
|
||||
model = nn.Sequential(
|
||||
nn.Linear(10, 128),
|
||||
nn.ReLU(),
|
||||
@@ -84,51 +147,66 @@ def run_training_node_func(rank: int, world_size: int):
|
||||
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
|
||||
loss_fn = nn.MSELoss()
|
||||
|
||||
# ── 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)
|
||||
|
||||
x = torch.randn(32, 10, device=device)
|
||||
y = torch.randn(32, 10, device=device)
|
||||
optimizer.zero_grad()
|
||||
out = ddp_model(x)
|
||||
loss = loss_fn(out, y)
|
||||
loss = loss_fn(ddp_model(x), y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
print(f"[{rank}] epoch={epoch} loss={loss.item():.4f}")
|
||||
|
||||
dist.destroy_process_group()
|
||||
print(f"[{rank}] DONE")
|
||||
return {"rank": rank, "status": "ok"}
|
||||
print(f"[rank{rank}] epoch={epoch} loss={loss.item():.4f}")
|
||||
|
||||
|
||||
# ── DAG ───────────────────────────────────────────────────────────────────────
|
||||
with DAG(
|
||||
dag_id="ddp_ib_test",
|
||||
dag_id="ddp_ib_xcom",
|
||||
start_date=pendulum.today("UTC").add(days=-1),
|
||||
schedule=None,
|
||||
catchup=False,
|
||||
max_active_runs=1,
|
||||
max_active_tasks=WORLD_SIZE,
|
||||
max_active_tasks=2,
|
||||
default_args=default_args,
|
||||
) as dag:
|
||||
|
||||
tasks = []
|
||||
for r in range(WORLD_SIZE):
|
||||
t = PythonOperator(
|
||||
task_id=f"train_rank_{r}",
|
||||
python_callable=run_training_node_func,
|
||||
op_kwargs={"rank": r, "world_size": WORLD_SIZE},
|
||||
rank0 = PythonOperator(
|
||||
task_id="train_rank_0",
|
||||
python_callable=run_rank_0,
|
||||
queue="gpu",
|
||||
executor_config={
|
||||
"pod_override": {
|
||||
"spec": {
|
||||
"containers": [{
|
||||
"name": "base",
|
||||
"securityContext": {
|
||||
"capabilities": {
|
||||
"add": ["IPC_LOCK", "SYS_RESOURCE"]
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
tasks.append(t)
|
||||
|
||||
# запускаем параллельно — без зависимостей между рангами
|
||||
# (prep/cleanup убраны, это минимальный тест)
|
||||
rank1 = PythonOperator(
|
||||
task_id="train_rank_1",
|
||||
python_callable=run_rank_1,
|
||||
queue="gpu",
|
||||
executor_config={
|
||||
"pod_override": {
|
||||
"spec": {
|
||||
"containers": [{
|
||||
"name": "base",
|
||||
"securityContext": {
|
||||
"capabilities": {
|
||||
"add": ["IPC_LOCK", "SYS_RESOURCE"]
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
# rank1 стартует сразу, но ждёт master_addr через XCom polling
|
||||
[rank0, rank1]
|
||||
|
||||
Reference in New Issue
Block a user