added KFP adapter for new torchx.specs.api interface (#10)

Summary:
This creates a KFP adapter for the new specs interface.

```
import torchx.specs.api as torchx
from kfp import compiler, components, dsl

app = torchx.Application(...)
kfp_copy: Callable = component_from_app(app)

def pipeline() -> dsl.PipelineParam:
    a = kfp_copy()
    b = kfp_copy()
    b.after(a)
    return b

compiler.Compiler().compile(pipeline, "pipeline.zip")
```

Pull Request resolved: pytorch#10

Reviewed By: kiukchung

Differential Revision: D28658161

Pulled By: d4l3k

fbshipit-source-id: 39969b6da9ab651c71dd348db67f610dbf981c9a

Author: d4l3k

torchx/pipelines/kfp/adapter.py

@@ -7,14 +7,16 @@
 
 import copy
 import os
-from typing import Type, Callable, List, Optional, Dict
+from typing import Callable, Dict, List, Optional, Type
 
 import yaml
 from kfp import components, dsl
 from torchx.runtime.component import Component, is_optional
+from torchx.specs import api
 
 from .version import __version__ as __version__  # noqa F401
 
+
 TORCHX_CONTAINER_ENV: str = "TORCHX_CONTAINER"
 TORCHX_CONTAINER: str = os.getenv(
     TORCHX_CONTAINER_ENV,
@@ -111,3 +113,41 @@ def outputs(self) -> Dict[str, dsl.PipelineParam]:
     @property
     def output(self) -> dsl.PipelineParam:
         ...
+
+
+def component_spec_from_app(app: api.Application) -> str:
+    assert len(app.roles) == 1, f"KFP adapter only support one role, got {app.roles}"
+
+    role = app.roles[0]
+    assert (
+        role.num_replicas == 1
+    ), f"KFP adapter only supports one replica, got {app.num_replicas}"
+    assert role.container != api.NULL_CONTAINER, "missing container for KFP"
+
+    container = role.container
+    assert container.base_image is None, "KFP adapter does not support base_image"
+    assert (
+        container.resources == api.NULL_RESOURCE
+    ), "KFP adapter requires you to specify resources in the pipeline"
+    assert len(container.port_map) == 0, "KFP adapter does not support port_map"
+
+    command = [role.entrypoint, *role.args]
+
+    spec = {
+        "name": f"{app.name}-{role.name}",
+        "description": f"KFP wrapper for TorchX component {app.name}, role {role.name}",
+        "implementation": {
+            "container": {
+                "image": container.image,
+                "command": command,
+                "env": role.env,
+            }
+        },
+    }
+    return yaml.dump(spec)
+
+
+# pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
+def component_from_app(app: api.Application) -> Callable:
+    spec = component_spec_from_app(app)
+    return components.load_component_from_text(spec)

torchx/pipelines/kfp/test/adapter_test.py

@@ -8,12 +8,18 @@
 import os.path
 import tempfile
 import unittest
-from typing import TypedDict, Optional
+from typing import Callable, Optional, TypedDict
 
 from kfp import compiler, components, dsl
 from torchx.apps.io.copy import Copy
-from torchx.pipelines.kfp.adapter import component_spec, TorchXComponent
+from torchx.pipelines.kfp.adapter import (
+    TorchXComponent,
+    component_from_app,
+    component_spec,
+    component_spec_from_app,
+)
 from torchx.runtime.component import Component
+from torchx.specs import api
 
 
 class Config(TypedDict):
@@ -103,3 +109,60 @@ class KFPCopy(TorchXComponent, component=Copy, image="foo"):
         print(copy)
         # pyre-fixme[16]: `KFPCopy` has no attribute `component_ref`.
         self.assertEqual(copy.component_ref.spec.implementation.container.image, "foo")
+
+
+class KFPSpecsTest(unittest.TestCase):
+    """
+    tests KFP components using torchx.specs.api
+    """
+
+    def _test_app(self) -> api.Application:
+        container = api.Container(image="pytorch/torchx:latest")
+        trainer_role = (
+            api.Role(name="trainer")
+            .runs(
+                "main",
+                "--output-path",
+                "blah",
+                FOO="bar",
+            )
+            .on(container)
+            .replicas(1)
+        )
+
+        return api.Application("test").of(trainer_role)
+
+    def test_component_spec_from_app(self) -> None:
+        app = self._test_app()
+
+        spec = component_spec_from_app(app)
+        self.assertIsNotNone(components.load_component_from_text(spec))
+        self.assertEqual(
+            spec,
+            """description: KFP wrapper for TorchX component test, role trainer
+implementation:
+  container:
+    command:
+    - main
+    - --output-path
+    - blah
+    env:
+      FOO: bar
+    image: pytorch/torchx:latest
+name: test-trainer
+""",
+        )
+
+    def test_pipeline(self) -> None:
+        app = self._test_app()
+        # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
+        kfp_copy: Callable = component_from_app(app)
+
+        def pipeline() -> dsl.PipelineParam:
+            a = kfp_copy()
+            b = kfp_copy()
+            b.after(a)
+            return b
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            compiler.Compiler().compile(pipeline, os.path.join(tmpdir, "pipeline.zip"))