airflow % kubectl get pod -n airflow
NAME                                 READY   STATUS    RESTARTS   AGE
airflow-postgresql-0                 1/1     Running   0          18h
airflow-scheduler-54ff57bcb6-tgtqs   2/2     Running   0          5h54m
airflow-statsd-769b757665-7pqjr      1/1     Running   0          18h
airflow-triggerer-0                  2/2     Running   0          5h53m
airflow-webserver-75b749479b-97llm   1/1     Running   0          5h54m

argo % kubectl get pod -n argo
NAME                                       READY   STATUS    RESTARTS   AGE
argo-server-75d76d9996-sqx9z           1/1     Running   0          34s
workflow-controller-798c4b99bb-vl8rx   1/1     Running   0          34s

# The following workflow executes a diamond workflow
# 
#   A
#  / \
# B   C
#  \ /
#   D

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator

# 定义默认参数
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2025, 5, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# 创建DAG对象
dag = DAG(
    'dag_diamond',
    default_args=default_args,
    description='A simple diamond-shaped workflow',
    schedule_interval=timedelta(days=1),  # 可根据需要调整调度间隔
)

# 定义一个简单的打印函数
def print_message(message):
    print(f"Message: {message}")

# 创建任务A
task_a = PythonOperator(
    task_id='A',
    python_callable=print_message,
    op_kwargs={'message': 'A'},
    dag=dag,
)

# 创建任务B
task_b = PythonOperator(
    task_id='B',
    python_callable=print_message,
    op_kwargs={'message': 'B'},
    dag=dag,
)

# 创建任务C
task_c = PythonOperator(
    task_id='C',
    python_callable=print_message,
    op_kwargs={'message': 'C'},
    dag=dag,
)

# 创建任务D
task_d = PythonOperator(
    task_id='D',
    python_callable=print_message,
    op_kwargs={'message': 'D'},
    dag=dag,
)

# 设置任务依赖关系
task_a >> [task_b, task_c]  # A完成后同时执行B和C
[task_b, task_c] >> task_d  # B和C都完成后才执行D

from datetime import timedelta
import os
from airflow import DAG
from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator
from airflow.utils.dates import days_ago
from airflow.models import Variable

default_args = {
    'owner': 'Howdy',
    'retries': 3,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'spark_on_k8s_airflow',
    start_date=days_ago(1),
    catchup=False,
    schedule_interval=timedelta(days=1),
    template_searchpath='/opt/airflow/dags/repo/dags/spark8s/'
)

spark_k8s_task = SparkKubernetesOperator(
    task_id='n-spark-on-k8s-airflow',
    trigger_rule="all_success",
    depends_on_past=False,
    retries=0,
    application_file='spark-k8s.yaml',
    namespace="spark-apps",
    kubernetes_conn_id="kubernetes_default",
    do_xcom_push=True,
    dag=dag
)

spark_k8s_task

apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: spark-pi-demo-
spec:
  entrypoint: spark-pi
  templates:
    - name: spark-pi
      plugin:
        spark:
          # SparkApplication definition (Spark Operator must be installed in advance)
          apiVersion: "sparkoperator.k8s.io/v1beta2"
          kind: SparkApplication
          metadata:
            name: spark-pi-demo
            namespace: argo
          spec:
            type: Scala
            mode: cluster
            image: "gcr.io/spark-operator/spark:v3.3.1"
            mainClass: org.apache.spark.examples.SparkPi
            mainApplicationFile: "local:///opt/spark/examples/jars/spark-examples_2.12-3.3.1.jar"
            sparkVersion: "3.3.1"
            restartPolicy:
              type: Never
            driver:
              cores: 1
              memory: "2g"
              serviceAccount: spark-sa
              labels:
                version: 3.3.1
            executor:
              cores: 2
              instances: 2
              memory: "4g"
              labels:
                version: 3.3.1
            arguments:
              - "1000" 

import datetime
import pendulum

from airflow.utils.state import DagRunState, TaskInstanceState
from airflow.utils.types import DagRunType
from airflow.models.dag import DAG
from airflow.decorators import task


DATA_INTERVAL_START = pendulum.datetime(2021, 9, 13, tz="UTC")
DATA_INTERVAL_END = DATA_INTERVAL_START + datetime.timedelta(days=1)
from torchx.schedulers.ids import make_unique

