气流¶
对于支持基于 Python 的执行的管道，您可以直接使用 TorchX API。TorchX 旨在通过编程 API 轻松集成到其他应用程序中。不需要特殊的 Airflow 集成。
借助 TorchX，您可以使用 Airflow 进行管道编排，并在远程 GPU 集群上运行 PyTorch 应用程序（即分布式训练）。
[1]:
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
[2]:
@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 中，并像往常一样运行它。
[3]:
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_4040/454499020.py:3 RemovedInAirflow3Warning: Param `schedule_interval` is deprecated and will be removed in a future release. Please use `schedule` instead.
[2023-10-19T01:33:50.915+0000] {taskinstance.py:1159} INFO - Dependencies all met for dep_context=non-requeueable deps ti=<TaskInstance: example_python_operator-kw25q9fcp19sdd.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2023-10-19T01:33:50.920+0000] {taskinstance.py:1159} INFO - Dependencies all met for dep_context=requeueable deps ti=<TaskInstance: example_python_operator-kw25q9fcp19sdd.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2023-10-19T01:33:50.921+0000] {taskinstance.py:1361} INFO - Starting attempt 1 of 1
[2023-10-19T01:33:50.922+0000] {taskinstance.py:1430} WARNING - cannot record queued_duration for task hello_torchx because previous state change time has not been saved
[2023-10-19T01:33:50.931+0000] {taskinstance.py:1382} INFO - Executing <Task(_PythonDecoratedOperator): hello_torchx> on 2021-09-13 00:00:00+00:00
[2023-10-19T01:33:51.168+0000] {taskinstance.py:1662} INFO - Exporting env vars: AIRFLOW_CTX_DAG_OWNER='airflow' AIRFLOW_CTX_DAG_ID='example_python_operator-kw25q9fcp19sdd' 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'
[2023-10-19T01:33:51.787+0000] {api.py:70} INFO - Tracker configurations: {}
[2023-10-19T01:33:51.791+0000] {local_scheduler.py:716} 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
[2023-10-19T01:33:51.791+0000] {local_scheduler.py:722} INFO - Log directory is: /tmp/torchx_okvtgkb3
Hello, TorchX!
[2023-10-19T01:33:51.898+0000] {python.py:194} INFO - Done. Returned value was: None
[2023-10-19T01:33:51.903+0000] {taskinstance.py:1400} INFO - Marking task as SUCCESS. dag_id=example_python_operator-kw25q9fcp19sdd, task_id=hello_torchx, execution_date=20210913T000000, start_date=20231019T013350, end_date=20231019T013351
如果一切顺利，您应该会看到上面打印的字样。Hello, TorchX!
后续步骤¶
查看运行器 API 文档，了解有关 TorchX 的编程用法的更多信息
浏览可在 Airflow 管道中使用的内置组件集合