component_integration_tests: add aws_batch, ray + more interpretable (#403)

Summary:
This adds aws_batch and ray to the integration tests and refactors them to run per scheduler in parallel to make it clearer what's running.

Pull Request resolved: #403

Test Plan: CI + pyre

Reviewed By: kiukchung

Differential Revision: D34568682

Pulled By: d4l3k

fbshipit-source-id: e2d18fc8b5f654fab994fafec46df5943da30c3f

Author: d4l3k

.github/workflows/components-integration-tests.yaml

@@ -9,6 +9,15 @@ on:
 jobs:
   components-launch:
     runs-on: ubuntu-18.04
+    strategy:
+      matrix:
+       include:
+         - scheduler: "aws_batch"
+         - scheduler: "kubernetes"
+         - scheduler: "local_cwd"
+         - scheduler: "local_docker"
+         - scheduler: "ray"
+      fail-fast: false
     permissions:
       id-token: write
       contents: read
@@ -48,8 +57,12 @@ jobs:
           set -eux
           pip install -r dev-requirements.txt
           pip install -e .[kubernetes]
+
+      - name: Start Ray
+        if: ${{ matrix.scheduler == 'ray' }}
+        run: ray start --head
       - name: Run Components Integration Tests
         env:
           INTEGRATION_TEST_STORAGE: ${{ secrets.INTEGRATION_TEST_STORAGE }}
           CONTAINER_REPO: ${{ secrets.CONTAINER_REPO }}
-        run: scripts/component_integration_tests.py
+        run: scripts/component_integration_tests.py --scheduler ${{ matrix.scheduler }}

scripts/component_integration_tests.py

@@ -8,6 +8,7 @@
 """
 Kubernetes integration tests.
 """
+import argparse
 import os
 
 import example_app_defs as examples_app_defs_providers
@@ -17,6 +18,7 @@
     IntegComponentTest,
     SchedulerInfo,
 )
+from torchx.schedulers import get_scheduler_factories
 
 
 # pyre-ignore-all-errors[21] # Cannot find module utils
@@ -44,54 +46,68 @@ def get_ray_sched_info(image: str) -> SchedulerInfo:
     return SchedulerInfo(name="ray", image=image, cfg=cfg)
 
 
-def get_local_cwd_sched_info(image: str) -> SchedulerInfo:
-    return SchedulerInfo(name="local_cwd", image=image)
-
-
-def get_local_docker_sched_info(image: str) -> SchedulerInfo:
-    return SchedulerInfo(name="local_docker", image=image)
+def argparser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Process some integers.")
+    choices = list(get_scheduler_factories().keys())
+    parser.add_argument("--scheduler", required=True, choices=choices)
+    return parser
 
 
 def main() -> None:
+    args = argparser().parse_args()
+    scheduler = args.scheduler
+
     print("Starting components integration tests")
     torchx_image = "dummy_image"
     dryrun: bool = False
-    try:
-        build = build_and_push_image()
-        torchx_image = build.torchx_image
-    except MissingEnvError:
-        dryrun = True
-        print("Skip runnig tests, executed only docker buid step")
-    test_suite = IntegComponentTest(timeout=600)  # 10 minutes
-    test_suite.run_components(
-        component_provider,
-        scheduler_infos=[
-            get_local_cwd_sched_info(os.getcwd()),
-            get_k8s_sched_info(torchx_image),
-        ],
-        dryrun=dryrun,
-    )
-
-    # Run components on `local_docker`  scheduler in sequence due to
-    # docker APIs are not atomic. Some of the APIs, e.g. `create_network`
-    # cause a race condition, making several networks with the same name to be created.
-    test_suite.run_components(
-        component_provider,
-        scheduler_infos=[
-            get_local_docker_sched_info(torchx_image),
-        ],
-        dryrun=dryrun,
-        run_in_parallel=False,
-    )
-
-    test_suite.run_components(
-        examples_app_defs_providers,
-        scheduler_infos=[
-            get_local_docker_sched_info(torchx_image),
-            get_k8s_sched_info(torchx_image),
-        ],
-        dryrun=dryrun,
-    )
+    if scheduler in ("kubernetes", "local_docker", "aws_batch"):
+        try:
+            build = build_and_push_image()
+            torchx_image = build.torchx_image
+        except MissingEnvError:
+            dryrun = True
+            print("Skip running tests, executed only docker build step")
+    test_suite: IntegComponentTest = IntegComponentTest(timeout=15 * 60)  # 15 minutes
+
+    def run_components(info: SchedulerInfo) -> None:
+        test_suite.run_components(
+            component_provider,
+            scheduler_infos=[info],
+            dryrun=dryrun,
+        )
+
+    def run_examples(info: SchedulerInfo) -> None:
+        test_suite.run_components(
+            examples_app_defs_providers,
+            scheduler_infos=[info],
+            dryrun=dryrun,
+        )
+
+    if scheduler == "kubernetes":
+        info = get_k8s_sched_info(torchx_image)
+        run_components(info)
+        run_examples(info)
+    elif scheduler == "local_cwd":
+        info = SchedulerInfo(name=scheduler, image=os.getcwd())
+        run_components(info)
+    elif scheduler == "local_docker":
+        info = SchedulerInfo(name=scheduler, image=torchx_image)
+        run_components(info)
+        run_examples(info)
+    elif scheduler == "aws_batch":
+        info = SchedulerInfo(
+            name=scheduler,
+            image=torchx_image,
+            cfg={
+                "queue": "torchx",
+            },
+        )
+        run_components(info)
+    elif scheduler == "ray":
+        info = get_ray_sched_info(torchx_image)
+        run_components(info)
+    else:
+        raise ValueError(f"component tests missing support for {scheduler}")
 
 
 if __name__ == "__main__":

torchx/components/component_test_base.py

@@ -57,6 +57,10 @@ def run_appdef_on_scheduler(
             app_status = runner.wait(app_handle)
             print(f"Final status: {app_status}")
             if none_throws(app_status).state != AppState.SUCCEEDED:
+                for line in runner.log_lines(
+                    app_handle, role_name=app_def.roles[0].name
+                ):
+                    print(f"{app_handle}: {line}")
                 raise AssertionError(
                     f"App {app_handle} failed with status: {app_status}"
                 )

torchx/components/dist.py

@@ -142,7 +142,7 @@ def ddp(
     memMB: int = 1024,
     j: str = "1x2",
     env: Optional[Dict[str, str]] = None,
-    max_restarts: Optional[int] = None,
+    max_retries: int = 0,
     rdzv_backend: str = "c10d",
     rdzv_endpoint: Optional[str] = None,
 ) -> specs.AppDef:
@@ -167,7 +167,7 @@ def ddp(
         h: a registered named resource (if specified takes precedence over cpu, gpu, memMB)
         j: {nnodes}x{nproc_per_node}, for gpu hosts, nproc_per_node must not exceed num gpus
         env: environment varibles to be passed to the run (e.g. ENV1=v1,ENV2=v2,ENV3=v3)
-        max_restarts: the number of restarts allowed
+        max_retries: the number of scheduler retries allowed
         rdzv_backend: rendezvous backend (only matters when nnodes > 1)
         rdzv_endpoint: rendezvous server endpoint (only matters when nnodes > 1), defaults to rank0 host for schedulers that support it
     """
@@ -219,8 +219,6 @@ def ddp(
         "--nproc_per_node",
         str(nproc_per_node),
     ]
-    if max_restarts is not None:
-        cmd += ["--max_restarts", str(max_restarts)]
     if script is not None:
         cmd += [script]
     elif m is not None:
@@ -240,6 +238,7 @@ def ddp(
                 port_map={
                     "c10d": 29500,
                 },
+                max_retries=max_retries,
             )
         ],
     )

torchx/components/integration_tests/component_provider.py

@@ -40,17 +40,18 @@ def get_app_def(self) -> AppDef:
             script="torchx/components/integration_tests/test/dummy_app.py",
             name="ddp-trainer",
             image=self._image,
+            cpu=1,
             j="2x2",
-            max_restarts=3,
+            max_retries=3,
         )
 
 
 class ServeComponentProvider(ComponentProvider):
     # TODO(aivanou): Remove dryrun and test e2e serve component+app
     def get_app_def(self) -> AppDef:
         return serve_components.torchserve(
-            model_path="",
-            management_api="",
+            model_path="dummy_path",
+            management_api="dummy_api",
             image=self._image,
             dryrun=True,
         )

torchx/schedulers/aws_batch_scheduler.py

@@ -35,6 +35,7 @@
 for how to create a image repository.
 """
 
+import threading
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Dict, Iterable, Mapping, Optional, Any, TYPE_CHECKING, Tuple
@@ -120,6 +121,19 @@ def __repr__(self) -> str:
         return str(self)
 
 
+def _thread_local_session() -> "boto3.session.Session":
+    KEY = "torchx_aws_batch_session"
+    local = threading.local()
+    if hasattr(local, KEY):
+        # pyre-ignore[16]
+        return getattr(local, KEY)
+    import boto3.session
+
+    session = boto3.session.Session()
+    setattr(local, KEY, session)
+    return session
+
+
 class AWSBatchScheduler(Scheduler, DockerWorkspace):
     """
     AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.
@@ -174,20 +188,16 @@ def __init__(
     @property
     # pyre-fixme[3]: Return annotation cannot be `Any`.
     def _client(self) -> Any:
-        if self.__client is None:
-            import boto3
-
-            self.__client = boto3.client("batch")
-        return self.__client
+        if self.__client:
+            return self.__client
+        return _thread_local_session().client("batch")
 
     @property
     # pyre-fixme[3]: Return annotation cannot be `Any`.
     def _log_client(self) -> Any:
-        if self.__log_client is None:
-            import boto3
-
-            self.__log_client = boto3.client("logs")
-        return self.__log_client
+        if self.__log_client:
+            return self.__log_client
+        return _thread_local_session().client("logs")
 
     def schedule(self, dryrun_info: AppDryRunInfo[BatchJob]) -> str:
         cfg = dryrun_info._cfg

torchx/schedulers/kubernetes_scheduler.py

@@ -291,7 +291,8 @@ def app_to_resource(app: AppDef, queue: str) -> Dict[str, object]:
             "tasks": tasks,
             "maxRetry": job_retries,
             "plugins": {
-                "svc": [],
+                # https://github.com/volcano-sh/volcano/issues/533
+                "svc": ["--publish-not-ready-addresses"],
                 "env": [],
             },
         },

torchx/schedulers/test/kubernetes_scheduler_test.py

@@ -193,7 +193,8 @@ def test_submit_dryrun(self) -> None:
   maxRetry: 3
   plugins:
     env: []
-    svc: []
+    svc:
+    - --publish-not-ready-addresses
   queue: testqueue
   schedulerName: volcano
   tasks: