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