Fixed VPC network

.cloud-build/execute_changed_notebooks_cli.py

@@ -51,7 +51,7 @@ parser.add_argument(
     "--build_id",
     type=str,
     help="The build id (which may be a Cloud Build job specific or user explicit.",
-    required=True
+    required=True,
 )
 parser.add_argument(
     "--base_branch",
@@ -147,16 +147,24 @@ results_bucket = f"{args.artifacts_bucket}"
 # artifacts_bucket may get set by trigger to a full gs:// folder path
 if results_bucket.startswith("gs://"):
     results_bucket = results_bucket[5:]
-results_bucket = results_bucket.split('/')[0]
+results_bucket = results_bucket.split("/")[0]
 results_file = f"build_results/{args.build_id}.json"
 
 if args.test_percent == 100:
     notebooks = changed_notebooks
     accumulative_results = {}
 else:
-    accumulative_results = execute_changed_notebooks_helper.load_results(results_bucket, results_file)
+    accumulative_results = execute_changed_notebooks_helper.load_results(
+        results_bucket, results_file
+    )
 
-    notebooks = [changed_notebook for changed_notebook in changed_notebooks if execute_changed_notebooks_helper.select_notebook(changed_notebook, accumulative_results, args.test_percent)]
+    notebooks = [
+        changed_notebook
+        for changed_notebook in changed_notebooks
+        if execute_changed_notebooks_helper.select_notebook(
+            changed_notebook, accumulative_results, args.test_percent
+        )
+    ]
 
 if args.dry_run:
     print("Dry run ...\n")
@@ -177,4 +185,4 @@ else:
         variable_vpc_network=args.variable_vpc_network,
         private_pool_id=args.private_pool_id,
         concurrent_notebooks=args.concurrent_notebooks,
-)
+    )

.cloud-build/execute_changed_notebooks_helper.py

@@ -47,6 +47,22 @@ PYTHON_VERSION = "3.9"  # Set default python version
 MAX_RESULTS_AGE_SECONDS: int = (60 * 60) * 24 * 60  # 60 days
 
 
+import re
+import nbformat
+
+
+def check_regex_in_notebook(notebook_path: str, regex: str):
+    """Checks if the given regex exists in the notebook."""
+    nb = nbformat.read(notebook_path, as_version=4)
+    cells = nb["cells"]
+    for cell in cells:
+        if cell["cell_type"] == "code":
+            match = re.search(regex, cell["source"])
+            if match:
+                return True
+    return False
+
+
 def format_timedelta(delta: datetime.timedelta) -> str:
     """Formats a timedelta duration to [N days] %H:%M:%S format"""
     seconds = int(delta.total_seconds())
@@ -89,11 +105,10 @@ 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 = {}
@@ -105,10 +120,14 @@ def load_results(results_bucket: str,
         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)
@@ -116,8 +135,12 @@ def load_results(results_bucket: str,
                 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]
 
