update dags

dags/test-train-pytorch.py

@@ -1,47 +1,77 @@
 from airflow import DAG
 from airflow.operators.python import PythonOperator
-import pendulum
-import subprocess
-from datetime import timedelta
+from airflow.operators.bash import BashOperator
+from datetime import datetime
+import os
 
 
+# -----------------------------
+# CONFIG
+# -----------------------------
 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"}
+MASTER_PORT = 29500
 
 
 default_args = {
     "owner": "airflow",
-    "retries": 1,
-    "retry_delay": timedelta(minutes=2),
 }
 
 
+# -----------------------------
+# 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,
-    start_date=pendulum.today("UTC").add(days=-1),
     catchup=False,
-    max_active_runs=1,
-    tags=["ddp", "torchrun", "stable"],
+    tags=["ddp", "pytorch", "fixed"],
 ) as dag:
 
-    train = PythonOperator(
-        task_id="train_ddp",
-        python_callable=run_ddp_job,
-        queue="gpu",
+    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