from airflow import DAG
from airflow.operators.python import PythonOperator
import pendulum
import subprocess
from datetime import timedelta


WORLD_SIZE = 2


def run_ddp_job():
    cmd = [
        "torchrun",
        "--nproc_per_node=2",
        "--standalone",
        "train.py"
    ]

    print("Running:", " ".join(cmd))

    subprocess.run(cmd, check=True)

    return {"status": "success"}


default_args = {
    "owner": "airflow",
    "retries": 1,
    "retry_delay": timedelta(minutes=2),
}


with DAG(
    dag_id="pytorch_ddp_airflow_fixed_production",
    default_args=default_args,
    schedule=None,
    start_date=pendulum.today("UTC").add(days=-1),
    catchup=False,
    max_active_runs=1,
    tags=["ddp", "torchrun", "stable"],
) as dag:

    train = PythonOperator(
        task_id="train_ddp",
        python_callable=run_ddp_job,
        queue="gpu",
    )