update dags

dags/test-train-pytorch.py

@@ -1,77 +1,80 @@
 from airflow import DAG
-from airflow.operators.python import PythonOperator
+from airflow.providers.standard.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
-# -----------------------------
+# -----------------------
+# PREP
+# -----------------------
 def prepare_training():
-    print("Preparing dataset / env for DDP")
+    print("Preparing DDP training environment")
+
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
-    return True
+    os.environ["NCCL_DEBUG"] = "INFO"
+
+    return "prepared"
 
 
-# -----------------------------
-# CLEANUP TASK
-# -----------------------------
+# -----------------------
+# CLEANUP
+# -----------------------
 def cleanup():
-    print("Cleaning up training artifacts")
-    return True
+    print("Cleaning up after training")
+    return "cleaned"
 
 
-# -----------------------------
+# -----------------------
 # 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"],
+    default_args=default_args,
+    tags=["ddp", "pytorch", "gpu"],
 ) as dag:
 
-    prepare = PythonOperator(
+    prepare_training_task = PythonOperator(
         task_id="prepare_training",
         python_callable=prepare_training,
+        queue="gpu",   # ✅ FIX #1
     )
 
-    # ---------------------------------------------------------
-    # MAIN FIX: use torchrun instead of spawn inside Python
-    # ---------------------------------------------------------
-    train_ddp = BashOperator(
+    train_ddp_task = BashOperator(
         task_id="train_ddp",
-        bash_command=f"""
+        bash_command="""
+        set -e
+
         export NCCL_DEBUG=INFO
         export NCCL_ASYNC_ERROR_HANDLING=1
+        export NCCL_IB_DISABLE=1
 
-        export NCCL_IB_DISABLE=1   # safe default (enable later if needed)
+        echo "Starting torchrun DDP training"
 
         torchrun \
-            --nproc_per_node={WORLD_SIZE} \
-            --master_port={MASTER_PORT} \
-            train.py
-        """
+            --nproc_per_node=2 \
+            --master_port=29500 \
+            /opt/airflow/dags/repo/train.py
+
+        echo "Training finished"
+        """,
+        queue="gpu",   # ✅ FIX #1
     )
 
-    finish = PythonOperator(
+    cleanup_task = PythonOperator(
         task_id="cleanup",
         python_callable=cleanup,
+        queue="gpu",   # ✅ FIX #1
     )
 
-    prepare >> train_ddp >> finish
+
+    # DAG FLOW
+    prepare_training_task >> train_ddp_task >> cleanup_task