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