Parallelize manual_aq_tests/initialize.py

manual_qa_tests/initialize.py

@@ -1,4 +1,5 @@
 import argparse
+from multiprocessing import Process
 
 import deploy_example
 from aqueduct.constants.enums import NotificationLevel
@@ -47,6 +48,24 @@
 TEMP_NOTEBOOK_PATH = "temp.py"
 RUN_NOTEBOOK_SCRIPT = "examples/run_notebook.py"
 
+
+def deploy_example_notebook(deploy_fn, dir_path, notebook_name, api_key, address) -> None:
+    print(f"Deploying example notebooks {notebook_name}...")
+    deploy_fn(
+        dir_path,
+        notebook_name,
+        TEMP_NOTEBOOK_PATH,
+        address,
+        api_key,
+    )
+
+
+def deploy_flow(name, deploy_fn, api_key, address, data_integration) -> None:
+    print(f"Deploying {name}...")
+    client = aq.Client(api_key, address)
+    deploy_fn(client, data_integration)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--addr", default="localhost:8080")
@@ -60,6 +79,7 @@
     parser.add_argument("--slack-channel", default="")
     parser.add_argument("--notification-level", default="success")
     parser.add_argument("--wait-to-complete", action="store_true")
+    parser.add_argument("--single-threaded", action="store_true")
     args = parser.parse_args()
 
     api_key = args.api_key if args.api_key else aq.get_apikey()
@@ -73,25 +93,48 @@
             NotificationLevel(args.notification_level),
         )
 
+    # This is only populated in the default multi-process case.
+    processes = []
+
     if args.example_notebooks or args.example_notebooks_only or args.demo_container_notebooks_only:
         notebooks = DEMO_NOTEBOOKS_PATHS
         if not args.demo_container_notebooks_only:
             notebooks += ADDITIONAL_EXAMPLE_NOTEBOOKS_PATHS
 
         for example_path in notebooks:
-            print(f"Deploying example notebooks {example_path[1]}...")
-            deploy_example.deploy(
-                example_path[0],
-                example_path[1],
-                TEMP_NOTEBOOK_PATH,
-                args.addr,
-                api_key,
-            )
+            if args.single_threaded:
+                deploy_example_notebook(
+                    deploy_example.deploy, example_path[0], example_path[1], api_key, args.addr
+                )
+            else:
+                p = Process(
+                    target=deploy_example_notebook,
+                    args=(
+                        deploy_example.deploy,
+                        example_path[0],
+                        example_path[1],
+                        api_key,
+                        args.addr,
+                    ),
+                )
+                processes.append(p)
+                p.start()
 
     if not args.example_notebooks_only and not args.demo_container_notebooks_only:
         for pkg in WORKFLOW_PKGS:
-            print(f"Deploying {pkg.NAME}...")
-            pkg.deploy(client, args.data_integration)
+            if args.single_threaded:
+                deploy_flow(pkg.NAME, pkg.deploy, api_key, args.addr, args.data_integration)
+            else:
+                p = Process(
+                    target=deploy_flow,
+                    args=(pkg.NAME, pkg.deploy, api_key, args.addr, args.data_integration),
+                )
+                processes.append(p)
+                p.start()
+
+    if len(processes) > 0:
+        for p in processes:
+            p.join()
 
     if args.wait_to_complete:
         wait_for_all_flows_to_complete(client)