78 lines
1.8 KiB
Python
78 lines
1.8 KiB
Python
from airflow import DAG
|
|
from airflow.operators.python import PythonOperator
|
|
from airflow.operators.bash import BashOperator
|
|
from datetime import datetime
|
|
import os
|
|
|
|
|
|
# -----------------------------
|
|
# CONFIG
|
|
# -----------------------------
|
|
WORLD_SIZE = 2
|
|
MASTER_PORT = 29500
|
|
|
|
|
|
default_args = {
|
|
"owner": "airflow",
|
|
}
|
|
|
|
|
|
# -----------------------------
|
|
# PREP TASK
|
|
# -----------------------------
|
|
def prepare_training():
|
|
print("Preparing dataset / env for DDP")
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
return True
|
|
|
|
|
|
# -----------------------------
|
|
# CLEANUP TASK
|
|
# -----------------------------
|
|
def cleanup():
|
|
print("Cleaning up training artifacts")
|
|
return True
|
|
|
|
|
|
# -----------------------------
|
|
# DAG
|
|
# -----------------------------
|
|
with DAG(
|
|
dag_id="pytorch_ddp_airflow_fixed_production",
|
|
default_args=default_args,
|
|
start_date=datetime(2024, 1, 1),
|
|
schedule=None,
|
|
catchup=False,
|
|
tags=["ddp", "pytorch", "fixed"],
|
|
) as dag:
|
|
|
|
prepare = PythonOperator(
|
|
task_id="prepare_training",
|
|
python_callable=prepare_training,
|
|
)
|
|
|
|
# ---------------------------------------------------------
|
|
# MAIN FIX: use torchrun instead of spawn inside Python
|
|
# ---------------------------------------------------------
|
|
train_ddp = BashOperator(
|
|
task_id="train_ddp",
|
|
bash_command=f"""
|
|
export NCCL_DEBUG=INFO
|
|
export NCCL_ASYNC_ERROR_HANDLING=1
|
|
|
|
export NCCL_IB_DISABLE=1 # safe default (enable later if needed)
|
|
|
|
torchrun \
|
|
--nproc_per_node={WORLD_SIZE} \
|
|
--master_port={MASTER_PORT} \
|
|
train.py
|
|
"""
|
|
)
|
|
|
|
finish = PythonOperator(
|
|
task_id="cleanup",
|
|
python_callable=cleanup,
|
|
)
|
|
|
|
prepare >> train_ddp >> finish
|