48 lines
881 B
Python
48 lines
881 B
Python
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",
|
|
)
|