docs: add airflow example

Author: d4l3k

.github/workflows/doc-build.yaml

@@ -37,6 +37,11 @@ jobs:
       - name: Doc Build
         run: |
           cd docs
+          # start airflow
+          airflow standalone&
+          # wait 5 seconds for airflow to start
+          sleep 5
+
           make html
       - name: Papermill
         run: |

docs/requirements.txt

@@ -10,3 +10,4 @@ jupytext
 ipython_genutils
 # https://github.com/jupyter/nbconvert/issues/1736
 jinja2<=3.0.3
+apache-airflow

docs/source/index.rst

@@ -68,6 +68,7 @@ Works With
    :caption: Pipelines
 
    pipelines/kfp
+   pipelines/airflow.md
 
 
 Examples

docs/source/pipelines/airflow.md

@@ -0,0 +1,97 @@
+---
+jupyter:
+  jupytext:
+    text_representation:
+      extension: .md
+      format_name: markdown
+      format_version: '1.3'
+      jupytext_version: 1.13.7
+  kernelspec:
+    display_name: Python 3
+    language: python
+    name: python3
+---
+
+# Airflow
+
+For pipelines that support Python based execution you can directly use the
+TorchX API. TorchX is designed to be easily integrated in to other applications
+via the programmatic API. No special Airflow integrations are needed.
+
+With TorchX, you can use Airflow for the pipeline orchestration and run your
+PyTorch application (i.e. distributed training) on a remote GPU cluster.
+
+```python
+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)
+```
+
+To launch a TorchX job from Airflow you can create a Airflow Python task to
+import the runner, launch the job and wait for it to complete. If you're running
+on a remote cluster you may need to use the virtualenv task to install the
+`torchx` package.
+
+```python
+@task(task_id=f'hello_torchx')
+def run_torchx(echo):
+    """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", "Hello, TorchX!"],
+            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="")
+```
+
+Once we have the task defined we can put it into a Airflow DAG and run it like
+normal.
+
+```python
+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("foo")
+
+
+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
+```
+
+If all goes well you should see `Hello, TorchX!` printed above.

torchx/specs/api.py

@@ -477,6 +477,25 @@ def __repr__(self) -> str:
         app_status_dict["state"] = repr(app_status_dict["state"])
         return yaml.dump({"AppStatus": app_status_dict})
 
+    def raise_for_status(self) -> None:
+        """
+        raise_for_status will raise an AppStatusError if the state is not SUCCEEDED.
+        """
+        if self.state != AppState.SUCCEEDED:
+            raise AppStatusError(self, f"job did not succeed: {self}")
+
+
+class AppStatusError(Exception):
+    """
+    AppStatusError is raised when the job status is in an exceptional state i.e.
+    not SUCCEEDED.
+    """
+
+    def __init__(self, status: AppStatus, *args: object) -> None:
+        super().__init__(*args)
+
+        self.status = status
+
 
 # valid run cfg values; only support primitives (str, int, float, bool, List[str])
 # TODO(wilsonhong): python 3.9+ supports list[T] in typing, which can be used directly

torchx/specs/test/api_test.py

@@ -13,22 +13,23 @@
 import torchx.specs.named_resources_aws as named_resources_aws
 from torchx.specs import named_resources, resource
 from torchx.specs.api import (
-    _TERMINAL_STATES,
     AppDef,
     AppDryRunInfo,
     AppState,
     AppStatus,
+    AppStatusError,
     CfgVal,
-    get_type_name,
     InvalidRunConfigException,
-    macros,
-    MalformedAppHandleException,
     MISSING,
+    MalformedAppHandleException,
     NULL_RESOURCE,
-    parse_app_handle,
     Resource,
     RetryPolicy,
     Role,
+    _TERMINAL_STATES,
+    get_type_name,
+    macros,
+    parse_app_handle,
     runopts,
 )
 
@@ -86,6 +87,24 @@ def test_serialize_embed_json(self) -> None:
 """,
         )
 
+    def test_raise_on_status(self) -> None:
+        AppStatus(state=AppState.SUCCEEDED).raise_for_status()
+
+        with self.assertRaisesRegex(
+            AppStatusError, r"(?s)job did not succeed:.*FAILED.*"
+        ):
+            AppStatus(state=AppState.FAILED).raise_for_status()
+
+        with self.assertRaisesRegex(
+            AppStatusError, r"(?s)job did not succeed:.*CANCELLED.*"
+        ):
+            AppStatus(state=AppState.CANCELLED).raise_for_status()
+
+        with self.assertRaisesRegex(
+            AppStatusError, r"(?s)job did not succeed:.*RUNNING.*"
+        ):
+            AppStatus(state=AppState.RUNNING).raise_for_status()
+
 
 class ResourceTest(unittest.TestCase):
     def test_copy_resource(self) -> None: