Move examples under torchx module (#224)

Summary:
Pull Request resolved: #224

The diff moves `examples` under `torchx` namespace, also removes examples Dockerfile, and makes torchx image to use dev-requirements

Reviewed By: kiukchung

Differential Revision: D31464358

fbshipit-source-id: c56a92ce1fbf5b2881156e57152aa9d40af9eac8

Author: aivanou

.coveragerc

@@ -4,3 +4,4 @@ source =
 
 omit =
     *test*
+    torchx/examples/*

dev-requirements.txt

@@ -14,4 +14,5 @@ torchserve>=0.4.0
 captum>=0.3.1
 importlib-metadata
 ax-platform[mysql]>=0.2.2
-fsspec>=2021.09.0
+fsspec[s3]>=2021.09.0
+torch-model-archiver>=0.4.2

docs/source/conf.py

@@ -324,8 +324,8 @@ def handle_item(fieldarg, content):
 
 sphinx_gallery_conf = {
     "examples_dirs": [
-        "../../examples/apps",
-        "../../examples/pipelines",
+        "../../torchx/examples/apps",
+        "../../torchx/examples/pipelines",
     ],
     "gallery_dirs": [
         "examples_apps",

examples/apps/Dockerfile

@@ -52,12 +52,10 @@ def get_local_docker_sched_info(image: str) -> SchedulerInfo:
 def main() -> None:
     print("Starting components integration tests")
     torchx_image = "dummy_image"
-    examples_image = "dummy_image"
     dryrun: bool = False
     try:
         build = build_and_push_image()
         torchx_image = build.torchx_image
-        examples_image = build.examples_image
     except MissingEnvError:
         dryrun = True
         print("Skip runnig tests, executed only docker buid step")
@@ -75,8 +73,8 @@ def main() -> None:
     test_suite.run_components(
         examples_app_defs_providers,
         scheduler_infos=[
-            get_local_docker_sched_info(examples_image),
-            get_k8s_sched_info(examples_image),
+            get_local_docker_sched_info(torchx_image),
+            get_k8s_sched_info(torchx_image),
         ],
         dryrun=dryrun,
     )

scripts/component_integration_tests.py

@@ -10,9 +10,9 @@
 """
 
 
-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
+import torchx.examples.apps.datapreproc.component as dp_component
+import torchx.examples.apps.dist_cifar.component as dist_cifar_component
+import torchx.examples.apps.lightning_classy_vision.component as cv_component
 from torchx.components.integration_tests.component_provider import ComponentProvider
 from torchx.specs import AppDef
 

scripts/example_app_defs.py

@@ -17,7 +17,6 @@
 class BuildInfo:
     id: str
     torchx_image: str
-    examples_image: str
 
 
 class MissingEnvError(AssertionError):
@@ -41,15 +40,6 @@ def run_in_bg(*args: str) -> "subprocess.Popen[str]":
     return subprocess.Popen(args)
 
 
-def build_examples_canary(id: str) -> str:
-    examples_tag = "torchx_examples_canary"
-
-    print(f"building {examples_tag}")
-    run("docker", "build", "-t", examples_tag, "examples/apps/")
-
-    return examples_tag
-
-
 def build_torchx_canary(id: str) -> str:
     torchx_tag = "torchx_canary"
 
@@ -65,30 +55,18 @@ def torchx_container_tag(id: str) -> str:
     return f"{CONTAINER_REPO}:canary_{id}_torchx"
 
 
-def examples_container_tag(id: str) -> str:
-    CONTAINER_REPO = getenv_asserts("CONTAINER_REPO")
-    return f"{CONTAINER_REPO}:canary_{id}_examples"
-
-
 def build_images() -> BuildInfo:
     id = f"{getuser()}_{random_id()}"
-    examples_image = build_examples_canary(id)
     torchx_image = build_torchx_canary(id)
     return BuildInfo(
         id=id,
         torchx_image=torchx_image,
-        examples_image=examples_image,
     )
 
 
 def push_images(build: BuildInfo) -> None:
-    examples_tag = examples_container_tag(build.id)
-    run("docker", "tag", build.examples_image, examples_tag)
-    build.examples_image = examples_tag
-
     torchx_tag = torchx_container_tag(build.id)
     run("docker", "tag", build.torchx_image, torchx_tag)
     build.torchx_image = torchx_tag
 
-    run("docker", "push", examples_tag)
     run("docker", "push", torchx_tag)

scripts/integ_test_utils.py

@@ -24,7 +24,6 @@
 ```
 export KFP_NAMESPACE=<kfp namespace>
 export INTEGRATION_TEST_STORAGE=<cloud storage path>
-export EXAMPLES_CONTAINER_REPO=<docker repo>
 export TORCHX_CONTAINER_REPO=<docker repo>
 ```
 
@@ -141,7 +140,6 @@ def save_advanced_pipeline_spec(path: str, build: BuildInfo) -> None:
 
     id = build.id
     torchx_image = build.torchx_image
-    examples_image = build.examples_image
 
     STORAGE_PATH = os.getenv("INTEGRATION_TEST_STORAGE", "/tmp/storage")
     root = os.path.join(STORAGE_PATH, id)
@@ -153,8 +151,6 @@ def save_advanced_pipeline_spec(path: str, build: BuildInfo) -> None:
         "--output_path",
         output,
         "--image",
-        examples_image,
-        "--torchx_image",
         torchx_image,
         "--model_name",
         f"tiny_image_net_{id}",
@@ -165,7 +161,7 @@ def save_pipeline_spec(path: str, pipeline_file: str, *args: str) -> None:
     print(f"generating pipeline spec for {pipeline_file}")
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        run(os.path.join("examples/pipelines/kfp", pipeline_file), *args)
+        run(os.path.join("torchx/examples/pipelines/kfp", pipeline_file), *args)
         shutil.copy("pipeline.yaml", path)
 
 
@@ -202,7 +198,6 @@ def run_pipeline(build: BuildInfo, pipeline_file: str) -> object:
     except MaxRetryError:
         print(_connection_error_message())
         raise
-    namespace = getenv_asserts("KFP_NAMESPACE")
     resp = client.create_run_from_pipeline_package(
         pipeline_file,
         arguments={},

scripts/kfpint.py

@@ -12,8 +12,6 @@
 import argparse
 import os
 
-from examples.apps.dist_cifar.component import trainer
-
 # pyre-ignore-all-errors[21] # Cannot find module utils
 # pyre-ignore-all-errors[11]
 from integ_test_utils import (
@@ -23,6 +21,7 @@
     BuildInfo,
 )
 from pyre_extensions import none_throws
+from torchx.examples.apps.dist_cifar.component import trainer
 from torchx.runner import get_runner
 from torchx.specs import Resource, named_resources, RunConfig, AppState
 
@@ -48,7 +47,7 @@ def build_and_push_image() -> BuildInfo:
 def run_job(dryrun: bool = False) -> None:
     register_gpu_resource()
     build = build_and_push_image()
-    image = build.examples_image
+    image = build.torchx_image
     runner = get_runner("kubeflow-dist-runner")
 
     storage_path = os.getenv("INTEGRATION_TEST_STORAGE", "/tmp/storage")

scripts/kube_dist_trainer.py

@@ -140,6 +140,8 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> Optional[str]:
             return
         else:
             app_handle = cast(specs.AppHandle, result)
+            # do not delete this line. It is used by slurm tests to retrieve the app id
+            print(app_handle)
 
             if args.scheduler.startswith("local"):
                 self._wait_and_exit(runner, app_handle)

torchx/cli/cmd_run.py

@@ -20,7 +20,7 @@
 def data_preproc(
     image: str,
     output_path: str,
-    entrypoint: str = "examples/apps/datapreproc/datapreproc.py",
+    entrypoint: str = "torchx/examples/apps/datapreproc/datapreproc.py",
     input_path: str = "http://cs231n.stanford.edu/tiny-imagenet-200.zip",
     env: Optional[Dict[str, str]] = None,
     resource: Optional[str] = None,

examples/README.rst renamed to torchx/examples/README.rst

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import examples.apps.datapreproc.component as datapreproc
+import torchx.examples.apps.datapreproc.component as datapreproc
 from torchx.components.component_test_base import ComponentTestCase
 
 

examples/apps/__init__.py renamed to torchx/examples/__init__.py

@@ -54,7 +54,7 @@ def trainer(
     return ddp(
         *script_args,
         image=image,
-        entrypoint="examples/apps/dist_cifar/train.py",
+        entrypoint="torchx/examples/apps/dist_cifar/train.py",
         rdzv_backend=rdzv_backend,
         rdzv_endpoint=rdzv_endpoint,
         resource=resource,

examples/apps/README.rst renamed to torchx/examples/apps/README.rst

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import examples.apps.dist_cifar.component as dist_cifar
+import torchx.examples.apps.dist_cifar.component as dist_cifar
 from torchx.components.component_test_base import ComponentTestCase
 
 

examples/apps/datapreproc/__init__.py renamed to torchx/examples/apps/__init__.py

@@ -21,7 +21,7 @@ def trainer(
     image: str,
     output_path: str,
     data_path: Optional[str] = None,
-    entrypoint: str = "examples/apps/lightning_classy_vision/train.py",
+    entrypoint: str = "torchx/examples/apps/lightning_classy_vision/train.py",
     load_path: str = "",
     log_path: str = "/logs",
     resource: Optional[str] = None,
@@ -84,7 +84,7 @@ def interpret(
     data_path: str,
     output_path: str,
     resource: Optional[str] = None,
-    entrypoint: str = "examples/apps/lightning_classy_vision/interpret.py",
+    entrypoint: str = "torchx/examples/apps/lightning_classy_vision/interpret.py",
 ) -> torchx.AppDef:
     """Runs the model interpretability app on the model outputted by the training
     component.

examples/apps/datapreproc/README.rst renamed to torchx/examples/apps/datapreproc/README.rst

@@ -32,11 +32,11 @@
 # ensure data and module are on the path
 sys.path.append(".")
 
-from examples.apps.lightning_classy_vision.data import (
+from torchx.examples.apps.lightning_classy_vision.data import (
     TinyImageNetDataModule,
     download_data,
 )
-from examples.apps.lightning_classy_vision.model import TinyImageNetModel
+from torchx.examples.apps.lightning_classy_vision.model import TinyImageNetModel
 
 
 # FIXME: captum must be imported after torch otherwise it causes python to crash

examples/apps/dist_cifar/__init__.py renamed to torchx/examples/apps/datapreproc/__init__.py

@@ -109,7 +109,7 @@ def export_inference_model(
             "--model-name",
             "tiny_image_net",
             "--handler",
-            "examples/apps/lightning_classy_vision/handler/handler.py",
+            "torchx/examples/apps/lightning_classy_vision/handler/handler.py",
             "--version",
             "1",
             "--serialized-file",

examples/apps/datapreproc/component.py renamed to torchx/examples/apps/datapreproc/component.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import examples.apps.lightning_classy_vision.component as lightning_classy_vision
+import torchx.examples.apps.lightning_classy_vision.component as lightning_classy_vision
 from torchx.components.component_test_base import ComponentTestCase
 
 

examples/apps/datapreproc/datapreproc.py renamed to torchx/examples/apps/datapreproc/datapreproc.py

@@ -7,7 +7,7 @@
 import unittest
 
 import torch
-from examples.apps.lightning_classy_vision.model import (
+from torchx.examples.apps.lightning_classy_vision.model import (
     TinyImageNetModel,
 )
 

examples/apps/datapreproc/test/component_test.py renamed to torchx/examples/apps/datapreproc/test/component_test.py

@@ -30,16 +30,16 @@
 # ensure data and module are on the path
 sys.path.append(".")
 
-from examples.apps.lightning_classy_vision.data import (
+from torchx.examples.apps.lightning_classy_vision.data import (
     TinyImageNetDataModule,
     download_data,
     create_random_data,
 )
-from examples.apps.lightning_classy_vision.model import (
+from torchx.examples.apps.lightning_classy_vision.model import (
     TinyImageNetModel,
     export_inference_model,
 )
-from examples.apps.lightning_classy_vision.profiler import (
+from torchx.examples.apps.lightning_classy_vision.profiler import (
     SimpleLoggingProfiler,
 )
 

examples/apps/lightning_classy_vision/__init__.py renamed to torchx/examples/apps/dist_cifar/__init__.py

@@ -32,18 +32,12 @@
 # the standard built in apps. If you modify the torchx example code you'll
 # need to rebuild the container before launching it on KFP
 
-from torchx.version import TORCHX_IMAGE, EXAMPLES_IMAGE
+from torchx.version import TORCHX_IMAGE
 
 parser.add_argument(
     "--image",
     type=str,
     help="docker image to use for the examples apps",
-    default=EXAMPLES_IMAGE,
-)
-parser.add_argument(
-    "--torchx_image",
-    type=str,
-    help="docker image to use for the builtin torchx apps",
     default=TORCHX_IMAGE,
 )
 
@@ -113,7 +107,7 @@
 copy_app: specs.AppDef = copy(
     "http://cs231n.stanford.edu/tiny-imagenet-200.zip",
     data_path,
-    image=args.torchx_image,
+    image=args.image,
 )
 
 # %%
@@ -124,7 +118,7 @@
 # specified ahead of time so we have a fully static pipeline.
 
 
-from examples.apps.datapreproc.component import data_preproc
+from torchx.examples.apps.datapreproc.component import data_preproc
 
 processed_data_path: str = os.path.join(args.output_path, "processed")
 datapreproc_app: specs.AppDef = data_preproc(
@@ -144,7 +138,7 @@
 if "__file__" in globals():
     sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
 
-from examples.apps.lightning_classy_vision.component import trainer
+from torchx.examples.apps.lightning_classy_vision.component import trainer
 
 logs_path: str = os.path.join(args.output_path, "logs")
 models_path: str = os.path.join(args.output_path, "models")
@@ -185,7 +179,7 @@
 serve_app: specs.AppDef = torchserve(
     model_path=os.path.join(models_path, "model.mar"),
     management_api=args.management_api,
-    image=args.torchx_image,
+    image=args.image,
     params={
         "model_name": args.model_name,
         # set this to allocate a worker
@@ -198,7 +192,7 @@
 # own component file. This component takes in the output from datapreproc and
 # train components and produces images with integrated gradient results.
 
-from examples.apps.lightning_classy_vision.component import interpret
+from torchx.examples.apps.lightning_classy_vision.component import interpret
 
 interpret_path: str = os.path.join(args.output_path, "interpret")
 interpret_app: specs.AppDef = interpret(