气流¶
对于支持基于 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_3974/454499020.py:3 RemovedInAirflow3Warning: Param `schedule_interval` is deprecated and will be removed in a future release. Please use `schedule` instead.
[2022-12-29 22:58:49,311] {taskinstance.py:1087} INFO - Dependencies all met for <TaskInstance: example_python_operator-wm3n61xf6z7h5c.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2022-12-29 22:58:49,318] {taskinstance.py:1087} INFO - Dependencies all met for <TaskInstance: example_python_operator-wm3n61xf6z7h5c.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2022-12-29 22:58:49,319] {taskinstance.py:1283} INFO -
--------------------------------------------------------------------------------
[2022-12-29 22:58:49,320] {taskinstance.py:1284} INFO - Starting attempt 1 of 1
[2022-12-29 22:58:49,320] {taskinstance.py:1285} INFO -
--------------------------------------------------------------------------------
[2022-12-29 22:58:49,332] {taskinstance.py:1304} INFO - Executing <Task(_PythonDecoratedOperator): hello_torchx> on 2021-09-13 00:00:00+00:00
[2022-12-29 22:58:49,544] {taskinstance.py:1511} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=example_python_operator-wm3n61xf6z7h5c
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
[2022-12-29 22:58:50,071] {local_scheduler.py:715} 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
[2022-12-29 22:58:50,072] {local_scheduler.py:721} INFO - Log directory is: /tmp/torchx_h9rktfpg
Hello, TorchX!
[2022-12-29 22:58:50,183] {python.py:177} INFO - Done. Returned value was: None
[2022-12-29 22:58:50,191] {taskinstance.py:1322} INFO - Marking task as SUCCESS. dag_id=example_python_operator-wm3n61xf6z7h5c, task_id=hello_torchx, execution_date=20210913T000000, start_date=20221229T225849, end_date=20221229T225850
如果一切顺利，您应该会看到上面打印的字样。Hello, TorchX!
后续步骤¶
查看运行器 API 文档，了解有关 TorchX 的编程用法的更多信息
浏览可在 Airflow 管道中使用的内置组件集合