ci: update testing pipeline for notebooks

synthtool/gcp/templates/python_notebooks_testing_pipeline/.cloud-build/cleanup/cleanup-cloudbuild.yaml

@@ -5,4 +5,4 @@ steps:
     args:
     - -c
     - 'python3 -m pip install -U -r .cloud-build/cleanup/cleanup-requirements.txt && python3 .cloud-build/cleanup/cleanup.py'
-timeout: 86400s
+timeout: 86400s

synthtool/gcp/templates/python_notebooks_testing_pipeline/.cloud-build/cleanup/cleanup.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# # Copyright 2021 Google LLC
+# # Copyright 2022 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 from typing import List
 from resource_cleanup_manager import (
     ResourceCleanupManager,
@@ -36,15 +37,13 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
                 resource_name = manager.resource_name(resource)
                 print(f"Will delete '{type_name}': {resource_name}")
             else:
-                manager.delete(resource)
-
+                try:
+                    manager.delete(resource)
+                except Exception as exception:
+                    print(exception)
 
-def set_dry_run(dry_run_status: bool):
-    if dry_run_status is True:
-        return True
-    print("Starting cleanup in dry run mode...")
-    return False
 
+is_dry_run = False
 
 # List of all cleanup managers
 managers = [
@@ -53,4 +52,4 @@ def set_dry_run(dry_run_status: bool):
     ModelResourceCleanupManager(),
 ]
 
-run_cleanup_managers(managers=managers, is_dry_run=set_dry_run(False))
+run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)

synthtool/gcp/templates/python_notebooks_testing_pipeline/.cloud-build/cleanup/resource_cleanup_manager.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# # Copyright 2021 Google LLC
+# # Copyright 2022 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import abc
 from google.cloud import aiplatform
 from typing import Any

synthtool/gcp/templates/python_notebooks_testing_pipeline/.cloud-build/execute_changed_notebooks_cli.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A CLI to process changed notebooks and execute them on Google Cloud Build"""
+
+import argparse
+import pathlib
+import execute_changed_notebooks_helper
+
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("String value expected.")
+
+
+parser = argparse.ArgumentParser(description="Run changed notebooks.")
+parser.add_argument(
+    "--test_paths_file",
+    type=pathlib.Path,
+    help="The path to the file that has newline-limited folders of notebooks that should be tested.",
+    required=True,
+)
+parser.add_argument(
+    "--base_branch",
+    help="The base git branch to diff against to find changed files.",
+    required=False,
+)
+parser.add_argument(
+    "--container_uri",
+    type=str,
+    help="The container uri to run each notebook in.",
+    required=True,
+)
+parser.add_argument(
+    "--variable_project_id",
+    type=str,
+    help="The GCP project id. This is used to inject a variable value into the notebook before running.",
+    required=True,
+)
+parser.add_argument(
+    "--variable_region",
+    type=str,
+    help="The GCP region. This is used to inject a variable value into the notebook before running.",
+    required=True,
+)
+parser.add_argument(
+    "--staging_bucket",
+    type=str,
+    help="The GCS bucket for staging temporary files.",
+    required=True,
+)
+parser.add_argument(
+    "--artifacts_bucket",
+    type=str,
+    help="The GCP directory for storing executed notebooks.",
+    required=True,
+)
+parser.add_argument(
+    "--should_parallelize",
+    type=str2bool,
+    nargs="?",
+    const=True,
+    default=True,
+    help="Should run notebooks in parallel.",
+)
+
+args = parser.parse_args()
+
+notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
+    test_paths_file=args.test_paths_file, base_branch=args.base_branch,
+)
+
+execute_changed_notebooks_helper.process_and_execute_notebooks(
+    notebooks=notebooks,
+    container_uri=args.container_uri,
+    staging_bucket=args.staging_bucket,
+    artifacts_bucket=args.artifacts_bucket,
+    variable_project_id=args.variable_project_id,
+    variable_region=args.variable_region,
+    should_parallelize=args.should_parallelize,
+)