docs: add airflow example (#519)

Summary:
This adds an example for how to use Airflow with TorchX. This example requires airflow server to be running on the local machine during the docs build process.

This also adds in a `AppStatus.raise_for_status()` method to match the requests/urllib3 behavior which just makes checking for whether or not a job failed easier.

Build is working locally -- might need to fiddle with CI to get airflow running there correctly still

Pull Request resolved: #519

Test Plan:
```
cd docs && make html
```
![Screenshot 2022-06-14 at 18-01-37 Airflow — PyTorch_TorchX main documentation](https://user-images.githubusercontent.com/909104/173714270-c3429671-b0e8-4bc1-83ca-9a487e2381b5.png)

Reviewed By: priyaramani

Differential Revision: D37161851

Pulled By: d4l3k

fbshipit-source-id: f7ee74741a25a21f7d4f376e8ab07070ae3b233d

Author: d4l3k

.github/workflows/doc-build.yaml

@@ -26,6 +26,12 @@ jobs:
       - name: Install TorchX
         run: |
           python setup.py develop
+      - name: Start Airflow
+        run: |
+          # start airflow in background
+          airflow standalone &
+          # wait 5 seconds for airflow to start
+          sleep 5
       - name: Doc Test
         run: |
           cd docs

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,104 @@
+---
+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(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="")
+```
+
+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("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
+```
+
+If all goes well you should see `Hello, TorchX!` printed above.
+
+## Next Steps
+
+* Checkout the [runner API documentation](../runner.rst) to learn more about
+  programmatic usage of TorchX
+* Browse through the collection of [builtin components](../components/overview.rst)
+  which can be used in your Airflow pipeline

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

@@ -18,6 +18,7 @@
     AppDryRunInfo,
     AppState,
     AppStatus,
+    AppStatusError,
     CfgVal,
     get_type_name,
     InvalidRunConfigException,
@@ -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: