feat: select notebook if failed on last run (#2420)

.cloud-build/execute_changed_notebooks_helper.py

@@ -89,10 +89,11 @@ class NotebookExecutionResult:
             return None
 
 
-def load_results(results_bucket: str, results_file: str) -> Dict[str, Any]:
-    """
+def load_results(results_bucket: str,
+                 results_file: str) -> Dict[str, Any]:
+    '''
     Load accumulated notebook test results
-    """
+    '''
 
     print("Loading existing accumulative results ...")
     accumulative_results = {}
@@ -104,14 +105,10 @@ def load_results(results_bucket: str, results_file: str) -> Dict[str, Any]:
         blobs = client.list_blobs(results_bucket, prefix=build_results_dir)
         for blob in blobs:
             time_created = blob.time_created.replace(tzinfo=None)
-            if (
-                datetime.datetime.now().replace(tzinfo=None) - time_created
-            ).total_seconds() > MAX_RESULTS_AGE_SECONDS:
+            if (datetime.datetime.now().replace(tzinfo=None) - time_created).total_seconds() > MAX_RESULTS_AGE_SECONDS:
                 continue
 
-            content = util.download_blob_into_memory(
-                results_bucket, blob.name, download_as_text=True
-            )
+            content = util.download_blob_into_memory(results_bucket, blob.name, download_as_text=True)
 
             try:
                 build_results = json.loads(content)
@@ -119,14 +116,11 @@ def load_results(results_bucket: str, results_file: str) -> Dict[str, Any]:
                 continue  # skip corrupted build results files
             for notebook in build_results:
                 if notebook in accumulative_results:
-                    accumulative_results[notebook]["passed"] += build_results[notebook][
-                        "passed"
-                    ]
-                    accumulative_results[notebook]["failed"] += build_results[notebook][
-                        "failed"
-                    ]
+                    accumulative_results[notebook]['passed'] += build_results[notebook]['passed']
+                    accumulative_results[notebook]['failed'] += build_results[notebook]['failed']
                 else:
                     accumulative_results[notebook] = build_results[notebook]
+                    accumulative_results[notebook]['failed_last_run'] = build_results[notebook]['failed']
 
         print(accumulative_results)
     except Exception as e:
@@ -135,22 +129,25 @@ def load_results(results_bucket: str, results_file: str) -> Dict[str, Any]:
     # If there are no accumulative results, an empty dict is returned
     return accumulative_results
 
-
-def select_notebook(
-    changed_notebook: str, accumulative_results: Dict[str, Any], test_percent: int
-) -> bool:
-    """
+def select_notebook(changed_notebook: str,
+                    accumulative_results: Dict[str, Any],
+                    test_percent: int) -> bool:
+    '''
     Algorithm to randomly select a notebook, but weight the propbability of selected based on past failures
-    """
+    '''
 
     if changed_notebook in accumulative_results:
-        pass_count = accumulative_results[changed_notebook]["passed"]
-        fail_count = accumulative_results[changed_notebook]["failed"]
+        pass_count = accumulative_results[changed_notebook]['passed']
+        fail_count = accumulative_results[changed_notebook]['failed']
     else:
         pass_count = 1
         fail_count = 0
 
-    inferred_failure_rate = fail_count / (pass_count + fail_count)
+    # if failed on last run, select the notebook
+    if accumulative_results['failed_last_run']:
+        inferred_failure_rate = 1
+    else:
+        inferred_failure_rate = fail_count / (pass_count + fail_count)
 
     # If failure rate is high, the chance of testing should be higher
     should_test_due_to_failure = random.uniform(0, 1) <= inferred_failure_rate
@@ -159,9 +156,7 @@ def select_notebook(
     should_test_due_to_random_subset = random.uniform(0, 1) <= (test_percent / 100)
 
     if should_test_due_to_failure or should_test_due_to_random_subset:
-        print(
-            f"Selected: {changed_notebook}, {should_test_due_to_failure}, {should_test_due_to_random_subset}"
-        )
+        print(f"Selected: {changed_notebook}, {should_test_due_to_failure}, {should_test_due_to_random_subset}")
         return True
     else:
         print(f"Not Selected: {changed_notebook}, pass {pass_count}, fail {fail_count}")
@@ -186,7 +181,7 @@ def _process_notebook(
             "PROJECT_ID": variable_project_id,
             "REGION": variable_region,
             "SERVICE_ACCOUNT": variable_service_account,
-            "VPC_NETWORK": variable_vpc_network or "",
+            "VPC_NETWORK": variable_vpc_network,
         },
     )
     unique_strings_preprocessor = NotebookProcessors.UniqueStringsPreprocessor()
@@ -243,6 +238,8 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
+
+
 def process_and_execute_notebook(
     container_uri: str,
     staging_bucket: str,
@@ -256,6 +253,7 @@ def process_and_execute_notebook(
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
+
     print(f"Running notebook: {notebook}")
 
     # Handle empty strings
@@ -420,11 +418,11 @@ def get_changed_notebooks(
 
     return notebooks
 
+def _save_results(results: List[NotebookExecutionResult],
+                  artifacts_bucket: str,
+                  results_file: str):
 
-def _save_results(
-    results: List[NotebookExecutionResult], artifacts_bucket: str, results_file: str
-):
-    artifacts_bucket = artifacts_bucket.replace("gs://", "").split("/")[0]
+    artifacts_bucket = artifacts_bucket.replace("gs://", "").split('/')[0]
 
     print("Updating build results ...")
     build_results = {}
@@ -436,10 +434,10 @@ def _save_results(
             pass_count = 0
             fail_count = 1
         build_results[result.path] = {
-            "duration": result.duration.total_seconds(),
-            "start_time": str(result.start_time),
-            "passed": pass_count,
-            "failed": fail_count,
+                'duration': result.duration.total_seconds(),
+                'start_time': str(result.start_time),
+                'passed': pass_count,
+                'failed': fail_count
         }
         print(f"adding {result.path}")
 
@@ -448,7 +446,8 @@ def _save_results(
 
     client = storage.Client()
     bucket = client.get_bucket(artifacts_bucket)
-    bucket.blob(str(results_file)).upload_from_string(content, "text/json")
+    bucket.blob(str(results_file)).upload_from_string(content, 'text/json')
+
 
 
 def process_and_execute_notebooks(
@@ -486,7 +485,7 @@ def process_and_execute_notebooks(
         artifacts_bucket (str):
             Required. The GCS staging bucket to write executed notebooks to.
         results_file (str):
-            Required: The path to the artifacts bucket to save results
+            Required: The path to the artifacts bucket to save results 
         variable_project_id (str):
             Required. The value for PROJECT_ID to inject into notebooks.
         variable_region (str):
@@ -513,9 +512,8 @@ def process_and_execute_notebooks(
                 "Running notebooks in parallel, so no logs will be displayed. Please wait..."
             )
 
-            with concurrent.futures.ThreadPoolExecutor(
-                max_workers=concurrent_notebooks
-            ) as executor:
+
+            with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_notebooks) as executor:
                 print(f"Max workers: {executor._max_workers}")
 
                 notebook_execution_results = list(
@@ -611,7 +609,9 @@ def process_and_execute_notebooks(
             else:
                 print(log_contents)
 
-        _save_results(results_sorted, artifacts_bucket, results_file)
+        _save_results(results_sorted, 
+                      artifacts_bucket, 
+                      results_file)
 
         print("\n=== END RESULTS===\n")