From a162066c947f7442e2894764f9ab6a313c4f08d9 Mon Sep 17 00:00:00 2001
From: Andrew Ferlitsch <aferlitsch@google.com>
Date: Wed, 29 Nov 2023 10:17:07 -0800
Subject: [PATCH] in-progress: SDK whl replacement (#2552)

* in-progress: SDK whl replacement

* add: cmd param
---
 .cloud-build/execute_changed_notebooks_cli.py |  8 ++++++
 .../execute_changed_notebooks_helper.py       |  2 ++
 .cloud-build/utils/NotebookProcessors.py      | 25 +++++++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/.cloud-build/execute_changed_notebooks_cli.py b/.cloud-build/execute_changed_notebooks_cli.py
index 182e2e5c3..befc61855 100755
--- a/.cloud-build/execute_changed_notebooks_cli.py
+++ b/.cloud-build/execute_changed_notebooks_cli.py
@@ -136,6 +136,13 @@ parser.add_argument(
     default=None,
     required=False,
 )
+parser.add_argument(
+    "--aiplatform_whl",
+    type=str,
+    help="The GCS path to a whl version google-cloud-aiplatform",
+    default=None,
+    required=False,
+)
 parser.add_argument(
     "--dry_run",
     type=str2bool,
@@ -208,4 +215,5 @@ else:
         variable_vpc_network=args.variable_vpc_network,
         private_pool_id=args.private_pool_id,
         concurrent_notebooks=args.concurrent_notebooks,
+        aiplatform_whl=args.aiplatform_whl
 )
diff --git a/.cloud-build/execute_changed_notebooks_helper.py b/.cloud-build/execute_changed_notebooks_helper.py
index 9c6751632..8394a1769 100755
--- a/.cloud-build/execute_changed_notebooks_helper.py
+++ b/.cloud-build/execute_changed_notebooks_helper.py
@@ -508,6 +508,7 @@ def process_and_execute_notebooks(
     variable_vpc_network: Optional[str] = None,
     private_pool_id: Optional[str] = None,
     concurrent_notebooks: Optional[int] = 10,
+    aiplatform_whl: Optional[str] = None,
 ):
     """
     Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -539,6 +540,7 @@ def process_and_execute_notebooks(
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
         concurrent_notebooks (int): Max number of notebooks per minute to run in parallel.
+        aiplatform_whl: alternate whl version of Vertex AI SDK to install
     """
 
     # Calculate deadline
diff --git a/.cloud-build/utils/NotebookProcessors.py b/.cloud-build/utils/NotebookProcessors.py
index 29d24c723..1f39fd361 100644
--- a/.cloud-build/utils/NotebookProcessors.py
+++ b/.cloud-build/utils/NotebookProcessors.py
@@ -98,3 +98,28 @@ class UniqueStringsPreprocessor(Preprocessor):
             executable_cells.append(cell)
         notebook.cells = executable_cells
         return notebook, resources
+
+class VertexAIInstallProprocessor(Preprocessor):
+    def __init__(self, vertex_ai_wheel):
+        self.vertex_ai_wheel = vertex_ai_wheel
+
+    @staticmethod
+    def update_vertex_ai_install(content: str):
+        if "google-cloud-aiplatform" not in content:
+            return content
+        return (
+            f"gsutil cp {self.vertex_ai_wheel} google-cloud-aiplatform.whl\n" + 
+            content.replace("google-cloud-aiplatform\n", "google-cloud-aiplatform.whl\n")
+            .replace("google-cloud-aiplatform ", "google-cloud-aiplatform.whl ")
+        )
+
+    def preprocess(self, notebook, resources=None):
+        executable_cells = []
+        for cell in notebook.cells:
+            if cell.cell_type == "code":
+                cell.source = self.update_vertex_ai_install(
+                    content=cell.source,
+                )
+
+            executable_cells.append(cell)
+        notebook.cells = executable_cells