feat: Added test_notebook_vm.txt to only test fast-running notebooks. Also fix private_pool not being used for tests. (#248)

* feat: Added test_notebook_vm.txt

* Propagate private pool to child builds

* Fixed workerpool issue

* Removed private pool requirement

* Added default pool

* Fixed private pool region issue

* Added comment

* Made private_pool_id arg optionally present

* Fixed when private pool is optional

* Fix regional issue bug

* Fixed pandas requirement

* Removed flaky notebook

.cloud-build/execute_changed_notebooks_cli.py

@@ -73,6 +73,12 @@ parser.add_argument(
     help="The GCP directory for storing executed notebooks.",
     required=True,
 )
+parser.add_argument(
+    "--private_pool_id",
+    type=str,
+    help="The private pool id.",
+    required=False,
+)
 parser.add_argument(
     "--should_parallelize",
     type=str2bool,
@@ -96,5 +102,6 @@ execute_changed_notebooks_helper.process_and_execute_notebooks(
     artifacts_bucket=args.artifacts_bucket,
     variable_project_id=args.variable_project_id,
     variable_region=args.variable_region,
+    private_pool_id=args.private_pool_id if not "default" else None,
     should_parallelize=args.should_parallelize,
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -109,6 +109,7 @@ def process_and_execute_notebook(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    private_pool_id: Optional[str],
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
@@ -150,6 +151,8 @@ def process_and_execute_notebook(
             notebook_output_uri=notebook_output_uri,
             container_uri=container_uri,
             tag=tag,
+            region=variable_region,
+            private_pool_id=private_pool_id,
         )
 
         operation_metadata = BuildOperationMetadata(mapping=operation.metadata)
@@ -236,6 +239,7 @@ def process_and_execute_notebooks(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    private_pool_id: Optional[str],
     should_parallelize: bool,
 ):
     """
@@ -283,6 +287,7 @@ def process_and_execute_notebooks(
                             artifacts_bucket,
                             variable_project_id,
                             variable_region,
+                            private_pool_id,
                         ),
                         notebooks,
                     )
@@ -295,6 +300,7 @@ def process_and_execute_notebooks(
                     artifacts_bucket=artifacts_bucket,
                     variable_project_id=variable_project_id,
                     variable_region=variable_region,
+                    private_pool_id=private_pool_id,
                     notebook=notebook,
                 )
                 for notebook in notebooks
@@ -338,4 +344,4 @@ def process_and_execute_notebooks(
 
     # Raise error if any notebooks failed
     if not all([result.is_pass for result in results_sorted]):
-        raise RuntimeError("Notebook failures detected. See logs for details")
+        raise RuntimeError("Notebook failures detected. See logs for details")

.cloud-build/execute_notebook_remote.py

@@ -15,6 +15,7 @@
 
 """Methods to run a notebook on Google Cloud Build"""
 
+from re import sub
 from google.protobuf import duration_pb2
 from yaml.loader import FullLoader
 
@@ -26,10 +27,12 @@ from typing import Optional
 import yaml
 
 from google.cloud.aiplatform import utils
-from google.api_core import operation
+from google.api_core import operation, client_options
+
 
 CLOUD_BUILD_FILEPATH = ".cloud-build/notebook-execution-test-cloudbuild-single.yaml"
 TIMEOUT_IN_SECONDS = 86400
+SERVICE_BASE_PATH = "cloudbuild.googleapis.com"
 
 
 def execute_notebook_remote(
@@ -37,19 +40,12 @@ def execute_notebook_remote(
     notebook_uri: str,
     notebook_output_uri: str,
     container_uri: str,
+    region: str,
+    private_pool_id: Optional[str],
     tag: Optional[str],
 ) -> operation.Operation:
     """Create and execute a single notebook on Google Cloud Build"""
-
-    # Authorize the client with Google defaults
-    credentials, project_id = google.auth.default()
-    client = cloudbuild_v1.services.cloud_build.CloudBuildClient()
-
-    build = cloudbuild_v1.Build()
-
-    # The following build steps will output "hello world"
-    # For more information on build configuration, see
-    # https://cloud.google.com/build/docs/configuring-builds/create-basic-configuration
+    # Load build steps from YAML
     cloudbuild_config = yaml.load(open(CLOUD_BUILD_FILEPATH), Loader=FullLoader)
 
     substitutions = {
@@ -58,6 +54,23 @@ def execute_notebook_remote(
         "_NOTEBOOK_OUTPUT_GCS_URI": notebook_output_uri,
     }
 
+    build = cloudbuild_v1.Build()
+
+    options: Optional[client_options.ClientOptions] = None
+    if private_pool_id:
+        substitutions["_PRIVATE_POOL_NAME"] = private_pool_id
+        build.options = cloudbuild_config["options"]
+
+        # Switch to the regional endpoint of the pool
+        options = client_options.ClientOptions(
+            api_endpoint=f"{region}-{SERVICE_BASE_PATH}"
+        )
+
+    # Authorize the client with Google defaults
+    credentials, project_id = google.auth.default()
+
+    client = cloudbuild_v1.services.cloud_build.CloudBuildClient(client_options=options)
+
     (
         source_archived_file_gcs_bucket,
         source_archived_file_gcs_object,

.cloud-build/notebook-execution-test-cloudbuild.yaml

@@ -5,10 +5,6 @@ steps:
     args:
     - -c
     - 'gcloud config list'
-  # # Clone the Git repo
-  # - name: ${_PYTHON_IMAGE}
-  #   entrypoint: git
-  #   args: ['clone', "${_GIT_REPO}", "--branch", "${_GIT_BRANCH_NAME}", "."]    
   # Check the Python version
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
@@ -28,11 +24,12 @@ steps:
     - -c
     - 'python3 -m pip install -U pip && python3 -m pip install -U --user -r .cloud-build/requirements.txt'
   # Install Python dependencies and run testing script
+  # TODO: Only pass in private_pool_id if it is set
   - name: ${_PYTHON_IMAGE}
     entrypoint: /bin/sh
     args:
     - -c
-    - 'python3 -m pip install -U pip && python3 -m pip freeze && python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION}'
+    - 'python3 -m pip install -U pip && python3 -m pip freeze && python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`'
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

.cloud-build/requirements.txt

@@ -3,7 +3,7 @@ numpy
 jupyter
 nbconvert
 papermill
-panda
+pandas
 matplotlib
 tabulate
 google-cloud-aiplatform

.cloud-build/test_notebook_vm.txt

@@ -0,0 +1,8 @@
+notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
+notebooks/official/pipelines/pipelines_intro_kfp.ipynb
+notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb
+notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
+notebooks/official/matching_engine/intro-swivel.ipynb
+notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
+notebooks/official/pipelines/metrics_viz_run_compare_kfp.ipynb
+notebooks/official/pipelines/control_flow_kfp.ipynb