@@ -128,16 +151,17 @@ def load_results(results_bucket: str,
     # 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
@@ -151,7 +175,9 @@ def select_notebook(changed_notebook: str,
     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}")
@@ -233,8 +259,6 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
-
-
 def process_and_execute_notebook(
     container_uri: str,
     staging_bucket: str,
@@ -303,13 +327,28 @@ def process_and_execute_notebook(
             int((deadline - datetime.datetime.now()).total_seconds()), 1
         )
 
+        # Only use private pool if the notebook makes use of the private pool variable
+        is_private_pool_required = check_regex_in_notebook(
+            notebook_path=notebook,
+            regex=r"VPC_NETWORK =",
+        )
+
+        if is_private_pool_required and (
+            not private_pool_id or len(private_pool_id.strip()) == 0
+        ):
+            raise ValueError(
+                "Private pool is used in notebook but `private_pool_id` is None"
+            )
+
+        private_pool_id_if_used = private_pool_id if is_private_pool_required else None
+
         operation = execute_notebook_remote.execute_notebook_remote(
             code_archive_uri=code_archive_uri,
             notebook_uri=notebook,
             notebook_output_uri=notebook_output_uri,
             container_uri=container_uri,
             tag=tag,
-            private_pool_id=private_pool_id,
+            private_pool_id=private_pool_id_if_used,
             private_pool_region=variable_region,
             timeout_in_seconds=timeout_in_seconds,
             python_version=notebook_exec_python_version,
@@ -413,11 +452,12 @@ def get_changed_notebooks(
 
     return notebooks
 
-def _save_results(results: List[NotebookExecutionResult],
-                  artifacts_bucket: str,
-                  results_file: str):
 
-    artifacts_bucket = artifacts_bucket.replace("gs://", "").split('/')[0]
+def _save_results(
+    results: List[NotebookExecutionResult], artifacts_bucket: str, results_file: str
+):
+
+    artifacts_bucket = artifacts_bucket.replace("gs://", "").split("/")[0]
 
     print("Updating build results ...")
     build_results = {}
@@ -429,10 +469,10 @@ def _save_results(results: List[NotebookExecutionResult],
             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}")
 
@@ -441,8 +481,7 @@ def _save_results(results: List[NotebookExecutionResult],
 
     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(
@@ -480,7 +519,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):
@@ -507,8 +546,9 @@ 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(
@@ -604,9 +644,7 @@ 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")
 

.cloud-build/execute_notebook_remote.py

@@ -40,7 +40,7 @@ def execute_notebook_remote(
     private_pool_region: Optional[str],
     tag: Optional[str],
     timeout_in_seconds: Optional[int] = None,
-    python_version: Optional[str] = None
+    python_version: Optional[str] = None,
 ) -> operation.Operation:
     """Create and execute a single notebook on Google Cloud Build"""
     # Load build steps from YAML
@@ -51,17 +51,16 @@ def execute_notebook_remote(
         "_PYTHON_IMAGE": container_uri,
         "_NOTEBOOK_GCS_URI": notebook_uri,
         "_NOTEBOOK_OUTPUT_GCS_URI": notebook_output_uri,
-        "_PYTHON_VERSION" : f"python{python_version}"
+        "_PYTHON_VERSION": f"python{python_version}",
     }
 
     if python_version is not None:
-      substitutions["_PYTHON_VERSION"] = "python" + python_version
+        substitutions["_PYTHON_VERSION"] = "python" + python_version
 
     build = cloudbuild_v1.Build()
 
     options: Optional[client_options.ClientOptions] = None
     if private_pool_id and private_pool_region:
-        # substitutions["_PRIVATE_POOL_NAME"] = private_pool_id
         build.options = cloudbuild_config.get("options")
         build.options.pool = {"name": private_pool_id}
 

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

@@ -36,7 +36,7 @@ steps:
     - -c
     - |
       . workspace/env/bin/activate &&
-      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_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GPC_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT} --concurrent_notebooks=${_CONCURRENT_NOTEBOOKS}
+      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} --should_parallelize ${_SHOULD_PARALLELIZE} --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GCP_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT}
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb

@@ -54,6 +54,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "b0a74aaf1481"
@@ -63,12 +64,11 @@
         "\n",
         "This example demonstrates how to encode custom text embeddings using the StackOverflow dataset and the sentence-T5 model. These are uploaded to the Vertex AI Matching Engine service. It is a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it is a fully managed offering, further reducing operational overhead. It is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research.\n",
         "\n",
-        "**Pre-requisite**: This notebook requires you to already have a VPC network set up. See the \"Prepare a VPC network\" section in [Create Vertex AI Matching Engine index notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb).\n",
-        "\n",
         "Learn more about [Vertex AI Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview)."
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "34a4b245e795"
@@ -84,13 +84,17 @@
         "\n",
         "The steps performed include:\n",
         "\n",
-        "* Create ANN index\n",
-        "* Create an index endpoint with VPC Network\n",
-        "* Deploy ANN index\n",
-        "* Perform online query\n"
+        "* Convert a BigQuery dataset to embeddings\n",
+        "* Create an index\n",
+        "* Upload embeddings to the index\n",
+        "* Create an index endpoint\n",
+        "* Deploy the index to the index endpoint\n",
+        "* Perform an online query\n",
+        "* Add metadata to a Redis store\n"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -104,6 +108,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "f0f1bea346db"
@@ -129,6 +134,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "1ae34c2a9ce7"
@@ -174,6 +180,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "5b08ba354c6e"
@@ -198,6 +205,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dd28c9e4f067"
@@ -220,13 +228,14 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[YOUR-PROJECT-ID]\"\n",
+        "PROJECT_ID = \"python-docs-samples-tests\"\n",
         "\n",
         "# Set the project id\n",
         "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "4f4512bf63b3"
@@ -249,6 +258,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "949271bfebe3"
@@ -260,6 +270,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "b65b4ce80d9a"
@@ -270,6 +281,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "985cdbfe7372"
@@ -290,6 +302,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "79efab26ad02"
@@ -311,6 +324,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0c0a44fa330f"
@@ -321,6 +335,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "d3uj8x73nDX_"
@@ -330,30 +345,7 @@
       ]
     },
     {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "hhq5zEbGg0XX"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 5,
-      "metadata": {
-        "id": "EzrelQZ22IZj"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "zgPO1eR3CYjk"
@@ -372,10 +364,11 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = \"gs://your-bucket-name-unique\"  # @param {type:\"string\"}"
+        "BUCKET_URI = \"gs://your-bucket-name-1gtoypb3\"  # @param {type:\"string\"}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "-EcIXiGsCePi"
@@ -396,6 +389,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "lR6Wwv-hCCN-"
@@ -581,6 +575,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "1124422cc200"
@@ -614,6 +609,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "43088937e820"
@@ -674,13 +670,14 @@
       "outputs": [],
       "source": [
         "# Encode 500 questions\n",
-        "questions = df.title.tolist()[:500]\n",
+        "test_questions = df.title.tolist()[:500]\n",
         "question_embeddings = encode_text_to_embedding(\n",
-        "    text_encoder=encoder, sentences=questions\n",
+        "    text_encoder=encoder, sentences=test_questions\n",
         ")"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "d3761f56648b"
@@ -751,27 +748,28 @@
       "source": [
         "question_index = 0\n",
         "\n",
-        "print(f\"Query question = {questions[question_index]}\")\n",
+        "print(f\"Query question = {test_questions[question_index]}\")\n",
         "scores = np.dot(question_embeddings[question_index], question_embeddings.T)\n",
         "\n",
         "# Print top 20 matches\n",
         "for index, (question, score) in enumerate(\n",
-        "    sorted(zip(questions, scores), key=lambda x: x[1], reverse=True)[:20]\n",
+        "    sorted(zip(test_questions, scores), key=lambda x: x[1], reverse=True)[:20]\n",
         "):\n",
         "    print(f\"\\t{index}: {question}: {score}\")"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "aQIQSyF9GtSv"
       },
       "source": [
-        "#### Save the train split in JSONL format.\n",
+        "#### Save the embeddings in JSONL format\n",
         "\n",
-        "The data must be formatted in JSONL format, which means each embedding dictionary is written as a JSON string on its own line.\n",
+        "The data must be formatted in JSONL format, which means each embedding dictionary is written as an individual JSON object on its own line.\n",
         "\n",
-        "See more information in the docs at [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup#input-data-format)."
+        "See more information in the docs at [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure#data-file-formats)."
       ]
     },
     {
@@ -825,6 +823,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "QuVl8DrWG8NS"
@@ -847,6 +846,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mglUPwHpJH98"
@@ -856,6 +856,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "qhIBCQ7dDSbW"
@@ -876,17 +877,6 @@
         "DESCRIPTION = \"questions from stackoverflow\""
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "svLYiDf0OD2G"
-      },
-      "source": [
-        "Create the ANN index configuration:\n",
-        "\n",
-        "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup#input-data-format).\n"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": 40,
@@ -900,6 +890,18 @@
         "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "42b843f8c1ae"
+      },
+      "source": [
+        "#### Create the index configuration\n",
+        "\n",
+        "For information on configuration settings, see the [Manage Indexes documentation](https://cloud.google.com/vertex-ai/docs/matching-engine/create-manage-index)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -933,6 +935,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0f1a9fbecabb"
@@ -953,29 +956,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "qV2xjAnDDObD"
       },
       "source": [
-        "## Create an IndexEndpoint with VPC Network"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "BpZQoJyxDlbO"
-      },
-      "outputs": [],
-      "source": [
-        "# Retrieve the project number\n",
-        "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
-        "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
-        "\n",
-        "VPC_NETWORK = \"[your-network-name]\"\n",
-        "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
-        "VPC_NETWORK_FULL"
+        "## Create an IndexEndpoint"
       ]
     },
     {
@@ -989,11 +976,12 @@
         "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
         "    display_name=DISPLAY_NAME,\n",
         "    description=DISPLAY_NAME,\n",
-        "    network=VPC_NETWORK_FULL,\n",
+        "    public_endpoint_enabled=True,\n",
         ")"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "np2cgVuuIe9k"
@@ -1003,6 +991,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8Ew1UgcIIiJG"
@@ -1075,7 +1064,7 @@
         "# Test query\n",
         "NUM_NEIGHBOURS = 20\n",
         "\n",
-        "response = my_index_endpoint.match(\n",
+        "response = my_index_endpoint.find_neighbors(\n",
         "    deployed_index_id=DEPLOYED_INDEX_ID,\n",
         "    queries=[test_embeddings.tolist()],\n",
         "    num_neighbors=NUM_NEIGHBOURS,\n",
@@ -1085,6 +1074,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ce2cf0297369"
@@ -1141,7 +1131,7 @@
         "REDIS_INSTANCE_NAME = \"stackoverflow-questions-unique\"\n",
         "\n",
         "# Create a Redis instance\n",
-        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=5 --region={REGION} --network={VPC_NETWORK_FULL} --connect-mode=private-service-access"
+        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=5 --region={REGION}"
       ]
     },
     {
@@ -1177,19 +1167,33 @@
         "# Connect to the instance\n",
         "import redis\n",
         "\n",
-        "redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT)"
+        "redis_client = redis.StrictRedis(\n",
+        "    host=REDIS_HOST,\n",
+        "    port=REDIS_PORT,\n",
+        "    socket_timeout=60 * 60 * 24,\n",
+        "    socket_connect_timeout=60 * 60 * 24,\n",
+        ")"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "f000f5432d13"
+        "id": "f68fe5d5b5f5"
       },
       "outputs": [],
       "source": [
-        "# Convert the id -> title relationship into a dict and write to redis\n",
-        "redis_client.mset({str(id): str(title) for id, title in zip(df.id, df.title)})"
+        "import os\n",
+        "\n",
+        "# Create a pipeline\n",
+        "pipe = redis_client.pipeline()\n",
+        "\n",
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Convert the id -> title relationship into a dict and write to redis\n",
+        "    pipe.mset({str(id): str(title) for id, title in zip(df.id, df.title)})\n",
+        "\n",
+        "    # Execute the pipeline\n",
+        "    pipe.execute()"
       ]
     },
     {
@@ -1200,14 +1204,16 @@
       },
       "outputs": [],
       "source": [
-        "# Verify that redis can retrieve the correct information\n",
-        "[\n",
-        "    f\"Actual = {title}, Retrieved = {redis_client.get(str(id))}\"\n",
-        "    for id, title in list(zip(df.id, df.title))[:10]\n",
-        "]"
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Verify that redis can retrieve the correct information\n",
+        "    [\n",
+        "        f\"Actual = {title}, Retrieved = {redis_client.get(str(id))}\"\n",
+        "        for id, title in list(zip(df.id, df.title))[:10]\n",
+        "    ]"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TpV-iwP9qw9c"
@@ -1229,31 +1235,18 @@
       "outputs": [],
       "source": [
         "# Force undeployment of indexes and delete endpoint\n",
-        "my_index_endpoint.delete(force=True)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "omj7N9iWv-Tq"
-      },
-      "outputs": [],
-      "source": [
+        "my_index_endpoint.delete(force=True)\n",
+        "\n",
         "# Delete indexes\n",
-        "tree_ah_index.delete()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d2fcf9468031"
-      },
-      "outputs": [],
-      "source": [
+        "tree_ah_index.delete()\n",
+        "\n",
         "# Delete redis instance\n",
-        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {REGION} --quiet"
+        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {REGION} --quiet\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings_vertex.ipynb

@@ -95,6 +95,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -188,7 +189,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
       "metadata": {
         "id": "80c0215f05a0"
       },
@@ -214,7 +215,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": null,
       "metadata": {
         "id": "474be5183c27"
       },
@@ -224,6 +225,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "949271bfebe3"
@@ -235,6 +237,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "b65b4ce80d9a"
@@ -245,6 +248,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "985cdbfe7372"
@@ -255,7 +259,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": null,
       "metadata": {
         "id": "fbc9cd30cc4b"
       },
@@ -265,6 +269,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "79efab26ad02"
@@ -275,7 +280,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": null,
       "metadata": {
         "id": "a336a05c6149"
       },
@@ -286,6 +291,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0c0a44fa330f"
@@ -296,6 +302,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "d3uj8x73nDX_"
@@ -306,30 +313,6 @@
     },
     {
       "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "hhq5zEbGg0XX"
-      },
-      "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 8,
-      "metadata": {
-        "id": "EzrelQZ22IZj"
-      },
-      "outputs": [],
-      "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
-        "\n",
-        "# app = IPython.Application.instance()\n",
-        "# app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
       "cell_type": "markdown",
       "metadata": {
         "id": "zgPO1eR3CYjk"
@@ -342,7 +325,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 27,
+      "execution_count": null,
       "metadata": {
         "id": "MzGDU7TWdts_"
       },
@@ -352,6 +335,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "-EcIXiGsCePi"
@@ -362,7 +346,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": null,
       "metadata": {
         "id": "NIq7R4HZCfIc"
       },
@@ -389,7 +373,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": null,
       "metadata": {
         "id": "ed1b3f87c475"
       },
@@ -423,7 +407,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 176,
+      "execution_count": null,
       "metadata": {
         "id": "b43937b6065d"
       },
@@ -551,6 +535,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "43088937e820"
@@ -561,6 +546,39 @@
         "Define a function to be used later that will take sentences and convert them to embeddings."
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3fb4ebbb4afa"
+      },
+      "source": [
+        "First, initialize aiplatform."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "26e0621c87dd"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d68e62511440"
+      },
+      "source": [
+        "Define the embedding function."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 4,
@@ -678,16 +696,17 @@
       "outputs": [],
       "source": [
         "# Encode a subset of questions for validation\n",
-        "questions = df.title.tolist()[:500]\n",
+        "test_questions = df.title.tolist()[:500]\n",
         "is_successful, question_embeddings = encode_text_to_embedding_batched(\n",
-        "    sentences=df.title.tolist()[:500]\n",
+        "    sentences=test_questions\n",
         ")\n",
         "\n",
         "# Filter for successfully embedded sentences\n",
-        "questions = np.array(questions)[is_successful]"
+        "test_questions = np.array(test_questions)[is_successful]"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "d3761f56648b"
@@ -776,14 +795,14 @@
         "\n",
         "question_index = random.randint(0, 99)\n",
         "\n",
-        "print(f\"Query question = {questions[question_index]}\")\n",
+        "print(f\"Query question = {test_questions[question_index]}\")\n",
         "\n",
         "# Get similarity scores for each embedding by using dot-product.\n",
         "scores = np.dot(question_embeddings[question_index], question_embeddings.T)\n",
         "\n",
         "# Print top 20 matches\n",
         "for index, (question, score) in enumerate(\n",
-        "    sorted(zip(questions, scores), key=lambda x: x[1], reverse=True)[:20]\n",
+        "    sorted(zip(test_questions, scores), key=lambda x: x[1], reverse=True)[:20]\n",
         "):\n",
         "    print(f\"\\t{index}: {question}: {score}\")"
       ]
@@ -903,6 +922,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "QuVl8DrWG8NS"
@@ -924,6 +944,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mglUPwHpJH98"
@@ -932,15 +953,6 @@
         "## Create Indexes\n"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "qhIBCQ7dDSbW"
-      },
-      "source": [
-        "### Create ANN Index (for Production Usage)"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": 231,
@@ -953,16 +965,6 @@
         "DESCRIPTION = \"question titles and bodies from stackoverflow\""
       ]
     },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "svLYiDf0OD2G"
