adapters/kfp: support distributed training (#109)

Summary:
This adds a new `resource_to_app` KFP adapter that allows adapting an app to a kfp ResourceOp that launches the operator using the Volcano scheduler. This reuses the same code that creates the resources for the kubernetes scheduler and embeds the resource inside a KFP pipeline.

This isn't supported under KFP v2 since it interacts directly with kubernetes resources/volcano. This also requires volcano to be installed on the cluster to use which is why it's a new adapter instead of automatically being used.

This is still fairly experimental and once KFP has better distributed support we likely want to rely on that instead since this has some less than ideal UX. You need to use the CLI to access the individual worker logs and there isn't any support for UI metadata yet.

UI metadata I think can be added by providing an output annotation for argo as part of the resource but I haven't looked into it.

Pull Request resolved: #109

Test Plan:
```
pyre
pytest
python dist_pipeline.py
```
http://5ab6bab9-istiosystem-istio-2af2-1926929629.us-west-2.elb.amazonaws.com/_/pipeline/#/runs/details/27707de9-bc67-42da-ab86-af2127ee54d1

![20210726_14h12m04s_grim](https://user-images.githubusercontent.com/909104/127059928-b4787429-e895-4b97-b53e-c6262e99c52b.png)

Reviewed By: kiukchung

Differential Revision: D29921246

Pulled By: d4l3k

fbshipit-source-id: ead6d0d4265cd22dd46fb89621a938c7281d7f5b

Author: d4l3k

.gitignore

@@ -138,3 +138,4 @@ dmypy.json
 cython_debug/
 
 wordlist.dic
+pipeline.yaml

docs/papermill.sh

@@ -7,8 +7,15 @@
 
 WORK_DIR=/tmp/papermill
 
-set -e
+set -ex
 mkdir -p "$WORK_DIR"
+
+# create empty master.tar.gz file and setup symlinks instead of pulling from
+# master so we can handle local changes
+tar -cJf "$WORK_DIR/master.tar.gz" -T /dev/null
+ROOT="$(pwd)/.."
+(cd "$WORK_DIR" && ln -s "$ROOT/torchx" . && ln -s "$ROOT/examples" .)
+
 files="$(find "$(pwd)"/build -name '*.ipynb')"
 for file in $files
 do

docs/source/pipelines.rst

@@ -13,6 +13,7 @@ torchx.pipelines.kfp
 .. currentmodule:: torchx.pipelines.kfp.adapter
 
 .. autofunction:: container_from_app
+.. autofunction:: resource_from_app
 .. autofunction:: component_from_app
 .. autofunction:: component_spec_from_app
 

examples/pipelines/kfp/dist_pipeline.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Distributed KubeFlow Pipelines Example
+======================================
+
+This is an example KFP pipeline that uses resource_from_app to launch a
+distributed operator using the kubernetes/volcano job scheduler. This only works
+in Kubernetes KFP clusters with https://volcano.sh/en/docs/ installed on them.
+"""
+
+import kfp
+from torchx import specs
+from torchx.pipelines.kfp.adapter import resource_from_app
+
+
+def pipeline() -> None:
+    # First we define our AppDef for the component, we set
+    echo_app = specs.AppDef(
+        name="test-dist",
+        roles=[
+            specs.Role(
+                name="dist-echo",
+                image="alpine",
+                entrypoint="/bin/echo",
+                args=["hello dist!"],
+                num_replicas=3,
+            ),
+        ],
+    )
+
+    # To convert the TorchX AppDef into a KFP container we use
+    # the resource_from_app adapter. This takes generates a KFP Kubernetes
+    # resource operator definition from the TorchX app def and instantiates it.
+    echo_container: kfp.dsl.BaseOp = resource_from_app(echo_app, queue="test")
+
+
+# %%
+# To generate the pipeline definition file we need to call into the KFP compiler
+# with our pipeline function.
+
+kfp.compiler.Compiler().compile(
+    pipeline_func=pipeline,
+    package_path="pipeline.yaml",
+)
+
+with open("pipeline.yaml", "rt") as f:
+    print(f.read())
+
+# %%
+# Once this has all run you should have a pipeline file (typically
+# pipeline.yaml) that you can upload to your KFP cluster via the UI or
+# a kfp.Client.
+#
+# See the
+# `KFP SDK Examples <https://www.kubeflow.org/docs/components/pipelines/tutorials/sdk-examples/#examples>`_
+# for more info on launching KFP pipelines.
+
+# %%
+# See the :ref:`Advanced KubeFlow Pipelines Example` for how to chain multiple
+# components together and use builtin components.

torchx/pipelines/kfp/adapter.py

@@ -22,6 +22,7 @@
     V1VolumeMount,
     V1EmptyDirVolumeSource,
 )
+from torchx.schedulers.kubernetes_scheduler import app_to_resource
 from torchx.specs import api
 from typing_extensions import Protocol
 
@@ -229,3 +230,41 @@ def container_from_app(
     """
     factory = component_from_app(app, ui_metadata)
     return factory(*args, **kwargs)
+
+
+def resource_from_app(
+    app: api.AppDef,
+    queue: str,
+) -> dsl.ResourceOp:
+    """
+    resource_from_app generates a KFP ResourceOp from the provided app that uses
+    the Volcano job scheduler on Kubernetes to run distributed apps. See
+    https://volcano.sh/en/docs/ for more info on Volcano and how to install.
+
+    Args:
+        app: The torchx AppDef to adapt.
+        queue: the Volcano queue to schedule the operator in.
+
+    >>> import kfp
+    >>> from torchx import specs
+    >>> from torchx.pipelines.kfp.adapter import resource_from_app
+    >>> app_def = specs.AppDef(
+    ...     name="trainer",
+    ...     roles=[specs.Role("trainer", image="foo:latest", num_replicas=3)],
+    ... )
+    >>> def pipeline():
+    ...     trainer = resource_from_app(app_def, queue="test")
+    ...     print(trainer)
+    >>> kfp.compiler.Compiler().compile(
+    ...     pipeline_func=pipeline,
+    ...     package_path="/tmp/pipeline.yaml",
+    ... )
+    {'ResourceOp': {... 'name': 'trainer-0', ... 'name': 'trainer-1', ... 'name': 'trainer-2', ...}}
+    """
+    return dsl.ResourceOp(
+        name=app.name,
+        action="create",
+        success_condition="status.state.phase = Completed",
+        failure_condition="status.state.phase = Failed",
+        k8s_resource=app_to_resource(app, queue),
+    )

torchx/schedulers/kubernetes_scheduler.py

@@ -164,6 +164,56 @@ def role_to_pod(name: str, role: Role) -> "V1Pod":
     )
 
 
+def app_to_resource(app: AppDef, queue: str) -> Dict[str, object]:
+    """
+    app_to_resource creates a volcano job kubernetes resource definition from
+    the provided AppDef. The resource definition can be used to launch the
+    app on Kubernetes.
+
+    To support macros we generate one task per replica instead of using the
+    volcano `replicas` field since macros change the arguments on a per
+    replica basis.
+
+    Volcano has two levels of retries: one at the task level and one at the
+    job level. When using the APPLICATION retry policy, the job level retry
+    count is set to the minimum of the max_retries of the roles.
+    """
+    tasks = []
+    for i, role in enumerate(app.roles):
+        for replica_id in range(role.num_replicas):
+            values = macros.Values(
+                img_root="",
+                app_id=macros.app_id,
+                replica_id=str(replica_id),
+            )
+            name = f"{role.name}-{replica_id}"
+            replica_role = values.apply(role)
+            pod = role_to_pod(name, replica_role)
+            tasks.append(
+                {
+                    "replicas": 1,
+                    "name": name,
+                    "template": pod,
+                    "maxRetry": role.max_retries,
+                    "policies": RETRY_POLICIES[role.retry_policy],
+                }
+            )
+
+    job_retries = min(role.max_retries for role in app.roles)
+    resource: Dict[str, object] = {
+        "apiVersion": "batch.volcano.sh/v1alpha1",
+        "kind": "Job",
+        "metadata": {"generateName": f"{app.name}-"},
+        "spec": {
+            "schedulerName": "volcano",
+            "queue": queue,
+            "tasks": tasks,
+            "maxRetry": job_retries,
+        },
+    }
+    return resource
+
+
 @dataclass
 class KubernetesJob:
     resource: Dict[str, object]
@@ -237,39 +287,9 @@ def _submit_dryrun(
         self, app: AppDef, cfg: RunConfig
     ) -> AppDryRunInfo[KubernetesJob]:
         queue = cfg.get("queue")
-        tasks = []
-        for i, role in enumerate(app.roles):
-            for replica_id in range(role.num_replicas):
-                values = macros.Values(
-                    img_root="",
-                    app_id=macros.app_id,
-                    replica_id=str(replica_id),
-                )
-                name = f"{role.name}-{replica_id}"
-                replica_role = values.apply(role)
-                pod = role_to_pod(name, replica_role)
-                tasks.append(
-                    {
-                        "replicas": 1,
-                        "name": name,
-                        "template": pod,
-                        "maxRetry": role.max_retries,
-                        "policies": RETRY_POLICIES[role.retry_policy],
-                    }
-                )
-
-        job_retries = min(role.max_retries for role in app.roles)
-        resource: Dict[str, object] = {
-            "apiVersion": "batch.volcano.sh/v1alpha1",
-            "kind": "Job",
-            "metadata": {"generateName": f"{app.name}-"},
-            "spec": {
-                "schedulerName": "volcano",
-                "queue": queue,
-                "tasks": tasks,
-                "maxRetry": job_retries,
-            },
-        }
+        if not isinstance(queue, str):
+            raise TypeError(f"config value 'queue' must be a string, got {queue}")
+        resource = app_to_resource(app, queue)
         req = KubernetesJob(resource=resource)
         info = AppDryRunInfo(req, repr)
         info._app = app