gstykalin/test-dags
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update dags

Browse Source
This commit is contained in:
George Stykalin
2026-04-15 10:37:43 +03:00
parent c0b03e9ab1
commit be1c81817d
1 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
dags/test-train-pytorch.py
View File

@@ -15,7 +15,7 @@ default_args = {
# ── rank 0: стартует, пушит свой hostname в XCom ───────────────────────────── # ── rank 0: стартует, пушит свой hostname в XCom ─────────────────────────────
def run_rank_0(**context): def run_rank_0(**context):
import os, resource, socket, torch import os, socket, torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
@@ -27,9 +27,6 @@ def run_rank_0(**context):
# Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться # Пушим MASTER_ADDR чтобы rank1 знал куда коннектиться
context["ti"].xcom_push(key="master_addr", value=fqdn) 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({ os.environ.update({
"MASTER_ADDR": fqdn, "MASTER_ADDR": fqdn,
"MASTER_PORT": MASTER_PORT, "MASTER_PORT": MASTER_PORT,
@@ -38,6 +35,8 @@ def run_rank_0(**context):
"NCCL_SOCKET_IFNAME": "eth0", # bootstrap/rendezvous через eth "NCCL_SOCKET_IFNAME": "eth0", # bootstrap/rendezvous через eth
"NCCL_IB_DISABLE": "0", # трафик через IB "NCCL_IB_DISABLE": "0", # трафик через IB
"NCCL_NET": "IB", "NCCL_NET": "IB",
"NCCL_IB_USE_INLINE": "1", # без ibv_reg_mr pinned memory
"NCCL_BUFFSIZE": "1048576",
"NCCL_P2P_DISABLE": "1", # разные ноды — p2p не нужен "NCCL_P2P_DISABLE": "1", # разные ноды — p2p не нужен
"NCCL_SHM_DISABLE": "1", # shm только для локальных рангов "NCCL_SHM_DISABLE": "1", # shm только для локальных рангов
"NCCL_DEBUG": "INFO", "NCCL_DEBUG": "INFO",
@@ -66,7 +65,7 @@ def run_rank_0(**context):
# ── rank 1: ждёт XCom от rank0, потом коннектится ──────────────────────────── # ── rank 1: ждёт XCom от rank0, потом коннектится ────────────────────────────
def run_rank_1(**context): def run_rank_1(**context):
import os, resource, socket, time, torch import os, socket, time, torch
import torch.distributed as dist import torch.distributed as dist
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
@@ -88,8 +87,6 @@ def run_rank_1(**context):
print(f"[rank1] master_addr={master_addr}") print(f"[rank1] master_addr={master_addr}")
resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
os.environ.update({ os.environ.update({
"MASTER_ADDR": master_addr, "MASTER_ADDR": master_addr,
"MASTER_PORT": MASTER_PORT, "MASTER_PORT": MASTER_PORT,
@@ -98,6 +95,8 @@ def run_rank_1(**context):
"NCCL_SOCKET_IFNAME": "eth0", "NCCL_SOCKET_IFNAME": "eth0",
"NCCL_IB_DISABLE": "0", "NCCL_IB_DISABLE": "0",
"NCCL_NET": "IB", "NCCL_NET": "IB",
"NCCL_IB_USE_INLINE": "1", # без ibv_reg_mr pinned memory
"NCCL_BUFFSIZE": "1048576",
"NCCL_P2P_DISABLE": "1", "NCCL_P2P_DISABLE": "1",
"NCCL_SHM_DISABLE": "1", "NCCL_SHM_DISABLE": "1",
"NCCL_DEBUG": "INFO", "NCCL_DEBUG": "INFO",