81 lines
1.7 KiB
Python
81 lines
1.7 KiB
Python
from airflow import DAG
|
|
from airflow.providers.standard.operators.python import PythonOperator
|
|
from airflow.operators.bash import BashOperator
|
|
from datetime import datetime
|
|
import os
|
|
|
|
|
|
default_args = {
|
|
"owner": "airflow",
|
|
}
|
|
|
|
|
|
# -----------------------
|
|
# PREP
|
|
# -----------------------
|
|
def prepare_training():
|
|
print("Preparing DDP training environment")
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
os.environ["NCCL_DEBUG"] = "INFO"
|
|
|
|
return "prepared"
|
|
|
|
|
|
# -----------------------
|
|
# CLEANUP
|
|
# -----------------------
|
|
def cleanup():
|
|
print("Cleaning up after training")
|
|
return "cleaned"
|
|
|
|
|
|
# -----------------------
|
|
# DAG
|
|
# -----------------------
|
|
with DAG(
|
|
dag_id="pytorch_ddp_airflow_fixed_production",
|
|
start_date=datetime(2024, 1, 1),
|
|
schedule=None,
|
|
catchup=False,
|
|
default_args=default_args,
|
|
tags=["ddp", "pytorch", "gpu"],
|
|
) as dag:
|
|
|
|
prepare_training_task = PythonOperator(
|
|
task_id="prepare_training",
|
|
python_callable=prepare_training,
|
|
queue="gpu", # ✅ FIX #1
|
|
)
|
|
|
|
train_ddp_task = BashOperator(
|
|
task_id="train_ddp",
|
|
bash_command="""
|
|
set -e
|
|
|
|
export NCCL_DEBUG=INFO
|
|
export NCCL_ASYNC_ERROR_HANDLING=1
|
|
export NCCL_IB_DISABLE=1
|
|
|
|
echo "Starting torchrun DDP training"
|
|
|
|
torchrun \
|
|
--nproc_per_node=2 \
|
|
--master_port=29500 \
|
|
/opt/airflow/dags/repo/train.py
|
|
|
|
echo "Training finished"
|
|
""",
|
|
queue="gpu", # ✅ FIX #1
|
|
)
|
|
|
|
cleanup_task = PythonOperator(
|
|
task_id="cleanup",
|
|
python_callable=cleanup,
|
|
queue="gpu", # ✅ FIX #1
|
|
)
|
|
|
|
|
|
# DAG FLOW
|
|
prepare_training_task >> train_ddp_task >> cleanup_task
|