Merge branch 'main' into imkc--vpc-network-arg

.cloud-build/cleanup/cleanup.py

@@ -1,45 +1,45 @@
 from typing import List
 from resource_cleanup_manager import (
-	ResourceCleanupManager,
-	DatasetResourceCleanupManager,
-	EndpointResourceCleanupManager,
-	ModelResourceCleanupManager,
+    ResourceCleanupManager,
+    DatasetResourceCleanupManager,
+    EndpointResourceCleanupManager,
+    ModelResourceCleanupManager,
 )
 
 
 def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: bool):
-	for manager in managers:
-		type_name = manager.type_name
+    for manager in managers:
+        type_name = manager.type_name
 
-		print(f"Fetching {type_name}'s...")
-		resources = manager.list()
-		print(f"Found {len(resources)} {type_name}'s")
-		for resource in resources:
-			if not manager.is_deletable(resource):
-				continue
+        print(f"Fetching {type_name}'s...")
+        resources = manager.list()
+        print(f"Found {len(resources)} {type_name}'s")
+        for resource in resources:
+            if not manager.is_deletable(resource):
+                continue
 
-			if is_dry_run:
-				resource_name = manager.resource_name(resource)
-				print(f"Will delete '{type_name}': {resource_name}")
-			else:
-				try:
-					manager.delete(resource)
-				except Exception as exception:
-					print(exception)
+            if is_dry_run:
+                resource_name = manager.resource_name(resource)
+                print(f"Will delete '{type_name}': {resource_name}")
+            else:
+                try:
+                    manager.delete(resource)
+                except Exception as exception:
+                    print(exception)
 
-		print("")
+        print("")
 
 
 is_dry_run = False
 
 if is_dry_run:
-	print("Starting cleanup in dry run mode...")
+    print("Starting cleanup in dry run mode...")
 
 # List of all cleanup managers
 managers = [
-	DatasetResourceCleanupManager(),
-	EndpointResourceCleanupManager(),
-	ModelResourceCleanupManager(),
+    DatasetResourceCleanupManager(),
+    EndpointResourceCleanupManager(),
+    ModelResourceCleanupManager(),
 ]
 
 run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)

.cloud-build/execute_changed_notebooks_cli.py

@@ -38,11 +38,6 @@ parser.add_argument(
     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,
@@ -73,12 +68,6 @@ 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,
@@ -87,6 +76,23 @@ parser.add_argument(
     default=True,
     help="Should run notebooks in parallel.",
 )
+parser.add_argument(
+    "--base_branch",
+    help="The base git branch to diff against to find changed files.",
+    required=False,
+)
+parser.add_argument(
+    "--variable_vpc_network",
+    type=str,
+    help="The full VPC network name. See https://cloud.google.com/compute/docs/networks-and-firewalls#networks. Format is projects/{project}/global/networks/{network}, where {project} is a project number, as in '12345', and {network} is network name. See <https://cloud.google.com/compute/docs/reference/rest/v1/networks/insert> for details. This is used to inject a variable value into the notebook before running.",
+    required=False,
+)
+parser.add_argument(
+    "--private_pool_id",
+    type=str,
+    help="The private pool id.",
+    required=False,
+)
 
 args = parser.parse_args()
 
@@ -102,6 +108,7 @@ 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,
+    variable_vpc_network=args.variable_vpc_network,
+    private_pool_id=args.private_pool_id,
     should_parallelize=args.should_parallelize,
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -67,6 +67,7 @@ def _process_notebook(
     notebook_path: str,
     variable_project_id: str,
     variable_region: str,
+    variable_vpc_network: Optional[str],
 ):
     # Read notebook
     with open(notebook_path) as f:
@@ -78,6 +79,7 @@ def _process_notebook(
         replacement_map={
             "PROJECT_ID": variable_project_id,
             "REGION": variable_region,
+            "VPC_NETWORK": variable_vpc_network,
         },
     )
 
@@ -109,12 +111,20 @@ def process_and_execute_notebook(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    variable_vpc_network: Optional[str],
     private_pool_id: Optional[str],
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
     print(f"Running notebook: {notebook}")
 
+    # Handle empty strings
+    if not variable_vpc_network:
+        variable_vpc_network = None
+
+    if not private_pool_id:
+        private_pool_id = None
+
     # Create paths
     notebook_output_uri = "/".join([artifacts_bucket, pathlib.Path(notebook).name])
 
@@ -140,6 +150,7 @@ def process_and_execute_notebook(
             notebook_path=notebook,
             variable_project_id=variable_project_id,
             variable_region=variable_region,
+            variable_vpc_network=variable_vpc_network,
         )
 
         # Upload the pre-processed code to a GCS bucket
@@ -239,6 +250,7 @@ def process_and_execute_notebooks(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    variable_vpc_network: Optional[str],
     private_pool_id: Optional[str],
     should_parallelize: bool,
 ):
@@ -287,6 +299,7 @@ def process_and_execute_notebooks(
                             artifacts_bucket,
                             variable_project_id,
                             variable_region,
+                            variable_vpc_network,
                             private_pool_id,
                         ),
                         notebooks,
@@ -300,6 +313,7 @@ def process_and_execute_notebooks(
                     artifacts_bucket=artifacts_bucket,
                     variable_project_id=variable_project_id,
                     variable_region=variable_region,
+                    variable_vpc_network=variable_vpc_network,
                     private_pool_id=private_pool_id,
                     notebook=notebook,
                 )

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

@@ -29,7 +29,7 @@ steps:
     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} `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`'
+    - '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} --variable_vpc_network "${_VPC_NETWORK_NAME}" --private_pool_id "${_PRIVATE_POOL_NAME}"'
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

.cloud-build/utils/NotebookProcessors.py

@@ -41,11 +41,12 @@ class UpdateVariablesPreprocessor(Preprocessor):
         # VARIABLE_NAME = '[description]'
 
         for variable_name, variable_value in replacement_map.items():
-            content = update_notebook_variables.get_updated_value(
-                content=content,
-                variable_name=variable_name,
-                variable_value=variable_value,
-            )
+            if variable_value is not None:
+                content = update_notebook_variables.get_updated_value(
+                    content=content,
+                    variable_name=variable_name,
+                    variable_value=variable_value,
+                )
 
         return content
 
@@ -60,4 +61,4 @@ class UpdateVariablesPreprocessor(Preprocessor):
 
             executable_cells.append(cell)
         notebook.cells = executable_cells
-        return notebook, resources
+        return notebook, resources

.cloud-build/utils/UpdateNotebookVariables.py

@@ -78,4 +78,4 @@ def test_region():
         variable_name="REGION",
         variable_value="us-central1",
     )
-    assert new_content == 'REGION = "us-central1"  # @param {type:"string"}'
+    assert new_content == 'REGION = "us-central1"  # @param {type:"string"}'

.cloud-build/utils/util.py

@@ -57,4 +57,4 @@ def archive_code_and_upload(staging_bucket: str):
 
     print(f"Uploaded source code archive to {source_archived_file_gcs}")
 
-    return source_archived_file_gcs
+    return source_archived_file_gcs