Merge pull request #151 from arup-group/faster_notebook_smoketest

Use multiprocessing in the notebooks smoketest
arup-group · Oct 13, 2022 · 74ef121 · 74ef121
2 parents 27f95bb + d613dc6
commit 74ef121
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/.gitignore b/.gitignore
@@ -133,3 +133,4 @@ scripts/secrets-mgr-exploration.py
 example_data/pt2matsim_network/genet_output/*
 genet_output/*
 example_data/output_*/*
+example_data/api_requests_send.json
diff --git a/notebooks/notebook_smoke_tests.py b/notebooks/notebook_smoke_tests.py
@@ -1,8 +1,11 @@
 import argparse
 import glob
+import multiprocessing
 import subprocess
 import sys
 from datetime import datetime
+from itertools import repeat
+from multiprocessing import Manager
 from pprint import pprint
 
 from colorama import Fore, Style
@@ -64,23 +67,30 @@ def install_jupyter_kernel(kernel_name):
     return run_shell_command(kernel_install_cmd)
 
 
-def execute_notebook(notebook_path):
+def execute_notebook(notebook_path, notebook_results_dict):
     print("Executing notebook '{}{}{}'...".format(Fore.YELLOW, notebook_path, Style.RESET_ALL))
     cmd = [
         'jupyter', 'nbconvert',
         '--to', 'notebook',
         '--execute', '"{}"'.format(notebook_path),
         '--output-dir=/tmp'
     ]
-    return run_shell_command(cmd)
+    return_code, cmd, run_time = run_shell_command(cmd)
+    notebook_results_dict[notebook_path] = (return_code, run_time)
 
 
 def run_shell_command(shell_cmd):
     print(Fore.BLUE + ' '.join(shell_cmd))
     start_time = datetime.now()
     rc = subprocess.call(' '.join(shell_cmd), shell=True)
     running_time = datetime.now() - start_time
-    print("{}Shell process return value was {}{}{}".format(Style.RESET_ALL, Fore.YELLOW, rc, Style.RESET_ALL))
+    print("{}Shell process return value for {}{}{} was {}{}{}".format(Style.RESET_ALL,
+                                                                      Fore.YELLOW,
+                                                                      shell_cmd,
+                                                                      Style.RESET_ALL,
+                                                                      Fore.YELLOW,
+                                                                      rc,
+                                                                      Style.RESET_ALL))
     return rc, ' '.join(shell_cmd), running_time
 
 
@@ -154,11 +164,9 @@ def print_summary(notebook_results_dict):
         print("{}Warning: Jupyter kernel installation shell command did not exit normally"
               " - this may cause problems later{}".format(Fore.RED, Style.RESET_ALL))
 
-    notebook_results = {}
-    for notebook in notebooks:
-        print('------------------------------------------------------')
-        return_code, cmd, run_time = execute_notebook(notebook)
-        notebook_results[notebook] = (return_code, run_time)
+    notebook_results = Manager().dict()
+    with multiprocessing.Pool() as pool:
+        pool.starmap(execute_notebook, zip(notebooks, repeat(notebook_results)))
 
     print('------------------------------------------------------')
     print("\nFinished the smoke test in {}{}{}".format(Fore.YELLOW, datetime.now() - start, Style.RESET_ALL))