diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 930b8e5..557c64e 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -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}, - queue="gpu", - ) - tasks.append(t) + 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"] + } + } + }] + } + } + }, + ) - # запускаем параллельно — без зависимостей между рангами - # (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]