Add components integration test (#182)

Summary:
Pull Request resolved: #182

The diff introduces an integ tests component framework. It consists of the following parts:

* `ComponentProvider` - an abstract class that users will use to implement the `AppDef` instantiation. Since many components have specific arguments that cannot be autofilled, e.g. `utils.copy` mast have  proper files as inputs, it does not make sense to auto-generate data

* Classes that extend `ComponentProvider`  and provide ability to retrieve scheduler-dependend app defs.

* Integraiton tests class. The class defines `run_components` method that allows users to run all component providers given a module.

* Components integration tests github workflow. The workflow uses Integration tests library to run the defined components.

Reviewed By: d4l3k

Differential Revision: D30980471

fbshipit-source-id: b77b5139303ba8943a5ae91333f51217f47f8811

Author: aivanou

.github/workflows/kubernetes-integration-tests.yaml renamed to .github/workflows/components-integration-tests.yaml

@@ -1,4 +1,4 @@
-name: Kubernetes Integration Tests
+name: Components Integration Tests
 
 on:
   push:
@@ -7,7 +7,7 @@ on:
   pull_request:
 
 jobs:
-  kubernetes-launch:
+  components-launch:
     runs-on: ubuntu-18.04
     steps:
       - name: Setup Python
@@ -26,28 +26,28 @@ jobs:
           if [ -n "$AWS_ACCESS_KEY_ID" ]; then
             aws eks update-kubeconfig --region=us-west-2 --name=${{ secrets.EKS_CLUSTER_NAME }}
           fi
+      - name: Configure Docker
+        env:
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: |
+          set -eux
+          if [ -n "$AWS_ACCESS_KEY_ID" ]; then
+            aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 495572122715.dkr.ecr.us-west-2.amazonaws.com
+          fi
       - name: Install dependencies
         run: |
           set -eux
           pip install -e .[kubernetes]
-      - name: Run Kubernetes Integration Tests
+      - name: Run Components Integration Tests
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
         run: |
-          set -eux
           if [ -z "$AWS_ACCESS_KEY_ID" ]; then
             # only dryrun if no secrets
             ARGS="--dryrun"
           else
             ARGS=
           fi
-
-          torchx runopts kubernetes
-          APP_ID="$(torchx run --wait $ARGS --scheduler kubernetes \
-            --scheduler_args queue=default utils.echo \
-            --image alpine:latest --num_replicas 3)"
-          if [ "$ARGS" != "--dryrun" ]; then
-            torchx status "$APP_ID"
-            torchx describe "$APP_ID"
-          fi
+          scripts/component_integration_tests.py $ARGS

examples/apps/datapreproc/component.py

@@ -36,6 +36,7 @@ def data_preproc(
         input_path: Url-like path to fetch the imagenet dataset
         env: Env variables to transfer to the user script
         resource: String representation of the resource
+        dryrun: Starts the app, but does not actually perform any work.
 
     Returns:
         specs.AppDef: Torchx AppDef

examples/apps/dist_cifar/component.py

@@ -54,16 +54,16 @@ def trainer(
         specs.AppDef: Torchx AppDef
     """
     return ddp(
-        image,
-        "examples/apps/dist_cifar/train.py",
-        rdzv_backend,
-        rdzv_endpoint,
-        resource,
-        nnodes,
-        nproc_per_node,
-        base_image,
-        "cifar-trainer",
-        "worker",
-        env,
         *script_args,
+        image=image,
+        entrypoint="examples/apps/dist_cifar/train.py",
+        rdzv_backend=rdzv_backend,
+        rdzv_endpoint=rdzv_endpoint,
+        resource=resource,
+        nnodes=nnodes,
+        nproc_per_node=nproc_per_node,
+        base_image=base_image,
+        name="cifar-trainer",
+        role="worker",
+        env=env,
     )

examples/apps/dist_cifar/train.py

@@ -151,6 +151,11 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Trainer that trains the last layer of pretrained resnet18 for cifar10 classification"
     )
+    parser.add_argument(
+        "--dryrun",
+        help=argparse.SUPPRESS,
+        action="store_true",
+    )
     parser.add_argument(
         "--epochs", type=int, default=1, help="number of epochs to train"
     )
@@ -173,6 +178,9 @@ def get_gpu_devices() -> int:
 
 def main() -> None:
     args = parse_args(sys.argv[1:])
+    if args.dryrun:
+        print("App dist_cifar.train started successfully")
+        return
     gpus = get_gpu_devices()
     batch_size = args.batch_size
     num_nodes = int(os.environ["GROUP_WORLD_SIZE"])

examples/apps/lightning_classy_vision/component.py

@@ -20,7 +20,7 @@
 def trainer(
     image: str,
     output_path: str,
-    data_path: str,
+    data_path: Optional[str] = None,
     entrypoint: str = "examples/apps/lightning_classy_vision/train.py",
     load_path: str = "",
     log_path: str = "/logs",
@@ -35,7 +35,8 @@ def trainer(
         image: image to run (e.g. foobar:latest)
         output_path: output path for model checkpoints (e.g. file:///foo/bar)
         load_path: path to load pretrained model from
-        data_path: path to the data to load
+        data_path: path to the data to load, if data_path is not provided,
+            auto generated test data will be used
         entrypoint: user script to launch.
         log_path: path to save tensorboard logs to
         resource: the resources to use
@@ -51,15 +52,17 @@ def trainer(
         load_path,
         "--log_path",
         log_path,
-        "--data_path",
-        data_path,
         "--epochs",
         str(epochs),
     ]
+    if data_path:
+        args += ["--data_path", data_path]
+    else:
+        args.append("--test")
     if skip_export:
         args.append("--skip_export")
     return torchx.AppDef(
-        name="examples-lightning_classy_vision-trainer",
+        name="cv-trainer",
         roles=[
             torchx.Role(
                 name="worker",
@@ -81,6 +84,7 @@ def interpret(
     data_path: str,
     output_path: str,
     resource: Optional[str] = None,
+    entrypoint: str = "examples/apps/lightning_classy_vision/interpret.py",
 ) -> torchx.AppDef:
     """Runs the model interpretability app on the model outputted by the training
     component.
@@ -91,13 +95,14 @@ def interpret(
         data_path: path to the data to load
         output_path: output path for model checkpoints (e.g. file:///foo/bar)
         resource: the resources to use
+        entrypoint: user script to launch.
     """
     return torchx.AppDef(
-        name="examples-lightning_classy_vision-interpret",
+        name="cv-interpret",
         roles=[
             torchx.Role(
                 name="worker",
-                entrypoint="examples/apps/lightning_classy_vision/interpret.py",
+                entrypoint=entrypoint,
                 args=[
                     "--load_path",
                     load_path,

examples/apps/lightning_classy_vision/train.py

@@ -120,6 +120,7 @@ def main(argv: List[str]) -> None:
 
         # Initialize a trainer
         trainer = pl.Trainer(
+            accelerator="ddp2",
             logger=logger,
             max_epochs=args.epochs,
             callbacks=[checkpoint_callback],

scripts/component_integration_tests.py

@@ -0,0 +1,92 @@
+#!/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.
+
+"""
+Kubernetes integration tests.
+"""
+import argparse
+import os
+
+# pyre-ignore-all-errors[21] # Cannot find module utils
+# pyre-ignore-all-errors[11]
+
+import example_app_defs as examples_app_defs_providers
+import torchx.components.integration_tests.component_provider as component_provider
+from integ_test_utils import (
+    MissingEnvError,
+    build_images,
+    push_images,
+    BuildInfo,
+)
+from torchx.components.integration_tests.integ_tests import (
+    IntegComponentTest,
+    SchedulerInfo,
+)
+from torchx.specs import RunConfig
+
+
+def build_and_push_image() -> BuildInfo:
+    build = build_images()
+    push_images(build)
+    return build
+
+
+def get_k8s_sched_info(image: str) -> SchedulerInfo:
+    cfg = RunConfig()
+    cfg.set("namespace", "torchx-dev")
+    cfg.set("queue", "default")
+    return SchedulerInfo(name="kubernetes", image=image, runconfig=cfg)
+
+
+def get_local_cwd_sched_info(image: str) -> SchedulerInfo:
+    return SchedulerInfo(name="local_cwd", image=image, runconfig=RunConfig())
+
+
+def get_local_docker_sched_info(image: str) -> SchedulerInfo:
+    return SchedulerInfo(name="local_docker", image=image, runconfig=RunConfig())
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="kubernetes integration test runner")
+    parser.add_argument(
+        "--dryrun",
+        action="store_true",
+        help="Does not actually submit the app," " just prints the scheduler request",
+    )
+    args = parser.parse_args()
+    print("Starting components integration tests")
+    torchx_image = "dummy_image"
+    examples_image = "dummy_image"
+    try:
+        build = build_and_push_image()
+        torchx_image = build.torchx_image
+        examples_image = build.examples_image
+    except MissingEnvError:
+        print("Skip runnig tests, executed only docker buid step")
+    test_suite = IntegComponentTest(timeout=900)  # 15 minutes
+    test_suite.run_components(
+        component_provider,
+        scheduler_infos=[
+            get_local_cwd_sched_info(os.getcwd()),
+            get_local_docker_sched_info(torchx_image),
+            get_k8s_sched_info(torchx_image),
+        ],
+        dryrun=args.dryrun,
+    )
+
+    test_suite.run_components(
+        examples_app_defs_providers,
+        scheduler_infos=[
+            get_local_docker_sched_info(examples_image),
+            get_k8s_sched_info(examples_image),
+        ],
+        dryrun=args.dryrun,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/example_app_defs.py

@@ -0,0 +1,44 @@
+#!/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.
+
+"""
+App Defs for integration tests.
+"""
+
+
+import examples.apps.datapreproc.component as dp_component
+import examples.apps.dist_cifar.component as dist_cifar_component
+import examples.apps.lightning_classy_vision.component as cv_component
+from torchx.components.integration_tests.component_provider import ComponentProvider
+from torchx.specs import AppDef
+
+
+class CvTrainerComponentProvider(ComponentProvider):
+    def get_app_def(self) -> AppDef:
+        return cv_component.trainer(
+            image=self._image,
+            output_path="/tmp",
+            skip_export=True,
+            log_path="/tmp",
+        )
+
+
+class DistCifarComponentProvider(ComponentProvider):
+    def get_app_def(self) -> AppDef:
+        args = ["--output_path", "/tmp", "--dryrun"]
+        return dist_cifar_component.trainer(
+            *args,
+            image=self._image,
+        )
+
+
+class DatapreprocComponentProvider(ComponentProvider):
+    def get_app_def(self) -> AppDef:
+        return dp_component.data_preproc(
+            image=self._image,
+            output_path="/tmp/test",
+        )

torchx/apps/serve/serve.py

@@ -21,7 +21,6 @@
 import fsspec
 import requests
 
-
 TORCHSERVE_PARAMS = (
     "model_name",
     "handler",
@@ -56,6 +55,12 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
         help="timeout for requests to management api",
         default=60,
     )
+    parser.add_argument(
+        "--dryrun",
+        action="store_true",
+        help=argparse.SUPPRESS,
+    )
+
     parser.add_argument(
         "--port",
         type=int,
@@ -96,6 +101,9 @@ def rand_id() -> str:
 
 def main(argv: List[str]) -> None:
     args = parse_args(argv)
+    if args.dryrun:
+        print("App serve started successfully")
+        return
     with tempfile.TemporaryDirectory() as tmpdir:
         model_name = args.model_name or "model"
         model_file = f"{model_name}_{rand_id()}.mar"