气流¶
对于支持基于 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_277834/454499020.py:3 RemovedInAirflow3Warning: Param `schedule_interval` is deprecated and will be removed in a future release. Please use `schedule` instead.
[2023-04-04 01:08:21,081] {taskinstance.py:1165} INFO - Dependencies all met for <TaskInstance: example_python_operator-xzg1ld7lds2m9.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2023-04-04 01:08:21,087] {taskinstance.py:1165} INFO - Dependencies all met for <TaskInstance: example_python_operator-xzg1ld7lds2m9.hello_torchx manual__2021-09-13T00:00:00+00:00 [None]>
[2023-04-04 01:08:21,088] {taskinstance.py:1362} INFO -
--------------------------------------------------------------------------------
[2023-04-04 01:08:21,089] {taskinstance.py:1363} INFO - Starting attempt 1 of 1
[2023-04-04 01:08:21,089] {taskinstance.py:1364} INFO -
--------------------------------------------------------------------------------
[2023-04-04 01:08:21,104] {taskinstance.py:1383} INFO - Executing <Task(_PythonDecoratedOperator): hello_torchx> on 2021-09-13 00:00:00+00:00
[2023-04-04 01:08:21,302] {taskinstance.py:1590} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=example_python_operator-xzg1ld7lds2m9
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-04-04 01:08:22,027] {api.py:70} INFO - Tracker configurations: {}
[2023-04-04 01:08:22,031] {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
[2023-04-04 01:08:22,032] {local_scheduler.py:721} INFO - Log directory is: /tmp/torchx_bag8ie9r
Hello, TorchX!
[2023-04-04 01:08:22,140] {python.py:177} INFO - Done. Returned value was: None
[2023-04-04 01:08:22,146] {taskinstance.py:1401} INFO - Marking task as SUCCESS. dag_id=example_python_operator-xzg1ld7lds2m9, task_id=hello_torchx, execution_date=20210913T000000, start_date=20230404T010821, end_date=20230404T010822
如果一切顺利，您应该会看到上面打印的字样。Hello, TorchX!
后续步骤¶
查看运行器 API 文档，了解有关 TorchX 的编程用法的更多信息
浏览可在 Airflow 管道中使用的内置组件集合