-      },
-      "source": [
-        "Create the index configuration\n"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": 232,
@@ -976,6 +978,18 @@
         "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2c6afd4d3e1"
+      },
+      "source": [
+        "#### Create the index configuration\n",
+        "\n",
+        "For information on configuration settings, see the [Manage Indexes documentation](https://cloud.google.com/vertex-ai/docs/matching-engine/create-manage-index)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 6,
@@ -1011,6 +1025,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0f1a9fbecabb"
@@ -1056,6 +1071,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "np2cgVuuIe9k"
@@ -1065,6 +1081,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8Ew1UgcIIiJG"
@@ -1113,6 +1130,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3cbfe4fd103a"
@@ -1276,10 +1294,10 @@
       },
       "outputs": [],
       "source": [
-        "REDIS_INSTANCE_NAME = \"stackoverflow-questions-vertex\"\n",
+        "REDIS_INSTANCE_NAME = \"stackoverflow-questions-vertex-unique\"\n",
         "\n",
         "# Create a Redis instance\n",
-        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=10 --region='{REGION}' --connect-mode=private-service-access"
+        "! gcloud redis instances create '{REDIS_INSTANCE_NAME}' --size=10 --region='{REGION}'"
       ]
     },
     {
@@ -1315,7 +1333,10 @@
         "# Connect to the instance\n",
         "import redis\n",
         "\n",
-        "redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT)"
+        "redis_client = redis.StrictRedis(\n",
+        "    host=REDIS_HOST,\n",
+        "    port=REDIS_PORT,\n",
+        ")"
       ]
     },
     {
@@ -1326,35 +1347,39 @@
       },
       "outputs": [],
       "source": [
-        "%%time\n",
-        "# Convert the id -> (title, body) relationship into a dict and write to Redis\n",
-        "for df in tqdm(\n",
-        "    query_bigquery_chunks(\n",
-        "        max_rows=BQ_NUM_ROWS, rows_per_chunk=BQ_CHUNK_SIZE, start_chunk=0\n",
-        "    ),\n",
-        "    total=BQ_NUM_CHUNKS,\n",
-        "    position=0,\n",
-        "    desc=\"Chunk of rows from BigQuery\",\n",
-        "):\n",
-        "    ids = df.id.tolist()\n",
-        "    titles = df.title.tolist()\n",
-        "    bodies = df.body.tolist()\n",
+        "import os\n",
         "\n",
-        "    # create a Redis pipeline\n",
-        "    pipe = redis_client.pipeline()\n",
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Convert the id -> (title, body) relationship into a dict and write to Redis\n",
+        "    for df in tqdm(\n",
+        "        query_bigquery_chunks(\n",
+        "            max_rows=BQ_NUM_ROWS, rows_per_chunk=BQ_CHUNK_SIZE, start_chunk=0\n",
+        "        ),\n",
+        "        total=BQ_NUM_CHUNKS,\n",
+        "        position=0,\n",
+        "        desc=\"Chunk of rows from BigQuery\",\n",
+        "    ):\n",
+        "        ids = df.id.tolist()\n",
+        "        titles = df.title.tolist()\n",
+        "        bodies = df.body.tolist()\n",
         "\n",
-        "    # iterate over the data and add hset commands to the pipeline\n",
-        "    for (id, title, body) in tqdm(zip(ids, titles, bodies), total=len(ids), position=1):\n",
-        "        pipe.hset(\n",
-        "            str(id),\n",
-        "            mapping={\n",
-        "                \"title\": str(title),\n",
-        "                \"body\": str(body[:100]),\n",
-        "            },\n",
-        "        )\n",
+        "        # create a Redis pipeline\n",
+        "        pipe = redis_client.pipeline()\n",
         "\n",
-        "    # execute the pipeline\n",
-        "    _ = pipe.execute()"
+        "        # iterate over the data and add hset commands to the pipeline\n",
+        "        for (id, title, body) in tqdm(\n",
+        "            zip(ids, titles, bodies), total=len(ids), position=1\n",
+        "        ):\n",
+        "            pipe.hset(\n",
+        "                str(id),\n",
+        "                mapping={\n",
+        "                    \"title\": str(title),\n",
+        "                    \"body\": str(body[:100]),\n",
+        "                },\n",
+        "            )\n",
+        "\n",
+        "        # execute the pipeline\n",
+        "        _ = pipe.execute()"
       ]
     },
     {
@@ -1365,16 +1390,18 @@
       },
       "outputs": [],
       "source": [
-        "# Verify that Redis can retrieve the correct information\n",
-        "df = next(query_bigquery_chunks(max_rows=10, rows_per_chunk=10))\n",
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Verify that Redis can retrieve the correct information\n",
+        "    df = next(query_bigquery_chunks(max_rows=10, rows_per_chunk=10))\n",
         "\n",
-        "[\n",
-        "    f\"Actual = {title}, Retrieved = {redis_client.hgetall(str(id))}\"\n",
-        "    for id, title in zip(df.id, df.title)\n",
-        "]"
+        "    [\n",
+        "        f\"Actual = {title}, Retrieved = {redis_client.hgetall(str(id))}\"\n",
+        "        for id, title in zip(df.id, df.title)\n",
+        "    ]"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TpV-iwP9qw9c"
@@ -1396,31 +1423,18 @@
       "outputs": [],
       "source": [
         "# Force undeployment of indexes and delete endpoint\n",
-        "my_index_endpoint.delete(force=True)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "omj7N9iWv-Tq"
-      },
-      "outputs": [],
-      "source": [
+        "my_index_endpoint.delete(force=True)\n",
+        "\n",
         "# Delete indexes\n",
-        "tree_ah_index.delete()"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d2fcf9468031"
-      },
-      "outputs": [],
-      "source": [
+        "tree_ah_index.delete()\n",
+        "\n",
         "# Delete redis instance\n",
-        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {REGION} --quiet"
+        "! gcloud redis instances delete '{REDIS_INSTANCE_NAME}' --region {REGION} --quiet\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
       ]
     }
   ],

notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb

@@ -173,7 +173,7 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "PROJECT_ID = \"[YOUR-PROJECT-ID]\"  # @param {type:\"string\"}\n",
         "\n",
         "# Set the project id\n",
         "! gcloud config set project {PROJECT_ID}"
@@ -293,7 +293,7 @@
       },
       "outputs": [],
       "source": [
-        "VPC_NETWORK = \"[your-vpc-network-name]\"  # @param {type:\"string\"}\n",
+        "VPC_NETWORK = \"[your-network-name]\"  # @param {type:\"string\"}\n",
         "\n",
         "PEERING_RANGE_NAME = \"ann-haystack-range\""
       ]
@@ -612,14 +612,15 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "svLYiDf0OD2G"
       },
       "source": [
-        "Create the ANN index configuration:\n",
+        "#### Create the index configuration\n",
         "\n",
-        "To learn more about configuring the index, see [Input data format and structure](https://cloud.google.com/vertex-ai/docs/matching-engine/match-eng-setup/format-structure).\n"
+        "For information on configuration settings, see the [Manage Indexes documentation](https://cloud.google.com/vertex-ai/docs/matching-engine/create-manage-index)"
       ]
     },
     {
@@ -870,7 +871,7 @@
         "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
         "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
         "\n",
-        "VPC_NETWORK = \"[your-network-name]\"\n",
+        "VPC_NETWORK = \"None\"\n",
         "VPC_NETWORK_FULL = \"projects/{}/global/networks/{}\".format(PROJECT_NUMBER, VPC_NETWORK)\n",
         "VPC_NETWORK_FULL"
       ]