@task(task_id=f'hello_torchx')
def run_torchx(message):
    """This is a function that will run within the DAG execution"""
    from torchx.runner import get_runner
    with get_runner() as runner:
        # Run the utils.sh component on the local_cwd scheduler.
        app_id = runner.run_component(
            "utils.sh",
            ["echo", message],
            scheduler="local_cwd",
        )

        # Wait for the the job to complete
        status = runner.wait(app_id, wait_interval=1)

        # Raise_for_status will raise an exception if the job didn't succeed
        status.raise_for_status()

        # Finally we can print all of the log lines from the TorchX job so it
        # will show up in the workflow logs.
        for line in runner.log_lines(app_id, "sh", k=0):
            print(line, end="")

with DAG(
    dag_id=make_unique('example_python_operator'),
    schedule_interval=None,
    start_date=DATA_INTERVAL_START,
    catchup=False,
    tags=['example'],
) as dag:
    run_job = run_torchx("Hello, TorchX!")

dagrun = dag.create_dagrun(
    state=DagRunState.RUNNING,
    execution_date=DATA_INTERVAL_START,
    data_interval=(DATA_INTERVAL_START, DATA_INTERVAL_END),
    start_date=DATA_INTERVAL_END,
    run_type=DagRunType.MANUAL,
)
ti = dagrun.get_task_instance(task_id="hello_torchx")
ti.task = dag.get_task(task_id="hello_torchx")
ti.run(ignore_ti_state=True)
assert ti.state == TaskInstanceState.SUCCESS      

apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: pytorch-deepspeed-demo-
spec:
  entrypoint: main
  arguments:
    parameters:
    - name: checkpoint-dir
      value: /root/deepspeed_data
    - name: instance-spec
      value: ecs.gn7i-c8g1.2xlarge
    - name: image
      value: mirrors-ssl.aliyuncs.com/kubeflow/pytorch-deepspeed-demo:latest

  templates:
  - name: main
    steps:
    - - name: run-pytorch-job
        template: pytorch-job-template
        arguments:
          parameters:
          - name: checkpoint-dir
            value: "{{workflow.parameters.checkpoint-dir}}"
    - - name: check-pytorch-job-status
        template: check-status
        arguments:
          parameters:
          - name: job-name
            value: "{{steps.run-pytorch-job.outputs.parameters.job-name}}"

  - name: pytorch-job-template
    inputs:
      parameters:
      - name: checkpoint-dir
    resource:
      action: create
      successCondition: status.replicaStatuses.Worker.succeeded = 1
      failureCondition: status.replicaStatuses.Worker.failed > 0
      manifest: |
        apiVersion: "kubeflow.org/v1"
        kind: PyTorchJob
        metadata:
          generateName: pytorch-deepspeed-demo-
        spec:
          pytorchReplicaSpecs:
            Master:
              replicas: 1
              restartPolicy: OnFailure
              template:
                metadata:
                  name: deepspeed-master
                spec:
                  containers:
                    - name: pytorch
                      image: "{{workflow.parameters.image}}"
                      command:
                        - torchrun
                        - /train_bert_ds.py
                        - --checkpoint_dir
                        - {{inputs.parameters.checkpoint-dir}}
                      resources:
                        limits:
                          nvidia.com/gpu: 1
            Worker:
              replicas: 1
              restartPolicy: OnFailure
              template:
                metadata:
                  name: deepspeed-worker
                spec:
                  containers:
                    - name: pytorch
                      image: "{{workflow.parameters.image}}"
                      command:
                        - torchrun
                        - /train_bert_ds.py
                        - --checkpoint_dir
                        - {{inputs.parameters.checkpoint-dir}}
                      resources:
                        limits:
                          nvidia.com/gpu: 1
    outputs:
      parameters:
      - name: job-name
        valueFrom:
          jsonPath: '{.metadata.name}'

  - name: check-status
    inputs:
      parameters:
      - name: job-name
    script:
      image: mirrors-ssl.aliyuncs.com/bitnami/kubectl:latest
      command: [bash]
      source: |
        #!/bin/bash
        set -ex

        JOB_NAME="{{inputs.parameters.job-name}}"

        STATUS=$(kubectl get pytorchjob $JOB_NAME -o jsonpath='{.status.conditions[?(@.type=="Succeeded")].status}')

        if [ "$STATUS" = "True" ]; then
          echo "PyTorch Job succeeded."
          exit 0
        else
          echo "PyTorch Job failed or is still running."
          exit 1
        fi