Airflow

对于支持基于 Python 执行的管道，您可以直接使用 TorchX API。TorchX 被设计为可以通过编程 API 轻松集成到其他应用程序中。无需特殊的 Airflow 集成。

使用 TorchX，您可以使用 Airflow 进行管道编排，并在远程 GPU 集群上运行您的 PyTorch 应用程序（例如分布式训练）。

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)

要从 Airflow 启动 TorchX 作业，您可以创建一个 Airflow Python 任务来导入运行程序、启动作业并等待作业完成。如果您在远程集群上运行，则可能需要使用 virtualenv 任务来安装 torchx 包。

@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="")

定义好任务后，我们可以将其放入 Airflow DAG 并像平常一样运行它。

from torchx.schedulers.ids import make_unique

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

/tmp/ipykernel_52183/454499020.py:3 RemovedInAirflow3Warning: Param `schedule_interval` is deprecated and will be removed in a future release. Please use `schedule` instead.

[2024-07-17T02:04:16.962+0000] {taskinstance.py:2076} INFO - Dependencies all met for dep_context=non-requeueable deps ti=<TaskInstance: example_python_operator-mz5brbw7kwqxs.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2024-07-17T02:04:16.968+0000] {taskinstance.py:2076} INFO - Dependencies all met for dep_context=requeueable deps ti=<TaskInstance: example_python_operator-mz5brbw7kwqxs.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2024-07-17T02:04:16.969+0000] {taskinstance.py:2306} INFO - Starting attempt 1 of 1
[2024-07-17T02:04:16.970+0000] {taskinstance.py:2388} WARNING - cannot record queued_duration for task hello_torchx because previous state change time has not been saved
[2024-07-17T02:04:16.981+0000] {taskinstance.py:2330} INFO - Executing <Task(_PythonDecoratedOperator): hello_torchx> on 2021-09-13 00:00:00+00:00
[2024-07-17T02:04:17.253+0000] {taskinstance.py:2648} INFO - Exporting env vars: AIRFLOW_CTX_DAG_OWNER='airflow' AIRFLOW_CTX_DAG_ID='example_python_operator-mz5brbw7kwqxs' AIRFLOW_CTX_TASK_ID='hello_torchx' AIRFLOW_CTX_EXECUTION_DATE='2021-09-13T00:00:00+00:00' AIRFLOW_CTX_TRY_NUMBER='1' AIRFLOW_CTX_DAG_RUN_ID='manual__2021-09-13T00:00:00+00:00'
[2024-07-17T02:04:17.255+0000] {taskinstance.py:430} INFO - ::endgroup::
[2024-07-17T02:04:17.953+0000] {api.py:72} INFO - Tracker configurations: {}
[2024-07-17T02:04:17.957+0000] {local_scheduler.py:771} INFO - Log directory not set in scheduler cfg. Creating a temporary log dir that will be deleted on exit. To preserve log directory set the `log_dir` cfg option
[2024-07-17T02:04:17.957+0000] {local_scheduler.py:777} INFO - Log directory is: /tmp/torchx_39w81v6k
Hello, TorchX!
[2024-07-17T02:04:18.062+0000] {python.py:237} INFO - Done. Returned value was: None
[2024-07-17T02:04:18.063+0000] {taskinstance.py:441} INFO - ::group::Post task execution logs
[2024-07-17T02:04:18.071+0000] {taskinstance.py:1206} INFO - Marking task as SUCCESS. dag_id=example_python_operator-mz5brbw7kwqxs, task_id=hello_torchx, run_id=manual__2021-09-13T00:00:00+00:00, execution_date=20210913T000000, start_date=20240717T020416, end_date=20240717T020418

如果一切顺利，您应该看到 Hello, TorchX! 打印在上面。

下一步

• 查看 运行程序 API 文档 以了解有关 TorchX 编程用法的更多信息

• 浏览 内置组件 集合，这些组件可用于您的 Airflow 管道